From 2979b787bdc8879a0bf105419d68786c17c5b17d Mon Sep 17 00:00:00 2001
From: Tschijnmo TSCHAU <tschijnmotschau@gmail.com>
Date: Sun, 19 Feb 2017 00:09:39 -0600
Subject: [PATCH] Add test for vector substitutions

---
 tests/free_algebra_test.py | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/tests/free_algebra_test.py b/tests/free_algebra_test.py
index bc810b3..2b6a805 100644
--- a/tests/free_algebra_test.py
+++ b/tests/free_algebra_test.py
@@ -317,7 +317,7 @@ def test_tensors_can_be_differentiated(free_alg):
 
 
 def test_tensors_can_be_substituted_scalars(free_alg):
-    """Test vector substitution facility for tensors."""
+    """Test scalar substitution facility for tensors."""
 
     dr = free_alg
     p = dr.names
@@ -341,6 +341,30 @@ def test_tensors_can_be_substituted_scalars(free_alg):
         assert res.simplify() == expected.simplify()
 
 
+def test_tensors_can_be_substituted_vectors(free_alg):
+    """Test vector substitution facility for tensors."""
+
+    dr = free_alg
+    p = dr.names
+
+    x = IndexedBase('x')
+    t = IndexedBase('t')
+    u = IndexedBase('u')
+    i, j = p.i, p.j
+    v = p.v
+    w = Vec('w')
+
+    orig = dr.einst(x[i] * v[i])
+    v_def = dr.einst(t[i, j] * w[j] + u[i, j] * w[j])
+
+    res = orig.subst(v[i], v_def).simplify()
+
+    expected = dr.einst(
+        x[i] * t[i, j] * w[j] + x[i] * u[i, j] * w[j]
+    ).simplify()
+    assert res == expected
+
+
 def test_tensor_method(free_alg):
     """Test tensor method can be injected."""