fix: interpolator type is now an enum

Changed interpolator type to an enum

LoopStructural/interpolators/__init__.py

@@ -2,17 +2,36 @@
 Interpolators and interpolation supports
 
 """
-from LoopStructural.interpolators.discrete_fold_interpolator import (
-    DiscreteFoldInterpolator,
-)
+from enum import IntEnum
+
+class InterpolatorType(IntEnum):
+    '''
+    Enum for the different interpolator types
+    
+    1-9 should cover interpolators with supports
+    9+ are data supported
+    '''
+    BASE = 0
+    BASE_DISCRETE = 1
+    FINITE_DIFFERENCE = 2
+    DISCRETE_FOLD = 3
+    PIECEWISE_LINEAR = 4
+    BASE_DATA_SUPPORTED = 10
+    SURFE = 11
+from LoopStructural.interpolators.geological_interpolator import GeologicalInterpolator
+from LoopStructural.interpolators.discrete_interpolator import DiscreteInterpolator
+from LoopStructural.interpolators.structured_tetra import TetMesh
+from LoopStructural.interpolators.unstructured_tetra import UnStructuredTetMesh
+from LoopStructural.interpolators.structured_grid import StructuredGrid
 from LoopStructural.interpolators.finite_difference_interpolator import (
     FiniteDifferenceInterpolator,
 )
 from LoopStructural.interpolators.piecewiselinear_interpolator import (
     PiecewiseLinearInterpolator,
 )
-from LoopStructural.interpolators.structured_tetra import TetMesh
-from LoopStructural.interpolators.unstructured_tetra import UnStructuredTetMesh
-from LoopStructural.interpolators.structured_grid import StructuredGrid
-from LoopStructural.interpolators.geological_interpolator import GeologicalInterpolator
-from LoopStructural.interpolators.discrete_interpolator import DiscreteInterpolator
+from LoopStructural.interpolators.discrete_fold_interpolator import (
+    DiscreteFoldInterpolator,
+)
+
+
+

LoopStructural/interpolators/discrete_fold_interpolator.py

@@ -6,8 +6,8 @@
 import numpy as np
 from LoopStructural.interpolators.cython.dsi_helper import fold_cg
 
-from LoopStructural.interpolators.piecewiselinear_interpolator import (
-    PiecewiseLinearInterpolator,
+from LoopStructural.interpolators import (
+    PiecewiseLinearInterpolator, InterpolatorType
 )
 
 from LoopStructural.utils import getLogger
@@ -31,7 +31,7 @@ def __init__(self, support, fold):
         """
 
         PiecewiseLinearInterpolator.__init__(self, support)
-        self.type = ["foldinterpolator"]
+        self.type = InterpolatorType.FINITE_DIFFERENCE
         self.fold = fold
 
     @classmethod

LoopStructural/interpolators/discrete_interpolator.py

@@ -8,8 +8,9 @@
 from scipy.sparse import linalg as sla
 from scipy.sparse.linalg import norm
 from sklearn.preprocessing import normalize
+from LoopStructural.interpolators import InterpolatorType
 
-from LoopStructural.interpolators.geological_interpolator import GeologicalInterpolator
+from LoopStructural.interpolators import GeologicalInterpolator
 from LoopStructural.utils import getLogger
 
 logger = getLogger(__name__)
@@ -31,30 +32,36 @@ def __init__(self, support):
         GeologicalInterpolator.__init__(self)
         self.B = []
         self.support = support
-        self.region_function = lambda xyz: np.ones(xyz.shape[0], dtype=int)
+        self.region_function = lambda xyz: np.ones(xyz.shape[0], dtype=bool)
         # self.region_map[self.region] = np.array(range(0,
         # len(self.region_map[self.region])))
         self.shape = "rectangular"
         if self.shape == "square":
             self.B = np.zeros(self.nx)
         self.c_ = 0
-        self.A = []  # sparse matrix storage coo format
-        self.col = []
-        self.row = []  # sparse matrix storage
-        self.w = []
+        # self.A = []  # sparse matrix storage coo format
+        # self.col = []
+        # self.row = []  # sparse matrix storage
+        # self.w = []
         self.solver = None
+
         self.eq_const_C = []
         self.eq_const_row = []
         self.eq_const_col = []
         self.eq_const_d = []
-        self.eq_const_c_ = 0
+
+        self.equal_constraints = {}
+        self.eq_const_c = 0
+        self.ineq_constraints = {}
+        self.ineq_const_c = 0
+
         self.non_linear_constraints = []
         self.constraints = {}
         self.interpolation_weights = {}
         logger.info(
             "Creating discrete interpolator with {} degrees of freedom".format(self.nx)
         )
-        self.type = "discrete"
+        self.type = InterpolatorType.BASE_DISCRETE
 
     @property
     def nx(self):
@@ -130,17 +137,7 @@ def reset(self):
 
         """
         logger.debug("Resetting interpolation constraints")
-        self.c_ = 0
-        self.A = []  # sparse matrix storage coo format
-        self.col = []
-        self.row = []  # sparse matrix storage
-        self.eq_const_C = []
-        self.eq_const_row = []
-        self.eq_const_col = []
-        self.eq_const_d = []
-        self.eq_const_c_ = 0
-        self.B = []
-        self.n_constraints = 0
+
 
     def add_constraints_to_least_squares(self, A, B, idc, w=1.0, name="undefined"):
         """
