|
| 1 | +import numpy as onp # as is "old numpy" |
| 2 | + |
| 3 | +from optimism.JaxConfig import * |
| 4 | +from optimism import EquationSolver |
| 5 | +from optimism import FunctionSpace |
| 6 | +from optimism import Interpolants |
| 7 | +from optimism.material import J2Plastic |
| 8 | +from optimism.material import Neohookean |
| 9 | +from optimism import Mechanics |
| 10 | +from optimism import Mesh |
| 11 | +from optimism import Objective |
| 12 | +from optimism import QuadratureRule |
| 13 | +from optimism import SparseMatrixAssembler |
| 14 | +from optimism import VTKWriter |
| 15 | + |
| 16 | + |
| 17 | +# set up the mesh |
| 18 | +w = 0.1 |
| 19 | +L = 1.0 |
| 20 | +nodesInX = 5 # must be odd |
| 21 | +nodesInY = 15 |
| 22 | +xRange = [0.0, w] |
| 23 | +yRange = [0.0, L] |
| 24 | +mesh = Mesh.construct_structured_mesh(nodesInX, nodesInY, xRange, yRange) |
| 25 | + |
| 26 | +def triangle_centroid(vertices): |
| 27 | + return np.average(vertices, axis=0) |
| 28 | + |
| 29 | +blockIds = -1*onp.ones(Mesh.num_elements(mesh), dtype=onp.int64) |
| 30 | +for e, t in enumerate(mesh.conns): |
| 31 | + vertices = mesh.coords[t,:] |
| 32 | + if triangle_centroid(vertices)[0] < w/2: |
| 33 | + blockIds[e] = 0 |
| 34 | + else: |
| 35 | + blockIds[e] = 1 |
| 36 | + |
| 37 | +blocks = {'soft': np.flatnonzero(np.array(blockIds) == 0), |
| 38 | + 'hard': np.flatnonzero(np.array(blockIds) == 1)} |
| 39 | + |
| 40 | +mesh = Mesh.mesh_with_blocks(mesh, blocks) |
| 41 | + |
| 42 | +nodeTol = 1e-8 |
| 43 | +nodeSets = {'left': np.flatnonzero(mesh.coords[:,0] < xRange[0] + nodeTol), |
| 44 | + 'right': np.flatnonzero(mesh.coords[:,0] > xRange[1] - nodeTol), |
| 45 | + 'bottom': np.flatnonzero(mesh.coords[:,1] < yRange[0] + nodeTol), |
| 46 | + 'top': np.flatnonzero(mesh.coords[:,1] > yRange[1] - nodeTol)} |
| 47 | + |
| 48 | +mesh = Mesh.mesh_with_nodesets(mesh, nodeSets) |
| 49 | + |
| 50 | +# set the essential boundary conditions and create the a DofManager to |
| 51 | +# handle the indexing between unknowns and degrees of freedom. |
| 52 | +EBCs = [Mesh.EssentialBC(nodeSet='left', field=0), |
| 53 | + Mesh.EssentialBC(nodeSet='bottom', field=1), |
| 54 | + Mesh.EssentialBC(nodeSet='top', field = 1)] |
| 55 | + |
| 56 | +fieldShape = mesh.coords.shape |
| 57 | +dofManager = Mesh.DofManager(mesh, fieldShape, EBCs) |
| 58 | + |
| 59 | +# write blocks and bcs to paraview output to check things are correct |
| 60 | +writer = VTKWriter.VTKWriter(mesh, baseFileName='check_problem_setup') |
| 61 | + |
| 62 | +writer.add_cell_field(name='block_id', cellData=blockIds, |
| 63 | + fieldType=VTKWriter.VTKFieldType.SCALARS, |
| 64 | + dataType=VTKWriter.VTKDataType.INT) |
| 65 | + |
| 66 | +bcs = np.array(dofManager.isBc, dtype=np.int64) |
| 67 | +writer.add_nodal_field(name='bcs', nodalData=bcs, fieldType=VTKWriter.VTKFieldType.VECTORS, |
| 68 | + dataType=VTKWriter.VTKDataType.INT) |
| 69 | + |
| 70 | +writer.write() |
| 71 | + |
| 72 | +# create the function space |
| 73 | +order = 2*mesh.masterElement.degree |
| 74 | +quadRule = QuadratureRule.create_quadrature_rule_on_triangle(degree=2*(order - 1)) |
| 75 | +fs = FunctionSpace.construct_function_space(mesh, quadRule) |
| 76 | + |
| 77 | +# create the material models |
| 78 | +ESoft = 5.0 |
| 79 | +nuSoft = 0.25 |
| 80 | +props = {'elastic modulus': ESoft, |
| 81 | + 'poisson ratio': nuSoft} |
| 82 | +softMaterialModel = Neohookean.create_material_model_functions(props) |
| 83 | + |
| 84 | +E = 10.0 |
| 85 | +nu = 0.25 |
| 86 | +Y0 = 0.01*E |
| 87 | +H = E/100 |
| 88 | +props = {'elastic modulus': E, |
| 89 | + 'poisson ratio': nu, |
| 90 | + 'yield strength': Y0, |
| 91 | + 'hardening model': 'linear', |
| 92 | + 'hardening modulus': H} |
| 93 | +hardMaterialModel = J2Plastic.create_material_model_functions(props) |
| 94 | + |
| 95 | +materialModels = {'soft': softMaterialModel, 'hard': hardMaterialModel} |
| 96 | + |
| 97 | +# mechanics functions |
| 98 | +mechanicsFunctions = Mechanics.create_multi_block_mechanics_functions(fs, mode2D="plane strain", materialModels=materialModels) |
| 99 | + |
| 100 | +# helper function to fill in nodal values of essential boundary conditions |
| 101 | +def get_ubcs(p): |
| 102 | + appliedDisp = p[0] |
| 103 | + EbcIndex = (mesh.nodeSets['top'], 1) |
| 104 | + V = np.zeros(fieldShape).at[EbcIndex].set(appliedDisp) |
| 105 | + return dofManager.get_bc_values(V) |
| 106 | + |
| 107 | +# helper function to go from unknowns to full DoF array |
| 108 | +def create_field(Uu, p): |
| 109 | + return dofManager.create_field(Uu, get_ubcs(p)) |
| 110 | + |
| 111 | +# write the energy to minimize |
| 112 | +def energy_function(Uu, p): |
| 113 | + U = create_field(Uu, p) |
| 114 | + internalVariables = p.state_data |
| 115 | + return mechanicsFunctions.compute_strain_energy(U, internalVariables) |
| 116 | + |
| 117 | +# Tell objective how to assemble preconditioner matrix |
| 118 | +def assemble_sparse_preconditioner_matrix(Uu, p): |
| 119 | + U = create_field(Uu, p) |
| 120 | + internalVariables = p.state_data |
| 121 | + elementStiffnesses = mechanicsFunctions.compute_element_stiffnesses(U, internalVariables) |
| 122 | + return SparseMatrixAssembler.assemble_sparse_stiffness_matrix( |
| 123 | + elementStiffnesses, mesh.conns, dofManager) |
| 124 | + |
| 125 | +# solver settings |
| 126 | +solverSettings = EquationSolver.get_settings(max_cumulative_cg_iters=100, |
| 127 | + max_trust_iters=1000, |
| 128 | + use_preconditioned_inner_product_for_cg=True) |
| 129 | + |
| 130 | +precondStrategy = Objective.PrecondStrategy(assemble_sparse_preconditioner_matrix) |
| 131 | + |
| 132 | +# initialize unknown displacements to zero |
| 133 | +Uu = dofManager.get_unknown_values(np.zeros(fieldShape)) |
| 134 | + |
| 135 | +# set initial values of parameters |
| 136 | +appliedDisp = 0.0 |
| 137 | +state = mechanicsFunctions.compute_initial_state() |
| 138 | +p = Objective.Params(appliedDisp, state) |
| 139 | + |
| 140 | +# Construct an objective object for the equation solver to work on |
| 141 | +objective = Objective.ScaledObjective(energy_function, Uu, p, precondStrategy=precondStrategy) |
| 142 | + |
| 143 | +# increment the applied displacement |
| 144 | +appliedDisp = L*0.003 |
| 145 | +p = Objective.param_index_update(p, 0, appliedDisp) |
| 146 | + |
| 147 | +# Find unknown displacements by minimizing the objective |
| 148 | +Uu = EquationSolver.nonlinear_equation_solve(objective, Uu, p, solverSettings) |
| 149 | + |
| 150 | +# update the state variables in the new equilibrium configuration |
| 151 | +U = create_field(Uu, p) |
| 152 | +state = mechanicsFunctions.compute_updated_internal_variables(U, p.state_data) |
| 153 | +p = Objective.param_index_update(p, 1, state) |
| 154 | + |
| 155 | +# write solution data to VTK file for post-processing |
| 156 | +writer = VTKWriter.VTKWriter(mesh, baseFileName='uniaxial_two_material') |
| 157 | + |
| 158 | +U = create_field(Uu, p) |
| 159 | +writer.add_nodal_field(name='displacement', nodalData=U, fieldType=VTKWriter.VTKFieldType.VECTORS) |
| 160 | + |
| 161 | +_, stresses = mechanicsFunctions.compute_output_energy_densities_and_stresses( |
| 162 | + U, state) |
| 163 | +cellStresses = FunctionSpace.project_quadrature_field_to_element_field(fs, stresses) |
| 164 | +writer.add_cell_field(name='stress', cellData=cellStresses, |
| 165 | + fieldType=VTKWriter.VTKFieldType.TENSORS) |
| 166 | + |
| 167 | +writer.write() |
0 commit comments