diff --git a/optimism/FunctionSpace.py b/optimism/FunctionSpace.py index d22dff3b..e7115414 100644 --- a/optimism/FunctionSpace.py +++ b/optimism/FunctionSpace.py @@ -3,11 +3,12 @@ from optimism import Interpolants from optimism import Mesh from optimism import QuadratureRule -from typing import Tuple +from typing import Tuple, List import equinox as eqx import jax import jax.numpy as np import numpy as onp +import scipy.sparse as sp class EssentialBC(eqx.Module): @@ -362,7 +363,6 @@ 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)) - class DofManager(eqx.Module): # TODO get type hints below correct # TODO this one could be moved to jax types if we move towards @@ -387,7 +387,7 @@ def __init__(self, functionSpace, dim, EssentialBCs): self.isUnknown = ~self.isBc self.ids = onp.arange(self.isBc.size).reshape(self.fieldShape) - + print(self.ids.shape) self.unknownIndices = self.ids[self.isUnknown] self.bcIndices = self.ids[self.isBc] @@ -466,3 +466,100 @@ def _make_hessian_bc_mask(self, conns): hessian_bc_mask[e,eFlag,:] = False hessian_bc_mask[e,:,eFlag] = False return hessian_bc_mask + +# DofManager for Multi-Point Constrained Problem +class DofManagerMPC(DofManager): + dim: int = eqx.field() + master_slave_pairs: dict = eqx.field() + C: np.ndarray = eqx.field() + C_reduced: np.ndarray = eqx.field() + s_reduced: np.ndarray = eqx.field() + s_tilde: np.ndarray = eqx.field() + T: np.ndarray = eqx.field() + isIndependent: np.ndarray = eqx.field() + isUnconstrained: np.ndarray = eqx.field() + is_slave: np.ndarray = eqx.field() + is_indep_dof: np.ndarray = eqx.field() + unconstrainedIndices: np.ndarray = eqx.field() + slaveIndices: np.ndarray = eqx.field() + bcIndices: np.ndarray = eqx.field() + dofToUnknown: np.ndarray = eqx.field() + + def __init__(self, functionSpace, dim, EssentialBCs, master_slave_pairs, C, s): + super().__init__(functionSpace, dim, EssentialBCs) + self.fieldShape = Mesh.num_nodes(functionSpace.mesh), dim + self.dim = dim + self.master_slave_pairs = master_slave_pairs + self.C = np.array(C) + self.s_tilde = np.array(s) + self.C_reduced = np.array(C) # Reduced constraint matrix + self.s_reduced = np.array(s) # Reduced shift vector + self.isIndependent = None + self.isUnconstrained = None + self.is_indep_dof = None + self.is_slave = None + self.slaveIndices = None + self.unconstrainedIndices = None + self.create_mpc_transformation() + + def create_mpc_transformation(self): + slave_nodes = np.array(list(self.master_slave_pairs.keys())) + master_nodes = np.array(list(self.master_slave_pairs.values())) + is_slave = onp.full(self.fieldShape, False, dtype=bool) + # self.is_slave = is_slave.at[slave_nodes, :].set(True) + + self.is_slave = onp.full(self.fieldShape, False, dtype=bool) + self.is_slave[slave_nodes,:] = True + + # self.is_slave = is_slave + self.isUnconstrained = ~self.is_slave + + # print("shape of isUnknown: ", self.isUnknown.shape) + + self.ids = onp.arange(self.is_slave.size).reshape(self.fieldShape) + print(self.ids.shape) + self.unconstrainedIndices = self.ids[self.isUnconstrained] + self.slaveIndices = self.ids[self.is_slave] + + T = np.zeros((self.is_slave.size, self.unconstrainedIndices.size)) + + for local_idx, global_idx in enumerate(self.unconstrainedIndices): + T = T.at[global_idx, local_idx].set(1.0) + + # Build s_tilde: global shift + s_tilde = np.zeros(self.is_slave.size) + + for i, (slave_node, master_node) in enumerate(self.master_slave_pairs.items()): + for d in range(self.dim): + slave_dof = slave_node * self.dim + d + master_dof = master_node * self.dim + d + + + # Find column index in reduced DOFs + reduced_master_index = np.where(self.unconstrainedIndices == master_dof)[0] + if reduced_master_index.size == 0: + raise ValueError(f"Master DOF {master_dof} is not independent") + + T = T.at[slave_dof, reduced_master_index[0]].set(1.0) + s_tilde = s_tilde.at[slave_dof].set(self.s_reduced[i * self.dim + d]) + + self.T = T + self.s_tilde = s_tilde + + # Track dof-to-unknown mapping + ones = onp.ones(self.is_slave.size, dtype=int) * -1 + dofToUnknown = ones + dofToUnknown[self.unconstrainedIndices] = onp.arange(self.unconstrainedIndices.size) + self.dofToUnknown = dofToUnknown + + def create_field(self, Uu, Ubc=0.0): + U_flat = np.matmul(self.T, Uu) + self.s_tilde + return U_flat.reshape(self.fieldShape) + + def get_unknown_values(self, U): + print("shape of U in get unconstrained values: ", U.shape) + print("shape of isUnconstrained: ", self.isUnconstrained.shape) + return U[self.isUnknown] + + def get_slave_values(self, U): + return U[self.is_slave] diff --git a/optimism/Objective.py b/optimism/Objective.py index 8f3d3f8b..a285e8d4 100644 --- a/optimism/Objective.py +++ b/optimism/Objective.py @@ -1,8 +1,12 @@ from optimism.JaxConfig import * from optimism.SparseCholesky import SparseCholesky import numpy as onp -from scipy.sparse import diags as sparse_diags +import jax.numpy as np +import jax +from jax import jit, grad, jacfwd, jvp, vjp from scipy.sparse import csc_matrix +from scipy.sparse import diags as sparse_diags + # static vs dynamics # differentiable vs undifferentiable @@ -265,4 +269,76 @@ def get_value(self, x): def get_residual(self, x): return self.gradient(self.scaling * x) - +# Objective class for handling Multi-Point Constraints +class ObjectiveMPC(Objective): + def __init__(self, objective_func, x_full, p, dofManagerMPC, precondStrategy=None): + """ + ObjectiveMPC: wraps the original energy functional to solve only for reduced DOFs. + - Applies u = T * ũ + s̃ for MPC condensation. + - Automatically condenses gradient and Hessian. + """ + self.dofManagerMPC = dofManagerMPC + self.T = dofManagerMPC.T + self.s_tilde = dofManagerMPC.s_tilde + self.p = p + self.precondStrategy = precondStrategy + + # Store full functional and derivatives + self.full_objective_func = jit(objective_func) + self.full_grad_x = jit(grad(objective_func, 0)) + self.full_hess = jit(jacfwd(self.full_grad_x, 0)) + + self.scaling = 1.0 + self.invScaling = 1.0 + + # Reduced initial guess from full vector + self.x_reduced0 = self.reduce_to_independent_dofs(x_full) + + # Define reduced-space objective and gradient + self.objective = jit(lambda x_red, p: self.full_objective_func(self.expand_to_full_dofs(x_red), p)) + self.grad_x = jit(lambda x_red, p: self.reduce_gradient(self.full_grad_x(self.expand_to_full_dofs(x_red), p))) + + self.precond = SparseCholesky() + + def reduce_to_independent_dofs(self, x_full): + """Extract reduced DOFs from full vector.""" + return self.dofManagerMPC.get_unknown_values(x_full) + + def expand_to_full_dofs(self, x_reduced): + """Reconstruct full DOF vector using u = T ũ + s̃.""" + return self.dofManagerMPC.create_field(x_reduced) + + def reduce_gradient(self, grad_full): + """Project full gradient to reduced space.""" + return self.dofManagerMPC.get_unknown_values(grad_full) + + def value(self, x_reduced): + """Return energy functional in reduced space.""" + return self.objective(x_reduced, self.p) + + def gradient(self, x_reduced): + """Return reduced gradient.""" + return self.grad_x(x_reduced, self.p) + + def hessian(self, x_reduced): + """Compute reduced Hessian H̃ = Tᵀ H T.""" + x_full = self.expand_to_full_dofs(self.scaling * x_reduced) + H_full = self.full_hess(x_full, self.p) + H_reduced = np.matmul(self.T.T, np.matmul(H_full, self.T)) + return H_reduced + + def update_precond(self, x_reduced): + """Update preconditioner from reduced Hessian.""" + print("Updating with condensed Hessian preconditioner.") + H_reduced = csc_matrix(self.hessian(x_reduced)) + self.precond.update(lambda attempt: H_reduced) + + def apply_precond(self, vx): + return self.precond.apply(vx) if self.precond else vx + + def multiply_by_approx_hessian(self, vx): + return self.precond.multiply_by_approximate(vx) if self.precond else vx + + def check_stability(self, x_reduced): + if self.precond: + self.precond.check_stability(x_reduced, self.p) diff --git a/optimism/test/MeshFixture.py b/optimism/test/MeshFixture.py index 341b4356..085d97bb 100644 --- a/optimism/test/MeshFixture.py +++ b/optimism/test/MeshFixture.py @@ -1,9 +1,15 @@ -# fea data structures +# ---------------- NEW MESHFIXTURE SCRIPT -------------------------- +# ------------------------------------------------------------------ +# Note - This script involves a function creating nodeset layers +# ------------------------------------------------------------------ + +import jax.numpy as jnp +import numpy as onp from optimism.Mesh import * from optimism import Surface # testing utils -from .TestFixture import * +from TestFixture import * d_kappa = 1.0 d_nu = 0.3 @@ -139,4 +145,80 @@ def is_edge_on_left(xyOnEdge): blocks = {'block': np.arange(conns.shape[0])} U = np.zeros(coords.shape) return construct_mesh_from_basic_data(coords, conns, None, nodeSets, sideSets), U - + + def create_mesh_disp_and_nodeset_layers(self, Nx, Ny, xRange, yRange, initial_disp_func, setNamePostFix=''): + coords, conns = create_structured_mesh_data(Nx, Ny, xRange, yRange) + tol = 1e-8 + nodeSets = {} + + # Predefined boundary node sets + nodeSets['left'+setNamePostFix] = jnp.flatnonzero(coords[:, 0] < xRange[0] + tol) + nodeSets['bottom'+setNamePostFix] = jnp.flatnonzero(coords[:, 1] < yRange[0] + tol) + nodeSets['right'+setNamePostFix] = jnp.flatnonzero(coords[:, 0] > xRange[1] - tol) + nodeSets['top'+setNamePostFix] = jnp.flatnonzero(coords[:, 1] > yRange[1] - tol) + nodeSets['all_boundary'+setNamePostFix] = jnp.flatnonzero( + (coords[:, 0] < xRange[0] + tol) | + (coords[:, 1] < yRange[0] + tol) | + (coords[:, 0] > xRange[1] - tol) | + (coords[:, 1] > yRange[1] - tol) + ) + + # Identify unique y-layers for nodes + unique_y_layers = sorted(onp.unique(coords[:, 1])) + # print("Unique y-layers identified:", unique_y_layers) + + # Ensure we have the expected number of layers + assert len(unique_y_layers) == Ny, f"Expected {Ny} y-layers, but found {len(unique_y_layers)}." + + # Initialize an empty list to store rows of nodes + y_layer_rows = [] + + # Map nodes to y_layers and construct rows + for i, y_val in enumerate(unique_y_layers): + nodes_in_layer = onp.flatnonzero(onp.abs(coords[:, 1] - y_val) < tol) + y_layer_rows.append(nodes_in_layer) + # print(f"Nodes in y-layer {i+1} (y = {y_val}):", nodes_in_layer) + + # Convert list of rows into a structured 2D JAX array, padding with -1 + max_nodes_per_layer = max(len(row) for row in y_layer_rows) + y_layers = -1 * jnp.ones((len(y_layer_rows), max_nodes_per_layer), dtype=int) # Initialize with -1 + + for i, row in enumerate(y_layer_rows): + y_layers = y_layers.at[i, :len(row)].set(row) # Fill each row with nodes from the layer + + # Print for debugging + # print("y_layers (2D array):", y_layers) + + def is_edge_on_left(xyOnEdge): + return np.all( xyOnEdge[:,0] < xRange[0] + tol ) + + def is_edge_on_bottom(xyOnEdge): + return np.all( xyOnEdge[:,1] < yRange[0] + tol ) + + def is_edge_on_right(xyOnEdge): + return np.all( xyOnEdge[:,0] > xRange[1] - tol ) + + def is_edge_on_top(xyOnEdge): + return np.all( xyOnEdge[:,1] > yRange[1] - tol ) + + sideSets = {} + sideSets['left'+setNamePostFix] = Surface.create_edges(coords, conns, is_edge_on_left) + sideSets['bottom'+setNamePostFix] = Surface.create_edges(coords, conns, is_edge_on_bottom) + sideSets['right'+setNamePostFix] = Surface.create_edges(coords, conns, is_edge_on_right) + sideSets['top'+setNamePostFix] = Surface.create_edges(coords, conns, is_edge_on_top) + + allBoundaryEdges = np.vstack([s for s in sideSets.values()]) + sideSets['all_boundary'+setNamePostFix] = allBoundaryEdges + + blocks = {'block'+setNamePostFix: np.arange(conns.shape[0])} + mesh = construct_mesh_from_basic_data(coords, conns, blocks, nodeSets, sideSets) + + return mesh, vmap(initial_disp_func)(mesh.coords), y_layers + + + + + + + + diff --git a/optimism/test/non_homogeneous_MPC.py b/optimism/test/non_homogeneous_MPC.py new file mode 100644 index 00000000..4c27e514 --- /dev/null +++ b/optimism/test/non_homogeneous_MPC.py @@ -0,0 +1,200 @@ +# -------------------------------------------------------------------------- +# TEST CASE - Compression Test for implementation of MPCs +# -------------------------------------------------------------------------- +# Note - Using Cubit Coreform for meshing, solved using NON-LINEAR SOLVER +# -------------------------------------------------------------------------- +import sys +sys.path.insert(0, "/home/sarvesh/Documents/Github/optimism") + +from jax import jit +from optimism import VTKWriter +from optimism import EquationSolver +from optimism import FunctionSpace +from optimism import Mechanics +from optimism import Mesh +from optimism import Objective +from optimism import QuadratureRule +from optimism import ReadExodusMesh +from optimism.FunctionSpace import DofManagerMPC +from optimism.FunctionSpace import EssentialBC +from optimism.material import Neohookean + +import jax.numpy as np +import numpy as onp +from scipy.spatial import cKDTree + + +class CompressionTest: + + def __init__(self): + self.writeOutput = True + + self.quadRule = QuadratureRule.create_quadrature_rule_on_triangle(degree=2) + self.lineQuadRule = QuadratureRule.create_quadrature_rule_1D(degree=2) + + self.EBCs = [ + EssentialBC(nodeSet='bottom', component=0), + EssentialBC(nodeSet='bottom', component=1) + ] + + shearModulus = 0.855 # MPa + bulkModulus = 1000*shearModulus # MPa + youngModulus = 9.0*bulkModulus*shearModulus / (3.0*bulkModulus + shearModulus) + poissonRatio = (3.0*bulkModulus - 2.0*shearModulus) / 2.0 / (3.0*bulkModulus + shearModulus) + props = { + 'elastic modulus': youngModulus, + 'poisson ratio': poissonRatio, + 'version': 'coupled' + } + + self.matModel = Neohookean.create_material_model_functions(props) + + self.input_mesh = 'square_mesh_compression.exo' + + self.maxForce = -0.2 + self.steps = 50 + + def create_field(self, Uu): + def get_ubcs(): + V = np.zeros(self.mesh.coords.shape) + return self.dofManager.get_bc_values(V) + + return self.dofManager.create_field(Uu, get_ubcs()) + + def reload_mesh(self): + origMesh = ReadExodusMesh.read_exodus_mesh(self.input_mesh) + self.mesh = Mesh.create_higher_order_mesh_from_simplex_mesh(origMesh, order=2, createNodeSetsFromSideSets=True) + + self.funcSpace = FunctionSpace.construct_function_space(self.mesh, self.quadRule) # This is where error is showcased + + surfaceXCoords = self.mesh.coords[self.mesh.nodeSets['top']][:,0] + self.tractionArea = np.max(surfaceXCoords) - np.min(surfaceXCoords) + + def master_slave_mapping(self): + coords = self.mesh.coords + tol = 1e-8 + + # Creating nodesets for master and slave + self.mesh.nodeSets['master'] = np.flatnonzero(coords[:,0] < -0.5 + tol) + self.mesh.nodeSets['slave'] = np.flatnonzero(coords[:,0] > 0.5 - tol) + + # Get master and slave node coordinates + master_coords = coords[self.mesh.nodeSets['master']] + slave_coords = coords[self.mesh.nodeSets['slave']] + + # Sort nodes by y-coordinate to ensure correct pairing + master_sorted_idx = np.argsort(master_coords[:, 1]) + slave_sorted_idx = np.argsort(slave_coords[:, 1]) + + # Ensure one-to-one pairing by iterating through sorted nodes + self.master_slave_pairs = { + int(self.mesh.nodeSets['master'][master_sorted_idx[i]]): + int(self.mesh.nodeSets['slave'][slave_sorted_idx[i]]) + for i in range(min(len(master_sorted_idx), len(slave_sorted_idx))) + } + + print("Master-Slave Pairs:", self.master_slave_pairs) + + def implement_constraints(self, alpha=2.0, shift_type="linear"): + coords = self.mesh.coords + + num_constraints = len(self.master_slave_pairs) + print("Number of Constraints: ", num_constraints) + + # Define Ap (Master Constraints) and Ac (Slave Constraints) + Ap = onp.ones((num_constraints, len(self.mesh.coords))) + Ac = onp.eye(num_constraints)*alpha + + b = onp.ones(num_constraints) + + for idx, (master, slave) in enumerate(self.master_slave_pairs.items()): + Ap[idx, master] = 1 # Master DOF stays unchanged + + # Compute nonzero shift `b[idx]` based on shift type + if shift_type == "linear": + b[idx] = coords[slave, 0] + coords[slave, 1] # Linear shift + elif shift_type == "sinusoidal": + b[idx] = onp.sin(coords[slave, 1]) # Sinusoidal variation + elif shift_type == "random": + b[idx] = onp.random.uniform(-0.1, 0.1) # Small random perturbation + else: + b[idx] = 0 # Default to zero shift + + # Getting C and s + Ac_inv = onp.linalg.inv(Ac) + C = -Ac_inv @ Ap + s = Ac_inv @ b + + print("Constraint Matrix:\n", C) + print("Shift Vector:\n", s) + + def run(self): + funcSpace = FunctionSpace.construct_function_space(self.mesh, self.quadRule) + mechFuncs = Mechanics.create_mechanics_functions(funcSpace, mode2D='plane strain', materialModel=self.matModel) + + def energy_function(Uu, p): + U = self.create_field(Uu) + internal_variables = p.state_data + strainEnergy = mechFuncs.compute_strain_energy(U, internal_variables) + + F = p.bc_data + def force_function(x, n): + return np.array([0.0, F/self.tractionArea]) + + loadPotential = Mechanics.compute_traction_potential_energy( + funcSpace, U, self.lineQuadRule, self.mesh.sideSets['top'], + force_function) + + return strainEnergy + loadPotential + + def write_output(Uu, p, step): + U = self.create_field(Uu) + plotName = 'results/' + 'output-' + str(step).zfill(3) + writer = VTKWriter.VTKWriter(self.mesh, baseFileName=plotName) + + writer.add_nodal_field(name='disp', nodalData=U, fieldType=VTKWriter.VTKFieldType.VECTORS) + energyDensities, stresses = mechFuncs.compute_output_energy_densities_and_stresses(U, p.state_data) + cellEnergyDensities = FunctionSpace.project_quadrature_field_to_element_field(funcSpace, energyDensities) + cellStresses = FunctionSpace.project_quadrature_field_to_element_field(funcSpace, stresses) + writer.add_cell_field(name='strain_energy_density', + cellData=cellEnergyDensities, + fieldType=VTKWriter.VTKFieldType.SCALARS) + writer.add_cell_field(name='cauchy_stress', + cellData=cellStresses, + fieldType=VTKWriter.VTKFieldType.TENSORS) + writer.write() + + # Problem Set Up + Uu = self.dofManager.get_unknown_values(np.zeros(self.mesh.coords.shape)) + ivs = mechFuncs.compute_initial_state() + p = Objective.Params(bc_data=0., state_data=ivs) + self.objective = Objective.ObjectiveMPC(energy_function, Uu, p, self.dofManager, precondStrategy=None) + + # Load over the steps + force = 0. + + force_inc = self.maxForce / self.steps + + + for step in range(1, self.steps): + print('------------------------------------') + print('LOAD STEP', step) + force += force_inc + p = Objective.param_index_update(p, 0, force) + Uu, solverSuccess = EquationSolver.nonlinear_equation_solve(self.objective, Uu, p, EquationSolver.get_settings()) + + if self.writeOutput: + write_output(Uu, p, step) + + + + +if __name__ == '__main__': + app = CompressionTest() + app.reload_mesh() + app.run() + + + + + diff --git a/optimism/test/square_mesh_compression.exo b/optimism/test/square_mesh_compression.exo new file mode 100644 index 00000000..5b793757 Binary files /dev/null and b/optimism/test/square_mesh_compression.exo differ diff --git a/optimism/test/test_compression.py b/optimism/test/test_compression.py new file mode 100644 index 00000000..cb88387a --- /dev/null +++ b/optimism/test/test_compression.py @@ -0,0 +1,145 @@ +# -------------------------------------------------------------------------- +# TEST CASE - Compression Test for implementation of MPCs +# -------------------------------------------------------------------------- +# Note - Using Cubit Coreform for meshing, solved using NON-LINEAR SOLVER +# -------------------------------------------------------------------------- + +from jax import jit +from jax.scipy.linalg import solve +from optimism import VTKWriter +from optimism import FunctionSpace +from optimism import EquationSolver +from optimism import Mechanics +from optimism import Mesh +from optimism import Objective +from optimism import QuadratureRule +from optimism import ReadExodusMesh +from optimism.FunctionSpace import DofManager +from optimism.FunctionSpace import EssentialBC +from optimism.material import Neohookean + +import jax.numpy as np +import numpy as onp + + + +class CompressionTest: + + def __init__(self): + self.writeOutput = True + + self.quadRule = QuadratureRule.create_quadrature_rule_on_triangle(degree=2) + self.lineQuadRule = QuadratureRule.create_quadrature_rule_1D(degree=2) + + self.EBCs = [ + EssentialBC(nodeSet='bottom', component=0), + EssentialBC(nodeSet='bottom', component=1) + ] + + shearModulus = 0.855 # MPa + bulkModulus = 1000*shearModulus # MPa + youngModulus = 9.0*bulkModulus*shearModulus / (3.0*bulkModulus + shearModulus) + poissonRatio = (3.0*bulkModulus - 2.0*shearModulus) / 2.0 / (3.0*bulkModulus + shearModulus) + props = { + 'elastic modulus': youngModulus, + 'poisson ratio': poissonRatio, + 'version': 'coupled' + } + + self.matModel = Neohookean.create_material_model_functions(props) + + self.input_mesh = 'square_mesh_compression.exo' + + self.maxForce = -0.2 + self.steps = 50 + + def create_field(self, Uu): + def get_ubcs(): + V = np.zeros(self.mesh.coords.shape) + return self.dofManager.get_bc_values(V) + + return self.dofManager.create_field(Uu, get_ubcs()) + + def reload_mesh(self): + origMesh = ReadExodusMesh.read_exodus_mesh(self.input_mesh) + self.mesh = Mesh.create_higher_order_mesh_from_simplex_mesh(origMesh, order=2, createNodeSetsFromSideSets=True) + + funcSpace = FunctionSpace.construct_function_space(self.mesh, self.quadRule) + self.dofManager = DofManager(funcSpace, 2, self.EBCs) + + surfaceXCoords = self.mesh.coords[self.mesh.nodeSets['top']][:,0] + self.tractionArea = np.max(surfaceXCoords) - np.min(surfaceXCoords) + + def run(self): + funcSpace = FunctionSpace.construct_function_space(self.mesh, self.quadRule) + mechFuncs = Mechanics.create_mechanics_functions(funcSpace, mode2D='plane strain', materialModel=self.matModel) + + def energy_function(Uu, p): + U = self.create_field(Uu) + internal_variables = p.state_data + strainEnergy = mechFuncs.compute_strain_energy(U, internal_variables) + + F = p.bc_data + def force_function(x, n): + return np.array([0.0, F/self.tractionArea]) + + loadPotential = Mechanics.compute_traction_potential_energy( + funcSpace, U, self.lineQuadRule, self.mesh.sideSets['top'], + force_function) + + return strainEnergy + loadPotential + + def write_output(Uu, p, step): + U = self.create_field(Uu) + plotName = 'results/' + 'output-' + str(step).zfill(3) + writer = VTKWriter.VTKWriter(self.mesh, baseFileName=plotName) + + writer.add_nodal_field(name='disp', nodalData=U, fieldType=VTKWriter.VTKFieldType.VECTORS) + energyDensities, stresses = mechFuncs.compute_output_energy_densities_and_stresses(U, p.state_data) + cellEnergyDensities = FunctionSpace.project_quadrature_field_to_element_field(funcSpace, energyDensities) + cellStresses = FunctionSpace.project_quadrature_field_to_element_field(funcSpace, stresses) + writer.add_cell_field(name='strain_energy_density', + cellData=cellEnergyDensities, + fieldType=VTKWriter.VTKFieldType.SCALARS) + writer.add_cell_field(name='cauchy_stress', + cellData=cellStresses, + fieldType=VTKWriter.VTKFieldType.TENSORS) + writer.write() + + # Problem Set Up + Uu = self.dofManager.get_unknown_values(np.zeros(self.mesh.coords.shape)) + ivs = mechFuncs.compute_initial_state() + p = Objective.Params(bc_data=0., state_data=ivs) + self.objective = Objective.Objective(energy_function, Uu, p) + + # Load over the steps + force = 0. + + force_inc = self.maxForce / self.steps + + # K = self.objective.hessian(Uu) + # print(f"Hessian first row: {K[0,:5]}") + # print(f"Hessian Shape: {K.shape}") + + for step in range(1, self.steps): + print('------------------------------------') + print('LOAD STEP', step) + force += force_inc + p = Objective.param_index_update(p, 0, force) + Uu, solverSuccess = EquationSolver.nonlinear_equation_solve(self.objective, Uu, p, EquationSolver.get_settings()) + + if self.writeOutput: + write_output(Uu, p, step) + + + + +if __name__ == '__main__': + app = CompressionTest() + app.reload_mesh() + app.run() + + + + + diff --git a/optimism/test/test_compression.txt b/optimism/test/test_compression.txt new file mode 100644 index 00000000..797beed0 --- /dev/null +++ b/optimism/test/test_compression.txt @@ -0,0 +1,149 @@ +# -------------------------------------------------------------------------- +# TEST CASE - Compression Test for implementation of MPCs +# -------------------------------------------------------------------------- +# Note - Using Cubit Coreform for meshing +# -------------------------------------------------------------------------- + +from jax import jit, grad +from jax.scipy.linalg import solve +from optimism import VTKWriter +from optimism import FunctionSpace +from optimism import Mechanics +from optimism import Mesh +from optimism import Objective +from optimism import QuadratureRule +from optimism import ReadExodusMesh +from optimism.FunctionSpace import DofManager +from optimism.FunctionSpace import EssentialBC +from optimism.material import Neohookean + +import jax.numpy as np +import numpy as onp + + + +class CompressionTest: + + def __init__(self): + self.writeOutput = True + + self.quadRule = QuadratureRule.create_quadrature_rule_on_triangle(degree=2) + self.lineQuadRule = QuadratureRule.create_quadrature_rule_1D(degree=2) + + self.EBCs = [ + EssentialBC(nodeSet='bottom', component=0), + EssentialBC(nodeSet='bottom', component=1) + ] + + shearModulus = 0.855 # MPa + bulkModulus = 1000*shearModulus # MPa + youngModulus = 9.0*bulkModulus*shearModulus / (3.0*bulkModulus + shearModulus) + poissonRatio = (3.0*bulkModulus - 2.0*shearModulus) / 2.0 / (3.0*bulkModulus + shearModulus) + props = { + 'elastic modulus': youngModulus, + 'poisson ratio': poissonRatio, + 'version': 'coupled' + } + + self.matModel = Neohookean.create_material_model_functions(props) + + self.input_mesh = 'square_mesh_compression.exo' + + self.maxForce = 0.1 + self.steps = 50 + + def create_field(self, Uu): + def get_ubcs(): + V = np.zeros(self.mesh.coords.shape) + return self.dofManager.get_bc_values(V) + + return self.dofManager.create_field(Uu, get_ubcs()) + + def reload_mesh(self): + origMesh = ReadExodusMesh.read_exodus_mesh(self.input_mesh) + self.mesh = Mesh.create_higher_order_mesh_from_simplex_mesh(origMesh, order=2, createNodeSetsFromSideSets=True) + + funcSpace = FunctionSpace.construct_function_space(self.mesh, self.quadRule) + self.dofManager = DofManager(funcSpace, 2, self.EBCs) + + surfaceXCoords = self.mesh.coords[self.mesh.nodeSets['top']][:,0] + self.tractionArea = np.max(surfaceXCoords) - np.min(surfaceXCoords) + + def run(self): + funcSpace = FunctionSpace.construct_function_space(self.mesh, self.quadRule) + mechFuncs = Mechanics.create_mechanics_functions(funcSpace, mode2D='plane strain', materialModel=self.matModel) + + def energy_function(Uu, p): + U = self.create_field(Uu) + internal_variables = p.state_data + strainEnergy = mechFuncs.compute_strain_energy(U, internal_variables) + + F = p.bc_data + def force_function(x, n): + return np.array([0.0, F/self.tractionArea]) + + loadPotential = Mechanics.compute_traction_potential_energy( + funcSpace, U, self.lineQuadRule, self.mesh.sideSets['top'], + force_function) + + return strainEnergy + loadPotential + + def write_output(Uu, p, step): + U = self.create_field(Uu) + plotName = 'results/' + 'output-' + str(step).zfill(3) + writer = VTKWriter.VTKWriter(self.mesh, baseFileName=plotName) + + writer.add_nodal_field(name='disp', nodalData=U, fieldType=VTKWriter.VTKFieldType.VECTORS) + energyDensities, stresses = mechFuncs.compute_output_energy_densities_and_stresses(U, p.state_data) + cellEnergyDensities = FunctionSpace.project_quadrature_field_to_element_field(funcSpace, energyDensities) + cellStresses = FunctionSpace.project_quadrature_field_to_element_field(funcSpace, stresses) + writer.add_cell_field(name='strain_energy_density', + cellData=cellEnergyDensities, + fieldType=VTKWriter.VTKFieldType.SCALARS) + writer.add_cell_field(name='cauchy_stress', + cellData=cellStresses, + fieldType=VTKWriter.VTKFieldType.TENSORS) + writer.write() + + # Problem Set Up + Uu = self.dofManager.get_unknown_values(np.zeros(self.mesh.coords.shape)) + ivs = mechFuncs.compute_initial_state() + p = Objective.Params(bc_data=0., state_data=ivs) + self.objective = Objective.Objective(energy_function, Uu, p) + + # Load over the steps + force = 0. + + force_inc = self.maxForce / self.steps + + K = self.objective.hessian(Uu) + # print(f"Hessian first row: {K[0,:5]}") + # print(f"Hessian Shape: {K.shape}") + + for step in range(1, self.steps): + print('------------------------------------') + print('LOAD STEP', step) + force += force_inc + p = Objective.param_index_update(p, 0, force) + print(f"Step {step}: Applied force = {force}, Updated p = {p.bc_data}") + print(f"Force inside energy function: {p.bc_data}") + + nodal_forces = jit(grad(energy_function, argnums=0)) + F = nodal_forces(Uu, p) + Uu = solve(K, F) + + if self.writeOutput: + write_output(Uu, p, step) + + + + +if __name__ == '__main__': + app = CompressionTest() + app.reload_mesh() + app.run() + + + + + diff --git a/optimism/test/test_compression_MPC.py b/optimism/test/test_compression_MPC.py new file mode 100644 index 00000000..5f85ef78 --- /dev/null +++ b/optimism/test/test_compression_MPC.py @@ -0,0 +1,199 @@ +# -------------------------------------------------------------------------- +# TEST CASE - Compression Test for implementation of MPCs +# -------------------------------------------------------------------------- +# Note - Using Cubit Coreform for meshing, solved using NON-LINEAR SOLVER +# -------------------------------------------------------------------------- +import sys +sys.path.insert(0, "/home/sarvesh/Documents/Github/optimism") + +from jax import jit +from optimism import VTKWriter +from optimism import EquationSolver +from optimism import FunctionSpace +from optimism import Mechanics +from optimism import Mesh +from optimism import Objective +from optimism import QuadratureRule +from optimism import ReadExodusMesh +from optimism.FunctionSpace import DofManagerMPC +from optimism.FunctionSpace import EssentialBC +from optimism.material import Neohookean + +import jax.numpy as np +import numpy as onp +from scipy.spatial import cKDTree + + +class CompressionTest: + + def __init__(self): + self.writeOutput = True + + self.quadRule = QuadratureRule.create_quadrature_rule_on_triangle(degree=2) + self.lineQuadRule = QuadratureRule.create_quadrature_rule_1D(degree=2) + + self.EBCs = [ + EssentialBC(nodeSet='bottom', component=0), + EssentialBC(nodeSet='bottom', component=1) + ] + + shearModulus = 0.855 # MPa + bulkModulus = 1000*shearModulus # MPa + youngModulus = 9.0*bulkModulus*shearModulus / (3.0*bulkModulus + shearModulus) + poissonRatio = (3.0*bulkModulus - 2.0*shearModulus) / 2.0 / (3.0*bulkModulus + shearModulus) + props = { + 'elastic modulus': youngModulus, + 'poisson ratio': poissonRatio, + 'version': 'coupled' + } + + self.matModel = Neohookean.create_material_model_functions(props) + + self.input_mesh = 'square_mesh_compression.exo' + + self.maxForce = -0.2 + self.steps = 50 + + def create_field(self, Uu): + def get_ubcs(): + V = np.zeros(self.mesh.coords.shape) + return self.dofManager.get_bc_values(V) + U = self.dofManager.create_field(Uu) + U.at[self.dofManager.isBc].set(get_ubcs()) + return U + # return self.dofManager.create_field(Uu, get_ubcs()) + + def create_C_and_s(self): + """ + Create constraint matrix C and shift vector s based on master-slave pairs (JAX-safe). + """ + num_pairs = len(self.master_slave_pairs) + total_dofs = self.mesh.coords.shape[0] * self.dim + + # Must initialize with jax.numpy (np here is jax.numpy in your imports) + C = np.zeros((num_pairs * self.dim, total_dofs)) + s = np.zeros((num_pairs * self.dim,)) + + master_nodes = list(self.master_slave_pairs.keys()) + slave_nodes = list(self.master_slave_pairs.values()) + + for i, (master, slave) in enumerate(zip(master_nodes, slave_nodes)): + for d in range(self.dim): + row_idx = i * self.dim + d + col_idx = master * self.dim + d + C = C.at[row_idx, col_idx].set(1.0) + + # s remains zeros unless you want shifts + return C, s + + def reload_mesh(self): + origMesh = ReadExodusMesh.read_exodus_mesh(self.input_mesh) + self.mesh = Mesh.create_higher_order_mesh_from_simplex_mesh(origMesh, order=2, createNodeSetsFromSideSets=True) + coords = self.mesh.coords + self.dim = coords.shape[1] # either 2 or 3 dimensions + tol = 1e-8 + + # Create node sets + self.mesh.nodeSets['master'] = np.flatnonzero(coords[:,0] < -0.5 + tol) + self.mesh.nodeSets['slave'] = np.flatnonzero(coords[:,0] > 0.5 - tol) + + # Sort by y-coordinate + master_coords = coords[self.mesh.nodeSets['master']] + slave_coords = coords[self.mesh.nodeSets['slave']] + + master_sorted_idx = np.argsort(master_coords[:, 1]) + slave_sorted_idx = np.argsort(slave_coords[:, 1]) + + self.master_slave_pairs = { + int(self.mesh.nodeSets['slave'][slave_sorted_idx[i]]): + int(self.mesh.nodeSets['master'][master_sorted_idx[i]]) + for i in range(min(len(master_sorted_idx), len(slave_sorted_idx))) + } + + print("Master-Slave Pairs:", self.master_slave_pairs) + + funcSpace = FunctionSpace.construct_function_space(self.mesh, self.quadRule) + + # USER CREATES C AND s HERE INSIDE TEST CASE + self.C, self.s = self.create_C_and_s() + + # PASS IT INTO DofManagerMPC + self.dofManager = DofManagerMPC(funcSpace, self.dim, self.EBCs, self.master_slave_pairs, self.C, self.s) + + # Top surface area + surfaceXCoords = self.mesh.coords[self.mesh.nodeSets['top']][:,0] + self.tractionArea = np.max(surfaceXCoords) - np.min(surfaceXCoords) + + def run(self): + funcSpace = FunctionSpace.construct_function_space(self.mesh, self.quadRule) + mechFuncs = Mechanics.create_mechanics_functions(funcSpace, mode2D='plane strain', materialModel=self.matModel) + + def energy_function(Uu, p): + U = self.create_field(Uu) + internal_variables = p.state_data + strainEnergy = mechFuncs.compute_strain_energy(U, internal_variables) + + F = p.bc_data + def force_function(x, n): + return np.array([0.0, F/self.tractionArea]) + + loadPotential = Mechanics.compute_traction_potential_energy( + funcSpace, U, self.lineQuadRule, self.mesh.sideSets['top'], + force_function) + + return strainEnergy + loadPotential + + def write_output(Uu, p, step): + U = self.create_field(Uu) + plotName = 'results/' + 'output-' + str(step).zfill(3) + writer = VTKWriter.VTKWriter(self.mesh, baseFileName=plotName) + + writer.add_nodal_field(name='disp', nodalData=U, fieldType=VTKWriter.VTKFieldType.VECTORS) + energyDensities, stresses = mechFuncs.compute_output_energy_densities_and_stresses(U, p.state_data) + cellEnergyDensities = FunctionSpace.project_quadrature_field_to_element_field(funcSpace, energyDensities) + cellStresses = FunctionSpace.project_quadrature_field_to_element_field(funcSpace, stresses) + writer.add_cell_field(name='strain_energy_density', + cellData=cellEnergyDensities, + fieldType=VTKWriter.VTKFieldType.SCALARS) + writer.add_cell_field(name='cauchy_stress', + cellData=cellStresses, + fieldType=VTKWriter.VTKFieldType.TENSORS) + writer.write() + + # Problem Set Up + Uu = self.dofManager.get_unknown_values(np.zeros(self.mesh.coords.shape)) + print("Mesh coords: ", self.mesh.coords.shape) + print("Shape of Uu: ", Uu.shape) + # Uu = self.dofManager.get_unknown_values(np.zeros((self.mesh.coords.shape[0] * self.dim))) + ivs = mechFuncs.compute_initial_state() + p = Objective.Params(bc_data=0., state_data=ivs) + self.objective = Objective.Objective(energy_function, Uu, p, precondStrategy=None) + + # Load over the steps + force = 0. + + force_inc = self.maxForce / self.steps + + + for step in range(1, self.steps): + print('------------------------------------') + print('LOAD STEP', step) + force += force_inc + p = Objective.param_index_update(p, 0, force) + Uu, solverSuccess = EquationSolver.nonlinear_equation_solve(self.objective, Uu, p, EquationSolver.get_settings()) + + if self.writeOutput: + write_output(Uu, p, step) + + + + +if __name__ == '__main__': + app = CompressionTest() + app.reload_mesh() + app.run() + + + + + diff --git a/optimism/test/test_shear.py b/optimism/test/test_shear.py new file mode 100644 index 00000000..82c18026 --- /dev/null +++ b/optimism/test/test_shear.py @@ -0,0 +1,161 @@ +# Test Case - Displacement based Shear Test +# --------------------------------------------------- +# Note - This approach involves segregation of master and slave DOFs +# --------------------------------------------------- + +import sys +sys.path.insert(0, "/home/sarvesh/Documents/Github/optimism") + + +import jax.numpy as np +from jax import jit, grad + +from optimism import EquationSolver as EqSolver +from optimism import QuadratureRule +from optimism import Mechanics +from optimism import Mesh +from optimism import Objective +from optimism import FunctionSpace +from optimism.Mesh import create_higher_order_mesh_from_simplex_mesh +from optimism.material import Neohookean +from optimism.FunctionSpace import DofManagerMPC + +from MeshFixture import * +import matplotlib.pyplot as plt + + + +class ShearTest(MeshFixture): + + def setUp(self): + dummyDispGrad = np.eye(2) + self.mesh = self.create_mesh_disp_and_nodeset_layers(5, 5, [0.,1.], [0.,1.], + lambda x: dummyDispGrad.dot(x))[0] + self.mesh = create_higher_order_mesh_from_simplex_mesh(self.mesh, order=1, createNodeSetsFromSideSets=True) + self.quadRule = QuadratureRule.create_quadrature_rule_on_triangle(degree=1) + self.fs = FunctionSpace.construct_function_space(self.mesh, self.quadRule) + + self.EBCs = [FunctionSpace.EssentialBC(nodeSet='top', component=0), + FunctionSpace.EssentialBC(nodeSet='top', component=1), + FunctionSpace.EssentialBC(nodeSet='bottom', component=0), + FunctionSpace.EssentialBC(nodeSet='bottom', component=1)] + + # dofManager = DofManagerMPC(self.fs, 2, self.EBCs, self.mesh) + + # self.master_array = dofManager.isMaster + # self.slave_array = dofManager.isSlave + # self.assembly = dofManager.dofAssembly + # print("Assembly: ", self.assembly) + + shearModulus = 0.855 # MPa + bulkModulus = 1000*shearModulus # MPa + youngModulus = 9.0*bulkModulus*shearModulus / (3.0*bulkModulus + shearModulus) + poissonRatio = (3.0*bulkModulus - 2.0*shearModulus) / 2.0 / (3.0*bulkModulus + shearModulus) + props = { + 'elastic modulus': youngModulus, + 'poisson ratio': poissonRatio, + 'version': 'coupled' + } + self.mat = Neohookean.create_material_model_functions(props) + + self.freq = 0.3 / 2.0 / np.pi + self.cycles = 1 + self.steps_per_cycle = 64 + self.steps = self.cycles*self.steps_per_cycle + totalTime = self.cycles / self.freq + self.dt = totalTime / self.steps + self.maxDisp = 1.0 + + def run(self): + mechFuncs = Mechanics.create_mechanics_functions(self.fs, + "plane strain", + self.mat) + dofManager = DofManagerMPC(self.fs, 2, self.EBCs, self.mesh) + + def create_field(Uu, disp): + def get_ubcs(disp): + V = np.zeros(self.mesh.coords.shape) + index = (self.mesh.nodeSets['top'], 0) + V = V.at[index].set(disp) + return dofManager.get_bc_values(V) + + return dofManager.create_field(Uu, get_ubcs(disp)) + + def energy_function_all_dofs(U, p): + internalVariables = p.state_data + return mechFuncs.compute_strain_energy(U, internalVariables, self.dt) + + def compute_energy(Uu, p): + U = create_field(Uu, p.bc_data) + return energy_function_all_dofs(U, p) + + nodal_forces = jit(grad(energy_function_all_dofs, argnums=0)) + integrate_free_energy = jit(mechFuncs.compute_strain_energy) + + def write_output(Uu, p, step): + from optimism import VTKWriter + U = create_field(Uu, p.bc_data) + plotName = "../test_results/disp_based_shear_test/"+'output-'+str(step).zfill(3) + writer = VTKWriter.VTKWriter(self.mesh, baseFileName=plotName) + + writer.add_nodal_field(name='displ', nodalData=U, fieldType=VTKWriter.VTKFieldType.VECTORS) + + energyDensities, stresses = mechFuncs.\ + compute_output_energy_densities_and_stresses(U, p.state_data, self.dt) + cellEnergyDensities = FunctionSpace.project_quadrature_field_to_element_field(self.fs, energyDensities) + cellStresses = FunctionSpace.project_quadrature_field_to_element_field(self.fs, stresses) + writer.add_cell_field(name='strain_energy_density', + cellData=cellEnergyDensities, + fieldType=VTKWriter.VTKFieldType.SCALARS) + writer.add_cell_field(name='stress', + cellData=cellStresses, + fieldType=VTKWriter.VTKFieldType.TENSORS) + writer.write() + + Uu = dofManager.get_unknown_values(np.zeros(self.mesh.coords.shape)) + ivs = mechFuncs.compute_initial_state() + p = Objective.Params(bc_data=0., state_data=ivs) + U = create_field(Uu, p.bc_data) + self.objective = Objective.ObjectiveMPC(compute_energy, Uu, p, dofManager) + + index = (self.mesh.nodeSets['top'], 0) + + time = 0.0 + times = [] + externalWorkStore = [] + incrementalPotentialStore = [] + forceHistory = [] + dispHistory = [] + for step in range(1, self.steps+1): + force_prev = np.array(nodal_forces(U, p).at[index].get()) + applied_disp_prev = U.at[index].get() + + # disp = self.maxDisp * np.sin(2.0 * np.pi * self.freq * time) + disp = self.maxDisp * self.freq * time + print("Displacement in this load step: ", disp) + + p = Objective.param_index_update(p, 0, disp) + Uu, solverSuccess = EqSolver.nonlinear_equation_solve(self.objective, Uu, p, EqSolver.get_settings(), useWarmStart=False) + U = create_field(Uu, p.bc_data) + ivs = mechFuncs.compute_updated_internal_variables(U, p.state_data, self.dt) + p = Objective.param_index_update(p, 1, ivs) + + write_output(Uu, p, step) + + force = np.array(nodal_forces(U, p).at[index].get()) + applied_disp = U.at[index].get() + externalWorkStore.append( 0.5*np.tensordot((force + force_prev),(applied_disp - applied_disp_prev), axes=1) ) + incrementalPotentialStore.append(integrate_free_energy(U, ivs, self.dt)) + + forceHistory.append( np.sum(force) ) + dispHistory.append(disp) + + times.append(time) + time += self.dt + + + + +app = ShearTest() +app.setUp() +app.run() \ No newline at end of file