fix: adding presolve and postsolve calls for interpolator

Author: lachlangrose

LoopStructural/interpolators/_discrete_interpolator.py

@@ -132,6 +132,28 @@ def set_interpolation_weights(self, weights):
             self.up_to_date = False
             self.interpolation_weights[key] = weights[key]
 
+    def _pre_solve(self):
+        """
+        Pre solve function to be run before solving the interpolation
+        """
+        self.c = np.zeros(self.support.n_nodes)
+        self.c[:] = np.nan
+        return True
+
+    def _post_solve(self):
+        """Post solve function(s) to be run after the solver has been called"""
+        self.clear_constraints()
+        return True
+
+    def clear_constraints(self):
+        """
+        Clear the constraints from the interpolator, this makes sure we are not storing
+        the constraints after the solver has been run
+        """
+        self.constraints = {}
+        self.ineq_constraints = {}
+        self.equal_constraints = {}
+
     def reset(self):
         """
         Reset the interpolation constraints
@@ -540,43 +562,37 @@ def solve_system(
 
         """
         starttime = time()
-        self.c = np.zeros(self.support.n_nodes)
-        self.c[:] = np.nan
+        if not self._pre_solve():
+            raise ValueError("Pre solve failed")
+
         A, b = self.build_matrix()
         Q, bounds = self.build_inequality_matrix()
         if callable(solver):
             logger.warning('Using custom solver')
             self.c = solver(A.tocsr(), b)
             self.up_to_date = True
-
-            return True
-        ## solve with lsmr
-        if isinstance(solver, str):
+        elif isinstance(solver, str) or solver is None:
             if solver not in ['cg', 'lsmr', 'admm']:
                 logger.warning(
                     f'Unknown solver {solver} using cg. \n Available solvers are cg and lsmr or a custom solver as a callable function'
                 )
                 solver = 'cg'
-        if solver is None:
-            solver = 'cg'
         if solver == 'cg':
             logger.info("Solving using cg")
             if 'atol' not in solver_kwargs or 'rtol' not in solver_kwargs:
                 if tol is not None:
                     solver_kwargs['atol'] = tol
 
-            ATA = A.T @ A
-            ATB = A.T @ b
             logger.info(f"Solver kwargs: {solver_kwargs}")
 
-            res = sparse.linalg.cg(ATA, ATB, **solver_kwargs)
+            res = sparse.linalg.cg(A.T @ A, A.T @ b, **solver_kwargs)
             if res[1] > 0:
                 logger.warning(
                     f'CG reached iteration limit ({res[1]})and did not converge, check input data. Setting solution to last iteration'
                 )
             self.c = res[0]
             self.up_to_date = True
-            return True
+
         elif solver == 'lsmr':
             logger.info("Solving using lsmr")
             if 'atol' not in solver_kwargs:
@@ -600,8 +616,7 @@ def solve_system(
                 )
                 self.c = res[0]
             self.up_to_date = True
-            logger.info("Interpolation took %f seconds" % (time() - starttime))
-            return True
+
         elif solver == 'admm':
             logger.info("Solving using admm")
 
@@ -616,29 +631,35 @@ def solve_system(
 
             try:
                 from loopsolver import admm_solve
+
+                try:
+                    linsys_solver = solver_kwargs.pop('linsys_solver', 'lsmr')
+                    res = admm_solve(
+                        A,
+                        b,
+                        Q,
+                        bounds,
+                        x0=x0,
+                        admm_weight=solver_kwargs.pop('admm_weight', 0.01),
+                        nmajor=solver_kwargs.pop('nmajor', 200),
+                        linsys_solver_kwargs=solver_kwargs,
+                        linsys_solver=linsys_solver,
+                    )
+                    self.c = res
+                    self.up_to_date = True
+                except ValueError as e:
+                    logger.error(f"ADMM solver failed: {e}")
+                    self.up_to_date = False
             except ImportError:
                 logger.warning(
                     "Cannot import admm solver. Please install loopsolver or use lsmr or cg"
                 )
-                return False
-            try:
-                res = admm_solve(
-                    A,
-                    b,
-                    Q,
-                    bounds,
-                    x0=x0,
-                    admm_weight=solver_kwargs.pop('admm_weight', 0.01),
-                    nmajor=solver_kwargs.pop('nmajor', 200),
-                    linsys_solver_kwargs=solver_kwargs,
-                )
-                self.c = res
-                self.up_to_date = True
-            except ValueError as e:
-                logger.error(f"ADMM solver failed: {e}")
-                return False
-        logger.info(f"{solver} not recognised")
-        return False
+                self.up_to_date = False
+        else:
+            logger.error(f"Unknown solver {solver}")
+            self.up_to_date = False
+        # self._post_solve()
+        return self.up_to_date
 
     def update(self) -> bool:
         """