TL: extended vectorized functions

Parallel-in-Time · Oct 20, 2023 · c71c186 · c71c186
1 parent 6048c20
commit c71c186
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 15 deletions.
diff --git a/blockops/schemes/__init__.py b/blockops/schemes/__init__.py
@@ -24,7 +24,7 @@ def getTransferMatrices(nodesFine, nodesCoarse, vectorized=False):
 
 @setParams(
     nPoints=PositiveInteger(latexName='$M$'),
-    ptsType=MultipleChoices(*NODE_TYPES, latexName="Point Distribution"), 
+    ptsType=MultipleChoices(*NODE_TYPES, latexName="Point Distribution"),
     quadType=MultipleChoices(*QUAD_TYPES, latexName="Quadrature Type"),
     form=MultipleChoices('Z2N', 'N2N')
 )
@@ -115,7 +115,7 @@ def getBlockOperators(self, lamDt, phiName, chiName) -> [BlockOperator, BlockOpe
             The value of :math:`\lambda\Delta{T}` for the block.
         phiName : str
             The symbol name for the :math:`\phi` operator.
-        chiName : Tstr
+        chiName : str
             The symbol name for the :math:`\chi` operator.
 
         Returns

diff --git a/blockops/tests/test_vectorize.py b/blockops/tests/test_vectorize.py
@@ -23,37 +23,47 @@ def testMatVecMul():
     """Test vectorized Matrix Vector Multiplication (matVecMul)"""
     mat, u = generate(M, nDOF)
 
+    # -- matVecMul for (nDOF, M, M), (nDOF, M)
     out = matVecMul(mat, u)
-
     for i in range(nDOF):
         assert np.allclose(out[i], mat[i] @ u[i])
 
+    # -- matVecMul for (M, M), (nDOF, M)
     out = matVecMul(mat[0], u)
-
     for i in range(nDOF):
         assert np.allclose(out[i], mat[0] @ u[i])
 
-    out = matVecMul(mat[0], u[0])
+    # -- matVecMul for (nDOF, M, M), (M)
+    out = matVecMul(mat, u[0])
+    for i in range(nDOF):
+        assert np.allclose(out[i], mat[i] @ u[0])
 
+    # -- matVecMul for (M, M), (M,)
+    out = matVecMul(mat[0], u[0])
     assert np.allclose(out, mat[0] @ u[0])
 
 
 def testMatVecInv():
     """Test vectorized Matrix Vector Inversion (matVecInv)"""
     mat, u = generate(M, nDOF)
 
+    # -- matVecInv for (nDOF, M, M), (nDOF, M)
     out = matVecInv(mat, u)
-
     for i in range(nDOF):
         assert np.allclose(out[i], np.linalg.solve(mat[i], u[i]))
 
+    # -- matVecInv for (M, M), (nDOF, M)
     out = matVecInv(mat[0], u)
-
     for i in range(nDOF):
         assert np.allclose(out[i], np.linalg.solve(mat[0], u[i]))
 
-    out = matVecInv(mat[0], u[0])
+    # -- matVecInv for (nDOF, M, M), (M,)
+    out = matVecInv(mat, u[0])
+    for i in range(nDOF):
+        assert np.allclose(out[i], np.linalg.solve(mat[i], u[0]))
 
+    # -- matVecInv for (M, M), (M,)
+    out = matVecInv(mat[0], u[0])
     assert np.allclose(out, np.linalg.solve(mat[0], u[0]))
 
 

diff --git a/blockops/utils/vectorize.py b/blockops/utils/vectorize.py
@@ -26,9 +26,10 @@ def matVecMul(mat, u):
 
     Notes
     -----
-    - matVecMul((nDOF, M, M), (nDOF, M)) -> (nDOF, M) <=> (M, M) @ (M,) for each nDOF
-    - matVecMul((M, M), (nDOF, M)) -> (nDOF, M) <=> (M, M) @ (M,) for each nDOF
-    - matVecMul((M, M), (M,)) -> (M,) <=> (M, M) @ (M,)
+    - matVecMul for (nDOF, M, M), (nDOF, M) -> (nDOF, M) <=> (M, M) @ (M,) for each nDOF
+    - matVecMul for (M, M), (nDOF, M) -> (nDOF, M) <=> (M, M) @ (M,) for each nDOF
+    - matVecMul for (nDOF, M, M), (M,) -> (nDOF, M) <=> (M, M) @ (M,) for each nDOF
+    - matVecMul for (M, M), (M,) -> (M,) <=> (M, M) @ (M,)
     """
     return np.matmul(mat, u[..., None]).squeeze(axis=-1)
 
@@ -51,14 +52,18 @@ def matVecInv(mat, u):
 
     Notes
     -----
-    - matVecInv((nDOF, M, M), (nDOF, M)) -> (nDOF, M) <=> (M, M) \ (M,) for each nDOF
-    - matVecInv((M, M), (nDOF, M)) -> (nDOF, M) <=> (M, M) \ (M,) for each nDOF
-    - matVecInv((M, M), (M,)) -> (M,) <=> (M, M) \ (M,)
+    - matVecInv for (nDOF, M, M), (nDOF, M) -> (nDOF, M) <=> (M, M) \ (M,) for each nDOF
+    - matVecInv for (M, M), (nDOF, M) -> (nDOF, M) <=> (M, M) \ (M,) for each nDOF
+    - matVecInv for (nDOF, M, M), (M,) -> (nDOF, M) <=> (M, M) \ (M,) for each nDOF
+    - matVecInv for (M, M), (M,)) -> (M,) <=> (M, M) \ (M,)
     """
     try:
         return np.linalg.solve(mat, u)
     except ValueError:
-        return np.linalg.solve(mat[None, ...], u)
+        try:
+            return np.linalg.solve(mat[None, ...], u)
+        except ValueError:
+            return np.linalg.solve(mat, u[None, ...])
 
 
 def matMatMul(m1, m2):