finished up the failing newmark tests and think this looks good for now.

examples/window_pain/window_pain_new.py

@@ -1,8 +1,5 @@
-from optimism.material import Neohookean as MatModel
-from optimism.Domain import Domain, Problem
-from optimism.EquationSolver import get_settings
-from optimism.FunctionSpace import EssentialBC
-from optimism.ReadExodusMesh import read_exodus_mesh
+from optimism import *
+# from optimism.material import Neohookean as MatModel
 
 mesh = read_exodus_mesh('./geometry.g')
 
@@ -18,7 +15,9 @@
   'poisson ratio'  : 0.3,
   'version'        : 'coupled'
 }
-mat_model = MatModel.create_material_model_functions(props)
+mat_models = {
+  'block_1': Neohookean.create_material_model_functions(props)
+}
 
 eq_settings = get_settings(
   use_incremental_objective=False,
@@ -28,7 +27,7 @@
   tol=5e-8
 )
 
-domain = Domain(mesh, ebcs, mat_model, 'plane strain')
+domain = Domain(mesh, ebcs, mat_models, 'plane strain', disp_nset='nset_outer_top')
 problem = Problem(domain, eq_settings)
 
 # setup
@@ -39,6 +38,6 @@
 n_steps = 20
 ddisp = -0.2 / n_steps
 for n in range(n_steps):
-  print(f'Load step {n}')
+  print(f'Load step {n + 1}')
   disp += ddisp
-  Uu, p = problem.solve_load_step(Uu, p, n, disp)
+  Uu, p = problem.solve_load_step(Uu, p, n + 1, disp)

optimism/Mechanics.py

@@ -1,100 +1,15 @@
-from collections import namedtuple
 from optimism.JaxConfig import *
 from optimism import SparseMatrixAssembler
 from optimism import FunctionSpace
 from optimism import Mesh
-from optimism.TensorMath import tensor_2D_to_3D
 from optimism import QuadratureRule
 from optimism import Interpolants
+from optimism.TensorMath import tensor_2D_to_3D
 from typing import Callable, Dict, List, Optional
 import equinox as eqx
 import functools
 import jax
 
-# # TODO add types to Callable's for improved autodocing
-# # TODO once all the equinox stuff is hooked up reconsider jitting by default here
-# class MechanicsFunctions(eqx.Module):
-#     fspace: FunctionSpace.FunctionSpace
-#     mat_models: List[any] # TODO type this
-#     modify_element_gradient: Callable
-#     element_hess_func: Callable
-#     L: Callable
-#     output_constitutive: Callable
-
-#     def __init__(self, fspace, mat_models, modify_element_gradient):
-#         self.fspace = fspace
-#         self.mat_models = mat_models
-#         self.modify_element_gradient = modify_element_gradient
-#         self.element_hess_func = hessian(FunctionSpace.integrate_element_from_local_field)
-#         self.L = strain_energy_density_to_lagrangian_density(mat_models[0].compute_energy_density)
-#         self.output_constitutive = jax.value_and_grad(self.L, 1)
-
-#     # public methods
-#     @eqx.filter_jit
-#     def compute_element_stiffnesses(self, U, state, dt=0.0):
-#         return self._compute_element_stiffnesses(U, state, dt)
-
-#     def compute_initial_state(self):
-#         shape = self.fspace.mesh.num_elements, QuadratureRule.len(self.fspace.quadratureRule), 1
-#         return np.tile(self.mat_models[0].compute_initial_state(), shape)
-
-#     @eqx.filter_jit
-#     def compute_output_energy_densities_and_stresses(self, U, state, dt=0.0):
-#         return FunctionSpace.evaluate_on_block(fs, U, state, dt, self.output_constitutive, slice(None), modify_element_gradient=modify_element_gradient)
-
-#     def compute_strain_energy(self, U, state, dt=0.0):
-#         L = strain_energy_density_to_lagrangian_density(self.mat_models[0].compute_energy_density)
-#         return FunctionSpace.integrate_over_block(
-#             self.fspace, U, state, dt, L,
-#             slice(None),
-#             modify_element_gradient=self.modify_element_gradient
-#         )
-
-#     @eqx.filter_jit
-#     def compute_updated_internal_variables(self, U, state, dt=0.0):
-#         dispGrads = FunctionSpace.compute_field_gradient(self.fspace, U, self.modify_element_gradient)
-#         dgQuadPointRavel = dispGrads.reshape(dispGrads.shape[0]*dispGrads.shape[1],*dispGrads.shape[2:])
-#         stQuadPointRavel = state.reshape(state.shape[0]*state.shape[1],*state.shape[2:])
-#         statesNew = vmap(self.mat_models[0].compute_state_new, (0, 0, None))(dgQuadPointRavel, stQuadPointRavel, dt)
-#         return statesNew.reshape(state.shape)
-
-#     @eqx.filter_jit
-#     def compute_output_material_qoi(self, U, state, dt=0.0):
-#         return FunctionSpace.evaluate_on_block(
-#             self.fspace, U, state, dt, 
-#             self._lagrangian_qoi, 
-#             slice(None), 
-#             modify_element_gradient=self.modify_element_gradient
-#         )
-
-#     def integrated_material_qoi(self, U, state, dt=0.0):
-#         return FunctionSpace.integrate_over_block(
-#             self.fspace, U, state, dt, 
-#             self._lagrangian_qoi,
-#             slice(None),
-#             modify_element_gradient=self.modify_element_gradient
-#         )
-
-#     # private methods
-#     # TODO there's probably some cleanup that can be done
-#     def _compute_element_stiffnesses(self, U, state, dt):
-#         f =  jax.vmap(self._compute_element_stiffness_from_global_fields,
-#                 (None, None, 0, None, 0, 0, 0, 0, None))
-
-#                 # (None, None, 0, None, 0, 0, 0, 0, None, None))
-#         fs = self.fspace
-#         return f(U, fs.mesh.coords, state, dt, fs.mesh.conns, fs.shapes, fs.shapeGrads, fs.vols, self.L)#,
-#                 # self.L, self.modify_element_gradient)
-
-#     def _compute_element_stiffness_from_global_fields(self, U, coords, elInternals, dt, elConn, elShapes, elShapeGrads, elVols, func):#, lagrangian_density, modify_element_gradient):
-#         elDisp = U[elConn,:]
-#         elCoords = coords[elConn,:]
-#         return self.element_hess_func(elDisp, elCoords, elInternals, dt, elShapes, elShapeGrads,
-#                                 elVols, func, self.modify_element_gradient)
-
-#     def _lagrangian_qoi(self, U, gradU, Q, X, dt):
-#         return self.mat_models[0].compute_material_qoi(gradU, Q, dt)
-
 
 # TODO eventually deprecate this and make multi blocks the default
 class MechanicsFunctions(eqx.Module):
@@ -257,17 +172,15 @@ def __init__(
 
     # public methods
     def compute_algorithmic_energy(self, U, UPredicted, stateVariables, dt):
-        # return self.compute_newmark_lagrangian(
-        #     U, UPredicted, stateVariables,
-        #     self.mat_models[0].density, dt, self.newmark_parameters.beta
-        # )
         fs, blockModels = self.fspace, self.mat_models
         L = 0.0
         for blockKey, blockModel in blockModels.items():
             elemIds = fs.mesh.blocks[blockKey]
             L = L + self.compute_newmark_lagrangian(
-
+                U, UPredicted, stateVariables, blockModel, dt, self.newmark_parameters.beta,
+                elemIds
             )
+        return L
 
     def compute_element_hessians(self, U, UPredicted, stateVariables, dt):
         fs, blockModels = self.fspace, self.mat_models
@@ -288,27 +201,42 @@ def compute_element_hessians(self, U, UPredicted, stateVariables, dt):
 
     def compute_element_masses(self):
         V = np.zeros_like(self.fspace.mesh.coords)
-        stateVariables = np.zeros((self.fspace.mesh.num_elements, QuadratureRule.len(self.fspace.quadratureRule)))
-        return self._compute_element_masses(V, stateVariables, self.mat_models[0].density, self.modify_element_gradient)
+        stateVariables = np.zeros((
+            self.fspace.mesh.num_elements, QuadratureRule.len(self.fspace.quadratureRule))
+        )
+        fs, blockModels = self.fspace, self.mat_models
+        nen = Interpolants.num_nodes(fs.mesh.parentElement)
+        elementMasses = np.zeros((fs.mesh.num_elements, nen, 2, nen, 2))
+
+        for blockKey, blockModel in blockModels.items():
+            elemIds = fs.mesh.blocks[blockKey]
+            elementMasses = self._compute_element_masses(
+                V, stateVariables, blockModel.density, elemIds
+            )
 
-    def compute_newmark_lagrangian(self, U, UPredicted, internals, density, dt, newmarkBeta):
+    def compute_newmark_lagrangian(
+        self, 
+        U, UPredicted, internals, blockModel, dt, newmarkBeta, elemIds
+    ):
         # We can't quite fuse these kernels because KE uses the velocity field and
         # the strain energy uses the displacements. If profiling suggests fusing
         # is beneficial, we could add the time derivative field to the Lagrangian
         # density definition.
 
         def lagrangian_density(W, gradW, Q, X, dtime):
-            return self._kinetic_energy_density(W, density)
+            return self._kinetic_energy_density(W, blockModel.density)
         KE = FunctionSpace.integrate_over_block(
             self.fspace, U - UPredicted, internals, dt,
-            lagrangian_density, slice(None)
+            lagrangian_density, elemIds
         )
         KE *= 1 / (newmarkBeta*dt**2)
 
-        lagrangian_density = strain_energy_density_to_lagrangian_density(self.mat_models[0].compute_energy_density)
+        lagrangian_density = strain_energy_density_to_lagrangian_density(
+            blockModel.compute_energy_density
+        )
         SE = FunctionSpace.integrate_over_block(
             self.fspace, U, internals, dt, lagrangian_density,
-            slice(None), modify_element_gradient=self.modify_element_gradient
+            elemIds, modify_element_gradient=self.modify_element_gradient
         )
         return SE + KE
 
@@ -346,14 +274,20 @@ def lagrangian_density(U, gradU, Q, X, dt):
         unused = 0.0
         return FunctionSpace.integrate_over_block(self.fspace, V, internals, unused, lagrangian_density, slice(None))
 
-    def _compute_element_masses(self, U, internals, density):
+    def _compute_element_masses(self, U, internals, density, elemIds):
         def lagrangian_density(V, gradV, Q, X, dt):
             return self._kinetic_energy_density(V, density)
-        f = vmap(self._compute_element_stiffness_from_global_fields, (None, None, 0, None, 0 ,0 ,0 ,0, None))
+        f = vmap(
+            self._compute_element_stiffness_from_global_fields, 
+            (None, None, 0, None, 0 ,0 ,0 ,0, None)
+        )
         fs = self.fspace
         unusedDt = 0.0
-        return f(U, fs.mesh.coords, internals, unusedDt, fs.mesh.conns, fs.shapes, fs.shapeGrads,
-                fs.vols, lagrangian_density)
+        return f(
+            U, fs.mesh.coords, internals, unusedDt, 
+            fs.mesh.conns[elemIds], fs.shapes[elemIds], fs.shapeGrads[elemIds], fs.vols[elemIds], 
+            lagrangian_density
+        )
 
     def _compute_newmark_element_hessians(self, U, UPredicted, internals, density, dt, newmarkBeta, strain_energy_density, modify_element_gradient, elemIds):
         def lagrangian_density(W, gradW, Q, X, dtime):
@@ -373,11 +307,13 @@ def _kinetic_energy_density(self, V, density):
         # assumes spatially homogeneous density
         return 0.5*density*np.dot(V, V)
 
-# NewmarkParameters = namedtuple('NewmarkParameters', ['gamma', 'beta'], defaults=[0.5, 0.25])
 class NewmarkParameters(eqx.Module):
     gamma: Optional[float] = 0.5
     beta: Optional[float] = 0.25
 
+
+
+
 # gradient transformations
 def plane_strain_gradient_transformation(elemDispGrads, elemShapes, elemVols, elemNodalDisps, elemNodalCoords):
     return vmap(tensor_2D_to_3D)(elemDispGrads)
@@ -488,13 +424,6 @@ def create_mechanics_functions(functionSpace, mode2D, materialModel, pressurePro
 
 def create_multi_block_mechanics_functions(functionSpace, mode2D, materialModels, pressureProjectionDegree=None):
     fs = functionSpace
-
-    # if mode2D == 'plane strain':
-    #     grad_2D_to_3D = plane_strain_gradient_transformation
-    # elif mode2D == 'axisymmetric':
-    #     grad_2D_to_3D = axisymmetric_element_gradient_transformation
-    # else:
-    #     raise NotImplementedError
     grad_2D_to_3D = parse_2D_to_3D_gradient_transformation(mode2D)
 
     modify_element_gradient = grad_2D_to_3D

optimism/__init__.py

@@ -13,3 +13,9 @@
 #config.update("jax_disable_jit", True)
 
 del config
+
+from .material import *
+from .Domain import Domain, Problem
+from .EquationSolver import get_settings
+from .FunctionSpace import EssentialBC
+from .ReadExodusMesh import read_exodus_mesh

optimism/material/__init__.py

@@ -0,0 +1,5 @@
+from . import Gent
+from . import HyperViscoelastic
+from . import J2Plastic
+from . import LinearElastic
+from . import Neohookean

test/test_Newmark.py

@@ -252,14 +252,18 @@ def setUp(self):
                  'density': 20.0}
         # density chosen to make kinetic energy and strain energy comparable
         # Want bugs in either to show up appreciably in error
-        materialModel = Neohookean.create_material_model_functions(props)
+        materialModels = {
+            'block': Neohookean.create_material_model_functions(props)
+        }
 
         newmarkParams = Mechanics.NewmarkParameters()
 
-        self.dynamics = Mechanics.create_dynamics_functions(self.fs, 'plane strain', materialModel, newmarkParams)
-
-        self.compute_stress = jax.jacfwd(materialModel.compute_energy_density)
+        self.dynamics = Mechanics.create_dynamics_functions(self.fs, 'plane strain', materialModels, newmarkParams)
 
+        # self.compute_stress = jax.jacfwd(materialModel.compute_energy_density)
+        self.compute_stresses = dict()
+        for blockKey, blockModel in materialModels.items():
+            self.compute_stresses[blockKey] = jax.jacfwd(blockModel.compute_energy_density)
 
     def test_patch_test(self):
         dt = 1.0
@@ -329,7 +333,7 @@ def traction(X, N):
                 dispGrad3D = np.zeros((3,3)).at[:2, :2].set(dispGrad)
                 q = np.array([0.0])
                 dt = 0.0
-                P = self.compute_stress(dispGrad3D, q, dt)[:2, :2]
+                P = self.compute_stresses['block'](dispGrad3D, q, dt)[:2, :2]
                 return np.dot(P, N)
 
             dt = t - p.time[1]