Merge pull request #51 from HERA-Team/avoid-build-solver

Add option to save memory when using LinProductSolver infrastructure but not its solvers

Author: tyler-a-cox

.gitignore

@@ -8,3 +8,4 @@ dist/*
 coverage.xml
 *.DS_Store
 */_version.py
+src/linsolve/_version.py

src/linsolve/linsolve.py

@@ -856,7 +856,7 @@ def taylor_expand(terms, consts=None, prepend="d"):
 # XXX make a version of linproductsolver that taylor expands in e^{a+bi} form
 # see https://github.com/HERA-Team/linsolve/issues/15
 class LinProductSolver:
-    def __init__(self, data, sol0, wgts={}, sparse=False, **kwargs):
+    def __init__(self, data, sol0, wgts={}, sparse=False, build_solver=True, **kwargs):
         """Set up a nonlinear system of equations of the form a*b + c*d = 1.0.
 
         Linearize via Taylor expansion and solve iteratively using the Gauss-Newton
@@ -882,6 +882,9 @@ def __init__(self, data, sol0, wgts={}, sparse=False, **kwargs):
         sparse : bool
             If True, represents A matrix sparsely (though AtA, Aty end up dense)
             May be faster for certain systems of equations.
+        build_solver : bool
+            Advanced users can turn this off to save memory when using only LinProductSolver
+            infrastructure but not solve() or solve_iteratively(), as in omnical.
         **kwargs: keyword arguments of constants (python variables in keys of data that
             are not to be solved for)
         """
@@ -894,8 +897,14 @@ def __init__(self, data, sol0, wgts={}, sparse=False, **kwargs):
         self.init_kwargs, self.sols_kwargs = constants, deepcopy(constants)
         self.sols_kwargs.update(sol0)
         self.all_terms, self.taylors, self.taylor_keys = self.gen_taylors()
-        self.build_solver(sol0)
-        self.dtype = self.ls.dtype
+        if build_solver:
+            self.build_solver(sol0)
+            self.dtype = self.ls.dtype
+        else:
+            self.sol0 = sol0
+            self.dtype = infer_dtype(list(self.data.values())
+                                     + list(self.sol0.values())
+                                     + list(self.wgts.values()))
 
     def gen_taylors(self, keys=None):
         """Perform Taylor expansion, and map eq. keys to taylor expansion keys."""

tests/test_linsolve.py

@@ -490,6 +490,10 @@ def test_init(self):
             np.testing.assert_almost_equal(eval(k), 0.002)
         assert len(ls.ls.prms) == 3
 
+        ls = linsolve.LinProductSolver(d, sol0, w, sparse=self.sparse, build_solver=False)
+        assert not hasattr(ls, "ls")
+        assert ls.dtype == np.complex64
+
     def test_real_solve(self):
         x, y, z = 1.0, 2.0, 3.0
         keys = ["x*y", "x*z", "y*z"]