From c2a0cfd91a62a0c2b10ed1081d7cfbc4d2dcf350 Mon Sep 17 00:00:00 2001 From: ralberd Date: Tue, 31 Oct 2023 16:48:41 -0600 Subject: [PATCH] adding finite difference check test harness; adding FD checks for hyperelastic objectives; adding FD check for J2 objective (is failing) --- .../inverse/test/FiniteDifferenceFixture.py | 50 ++++ .../test/adjoint_problem_function_space.py | 0 .../inverse/test/test_Hyperelastic_inverse.py | 213 +++++++++++++++ .../test/test_J2Plastic_inverse.py | 256 ++++++++++++------ 4 files changed, 437 insertions(+), 82 deletions(-) create mode 100644 optimism/inverse/test/FiniteDifferenceFixture.py rename optimism/{material => inverse}/test/adjoint_problem_function_space.py (100%) create mode 100644 optimism/inverse/test/test_Hyperelastic_inverse.py rename optimism/{material => inverse}/test/test_J2Plastic_inverse.py (61%) diff --git a/optimism/inverse/test/FiniteDifferenceFixture.py b/optimism/inverse/test/FiniteDifferenceFixture.py new file mode 100644 index 0000000..be8ca05 --- /dev/null +++ b/optimism/inverse/test/FiniteDifferenceFixture.py @@ -0,0 +1,50 @@ +from optimism.test.MeshFixture import * +from collections import namedtuple +import numpy + +class FiniteDifferenceFixture(MeshFixture): + def assertFiniteDifferenceCheckHasVShape(self, errors, tolerance=1e-6): + minError = min(errors) + self.assertLess(minError, tolerance, "Smallest finite difference error not less than tolerance.") + self.assertLess(minError, errors[0], "Finite difference error does not decrease from initial step size.") + self.assertLess(minError, errors[-1], "Finite difference error does not increase after reaching minimum. Try more finite difference steps.") + + def build_direction_vector(self, numDesignVars, seed=123): + + numpy.random.seed(seed) + directionVector = numpy.random.uniform(-1.0, 1.0, numDesignVars) + normVector = directionVector / numpy.linalg.norm(directionVector) + + return np.array(normVector) + + def compute_finite_difference_errors(self, stepSize, steps, initialParameters, printOutput=True): + storedState = self.forward_solve(initialParameters) + originalObjective = self.compute_objective_function(storedState, initialParameters) + gradient = self.compute_gradient(storedState, initialParameters) + + directionVector = self.build_direction_vector(initialParameters.shape[0]) + directionalDerivative = np.tensordot(directionVector, gradient, axes=1) + + fd_values = [] + errors = [] + for i in range(0, steps): + perturbedParameters = initialParameters + stepSize * directionVector + storedState = self.forward_solve(perturbedParameters) + perturbedObjective = self.compute_objective_function(storedState, perturbedParameters) + + fd_value = (perturbedObjective - originalObjective) / stepSize + fd_values.append(fd_value) + + error = abs(directionalDerivative - fd_value) + errors.append(error) + + stepSize *= 1e-1 + + if printOutput: + print("\n grad'*dir | FD approx | abs error") + print("--------------------------------------------------------------------------------") + for i in range(0, steps): + print(f" {directionalDerivative} | {fd_values[i]} | {errors[i]}") + + return errors + \ No newline at end of file diff --git a/optimism/material/test/adjoint_problem_function_space.py b/optimism/inverse/test/adjoint_problem_function_space.py similarity index 100% rename from optimism/material/test/adjoint_problem_function_space.py rename to optimism/inverse/test/adjoint_problem_function_space.py diff --git a/optimism/inverse/test/test_Hyperelastic_inverse.py b/optimism/inverse/test/test_Hyperelastic_inverse.py new file mode 100644 index 0000000..bbe7ce4 --- /dev/null +++ b/optimism/inverse/test/test_Hyperelastic_inverse.py @@ -0,0 +1,213 @@ +import unittest + +import jax +import jax.numpy as np +from matplotlib import pyplot as plt + +from optimism import EquationSolver as EqSolver +from optimism import QuadratureRule +from optimism import FunctionSpace +from optimism import Mechanics +from optimism import Objective +from optimism import Mesh +from optimism.material import Neohookean + +from optimism.inverse.test.FiniteDifferenceFixture import FiniteDifferenceFixture + +# for adjoint +import numpy as onp +from scipy.sparse import linalg + +# misc. modules in test directory for now while I test +import adjoint_problem_function_space as AdjointFunctionSpace + +class NeoHookeanGlobalMeshAdjointSolveFixture(FiniteDifferenceFixture): + def setUp(self): + dispGrad0 = np.array([[0.4, -0.2], + [-0.04, 0.68]]) + self.initialMesh, self.U = self.create_mesh_and_disp(4,4,[0.,1.],[0.,1.], + lambda x: dispGrad0@x) + + shearModulus = 0.855 # MPa + bulkModulus = 100*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.materialModel = Neohookean.create_material_model_functions(props) + + self.quadRule = QuadratureRule.create_quadrature_rule_on_triangle(degree=1) + + self.EBCs = [FunctionSpace.EssentialBC(nodeSet='all_boundary', component=0), + FunctionSpace.EssentialBC(nodeSet='all_boundary', component=1)] + + self.steps = 2 + + def forward_solve(self, parameters): + + self.mesh = Mesh.construct_mesh_from_basic_data(parameters.reshape(self.initialMesh.coords.shape),\ + self.initialMesh.conns, self.initialMesh.blocks,\ + self.initialMesh.nodeSets, self.initialMesh.sideSets) + + functionSpace = FunctionSpace.construct_function_space(self.mesh, self.quadRule) + mechFuncs = Mechanics.create_mechanics_functions(functionSpace, + "plane strain", + self.materialModel) + self.dofManager = FunctionSpace.DofManager(functionSpace, 2, self.EBCs) + Ubc = self.dofManager.get_bc_values(self.U) + + Ubc_inc = Ubc / self.steps + ivs = mechFuncs.compute_initial_state() + p = Objective.Params(bc_data=np.zeros(Ubc.shape), state_data=ivs) + + def compute_energy(Uu, p): + U = self.dofManager.create_field(Uu, p.bc_data) + internalVariables = p.state_data + return mechFuncs.compute_strain_energy(U, internalVariables) + + Uu = 0.0*self.dofManager.get_unknown_values(self.U) + self.objective = Objective.Objective(compute_energy, Uu, p) + + storedState = [] + storedState.append((Uu, p)) + + for step in range(1, self.steps+1): + p = Objective.param_index_update(p, 0, step*Ubc_inc) + Uu = EqSolver.nonlinear_equation_solve(self.objective, Uu, p, EqSolver.get_settings(), useWarmStart=False) + storedState.append((Uu, p)) + + return storedState + + def strain_energy_objective(self, storedState, parameters): + coords = parameters.reshape(self.mesh.coords.shape) + adjoint_func_space = AdjointFunctionSpace.construct_function_space_for_adjoint(coords, self.mesh, self.quadRule) + mech_funcs = Mechanics.create_mechanics_functions(adjoint_func_space, mode2D='plane strain', materialModel=self.materialModel) + + def energy_function_coords(Uu, p, coords): + U = self.dofManager.create_field(Uu, p.bc_data) + internal_vars = p[1] + return mech_funcs.compute_strain_energy(U, internal_vars) + + return energy_function_coords(storedState[-1][0], storedState[-1][1], parameters) + + def strain_energy_gradient(self, storedState, parameters): + return jax.grad(self.strain_energy_objective, argnums=1)(storedState, parameters) + + def total_work_increment(self, Uu, Uu_prev, p, p_prev, coordinates): + coords = coordinates.reshape(self.mesh.coords.shape) + adjoint_func_space = AdjointFunctionSpace.construct_function_space_for_adjoint(coords, self.mesh, self.quadRule) + mech_funcs = Mechanics.create_mechanics_functions(adjoint_func_space, mode2D='plane strain', materialModel=self.materialModel) + + def energy_function_all_dofs(U, p): + internal_variables = p[1] + return mech_funcs.compute_strain_energy(U, internal_variables) + + nodal_forces = jax.grad(energy_function_all_dofs, argnums=0) + + index = (self.mesh.nodeSets['left'], 1) # arbitrarily choosing left side nodeset for reaction force + + U = self.dofManager.create_field(Uu, p.bc_data) + force = np.array(nodal_forces(U, p).at[index].get()) + disp = U.at[index].get() + + U_prev = self.dofManager.create_field(Uu_prev, p_prev.bc_data) + force_prev = np.array(nodal_forces(U_prev, p_prev).at[index].get()) + disp_prev = U_prev.at[index].get() + + return 0.5*np.tensordot((force + force_prev),(disp - disp_prev), axes=1) + + def total_work_objective(self, storedState, parameters): + val = 0.0 + for step in range(1, self.steps+1): + val += self.total_work_increment(storedState[step][0], storedState[step-1][0], storedState[step][1], storedState[step-1][1], parameters) + + return val + + def total_work_gradient_just_jax(self, storedState, parameters): + return jax.grad(self.total_work_objective, argnums=1)(storedState, parameters) + + def total_work_gradient_with_adjoint(self, storedState, parameters): + compute_df = jax.grad(self.total_work_increment, (0, 1, 4)) + + def energy_function_coords(Uu, Uu_prev, p, coordinates): + coords = coordinates.reshape(self.mesh.coords.shape) + adjoint_func_space = AdjointFunctionSpace.construct_function_space_for_adjoint(coords, self.mesh, self.quadRule) + mech_funcs = Mechanics.create_mechanics_functions(adjoint_func_space, mode2D='plane strain', materialModel=self.materialModel) + U = self.dofManager.create_field(Uu, p.bc_data) + internal_vars = p[1] + return mech_funcs.compute_strain_energy(U, internal_vars) + + gradient = np.zeros(parameters.shape) + adjointLoad = np.zeros(storedState[0][0].shape) + + for step in reversed(range(1, self.steps+1)): + Uu = storedState[step][0] + p = storedState[step][1] + + Uu_prev = storedState[step-1][0] + p_prev = storedState[step-1][1] + + df_du, df_dun, df_dx = compute_df(Uu, Uu_prev, p, p_prev, parameters) + + adjointLoad -= df_du + + n = self.dofManager.get_unknown_size() + self.objective.p = p # have to update parameters to get precond to work + self.objective.update_precond(Uu) # update preconditioner for use in cg (will converge in 1 iteration as long as the preconditioner is not approximate) + dRdu = linalg.LinearOperator((n, n), lambda V: onp.asarray(self.objective.hessian_vec(Uu, V))) + dRdu_decomp = linalg.LinearOperator((n, n), lambda V: onp.asarray(self.objective.apply_precond(V))) + adjointVector = linalg.cg(dRdu, onp.array(adjointLoad, copy=False), tol=1e-10, atol=0.0, M=dRdu_decomp)[0] + + gradient += df_dx + # The action dRdX * lambda (the same as lambda^T * dRdX) + gradient += jax.vjp(lambda z: jax.grad(energy_function_coords, 0)(Uu, Uu_prev, p, z), parameters)[1](adjointVector)[0] + + adjointLoad = -df_dun + # The action dRdU * lambda (the same as lambda^T * dRdU) - For Hyperelastic the residual is not dependent on Uu_prev, so don't actually need this term + # adjointLoad -= jax.vjp(lambda z: jax.grad(energy_function_coords, 0)(Uu, z, p, parameters), Uu_prev)[1](adjointVector)[0] + + return gradient + + + + def test_self_adjoint_gradient(self): + self.compute_objective_function = self.strain_energy_objective + self.compute_gradient = self.strain_energy_gradient + + initialStepSize = 1e-5 + numSteps = 4 + + errors = self.compute_finite_difference_errors(initialStepSize, numSteps, self.initialMesh.coords.ravel()) + + self.assertFiniteDifferenceCheckHasVShape(errors) + + @unittest.expectedFailure + def test_non_self_adjoint_gradient_without_adjoint_solve(self): + self.compute_objective_function = self.total_work_objective + self.compute_gradient = self.total_work_gradient_just_jax + + initialStepSize = 1e-5 + numSteps = 4 + + errors = self.compute_finite_difference_errors(initialStepSize, numSteps, self.initialMesh.coords.ravel()) + + self.assertFiniteDifferenceCheckHasVShape(errors) + + def test_non_self_adjoint_gradient_with_adjoint_solve(self): + self.compute_objective_function = self.total_work_objective + self.compute_gradient = self.total_work_gradient_with_adjoint + + initialStepSize = 1e-5 + numSteps = 4 + + errors = self.compute_finite_difference_errors(initialStepSize, numSteps, self.initialMesh.coords.ravel()) + + self.assertFiniteDifferenceCheckHasVShape(errors) + + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/optimism/material/test/test_J2Plastic_inverse.py b/optimism/inverse/test/test_J2Plastic_inverse.py similarity index 61% rename from optimism/material/test/test_J2Plastic_inverse.py rename to optimism/inverse/test/test_J2Plastic_inverse.py index 9f4f5c3..989afd8 100644 --- a/optimism/material/test/test_J2Plastic_inverse.py +++ b/optimism/inverse/test/test_J2Plastic_inverse.py @@ -2,22 +2,22 @@ import jax import jax.numpy as np -from jax.scipy import linalg from matplotlib import pyplot as plt from optimism import EquationSolver as EqSolver from optimism import FunctionSpace from optimism.material import J2Plastic as J2 -from optimism.material import MaterialUniaxialSimulator from optimism import Mechanics from optimism import Mesh from optimism import Objective from optimism import QuadratureRule -from optimism.test.TestFixture import TestFixture -from optimism.test.MeshFixture import MeshFixture from optimism import TensorMath +from optimism.test.TestFixture import TestFixture +from optimism.test.MeshFixture import MeshFixture +from optimism.inverse.test.FiniteDifferenceFixture import FiniteDifferenceFixture + # for adjoint import numpy as onp from scipy.sparse import linalg @@ -133,12 +133,13 @@ def test_jax_computation_of_state_derivs_at_plastic_step(self): self.assertArrayNear(dc_dc_n[1:,0], dc_dc_n_gold[1:,0], 12) self.assertArrayNear(dc_dc_n[1:,1:].ravel(), dc_dc_n_gold[1:,1:].ravel() , 12) + class J2GlobalMeshUpdateGradsFixture(MeshFixture): def setUp(self): - self.dispGrad0 = np.array([[0.4, -0.2], + dispGrad0 = np.array([[0.4, -0.2], [-0.04, 0.68]]) self.mesh, self.U = self.create_mesh_and_disp(4,4,[0.,1.],[0.,1.], - lambda x: self.dispGrad0@x) + lambda x: dispGrad0@x) E = 100.0 poisson = 0.321 @@ -253,136 +254,227 @@ def compute_energy(Uu, p): self.assertArrayNear(dc_er_du[:,conn,:].ravel(), dc_duele_gold.reshape(10,3,2).ravel(), 10) -class J2GlobalMeshAdjointSolveFixture(MeshFixture): +class J2GlobalMeshAdjointSolveFixture(FiniteDifferenceFixture): def setUp(self): - self.dispGrad0 = np.array([[0.4, -0.2], + dispGrad0 = np.array([[0.4, -0.2], [-0.04, 0.68]]) - self.mesh, self.U = self.create_mesh_and_disp(4,4,[0.,1.],[0.,1.], - lambda x: self.dispGrad0@x) + self.initialMesh, self.U = self.create_mesh_and_disp(4,4,[0.,1.],[0.,1.], + lambda x: dispGrad0@x) E = 100.0 poisson = 0.321 H = 1e-2 * E Y0 = 0.3 * E - - self.props = {'elastic modulus': E, - 'poisson ratio': poisson, - 'yield strength': Y0, - 'kinematics': 'small deformations', - 'hardening model': 'linear', - 'hardening modulus': H} - - self.materialModel = J2.create_material_model_functions(self.props) + props = { + 'elastic modulus': E, + 'poisson ratio': poisson, + 'yield strength': Y0, + 'kinematics': 'small deformations', + 'hardening model': 'linear', + 'hardening modulus': H + } + self.materialModel = J2.create_material_model_functions(props) self.quadRule = QuadratureRule.create_quadrature_rule_on_triangle(degree=1) - self.fs = FunctionSpace.construct_function_space(self.mesh, self.quadRule) - - self.mechFuncs = Mechanics.create_mechanics_functions(self.fs, - "plane strain", - self.materialModel) - - EBCs = [FunctionSpace.EssentialBC(nodeSet='all_boundary', component=0), + self.EBCs = [FunctionSpace.EssentialBC(nodeSet='all_boundary', component=0), FunctionSpace.EssentialBC(nodeSet='all_boundary', component=1)] - self.dofManager = FunctionSpace.DofManager(self.fs, 2, EBCs) + self.steps = 2 - self.Ubc = self.dofManager.get_bc_values(self.U) - - def test_adjoint_solve(self): + def forward_solve(self, parameters): + self.mesh = Mesh.construct_mesh_from_basic_data(parameters.reshape(self.initialMesh.coords.shape),\ + self.initialMesh.conns, self.initialMesh.blocks,\ + self.initialMesh.nodeSets, self.initialMesh.sideSets) + + functionSpace = FunctionSpace.construct_function_space(self.mesh, self.quadRule) + mechFuncs = Mechanics.create_mechanics_functions(functionSpace, + "plane strain", + self.materialModel) + self.dofManager = FunctionSpace.DofManager(functionSpace, 2, self.EBCs) + Ubc = self.dofManager.get_bc_values(self.U) - # solve forward problem - steps = 2 - Ubc_inc = self.Ubc / steps - ivs = self.mechFuncs.compute_initial_state() - p = Objective.Params(bc_data=np.zeros(self.Ubc.shape), state_data=ivs) + Ubc_inc = Ubc / self.steps + ivs = mechFuncs.compute_initial_state() + p = Objective.Params(bc_data=np.zeros(Ubc.shape), state_data=ivs) def compute_energy(Uu, p): U = self.dofManager.create_field(Uu, p.bc_data) internalVariables = p.state_data - return self.mechFuncs.compute_strain_energy(U, internalVariables) + return mechFuncs.compute_strain_energy(U, internalVariables) Uu = 0.0*self.dofManager.get_unknown_values(self.U) - objective = Objective.Objective(compute_energy, Uu, p) + self.objective = Objective.Objective(compute_energy, Uu, p) storedState = [] storedState.append((Uu, p)) - for step in range(1, steps+1): + for step in range(1, self.steps+1): p = Objective.param_index_update(p, 0, step*Ubc_inc) - Uu = EqSolver.nonlinear_equation_solve(objective, Uu, p, EqSolver.get_settings(), useWarmStart=False) + Uu = EqSolver.nonlinear_equation_solve(self.objective, Uu, p, EqSolver.get_settings(), useWarmStart=False) U = self.dofManager.create_field(Uu, p.bc_data) - ivs = self.mechFuncs.compute_updated_internal_variables(U, p.state_data) + ivs = mechFuncs.compute_updated_internal_variables(U, p.state_data) p = Objective.param_index_update(p, 1, ivs) storedState.append((Uu, p)) - # define functions for gradients - internal_vars_shape = ivs.shape - coords_shape = self.mesh.coords.shape + return storedState - def update_internal_vars_test(Uu, ivs, coords): - U = self.dofManager.create_field(Uu, p.bc_data) - internal_vars = ivs.reshape(internal_vars_shape) - coordinates = coords.reshape(coords_shape) - adjoint_func_space = AdjointFunctionSpace.construct_function_space_for_adjoint(coordinates, self.mesh, self.quadRule) - mech_funcs = Mechanics.create_mechanics_functions(adjoint_func_space, mode2D='plane strain', materialModel=self.materialModel) - return mech_funcs.compute_updated_internal_variables(U, internal_vars).ravel() + def strain_energy_increment(self, Uu, ivs, p, coordinates): + coords = coordinates.reshape(self.mesh.coords.shape) + internal_vars = ivs.reshape(p.state_data.shape) + adjoint_func_space = AdjointFunctionSpace.construct_function_space_for_adjoint(coords, self.mesh, self.quadRule) + mech_funcs = Mechanics.create_mechanics_functions(adjoint_func_space, mode2D='plane strain', materialModel=self.materialModel) + U = self.dofManager.create_field(Uu, p.bc_data) - compute_dc = jax.jacfwd(update_internal_vars_test, (0, 1, 2)) + return mech_funcs.compute_strain_energy(U, internal_vars) - # for now just using sum of incremental internal energy as objective - def energy_function_coords(Uu, ivs, ivs_old, p, coords): - U = self.dofManager.create_field(Uu, p.bc_data) - internal_vars = ivs.reshape(internal_vars_shape) - coordinates = coords.reshape(coords_shape) - adjoint_func_space = AdjointFunctionSpace.construct_function_space_for_adjoint(coordinates, self.mesh, self.quadRule) + def strain_energy_objective(self, storedState, parameters): + val = 0.0 + for step in range(1, self.steps+1): + Uu = storedState[step][0] + p = storedState[step][1] + ivs = storedState[step][1].state_data.ravel() + val += self.strain_energy_increment(Uu, ivs, p, parameters) + + return val + + def strain_energy_gradient(self, storedState, parameters): + return jax.grad(self.strain_energy_objective, argnums=1)(storedState, parameters) + + # def total_work_increment(self, Uu, Uu_prev, ivs, ivs_prev, p, p_prev, coordinates): + # coords = coordinates.reshape(self.mesh.coords.shape) + # internal_vars = ivs.reshape(p.state_data.shape) + # internal_vars_prev = ivs_prev.reshape(p_prev.state_data.shape) + + # adjoint_func_space = AdjointFunctionSpace.construct_function_space_for_adjoint(coords, self.mesh, self.quadRule) + # mech_funcs = Mechanics.create_mechanics_functions(adjoint_func_space, mode2D='plane strain', materialModel=self.materialModel) + + # def energy_function_all_dofs(U, ivs): + # return mech_funcs.compute_strain_energy(U, ivs) + + # nodal_forces = jax.grad(energy_function_all_dofs, argnums=0) + + # index = (self.mesh.nodeSets['left'], 1) # arbitrarily choosing left side nodeset for reaction force + + # U = self.dofManager.create_field(Uu, p.bc_data) + # force = np.array(nodal_forces(U, internal_vars).at[index].get()) + # disp = U.at[index].get() + + # U_prev = self.dofManager.create_field(Uu_prev, p_prev.bc_data) + # force_prev = np.array(nodal_forces(U_prev, internal_vars_prev).at[index].get()) + # disp_prev = U_prev.at[index].get() + + # return 0.5*np.tensordot((force + force_prev),(disp - disp_prev), axes=1) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~TESTER~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~# + def total_work_increment(self, Uu, Uu_prev, ivs, ivs_prev, p, p_prev, coordinates): + coords = coordinates.reshape(self.mesh.coords.shape) + internal_vars = ivs.reshape(p.state_data.shape) + adjoint_func_space = AdjointFunctionSpace.construct_function_space_for_adjoint(coords, self.mesh, self.quadRule) + mech_funcs = Mechanics.create_mechanics_functions(adjoint_func_space, mode2D='plane strain', materialModel=self.materialModel) + U = self.dofManager.create_field(Uu, p.bc_data) + + def energy_function_all_dofs(U, ivs): + return mech_funcs.compute_strain_energy(U, ivs) + + nodal_forces = jax.grad(energy_function_all_dofs, argnums=0) + index = (self.mesh.nodeSets['left'], 1) # arbitrarily choosing left side nodeset for reaction force + + return np.sum(np.array(nodal_forces(U, internal_vars).at[index].get())) + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~# + + def total_work_objective(self, storedState, parameters): + val = 0.0 + for step in range(1, self.steps+1): + Uu = storedState[step][0] + Uu_prev = storedState[step-1][0] + p = storedState[step][1] + p_prev = storedState[step-1][1] + ivs = p.state_data.ravel() + ivs_prev = p_prev.state_data.ravel() + + val += self.total_work_increment(Uu, Uu_prev, ivs, ivs_prev, p, p_prev, parameters) + + return val + + def total_work_gradient_with_adjoint(self, storedState, parameters): + + def energy_function_coords(Uu, ivs_prev, p, coordinates): + coords = coordinates.reshape(self.mesh.coords.shape) + internal_vars = ivs_prev.reshape(p.state_data.shape) + adjoint_func_space = AdjointFunctionSpace.construct_function_space_for_adjoint(coords, self.mesh, self.quadRule) mech_funcs = Mechanics.create_mechanics_functions(adjoint_func_space, mode2D='plane strain', materialModel=self.materialModel) + U = self.dofManager.create_field(Uu, p.bc_data) return mech_funcs.compute_strain_energy(U, internal_vars) - compute_df = jax.grad(energy_function_coords, (0, 1, 2, 4)) + def update_internal_vars_test(Uu, ivs_prev, p, coordinates): + coords = coordinates.reshape(self.mesh.coords.shape) + internal_vars = ivs_prev.reshape(p.state_data.shape) + adjoint_func_space = AdjointFunctionSpace.construct_function_space_for_adjoint(coords, self.mesh, self.quadRule) + mech_funcs = Mechanics.create_mechanics_functions(adjoint_func_space, mode2D='plane strain', materialModel=self.materialModel) + U = self.dofManager.create_field(Uu, p.bc_data) + return mech_funcs.compute_updated_internal_variables(U, internal_vars).ravel() - # initialize - gradient = np.zeros(np.prod(np.array(coords_shape))) - mu = np.zeros(np.prod(np.array(internal_vars_shape))) - adjointLoad = np.zeros(Uu.shape) + compute_df = jax.grad(self.total_work_increment, (0, 1, 2, 3, 6)) + compute_dc = jax.jacfwd(update_internal_vars_test, (0, 1, 3)) - coords = self.mesh.coords.ravel() - for step in reversed(range(1, steps+1)): + gradient = np.zeros(parameters.shape) + mu = np.zeros(np.prod(np.array(storedState[0][1].state_data.shape))) + adjointLoad = np.zeros(storedState[0][0].shape) + + for step in reversed(range(1, self.steps+1)): Uu = storedState[step][0] + Uu_prev = storedState[step-1][0] p = storedState[step][1] + p_prev = storedState[step-1][1] ivs = p.state_data.ravel() - ivs_old = storedState[step-1][1].state_data.ravel() - - # partial derivatives of internal variables update - dc_du, dc_dcn, dc_dx = compute_dc(Uu, ivs_old, coords) + ivs_prev = p_prev.state_data.ravel() - # partial derivatives of objective increment - df_du, df_dc, df_dcn, df_dx = compute_df(Uu, ivs, ivs_old, p, coords) + dc_du, dc_dcn, dc_dx = compute_dc(Uu, ivs_prev, p, parameters) + df_du, df_dun, df_dc, df_dcn, df_dx = compute_df(Uu, Uu_prev, ivs, ivs_prev, p, p_prev, parameters) - # Compute adjoint load mu += df_dc adjointLoad -= df_du adjointLoad -= np.tensordot(mu, dc_du, axes=1) # mu^T dc/du - # Solve adjoint equation n = self.dofManager.get_unknown_size() - objective.p = p # have to update parameters to get precond to work - objective.update_precond(Uu) # update preconditioner for use in cg (will converge in 1 iteration as long as the preconditioner is not approximate) - dRdu = linalg.LinearOperator((n, n), lambda V: onp.asarray(objective.hessian_vec(Uu, V))) - dRdu_decomp = linalg.LinearOperator((n, n), lambda V: onp.asarray(objective.apply_precond(V))) + self.objective.p = p # have to update parameters to get precond to work + self.objective.update_precond(Uu) # update preconditioner for use in cg (will converge in 1 iteration as long as the preconditioner is not approximate) + dRdu = linalg.LinearOperator((n, n), lambda V: onp.asarray(self.objective.hessian_vec(Uu, V))) + dRdu_decomp = linalg.LinearOperator((n, n), lambda V: onp.asarray(self.objective.apply_precond(V))) adjointVector = linalg.cg(dRdu, onp.array(adjointLoad, copy=False), tol=1e-10, atol=0.0, M=dRdu_decomp)[0] - # Update gradient gradient += df_dx # The action dRdX * lambda (the same as lambda^T * dRdX) - gradient += jax.vjp(lambda z: jax.grad(energy_function_coords, 0)(Uu, ivs, ivs_old, p, z), coords)[1](adjointVector)[0] + gradient += jax.vjp(lambda z: jax.grad(energy_function_coords, 0)(Uu, ivs_prev, p, z), parameters)[1](adjointVector)[0] gradient += np.tensordot(mu, dc_dx, axes=1) # mu^T dc/dx - # Update mu mu = np.tensordot(mu, dc_dcn, axes=1) - # The action dRdcn * lambda (the same as lambda^T * dRdcn) - mu += jax.vjp(lambda z: jax.grad(energy_function_coords, 0)(Uu, ivs, z, p, coords), ivs_old)[1](adjointVector)[0] - mu += df_dcn + # The action dRdcn * lambda (the same as lambda^T * dRdcn) - Is the residual dependent on ivs_prev? Is this term needed? + mu += jax.vjp(lambda z: jax.grad(energy_function_coords, 0)(Uu, z, p, parameters), ivs_prev)[1](adjointVector)[0] + mu += df_dcn + + adjointLoad = -df_dun + + return gradient + + + + def test_gradient_with_adjoint_solve(self): + + # self.compute_objective_function = self.strain_energy_objective + # self.compute_gradient = self.strain_energy_gradient + + + self.compute_objective_function = self.total_work_objective + self.compute_gradient = self.total_work_gradient_with_adjoint + + initialStepSize = 1e-5 + numSteps = 4 + + errors = self.compute_finite_difference_errors(initialStepSize, numSteps, self.initialMesh.coords.ravel()) + + self.assertFiniteDifferenceCheckHasVShape(errors) if __name__ == '__main__':