Implements MPCs

optimism/FunctionSpace.py

@@ -362,6 +362,30 @@ def integrate_function_on_edges(functionSpace, func, U, quadRule, edges):
     integrate_on_edges = jax.vmap(integrate_function_on_edge, (None, None, None, None, 0))
     return np.sum(integrate_on_edges(functionSpace, func, U, quadRule, edges))
 
+def create_nodeset_layers(mesh):
+    coords = mesh.coords
+    tol = 1e-8
+    # Create unique layers of nodesets along the y-direction
+    unique_layers = sorted(onp.unique(coords[:,1]))
+    Ny = int(input("Enter the expected number of nodeset layers in y-direction: "))
+    assert len(unique_layers) == Ny, f"ERROR - Expected {Ny} layers, but found {len(unique_layers)}."
+
+    layer_rows = []
+
+    for i, y_val in enumerate(unique_layers):
+        nodes_in_layer = onp.flatnonzero(onp.abs(coords[:, 1] - y_val) < tol)
+        layer_rows.append(nodes_in_layer)
+
+    max_nodes_per_layer = max(len(row) for row in layer_rows)
+    y_layers = -1 * np.ones((len(layer_rows), max_nodes_per_layer), dtype=int)  # Initialize with -1
+
+    for i, row in enumerate(layer_rows):
+        y_layers = y_layers.at[i, :len(row)].set(row)  # Fill each row with nodes from the layer
+
+    # # For debugging
+    # print("Layers in y-direction: ", y_layers)
+    return y_layers
+
 
 class DofManager(eqx.Module):
     # TODO get type hints below correct
@@ -466,3 +490,181 @@ def _make_hessian_bc_mask(self, conns):
             hessian_bc_mask[e,eFlag,:] = False
             hessian_bc_mask[e,:,eFlag] = False
         return hessian_bc_mask
+
+
+# Different class for Multi-Point Constrained Problem
+from optimism import Mesh
+import equinox as eqx
+import jax.numpy as np
+import numpy as onp
+from typing import List, Tuple
+
+
+class DofManagerMPC(eqx.Module):
+    fieldShape: Tuple[int, int]
+    isBc: np.ndarray
+    isUnknown: np.ndarray
+    ids: np.ndarray
+    unknownIndices: np.ndarray
+    bcIndices: np.ndarray
+    dofToUnknown: np.ndarray
+    HessRowCoords: np.ndarray
+    HessColCoords: np.ndarray
+    hessian_bc_mask: np.ndarray
+    master_layer: int = eqx.static_field()
+    slave_layer: int = eqx.static_field()
+    master_array: np.ndarray
+    slave_array: np.ndarray
+    layers: List[List[int]] = eqx.static_field()
+
+    def __init__(self, functionSpace, dim, EssentialBCs, mesh):
+        self.fieldShape = (Mesh.num_nodes(functionSpace.mesh), dim)
+        self.isBc = onp.full(self.fieldShape, False, dtype=bool)
+        for ebc in EssentialBCs:
+            self.isBc[functionSpace.mesh.nodeSets[ebc.nodeSet], ebc.component] = True
+        self.isUnknown = ~self.isBc
+
+        self.ids = onp.arange(self.isBc.size).reshape(self.fieldShape)
+
+        # Create layers and assign master and slave rows
+        self.layers = create_nodeset_layers(mesh)
+        print("These are the layers created for the mesh: ", self.layers)
+        self.master_layer = int(input("Choose the row index for master nodes: "))
+        self.slave_layer = int(input("Choose the row index for slave nodes: "))
+
+        # Generate master and slave arrays
+        self.master_array = self._create_layer_array(self.master_layer)
+        self.slave_array = self._create_layer_array(self.slave_layer)
+
+        # Create DOF mappings for unknowns
+        self.unknownIndices = self.ids[self.isUnknown]
+        self.bcIndices = self.ids[self.isBc]
+
+        ones = np.ones(self.isBc.size, dtype=int) * -1
+        self.dofToUnknown = ones.at[self.unknownIndices].set(np.arange(self.unknownIndices.size))
+
+        # Compute Hessian-related data
+        self.HessRowCoords, self.HessColCoords = self._make_hessian_coordinates(functionSpace.mesh.conns)
+        self.hessian_bc_mask = self._make_hessian_bc_mask(functionSpace.mesh.conns)
+
+    def get_bc_size(self):
+        return np.sum(self.isBc).item()
+
+    def get_unknown_size(self):
+        return np.sum(self.isUnknown).item()
+
+    def create_field(self, Uu, Ubc=0.0):
+        U = np.zeros(self.isBc.shape).at[self.isBc].set(Ubc)
+        return U.at[self.isUnknown].set(Uu)
+
+    def get_bc_values(self, U):
+        return U[self.isBc]
+
+    def get_unknown_values(self, U):
+        return U[self.isUnknown]
+
+    def get_master_values(self, U):
+        return U[self.master_array[:, 1:]]
+
+    def get_slave_values(self, U):
+        return U[self.slave_array[:, 1:]]
+
+    def _create_layer_array(self, layer_index):
+        """Creates an array for the specified layer (master or slave) with DOF mappings."""
+        layer_nodes = self.layers[layer_index]
+        layer_array = []
+        for node in layer_nodes:
+            node_dofs = self.ids[node]
+            layer_array.append([node, *node_dofs])
+        return np.array(layer_array, dtype=int)
+
+    def _make_hessian_coordinates(self, conns):
+        """Creates row and column coordinates for the Hessian, considering master and slave nodes."""
+        nElUnknowns = onp.zeros(conns.shape[0], dtype=int)
+        nHessianEntries = 0
+        for e, eNodes in enumerate(conns):
+            elUnknownFlags = self.isUnknown[eNodes, :].ravel()
+            nElUnknowns[e] = onp.sum(elUnknownFlags)
+
+            # Include master and slave nodes in the size calculation
+            eNodes_list = onp.asarray(eNodes).tolist()
+            elMasterNodes = set(eNodes_list).intersection(set(onp.array(self.master_array[:, 0])))
+            elSlaveNodes = set(eNodes_list).intersection(set(onp.array(self.slave_array[:, 0])))
+
+            nElMasters = sum(len(self.master_array[onp.where(self.master_array[:, 0] == node)[0], 1:].ravel()) for node in elMasterNodes)
+            nElSlaves = sum(len(self.slave_array[onp.where(self.slave_array[:, 0] == node)[0], 1:].ravel()) for node in elSlaveNodes)
+
+            totalDOFs = nElUnknowns[e] + nElMasters + nElSlaves
+            nHessianEntries += totalDOFs ** 2
+
+        # Allocate sufficient space for Hessian coordinates
+        rowCoords = onp.zeros(nHessianEntries, dtype=int)
+        colCoords = onp.zeros(nHessianEntries, dtype=int)
+        rangeBegin = 0
+
+        for e, eNodes in enumerate(conns):
+            eNodes_list = onp.asarray(eNodes).tolist()
+            elDofs = self.ids[eNodes, :]
+            elUnknownFlags = self.isUnknown[eNodes, :]
+            elUnknowns = self.dofToUnknown[elDofs[elUnknownFlags]]
+
+            # Identify master and slave DOFs for the element
+            elMasterNodes = set(eNodes_list).intersection(set(onp.array(self.master_array[:, 0])))
+            elSlaveNodes = set(eNodes_list).intersection(set(onp.array(self.slave_array[:, 0])))
+
+            elMasterDofs = []
+            for node in elMasterNodes:
+                master_idx = onp.where(self.master_array[:, 0] == node)[0]
+                elMasterDofs.extend(self.master_array[master_idx, 1:].ravel())
+
+            elSlaveDofs = []
+            for node in elSlaveNodes:
+                slave_idx = onp.where(self.slave_array[:, 0] == node)[0]
+                elSlaveDofs.extend(self.slave_array[slave_idx, 1:].ravel())
+
+            # Combine unknowns, master, and slave DOFs
+            elAllUnknowns = onp.concatenate([elUnknowns, elMasterDofs, elSlaveDofs])
+            elHessCoords = onp.tile(elAllUnknowns, (len(elAllUnknowns), 1))
+
+            nElHessianEntries = len(elAllUnknowns) ** 2
+            rangeEnd = rangeBegin + nElHessianEntries
+
+            # Assign values to the Hessian coordinate arrays
+            rowCoords[rangeBegin:rangeEnd] = elHessCoords.ravel()
+            colCoords[rangeBegin:rangeEnd] = elHessCoords.T.ravel()
+
+            rangeBegin += nElHessianEntries
+
+        return rowCoords, colCoords
+
+    def _make_hessian_bc_mask(self, conns):
+        """Creates a mask for BCs in the Hessian, considering master and slave nodes."""
+        nElements, nNodesPerElement = conns.shape
+        nFields = self.ids.shape[1]
+        nDofPerElement = nNodesPerElement * nFields
+
+        hessian_bc_mask = onp.full((nElements, nDofPerElement, nDofPerElement), True, dtype=bool)
+        for e, eNodes in enumerate(conns):
+            # Get boundary condition flags for all DOFs
+            eFlag = self.isBc[eNodes, :].ravel()
+
+            # Identify master and slave nodes
+            eNodes_list = onp.asarray(eNodes).tolist()
+            masterFlag = onp.isin(eNodes_list, self.master_array[:, 0])
+            slaveFlag = onp.isin(eNodes_list, self.slave_array[:, 0])
+
+            # Expand master and slave flags to match the DOF shape
+            masterFlag_expanded = onp.repeat(masterFlag, nFields)
+            slaveFlag_expanded = onp.repeat(slaveFlag, nFields)
+
+            # Combine flags
+            combinedFlag = eFlag | masterFlag_expanded | slaveFlag_expanded
+
+            # Update Hessian mask
+            hessian_bc_mask[e, combinedFlag, :] = False
+            hessian_bc_mask[e, :, combinedFlag] = False
+
+        return hessian_bc_mask
+
+
+

optimism/Objective.py

@@ -3,6 +3,7 @@
 import numpy as onp
 from scipy.sparse import diags as sparse_diags
 from scipy.sparse import csc_matrix
+import jax.random as rand
 
 # static vs dynamics
 # differentiable vs undifferentiable
@@ -265,4 +266,68 @@ def get_value(self, x):
     def get_residual(self, x):
         return self.gradient(self.scaling * x)
 
-
+
+class ObjectiveMPC(Objective):
+    def __init__(self, f, x, p, dofManagerMPC, precondStrategy=None):
+        """
+        Initialize the Objective class for MPC with full DOFs and master-slave constraints.
+
+        Parameters:
+        - f: Objective function.
+        - x: Initial guess for the primal DOF vector.
+        - p: Parameters for the objective function.
+        - dofManagerMPC: Instance of DofManagerMPC with master-slave relationships.
+        - precondStrategy: Preconditioning strategy (optional).
+        """
+        super().__init__(f, x, p, precondStrategy)
+        self.dofManagerMPC = dofManagerMPC
+        self.master_dofs = dofManagerMPC.master_array[:, 1:]
+        self.slave_dofs = dofManagerMPC.slave_array[:, 1:]
+        self.slave_to_master_map = self.create_slave_to_master_map()
+
+    def create_slave_to_master_map(self):
+        """
+        Create the mapping to express slave DOFs in terms of master DOFs.
+        Returns:
+        - C: Constraint matrix (2D array).
+        - d: Constant vector (1D array).
+        """
+        num_slaves = len(self.slave_dofs)
+        num_masters = len(self.master_dofs)
+
+        key = rand.PRNGKey(0)
+        # Initialize the constraint matrix with the correct dimensions
+        C = rand.uniform(key, (num_slaves, num_masters))  # Replace with actual mapping logic
+
+        # Initialize the constant vector
+        d = np.zeros(num_slaves)  # Replace with actual constant values if needed
+
+        return C, d
+
+
+    def value(self, x_full):
+        # Ensure the constraints are satisfied
+        x_full = self.enforce_constraints(x_full)
+        return super().value(x_full)
+
+    def gradient(self, x_full):
+        x_full = self.enforce_constraints(x_full)
+        grad_full = super().gradient(x_full)
+        return grad_full
+
+    def hessian_vec(self, x_full, vx_full):
+        x_full = self.enforce_constraints(x_full)
+        hess_vec_full = super().hessian_vec(x_full, vx_full)
+        return hess_vec_full
+
+    def enforce_constraints(self, x_full):
+        if self.slave_to_master_map is not None:
+            C, d = self.slave_to_master_map
+            slave_values = C @ x_full[self.master_dofs] + d[:, None]
+            print("Slave values before assignment:", slave_values)
+            x_full = x_full.at[self.slave_dofs].set(slave_values)
+            print("x_full after enforcing constraints:", x_full)
+        return x_full
+
+
+