TL: fixed tests

Author: tlunet

blockops/taskPool.py

@@ -1,7 +1,6 @@
 import sympy as sy
 import re
 import warnings
-import time
 
 from blockops.run import PintRun
 
@@ -223,21 +222,22 @@ def addTask(self, ope, inp, dep, n, k, result=None):
         res : Expression
             The result of this task (full expression).
         """
-        tmpRes = ope * inp
-        notZero = type(tmpRes) != sy.core.numbers.Zero
+        tmpRes = sy.simplify(ope * inp)
 
-        if tmpRes in self.results and notZero:
+        # Check if result is already computed
+        notZero = type(tmpRes) != sy.core.numbers.Zero
+        if notZero and tmpRes in self.results:
             task = self.results[tmpRes]
             return task, tmpRes
 
-        check2 = -ope * inp
-        if check2 in self.results and notZero:
-            task = self.results[check2]
+        # Check if negative result is already computed
+        neg = ope * sy.expand(-inp, deep=False)
+        if notZero and neg in self.results:
+            task = self.results[neg]
             return task, tmpRes
-
-        check3 = ope * -inp
-        if check3 in self.results and notZero:
-            task = self.results[check3]
+        neg = sy.expand(-ope, deep=False) * inp
+        if notZero and neg in self.results:
+            task = self.results[neg]
             return task, tmpRes
 
         # Task not in pool, create and add it

blockops/tests/test_blockMethods.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from .data.methodData import results
+from blockops.tests.data.methodData import results
 
 from blockops.run import PintRun
 from blockops.block import BlockOperator, I
@@ -25,17 +25,16 @@
 
 
 def checkResults(method, run, pool):
-    i = 0
-    for key, value in run.facBlockRules.items():
+    for i, (key, value) in enumerate(run.facBlockRules.items()):
         assert results[method]['blockRules'][i][0] == str(value["result"])
         assert results[method]['blockRules'][i][1] == str(value["rule"])
-        i = i + 1
-
-    i = 0
-    for key, value in pool.pool.items():
-        assert results[method]['taskPool'][i][0] == str(key)
-        assert str(value.fullOP) in results[method]['taskPool'][i][1]
-        i = i + 1
+    for i, (key, value) in enumerate(pool.pool.items()):
+        refKey, refVal = results[method]['taskPool'][i]
+        assert refKey == str(key)
+        if isinstance(refVal, list):
+            assert str(value.fullOP) in refVal
+        else:
+            assert str(value.fullOP) == refVal
 
 
 class TestMethods:

blockops/utils/params.py

@@ -123,20 +123,17 @@ def extractParamDocs(cls, *names):
 
         if docs is None:
             raise ValueError(f'undocumented class {cls}')
-
         for name in names:
-            iStart = docs.find(f'\n    {name} :')
+            iStart = docs.find(f'{name} :')
             if iStart == -1:
-                iStart = docs.find(f'\n    **{name} :')
-                if iStart == -1:
-                    raise ValueError(f'{name} parameter not in {cls} docs')
-            docLines = docs[iStart:].splitlines()[2:]
+                raise ValueError(f'{name} parameter not in {cls} docs')
+            docLines = docs[iStart:].splitlines()[1:]
             descr = []
             for line in docLines:
-                if line.startswith(8 * ' '):
-                    descr.append(line.strip())
-                elif line.strip() == '':
+                if line.strip() == "":
                     continue
+                elif line.startswith(" "):
+                    descr.append(line.strip())
                 else:
                     break
             if len(descr) == 0:

blockops/utils/vectorize.py

@@ -34,15 +34,15 @@ def matVecMul(mat, u):
     return np.matmul(mat, u[..., None]).squeeze(axis=-1)
 
 
-def matVecInv(mat, u):
+def matVecInv(mat, vec):
     r"""
     Compute vectorized Matrix Vector Inversion :math:`A^{-1}x` (A / x)
 
     Parameters
     ----------
     mat : np.ndarray, size (nDOF, M, M) or (M, M)
         Matrix or array of matrices.
-    u : np.ndarray, size (nDOF, M) or (M,)
+    vec : np.ndarray, size (nDOF, M) or (M,)
         Vector or array of vectors.
 
     Returns
@@ -57,13 +57,10 @@ def matVecInv(mat, u):
     - 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:
-        try:
-            return np.linalg.solve(mat[None, ...], u)
-        except ValueError:
-            return np.linalg.solve(mat, u[None, ...])
+    mat, vec = np.asarray(mat), np.asarray(vec)
+    if mat.ndim > 2 and vec.ndim > 1:
+        assert mat.shape[0] == vec.shape[0], "different nDOF for mat and vec"
+    return np.linalg.solve(mat, vec[..., None]).squeeze()
 
 
 def matMatMul(m1, m2):