@@ -255,31 +252,9 @@ def remove_constraints_from_least_squares(
 
         """
 
-        if constraint_ids is None:
-            constraint_ids = self.constraints[name]
-        print(
-            "Removing {} {} constraints from least squares".format(
-                len(constraint_ids), name
-            )
-        )
-        A = np.array(self.A)
-        B = np.array(self.B)
-        col = np.array(self.col)
-        row = np.array(self.row)
-        mask = np.any((row[:, None] == constraint_ids[None, :]) == True, axis=1)
-        # np.any((numbers[:, None] == np.array([0, 10, 30])[None, :]) == True,
-        #        axis=1)
-        bmask = np.ones(B.shape, dtype=bool)
-        bmask[constraint_ids] = 0
-        self.A = A[~mask].tolist()
-        self.B = B[bmask]
-        self.col = col[~mask].tolist()
-        rowmax = np.max(row[mask])
-        rowrange = rowmax - np.min(row[mask])
-        # row[np.logical_and(~mask,row>rowmax)] -= rowrange
-        return row[~mask]
-
-    def add_equality_constraints(self, node_idx, values):
+        pass
+
+    def add_equality_constraints(self, node_idx, values,name="undefined"):
         """
         Adds hard constraints to the least squares system. For now this just
         sets
@@ -302,16 +277,60 @@ def add_equality_constraints(self, node_idx, values):
         gi[self.region] = np.arange(0, self.nx)
         idc = gi[node_idx]
         outside = ~(idc == -1)
-
-        self.eq_const_C.extend(np.ones(idc[outside].shape[0]).tolist())
-        self.eq_const_col.extend(idc[outside].tolist())
-        self.eq_const_row.extend((np.arange(0, idc[outside].shape[0])))
-        self.eq_const_d.extend(values[outside].tolist())
-        self.eq_const_c_ += idc[outside].shape[0]
+        self.equal_constraints[name] = {
+            "A": np.ones(idc[outside].shape[0]).tolist(),
+            "B": values[outside].tolist(),
+            "col": idc[outside].tolist(),
+            # "w": w,
+            "row": np.arange(self.eq_const_c, self.eq_const_c+idc[outside].shape[0])
+
+        }
+        self.eq_const_c += idc[outside].shape[0]
+        # ,'C':np.ones(idc[outside].shape[0]).tolist(),}
+        # self.eq_const_C.extend(np.ones(idc[outside].shape[0]).tolist())
+        # self.eq_const_col.extend(idc[outside].tolist())
+        # self.eq_const_row.extend((np.arange(0, idc[outside].shape[0])))
+        # self.eq_const_d.extend(values[outside].tolist())
+        # self.eq_const_c_ += idc[outside].shape[0]
 
     def add_non_linear_constraints(self, nonlinear_constraint):
         self.non_linear_constraints.append(nonlinear_constraint)
 
+    def add_inequality_constraints_to_matrix(self, A, l, u, idc, name="undefined"):
+        """Adds constraints for a matrix where the linear function
+        l < Ax > u constrains the objective function
+
+
+        Parameters
+        ----------
+        A : numpy array
+            matrix of coefficients
+        l : numpy array
+            lower bounds
+        u : numpy array
+            upper bounds
+        idc : numpy array
+            index of constraints
+        Returns
+        -------
+
+        """
+        # map from mesh node index to region node index
+        gi = np.zeros(self.support.n_nodes)
+        gi[:] = -1
+        gi[self.region] = np.arange(0, self.nx)
+        idc = gi[idc]
+        rows = np.arange(self.ineq_const_c, self.ineq_const_c+idc.shape[0])
+        rows =  np.tile(rows, (A.shape[-1], 1)).T
+        self.ineq_constraints[name] = {
+            "A": A,
+            "l": l,
+            "col": idc,
+            "u": u,
+            "row": rows
+        }
+        self.ineq_const_c+=idc.shape[0]
+
     def add_tangent_constraints(self, w=1.0):
         """
 
@@ -328,7 +347,7 @@ def add_tangent_constraints(self, w=1.0):
         if points.shape[0] > 1:
             self.add_gradient_orthogonal_constraints(points[:, :3], points[:, 3:6], w)
 
-    def build_matrix(self, square=True, damp=0.0):
+    def build_matrix(self, square=True, damp=0.0, ie=False):
         """
         Assemble constraints into interpolation matrix. Adds equaltiy
         constraints
@@ -344,7 +363,6 @@ def build_matrix(self, square=True, damp=0.0):
         """
 
         logger.info("Interpolation matrix is %i x %i" % (self.c_, self.nx))
-        cols = np.array(self.col)
         # To keep the solvers consistent for different model scales the range of the constraints should be similar.
         # We normalise the row vectors for the interpolation matrix
         # Each constraint can then be weighted separately for the least squares problem
@@ -381,7 +399,7 @@ def build_matrix(self, square=True, damp=0.0):
 
         A = coo_matrix(
             (np.array(a), (np.array(rows), cols)), shape=(self.c_, self.nx), dtype=float
-        )  # .tocsr()
+        ).tocsc()  # .tocsr()
 
         B = np.array(b)
         if not square:
@@ -392,7 +410,7 @@ def build_matrix(self, square=True, damp=0.0):
         # add a small number to the matrix diagonal to smooth the results
         # can help speed up solving, but might also introduce some errors
 
-        if self.eq_const_c_ > 0:
+        if len(self.equal_constraints) > 0:
             logger.info("Equality block is %i x %i" % (self.eq_const_c_, self.nx))
             # solving constrained least squares using
             # | ATA CT | |c| = b
@@ -404,16 +422,23 @@ def build_matrix(self, square=True, damp=0.0):
             # and d are the equality constraints
             # c are the node values and y are the
             # lagrange multipliers#
+            nc = 0
+            for c in self.equal_constraints.values():
+                aa = (c["A"]).flatten()
+                b.extend((c["B"] ).tolist())
+                mask = aa == 0
+                a.extend(aa[~mask].tolist())
+                rows.extend(c["row"].flatten()[~mask].tolist())
+                cols.extend(c["col"].flatten()[~mask].tolist())
             C = coo_matrix(
-                (
-                    np.array(self.eq_const_C),
-                    (np.array(self.eq_const_row), np.array(self.eq_const_col)),
-                ),
-                shape=(self.eq_const_c_, self.nx),
-            )
-            d = np.array(self.eq_const_d)
+
+                    (np.array(a), (np.array(rows), cols)), shape=(self.eq_const_c_, self.nx), dtype=float
+                ).tocsr()
+
+            d = np.array(b)
             ATA = bmat([[ATA, C.T], [C, None]])
             ATB = np.hstack([ATB, d])
+
         if isinstance(damp, bool):
             if damp == True:
                 damp = np.finfo("float").eps
@@ -422,7 +447,67 @@ def build_matrix(self, square=True, damp=0.0):
         if isinstance(damp, float):
             logger.info("Adding eps to matrix diagonal")
             ATA += eye(ATA.shape[0]) * damp
+        if len(self.ineq_constraints) > 0 and ie:
+            print('using inequality constraints')
+            a = []
+            l = []
+            u = []
+            rows = []
+            cols = []
+            for c in self.ineq_constraints.values():
+                aa = (c["A"]).flatten()
+                l.extend((c["l"]).tolist())
+                u.extend((c["u"]).tolist())
+
+                mask = aa == 0
+                a.extend(aa[~mask].tolist())
+                rows.extend(c["row"].flatten()[~mask].tolist())
+                cols.extend(c["col"].flatten()[~mask].tolist())
+            Aie = coo_matrix(
+            (np.array(a), (np.array(rows), cols)), shape=(self.ineq_const_c, self.nx), dtype=float
+            ).tocsc()    # .tocsr()
+
+            uie = np.array(u)
+            lie = np.array(l)
+
+            return ATA, ATB, Aie.T.dot(Aie), Aie.T.dot(uie), Aie.T.dot(lie)
         return ATA, ATB
+    def _solve_osqp(self, P, A, q, l, u):
+
+
+        import osqp
+        m = A.shape[0]
+        n = A.shape[1]
+        Ad = sparse.random(m, n, density=0.7, format='csc')
+        b = np.random.randn(m)
+
+        # OSQP data
+        P = sparse.block_diag([sparse.csc_matrix((n, n)), sparse.eye(m)], format='csc')
+        q = np.zeros(n+m)
+        A = sparse.vstack([
+                sparse.hstack([Ad, -sparse.eye(m)]),
+                sparse.hstack([sparse.eye(n), sparse.csc_matrix((n, m))])], format='csc')
+        l = np.hstack([b, np.zeros(n)])
+        u = np.hstack([b, np.ones(n)])
+
+        # Create an OSQP object
+        prob = osqp.OSQP()
+
+        # Setup workspace
+        prob.setup(P, q, A, l, u)
+
+        # Solve problem
+        res = prob.solve()
+
+
+        # Create an OSQP object
+        prob = osqp.OSQP()
+
+        # Setup workspace
+        # osqp likes csc matrices
+        prob.setup(P.tocsc(), np.array(q), A.tocsc(), np.array(u), np.array(l))
+        res = prob.solve()
+        return res.x
 
     def _solve_lu(self, A, B):
         """
@@ -588,6 +673,8 @@ def _solve(self, solver="cg", **kwargs):
             damp = True
         if solver == "lsqr":
             A, B = self.build_matrix(False)
+        elif solver =='osqp':
+            P, q, A, l, u = self.build_matrix(True,ie=True)
         else:
             A, B = self.build_matrix(damp=damp)
 
@@ -612,6 +699,8 @@ def _solve(self, solver="cg", **kwargs):
         if solver == "external":
             logger.warning("Using external solver")
             self.c[self.region] = kwargs["external"](A, B)[: self.nx]
+        if solver == 'osqp':
+            self.c[self.region] = self._solve_osqp(P, A, q, l, u)#, **kwargs)
         # check solution is not nan
         # self.support.properties[self.propertyname] = self.c
         if np.all(self.c == np.nan):