feat: another ns implementation to accelerate A^-1b

Benjamin-eecs · Benjamin-eecs · commit 180f62f2a5cb · 2022-10-13T16:58:14.000+08:00
diff --git a/torchopt/linalg.py b/torchopt/linalg.py
@@ -218,6 +218,45 @@ def cg(
 
 
 def ns(
+    A: Union[torch.Tensor, Callable[[TensorTree], TensorTree]],
+    b: TensorTree,
+    maxiter: Optional[int] = None,
+    *,
+    alpha: Optional[float] = None,
+) -> TensorTree:
+    """Use Neumann Series Inverse Matric Approximation to solve ``Ax = b``.
+
+    Args:
+        A: (tensor or tree of tensors or function)
+            2D array or function that calculates the linear map (matrix-vector product) ``Ax`` when
+            called like ``A(x)``. ``A`` must represent a hermitian, positive definite matrix, and
+            must return array(s) with the same structure and shape as its argument.
+        b: (tensor or tree of tensors)
+            Right hand side of the linear system representing a single vector. Can be stored as an
+            array or Python container of array(s) with any shape.
+        maxiter: (integer, optional)
+            Maximum number of iterations. Iteration will stop after maxiter steps even if the
+            specified tolerance has not been achieved.
+        alpha: (float, optional)
+            Decay coefficient.
+
+    Returns:
+        the Neumann Series (NS) matrix inversion approximation
+    """
+    if maxiter is None:
+        size = sum(_shapes(b))
+        maxiter = 10 * size
+    inv_A_hat_b = b
+    for rank in range(maxiter):
+        if alpha is not None:
+            b = pytree.tree_sub(b, pytree.tree_mul(alpha, A(b)))
+        else:
+            b = pytree.tree_sub(b, A(b))
+        inv_A_hat_b = pytree.tree_sub(inv_A_hat_b, b)
+    return inv_A_hat_b
+
+
+def ns_inv(
     A: TensorTree,
     n: int,
     maxiter: Optional[int] = None,
@@ -229,7 +268,7 @@ def ns(
     Args:
         A: (tensor or tree of tensors or function)
             2D array or function that calculates the linear map (matrix-vector product) ``Ax`` when
-            called like ``A(x)``. ``A`` must represent a Hermitian, positive definite matrix, and
+            called like ``A(x)``. ``A`` must represent a hermitian, positive definite matrix, and
             must return array(s) with the same structure and shape as its argument.
         n: (integer)
             Number of rows (and columns) in n x n output.
diff --git a/torchopt/pytree.py b/torchopt/pytree.py
@@ -13,6 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 """The PyTree utilities."""
+import functools
+import itertools
+import operator
 
 import optree
 import optree.typing as typing  # pylint: disable=unused-import
@@ -22,7 +25,21 @@
 from torchopt.typing import Future, PyTree, RRef, T
 
 
-__all__ = [*optree.__all__, 'tree_wait']
+__all__ = [*optree.__all__, 'tree_wait', 'tree_add', 'tree_sub', 'tree_mul', 'tree_div']
+
+tree_map = optree.tree_map
+
+tree_add = functools.partial(tree_map, operator.add)
+tree_add.__doc__ = "Tree addition."
+
+tree_sub = functools.partial(tree_map, operator.sub)
+tree_sub.__doc__ = "Tree subtraction."
+
+tree_mul = functools.partial(tree_map, operator.mul)
+tree_mul.__doc__ = "Tree multiplication."
+
+tree_div = functools.partial(tree_map, operator.truediv)
+tree_div.__doc__ = "Tree division."
 
 
 def tree_wait(future_tree: PyTree[Future[T]]) -> PyTree[T]:
