From 66caab8ee17dc358c0dc7813d6f4be6424a2d017 Mon Sep 17 00:00:00 2001 From: "Craig M. Hamel" Date: Thu, 1 Aug 2024 21:33:34 -0400 Subject: [PATCH] finished up the failing newmark tests and think this looks good for now. --- examples/window_pain/window_pain_new.py | 17 ++- optimism/Mechanics.py | 149 +++++++----------------- optimism/__init__.py | 6 + optimism/material/__init__.py | 5 + test/test_Newmark.py | 14 ++- 5 files changed, 67 insertions(+), 124 deletions(-) diff --git a/examples/window_pain/window_pain_new.py b/examples/window_pain/window_pain_new.py index a6d11a10..d5e4f104 100644 --- a/examples/window_pain/window_pain_new.py +++ b/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) \ No newline at end of file + Uu, p = problem.solve_load_step(Uu, p, n + 1, disp) diff --git a/optimism/Mechanics.py b/optimism/Mechanics.py index 6d3d4f86..bc510118 100644 --- a/optimism/Mechanics.py +++ b/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 diff --git a/optimism/__init__.py b/optimism/__init__.py index 4c2de879..0870c305 100644 --- a/optimism/__init__.py +++ b/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 diff --git a/optimism/material/__init__.py b/optimism/material/__init__.py index e69de29b..8ac5cf16 100644 --- a/optimism/material/__init__.py +++ b/optimism/material/__init__.py @@ -0,0 +1,5 @@ +from . import Gent +from . import HyperViscoelastic +from . import J2Plastic +from . import LinearElastic +from . import Neohookean diff --git a/test/test_Newmark.py b/test/test_Newmark.py index 03642faf..d1a6417d 100644 --- a/test/test_Newmark.py +++ b/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]