From a619d35bceb08833822381dffd7d17e864b3a61c Mon Sep 17 00:00:00 2001 From: ryany Date: Fri, 20 Dec 2024 00:42:37 -0500 Subject: [PATCH 1/5] implemented FEM for linear elasticity --- cmad/fem_utils/fem_functions.py | 216 ++++++++++++++ cmad/fem_utils/fem_problem.py | 196 ++++++++++++ cmad/fem_utils/interpolants.py | 201 +++++++++++++ cmad/fem_utils/mesh.py | 69 +++++ cmad/fem_utils/quadrature_rule.py | 346 ++++++++++++++++++++++ cmad/fem_utils/solve_fem.py | 70 +++++ cmad/fem_utils/tests/test_fem.py | 43 +++ cmad/fem_utils/tests/test_interpolants.py | 60 ++++ 8 files changed, 1201 insertions(+) create mode 100644 cmad/fem_utils/fem_functions.py create mode 100644 cmad/fem_utils/fem_problem.py create mode 100644 cmad/fem_utils/interpolants.py create mode 100644 cmad/fem_utils/mesh.py create mode 100644 cmad/fem_utils/quadrature_rule.py create mode 100644 cmad/fem_utils/solve_fem.py create mode 100644 cmad/fem_utils/tests/test_fem.py create mode 100644 cmad/fem_utils/tests/test_interpolants.py diff --git a/cmad/fem_utils/fem_functions.py b/cmad/fem_utils/fem_functions.py new file mode 100644 index 0000000..4436af7 --- /dev/null +++ b/cmad/fem_utils/fem_functions.py @@ -0,0 +1,216 @@ +import numpy as np +import time + +def initialize_equation(NUM_NODES, DOF_NODE, disp_node): + + eq_num = np.zeros((NUM_NODES, DOF_NODE), dtype = int) + for i, node in enumerate(disp_node): + node_number = node[0] + direction = node[1] + eq_num[node_number][direction - 1] = -(i + 1) + + NUM_FREE_DOF = 0 + for i in range(len(eq_num)): + for j in range(len(eq_num[i])): + if (eq_num[i, j] == 0): + NUM_FREE_DOF += 1 + eq_num[i, j] = NUM_FREE_DOF + NUM_PRES_DOF = NUM_NODES * DOF_NODE - NUM_FREE_DOF + + return eq_num, NUM_FREE_DOF, NUM_PRES_DOF + +def assemble_prescribed_displacement( + NUM_PRES_DOF, disp_node, disp_val, eq_num): + + UP = np.zeros(NUM_PRES_DOF) + for i, row in enumerate(disp_node): + node_number = row[0] + dof = row[1] + displacement = disp_val[i] + eqn_number = -eq_num[node_number][dof - 1] + UP[eqn_number - 1] = displacement + + return UP + +def calc_element_stiffness_matrix( + elem_num, volume_conn, nodal_coords, NUM_NODES_ELE, + DOF_NODE, E, nu, quad_rule_3D, shape_func_tetra): + + gauss_weights_3D = quad_rule_3D.wgauss + dshape_tetra = shape_func_tetra.gradients + + k = E / (3 * (1 - 2 * nu)) + mu = E / (2 * (1 + nu)) + + c1 = k + 4/3 * mu + c2 = k - 2/3 * mu + material_stiffness = np.array([[c1, c2, c2, 0, 0, 0], + [c2, c1, c2, 0, 0, 0], + [c2, c2, c1, 0, 0, 0], + [0, 0, 0, mu, 0, 0], + [0, 0, 0, 0, mu, 0], + [0, 0, 0, 0, 0, mu]]) + + elem_points = nodal_coords[volume_conn[elem_num], :] + + KEL = np.zeros((NUM_NODES_ELE * DOF_NODE, NUM_NODES_ELE * DOF_NODE)) + + for gaussPt3D in range(len(gauss_weights_3D)): + Nzeta = dshape_tetra[gaussPt3D, :, :] + + J = (Nzeta @ elem_points).T + + det_J = np.linalg.det(J) + + gradphiXYZ = np.linalg.inv(J).T @ Nzeta + + D_phi = np.zeros((6, NUM_NODES_ELE * DOF_NODE)) + D_phi[0, 0:NUM_NODES_ELE] = gradphiXYZ[0, :] + D_phi[1, NUM_NODES_ELE:NUM_NODES_ELE * 2] = gradphiXYZ[1, :] + D_phi[2, NUM_NODES_ELE * 2:NUM_NODES_ELE * 3] = gradphiXYZ[2, :] + D_phi[3, 0:NUM_NODES_ELE * 2] \ + = np.concatenate((gradphiXYZ[1, :], gradphiXYZ[0, :])) + D_phi[4, NUM_NODES_ELE:NUM_NODES_ELE * 3] \ + = np.concatenate((gradphiXYZ[2, :], gradphiXYZ[1, :])) + D_phi[5, 0:NUM_NODES_ELE] = gradphiXYZ[2, :] + D_phi[5, NUM_NODES_ELE * 2:NUM_NODES_ELE * 3] = gradphiXYZ[0, :] + + KEL += gauss_weights_3D[gaussPt3D] \ + * D_phi.T @ material_stiffness @ D_phi * det_J + + return KEL + +def calc_element_traction_vector( + surf_num, pres_surf, nodal_coords, NUM_NODES_SURF, DOF_NODE, + SURF_TRACTION_VECTOR, quad_rule_2D, shape_func_triangle): + + gauss_weights_2D = quad_rule_2D.wgauss + shape_triangle = shape_func_triangle.values + dshape_triangle = shape_func_triangle.gradients + + surf_points = nodal_coords[pres_surf[surf_num], :] + + PEL = np.zeros(NUM_NODES_SURF * DOF_NODE) + + for gaussPt2D in range(len(gauss_weights_2D)): + N = shape_triangle[gaussPt2D, :] + Nzeta = dshape_triangle[gaussPt2D, :, :] + + J = Nzeta @ surf_points + + Js = np.linalg.norm(np.cross(J[0, :], J[1, :])) + + PEL += gauss_weights_2D[gaussPt2D] \ + * (np.column_stack([N, N, N]) \ + * SURF_TRACTION_VECTOR).T.reshape(-1) * Js + + return PEL + +def assemble_global_stiffness( + KEL, volume_conn, eq_num, ELEM_NUM, KPP, KPF, KFF, KFP): + + elem_conn = volume_conn[ELEM_NUM] + elem_eq_num = eq_num[elem_conn, : ] + global_indices = elem_eq_num.T.reshape(-1) + + local_pres_indices = np.where(global_indices < 0)[0] + local_free_indices = np.where(global_indices > 0)[0] + + global_free_indices = global_indices[local_free_indices] - 1 + global_pres_indices = - global_indices[local_pres_indices] - 1 + + KFF[np.ix_(global_free_indices,global_free_indices)] \ + += KEL[np.ix_(local_free_indices,local_free_indices)] + KFP[np.ix_(global_free_indices,global_pres_indices)] \ + += KEL[np.ix_(local_free_indices,local_pres_indices)] + KPF[np.ix_(global_pres_indices,global_free_indices)] \ + += KEL[np.ix_(local_pres_indices,local_free_indices)] + KPP[np.ix_(global_pres_indices,global_pres_indices)] \ + += KEL[np.ix_(local_pres_indices,local_pres_indices)] + + +def assemble_global_traction_vector( + PEL, pres_surf, surf_num, eq_num, PF, PP): + + surf_conn = pres_surf[surf_num] + surf_eq_num = eq_num[surf_conn, : ] + global_indices = surf_eq_num.T.reshape(-1) + + local_pres_indices = np.where(global_indices < 0)[0] + local_free_indices = np.where(global_indices > 0)[0] + + global_free_indices = global_indices[local_free_indices] - 1 + global_pres_indices = - global_indices[local_pres_indices] - 1 + + PF[global_free_indices] += PEL[local_free_indices] + PP[global_pres_indices] += PEL[local_pres_indices] + + +def assemble_module( + NUM_PRES_DOF, NUM_FREE_DOF, NUM_ELE, NUM_NODES_ELE, DOF_NODE, + NUM_NODES_SURF, SURF_TRACTION_VECTOR, E, nu, disp_node, + disp_val, eq_num, volume_conn, nodal_coords, pres_surf, + quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle): + + # Initialize arrays that need to be returned (KPP, KPF, KFF, KFP, PP) + KPP = np.zeros((NUM_PRES_DOF, NUM_PRES_DOF)) + KPF = np.zeros((NUM_PRES_DOF, NUM_FREE_DOF)) + KFP = np.zeros((NUM_FREE_DOF, NUM_PRES_DOF)) + KFF = np.zeros((NUM_FREE_DOF, NUM_FREE_DOF)) + PP = np.zeros(NUM_PRES_DOF) + PF = np.zeros(NUM_FREE_DOF) + + ## Prescribe boundary conditions + UP = assemble_prescribed_displacement(NUM_PRES_DOF, disp_node, + disp_val, eq_num) + + start = time.time() + # assemble global stiffness and force + for elem_num in range(0, NUM_ELE): + + # get element stiffness matrix + KEL = calc_element_stiffness_matrix( + elem_num, volume_conn, nodal_coords, NUM_NODES_ELE, + DOF_NODE, E, nu, quad_rule_3D, shape_func_tetra) + + # assemble global stiffness + assemble_global_stiffness(KEL, volume_conn, eq_num, + elem_num, KPP, KPF, KFF, KFP) + end = time.time() + print("Time to assemble stiffness matrix: ", end - start) + + start = time.time() + for surf_num in range(len(pres_surf)): + + # get local traction vector + PEL = calc_element_traction_vector( + surf_num, pres_surf, nodal_coords, NUM_NODES_SURF, DOF_NODE, + SURF_TRACTION_VECTOR, quad_rule_2D, shape_func_triangle) + + # assemble traction vector + assemble_global_traction_vector(PEL, pres_surf, + surf_num, eq_num, PF, PP) + end = time.time() + print("Time to assemble load vector: ", end - start) + + return KPP, KPF, KFF, KFP, PF, PP, UP + +def assemble_global_displacement_field(eq_num, UF, UP): + + UUR = np.zeros(eq_num.shape) + for i,val in enumerate(eq_num): + for j, row in enumerate(val): + if row > 0: + UUR[i, j] = UF[row - 1] + else: + UUR[i, j] = UP[- row - 1] + return UUR + +def solve_module(KPP, KPF, KFF, KFP, PP, PF, UP, eq_num): + + UF = np.linalg.solve(KFF, PF - KFP @ UP) + R = KPP @ UP + KPF @ UF - PP + + UUR = assemble_global_displacement_field(eq_num, UF, UP) + + return UUR, UF, R \ No newline at end of file diff --git a/cmad/fem_utils/fem_problem.py b/cmad/fem_utils/fem_problem.py new file mode 100644 index 0000000..2e0cc6e --- /dev/null +++ b/cmad/fem_utils/fem_problem.py @@ -0,0 +1,196 @@ +import numpy as np +from cmad.fem_utils.mesh import Mesh +from cmad.fem_utils import interpolants +from cmad.fem_utils import quadrature_rule + +class fem_problem(): + def __init__(self, problem_type, order): + + # evaluate triangle basis functions at quadrature points + self._quad_rule_2D\ + = quadrature_rule.create_quadrature_rule_on_triangle(order) + gauss_pts_2D = self._quad_rule_2D.xigauss + self._shape_func_triangle = interpolants.shape_triangle(gauss_pts_2D) + + # evaluate tetrahedton basis functions at quadrature points + self._quad_rule_3D \ + = quadrature_rule.create_quadrature_rule_on_tetrahedron(order) + gauss_pts_3D = self._quad_rule_3D.xigauss + self._shape_func_tetra = interpolants.shape_tetrahedron(gauss_pts_3D) + + if problem_type == "hole_block_example": + + self._mesh = Mesh("hole_block") + + self._nodal_coords = self._mesh.get_nodal_coordinates() + self._colume_conn = self._mesh.get_volume_connectivity() + self._surface_conn = self._mesh.get_surface_connectivity() + + self._DOF_NODE = 3 + self._NUM_NODES = len(self._nodal_coords) + self._NUM_NODES_ELE = 4 + self._NUM_ELE = len(self._colume_conn) + self._NUM_NODES_SURF = 3 + + # fix all nodes on plane x = 0 + pres_nodes = [] + for i in range(self._NUM_NODES): + if self._nodal_coords[i][0] == 0.0: + pres_nodes.append(i) + NUM_PRES_NODES = len(pres_nodes) + + self._disp_node = np.zeros((NUM_PRES_NODES \ + * self._DOF_NODE, 2), dtype = int) + for i in range(NUM_PRES_NODES): + for j in range(self._DOF_NODE): + self._disp_node[i * self._DOF_NODE + j][0] \ + = pres_nodes[i] + self._disp_node[i * self._DOF_NODE + j][1] = j+1 + self._disp_val = np.zeros(NUM_PRES_NODES * self._DOF_NODE) + + # normal traction on plane x = 1 + self._SURF_TRACTION_VECTOR = np.array([1.0, 0.0, 0.0]) + pres_surf = [] + for surface in self._surface_conn: + surface_points = self._nodal_coords[surface, :] + if (surface_points[:, 0] == 1 * np.ones(3)).all(): + pres_surf.append(surface) + self._pres_surf = np.array(pres_surf) + + if problem_type == "uniaxial_stress": + + self._mesh = Mesh("bar") + + self._nodal_coords = self._mesh.get_nodal_coordinates() + self._colume_conn = self._mesh.get_volume_connectivity() + self._surface_conn = self._mesh.get_surface_connectivity() + + self._DOF_NODE = 3 + self._NUM_NODES = len(self._nodal_coords) + self._NUM_NODES_ELE = 4 + self._NUM_ELE = len(self._colume_conn) + self._NUM_NODES_SURF = 3 + + # prescribe ux = 0 on x = 0, uy = 0 on y = 0, and uz = 0 on z = 0 + disp_node = [] + for i in range(self._NUM_NODES): + if self._nodal_coords[i][0] == 0.0: + disp_node.append(np.array([i, 1], dtype = int)) + if self._nodal_coords[i][1] == 0.0: + disp_node.append(np.array([i, 2], dtype = int)) + if self._nodal_coords[i][2] == 0.0: + disp_node.append(np.array([i, 3], dtype = int)) + self._disp_node = np.array(disp_node, dtype = int) + self._disp_val = np.zeros(len(disp_node)) + + # normal traction on x = 2 + self._SURF_TRACTION_VECTOR = np.array([1.0, 0.0, 0.0]) + pres_surf = [] + for surface in self._surface_conn: + surface_points = self._nodal_coords[surface, :] + if (surface_points[:, 0] == 2 * np.ones(3)).all(): + pres_surf.append(surface) + self._pres_surf = np.array(pres_surf) + + # setup patch test form B + if problem_type == "patch_B": + + self._mesh = Mesh("cube") + + self._nodal_coords = self._mesh.get_nodal_coordinates() + self._colume_conn = self._mesh.get_volume_connectivity() + self._surface_conn = self._mesh.get_surface_connectivity() + + self._DOF_NODE = 3 + self._NUM_NODES = len(self._nodal_coords) + self._NUM_NODES_ELE = 4 + self._NUM_ELE = len(self._colume_conn) + self._NUM_NODES_SURF = 3 + + #impose linear displacement field on the boundary + disp_node = [] + disp_val = [] + self.UUR_true = np.zeros((self._NUM_NODES, self._DOF_NODE)) + + for i in range(self._NUM_NODES): + coord = self._nodal_coords[i] + x = coord[0] + y = coord[1] + z = coord[2] + + ux = (2 * x + y + z - 4) / 2 + uy = (x + 2 * y + z - 4) / 2 + uz = (x + y + 2 * z - 4) / 2 + + self.UUR_true[i, :] = np.array([ux, uy, uz]) + + if (x == 0.0 or x == 2.0 or y == 0.0 + or y == 2.0 or z == 0 or z == 2.0): + disp_node.append(np.array([i, 1], dtype = int)) + disp_node.append(np.array([i, 2], dtype = int)) + disp_node.append(np.array([i, 3], dtype = int)) + disp_val.append(ux) + disp_val.append(uy) + disp_val.append(uz) + + self._disp_node = np.array(disp_node, dtype = int) + self._disp_val = np.array(disp_val) + + # no surface tractions + self._SURF_TRACTION_VECTOR = np.array([0.0, 0.0, 0.0]) + self._pres_surf = np.array([]) + + if problem_type == "simple_shear": + + self._mesh = Mesh("cube") + + self._nodal_coords = self._mesh.get_nodal_coordinates() + self._colume_conn = self._mesh.get_volume_connectivity() + self._surface_conn = self._mesh.get_surface_connectivity() + + self._DOF_NODE = 3 + self._NUM_NODES = len(self._nodal_coords) + self._NUM_NODES_ELE = 4 + self._NUM_ELE = len(self._colume_conn) + self._NUM_NODES_SURF = 3 + + # fix all nodes on plane z = 0, and set uz = 0 on plane z = 2 + disp_node = [] + for i in range(self._NUM_NODES): + if self._nodal_coords[i][2] == 0.0: + disp_node.append(np.array([i, 1], dtype = int)) + disp_node.append(np.array([i, 2], dtype = int)) + disp_node.append(np.array([i, 3], dtype = int)) + if self._nodal_coords[i][2] == 2.0: + disp_node.append(np.array([i, 3], dtype = int)) + + self._disp_node = np.array(disp_node, dtype = int) + self._disp_val = np.zeros(len(disp_node)) + + # shear traction in x direction on plane z = 2 + self._SURF_TRACTION_VECTOR = np.array([1.0, 0.0, 0.0]) + pres_surf = [] + for surface in self._surface_conn: + surface_points = self._nodal_coords[surface, :] + if (surface_points[:, 2] == 2 * np.ones(3)).all(): + pres_surf.append(surface) + self._pres_surf = np.array(pres_surf) + + def get_2D_basis_functions(self): + return self._quad_rule_2D, self._shape_func_triangle + + def get_3D_basis_functions(self): + return self._quad_rule_3D, self._shape_func_tetra + + def get_mesh_properties(self): + return self._DOF_NODE, self._NUM_NODES, self._NUM_NODES_ELE, \ + self._NUM_ELE, self._NUM_NODES_SURF, self._nodal_coords, \ + self._colume_conn + + def get_boundary_conditions(self): + return self._disp_node, self._disp_val, \ + self._pres_surf, self._SURF_TRACTION_VECTOR + + def save_data(self, filename, data): + self._mesh.add_point_data(data) + self._mesh.save_mesh(filename) \ No newline at end of file diff --git a/cmad/fem_utils/interpolants.py b/cmad/fem_utils/interpolants.py new file mode 100644 index 0000000..e1fcb45 --- /dev/null +++ b/cmad/fem_utils/interpolants.py @@ -0,0 +1,201 @@ +from jaxtyping import Array, Float, Int +from scipy import special +import numpy as onp +import equinox as eqx +import jax.numpy as np + +class ShapeFunctions(eqx.Module): + """ + + Shape functions and shape function gradients (in the parametric space). + + Attributes: + values: Values of the shape functions at a discrete set of points. + Shape is ``(nPts, nNodes)``, where ``nPts`` is the number of + points at which the shame functinos are evaluated, and ``nNodes`` + is the number of nodes in the element (which is equal to the + number of shape functions). + gradients: Values of the parametric gradients of the shape functions. + Shape is ``(nPts, nDim, nNodes)``, where ``nDim`` is the number + of spatial dimensions. Line elements are an exception, which + have shape ``(nPts, nNdodes)``. + + """ + values: Float[Array, "nq nn"] + gradients: Float[Array, "nq nn nd"] + + def __iter__(self): + yield self.values + yield self.gradients + +def shape1d(evaluationPoints): + """ + + Evaluate 1D shape functions and derivatives at points in the master element. + + Args: + evaluationPoints: Array of points in the master element domain at + which to evaluate the shape functions and derivatives. + + Returns: + Shape function values and shape function derivatives at ``evaluationPoints``, + in a tuple (``shape``, ``dshape``). + shapes: [nEvalPoints, nNodes] + dshapes: [nEvalPoints, nNodes] + + """ + + shape = np.vstack(((1 - evaluationPoints[:,0])/2.0, (1 + evaluationPoints[:,0])/2.0)).T + dshape = np.vstack((-0.5 * np.ones(len(evaluationPoints)), 0.5 * np.ones(len(evaluationPoints)))).T + + return ShapeFunctions(shape, dshape) + +def shape_triangle(evaluationPoints): + """ + + Evaluate triangle shape functions and derivatives at points in the master element. + + Args: + evaluationPoints: Array of points in the master element domain at + which to evaluate the shape functions and derivatives. + + Returns: + Shape function values and shape function derivatives at ``evaluationPoints``, + in a tuple (``shape``, ``dshape``). + shapes: [nEvalPoints, nNodes] + dshapes: [nEvalPoints, nDims, nNodes] + + """ + nEvalPoints = len(evaluationPoints) + nDims = 2 + nNodes = 3 + + shape = np.vstack((1 - evaluationPoints[:,0] - evaluationPoints[:,1], + evaluationPoints[:,0], evaluationPoints[:,1])).T + + dshape = onp.zeros((nEvalPoints, nDims, nNodes)) + + for i in range(nEvalPoints): + dshape[i,:,:] = onp.array([[-1, 1, 0],[-1, 0, 1]]) + + return ShapeFunctions(shape, np.array(dshape)) + +def shape_quad(evaluationPoints): + """ + + Evaluate quad shape functions and derivatives at points in the master element. + + Args: + evaluationPoints: Array of points in the master element domain at + which to evaluate the shape functions and derivatives. + + Returns: + Shape function values and shape function derivatives at ``evaluationPoints``, + in a tuple (``shape``, ``dshape``). + shapes: [nEvalPoints, nNodes] + dshapes: [nEvalPoints, nDims, nNodes] + + """ + + nEvalPoints = len(evaluationPoints) + nDims = 2 + nNodes = 4 + + l0x = 1 - evaluationPoints[:,0] + l1x = 1 + evaluationPoints[:,0] + l0y = 1 - evaluationPoints[:,1] + l1y = 1 + evaluationPoints[:,1] + + shape = np.vstack((l0x * l0y / 4, l1x * l0y / 4, l1x * l1y / 4, l0x * l1y / 4)).T + dshape = onp.zeros((nEvalPoints, nDims, nNodes)) + + for i in range(nEvalPoints): + point = evaluationPoints[i] + l0x = 1 - point[0] + l1x = 1 + point[0] + l0y = 1 - point[1] + l1y = 1 + point[1] + dshape[i,:,:] = onp.array([[-l0y, l0y, l1y, -l1y],[-l0x, -l1x, l1x, l0x]]) + + return ShapeFunctions(shape, np.array(dshape)) + +def shape_tetrahedron(evaluationPoints): + """ + + Evaluate tetrahedron shape functions and derivatives at points in the master element. + + Args: + evaluationPoints: Array of points in the master element domain at + which to evaluate the shape functions and derivatives. + + Returns: + Shape function values and shape function derivatives at ``evaluationPoints``, + in a tuple (``shape``, ``dshape``). + shapes: [nEvalPoints, nNodes] + dshapes: [nEvalPoints, nDims, nNodes] + + """ + + nEvalPoints = len(evaluationPoints) + nDims = 3 + nNodes = 4 + + shape = np.vstack((1 - evaluationPoints[:,0] - evaluationPoints[:,1] - evaluationPoints[:,2], + evaluationPoints[:,0], evaluationPoints[:,1], evaluationPoints[:,2])).T + + dshape = onp.zeros((nEvalPoints, nDims, nNodes)) + + for i in range(nEvalPoints): + dshape[i,:,:] = onp.array([[-1, 1, 0, 0], [-1, 0, 1, 0], [-1, 0, 0, 1]]) + + return ShapeFunctions(shape, np.array(dshape)) + + +def shape_brick(evaluationPoints): + """ + + Evaluate brick shape functions and derivatives at points in the master element. + + Args: + evaluationPoints: Array of points in the master element domain at + which to evaluate the shape functions and derivatives. + + Returns: + Shape function values and shape function derivatives at ``evaluationPoints``, + in a tuple (``shape``, ``dshape``). + shapes: [nEvalPoints, nNodes] + dshapes: [nEvalPoints, nDims, nNodes] + + """ + + nEvalPoints = len(evaluationPoints) + nDims = 3 + nNodes = 8 + + m1 = 1 - evaluationPoints[:,0] + p1 = 1 + evaluationPoints[:,0] + m2 = 1 - evaluationPoints[:,1] + p2 = 1 + evaluationPoints[:,1] + m3 = 1 - evaluationPoints[:,2] + p3 = 1 + evaluationPoints[:,3] + + shape = np.vstack((m1 * m2 * m3 / 8, p1 * m2 * m3 / 8, p1 * p2 * m3 / 8, m1 * p2 * m3 / 8, + m1 * m2 * p3 / 8, p1 * m2 * p3 / 8, p1 * p2 * p3 / 8, m1 * p2 * p3 / 8)).T + + dshape = onp.zeros((nEvalPoints, nDims, nNodes)) + + for i in range(nEvalPoints): + point = evaluationPoints[i] + m1 = 1 - point[0] + p1 = 1 + point[0] + m2 = 1 - point[1] + p2 = 1 + point[1] + m3 = 1 - point[2] + p3 = 1 + point[2] + dshape[i,:,:] = onp.array([[-m2 * m3 / 8, m2 * m3 / 8, p2 * m3 / 8, -p2 * m3 / 8, -m2 * p3 / 8, m2 * p3 / 8, p2 * p3 / 8, -p2 * p3 / 8], + [-m1 * m3 / 8, -p1 * m3 / 8, p1 * m3 / 8, m1 * m3 / 8, -m1 * p3 / 8, -p1 * p3 / 8, p1 * p3 / 8, m1 * p3 / 8], + [-m1 * m2 / 8, -p1 * m2 / 8, -p1 * p2 / 8, -m1 * p2 / 8, m1 * m2 / 8, p1 * m2 / 8, p1 * p2 / 8, m1 * p2 / 8]]) + + return ShapeFunctions(shape, np.array(dshape)) + + diff --git a/cmad/fem_utils/mesh.py b/cmad/fem_utils/mesh.py new file mode 100644 index 0000000..1f2a18f --- /dev/null +++ b/cmad/fem_utils/mesh.py @@ -0,0 +1,69 @@ +import meshio +import pygmsh + +class Mesh(): + + def __init__(self, mesh_type): + + self._mesh_type = mesh_type + with pygmsh.occ.Geometry() as geom: + + # square prism with hole through center + if self._mesh_type == "hole_block": + + geom.characteristic_length_min = 0.1 + geom.characteristic_length_max = 0.1 + + rectangle = geom.add_rectangle([0.0, 0.0, 0.0], 1.0, 1.0) + + disk = geom.add_disk([0.5, 0.5, 0.0], 0.25) + + flat = geom.boolean_difference( + rectangle, + disk + ) + geom.extrude(flat, [0.0, 0.0, 0.25], num_layers = 5) + self._mesh = geom.generate_mesh() + + # thin beam + if self._mesh_type == "bar": + + geom.characteristic_length_min = 0.1 + geom.characteristic_length_max = 0.1 + + rectangle = geom.add_rectangle([0.0, 0.0, 0.0], 2, 0.5) + + geom.extrude(rectangle, [0.0, 0.0, 0.5], num_layers = 5) + self._mesh = geom.generate_mesh() + + # 2 x 2 x 2 cube + if self._mesh_type == "cube": + + geom.characteristic_length_min = 0.2 + geom.characteristic_length_max = 0.2 + + rectangle = geom.add_rectangle([0.0, 0.0, 0.0], 2, 2) + + geom.extrude(rectangle, [0.0, 0.0, 2], num_layers = 10) + self._mesh = geom.generate_mesh() + + self._points = self._mesh.points + self._cells = self._mesh.cells + self._volume_conn = self._cells[2].data + self._surface_conn = self._cells[1].data + + + def get_nodal_coordinates(self): + return self._points + + def get_volume_connectivity(self): + return self._volume_conn + + def get_surface_connectivity(self): + return self._surface_conn + + def add_point_data(self, data): + self._mesh.point_data = {"displacement_field": data} + + def save_mesh(self, filename): + self._mesh.write(filename + ".vtk") diff --git a/cmad/fem_utils/quadrature_rule.py b/cmad/fem_utils/quadrature_rule.py new file mode 100644 index 0000000..0e7afaf --- /dev/null +++ b/cmad/fem_utils/quadrature_rule.py @@ -0,0 +1,346 @@ +from jax.lax import switch +from jaxtyping import Array, Float +import equinox as eqx +import jax.numpy as np +import math +import numpy as onp +import scipy.special + +class QuadratureRule(eqx.Module): + """ + Quadrature rule points and weights. + An ``equinox`` ``Module`` containing ``xigauss``, a ``jax.numpy`` array of the + coordinates of the sample points in the reference domain, and + ``wgauss``, a ``jax.numpy`` array with the weights. + """ + xigauss: Float[Array, "nq 2"] + wgauss: Float[Array, "nq"] + + def __iter__(self): + yield self.xigauss + yield self.wgauss + + def __len__(self): + return self.xigauss.shape[0] + + +def create_quadrature_rule_1D(degree): + """Creates a Gauss-Legendre quadrature on the interval [-1, 1]. + + The rule can exactly integrate polynomials of degree up to + ``degree``. + + Parameters + ---------- + degree: Highest degree polynomial to be exactly integrated by the quadrature rule + + Returns + ------- + A ``QuadratureRule`` named tuple containing the quadrature point coordinates + and the weights. + """ + + n = math.ceil((degree + 1)/2) + xi, w = scipy.special.roots_sh_legendre(n) + xi = -1 + 2 * xi + + xi_matrix = onp.zeros((len(xi), 1)) + xi_matrix[:,0]=xi + return QuadratureRule(np.array(xi_matrix), w) + + +def create_quadrature_rule_on_triangle(degree): + """Creates a Gauss-Legendre quadrature on the unit triangle. + + The rule can exactly integrate 2D polynomials up to the value of + ``degree``. The domain is the triangle between the vertices + (0, 0)-(1, 0)-(0, 1). The rules here are guaranteed to be + cyclically symmetric in triangular coordinates and to have strictly + positive weights. + + Parameters + ---------- + degree: Highest degree polynomial to be exactly integrated by the quadrature rule + + Returns + ------- + A ``QuadratureRule`` named tuple containing the quadrature point coordinates + and the weights. + """ + if degree <= 1: + xi = np.array([[3.33333333333333333E-01, 3.33333333333333333E-01]]) + + w = np.array([ 5.00000000000000000E-01 ]) + elif degree == 2: + xi = np.array([[6.66666666666666667E-01, 1.66666666666666667E-01], + [1.66666666666666667E-01, 6.66666666666666667E-01], + [1.66666666666666667E-01, 1.66666666666666667E-01]]) + + w = np.array([1.66666666666666666E-01, + 1.66666666666666667E-01, + 1.66666666666666667E-01]) + elif degree <= 4: + xi = np.array([[1.081030181680700E-01, 4.459484909159650E-01], + [4.459484909159650E-01, 1.081030181680700E-01], + [4.459484909159650E-01, 4.459484909159650E-01], + [8.168475729804590E-01, 9.157621350977100E-02], + [9.157621350977100E-02, 8.168475729804590E-01], + [9.157621350977100E-02, 9.157621350977100E-02]]) + + w = np.array([1.116907948390055E-01, + 1.116907948390055E-01, + 1.116907948390055E-01, + 5.497587182766100E-02, + 5.497587182766100E-02, + 5.497587182766100E-02]) + elif degree <= 5: + xi = np.array([[3.33333333333333E-01, 3.33333333333333E-01], + [5.97158717897700E-02, 4.70142064105115E-01], + [4.70142064105115E-01, 5.97158717897700E-02], + [4.70142064105115E-01, 4.70142064105115E-01], + [7.97426985353087E-01, 1.01286507323456E-01], + [1.01286507323456E-01, 7.97426985353087E-01], + [1.01286507323456E-01, 1.01286507323456E-01]]) + + w = np.array([1.12500000000000E-01, + 6.61970763942530E-02, + 6.61970763942530E-02, + 6.61970763942530E-02, + 6.29695902724135E-02, + 6.29695902724135E-02, + 6.29695902724135E-02]) + elif degree <= 6: + xi = np.array([[5.01426509658179E-01, 2.49286745170910E-01], + [2.49286745170910E-01, 5.01426509658179E-01], + [2.49286745170910E-01, 2.49286745170910E-01], + [8.73821971016996E-01, 6.30890144915020E-02], + [6.30890144915020E-02, 8.73821971016996E-01], + [6.30890144915020E-02, 6.30890144915020E-02], + [5.31450498448170E-02, 3.10352451033784E-01], + [6.36502499121399E-01, 5.31450498448170E-02], + [3.10352451033784E-01, 6.36502499121399E-01], + [5.31450498448170E-02, 6.36502499121399E-01], + [6.36502499121399E-01, 3.10352451033784E-01], + [3.10352451033784E-01, 5.31450498448170E-02]]) + + w = np.array([5.83931378631895E-02, + 5.83931378631895E-02, + 5.83931378631895E-02, + 2.54224531851035E-02, + 2.54224531851035E-02, + 2.54224531851035E-02, + 4.14255378091870E-02, + 4.14255378091870E-02, + 4.14255378091870E-02, + 4.14255378091870E-02, + 4.14255378091870E-02, + 4.14255378091870E-02]) + elif degree <= 10: + xi = np.array([[0.33333333333333333E+00, 0.33333333333333333E+00], + [0.4269134091050342E-02, 0.49786543295447483E+00], + [0.49786543295447483E+00, 0.4269134091050342E-02], + [0.49786543295447483E+00, 0.49786543295447483E+00], + [0.14397510054188759E+00, 0.42801244972905617E+00], + [0.42801244972905617E+00, 0.14397510054188759E+00], + [0.42801244972905617E+00, 0.42801244972905617E+00], + [0.6304871745135507E+00, 0.18475641274322457E+00], + [0.18475641274322457E+00, 0.6304871745135507E+00], + [0.18475641274322457E+00, 0.18475641274322457E+00], + [0.9590375628566448E+00, 0.20481218571677562E-01], + [0.20481218571677562E-01, 0.9590375628566448E+00], + [0.20481218571677562E-01, 0.20481218571677562E-01], + [0.3500298989727196E-01, 0.1365735762560334E+00], + [0.1365735762560334E+00, 0.8284234338466947E+00], + [0.8284234338466947E+00, 0.3500298989727196E-01], + [0.1365735762560334E+00, 0.3500298989727196E-01], + [0.8284234338466947E+00, 0.1365735762560334E+00], + [0.3500298989727196E-01, 0.8284234338466947E+00], + [0.37549070258442674E-01, 0.3327436005886386E+00], + [0.3327436005886386E+00, 0.6297073291529187E+00], + [0.6297073291529187E+00, 0.37549070258442674E-01], + [0.3327436005886386E+00, 0.37549070258442674E-01], + [0.6297073291529187E+00, 0.3327436005886386E+00], + [0.37549070258442674E-01, 0.6297073291529187E+00]]) + + w = np.array([0.4176169990259819E-01, + 0.36149252960283717E-02, + 0.36149252960283717E-02, + 0.36149252960283717E-02, + 0.3724608896049025E-01, + 0.3724608896049025E-01, + 0.3724608896049025E-01, + 0.39323236701554264E-01, + 0.39323236701554264E-01, + 0.39323236701554264E-01, + 0.3464161543553752E-02, + 0.3464161543553752E-02, + 0.3464161543553752E-02, + 0.147591601673897E-01, + 0.147591601673897E-01, + 0.147591601673897E-01, + 0.147591601673897E-01, + 0.147591601673897E-01, + 0.147591601673897E-01, + 0.1978968359803062E-01, + 0.1978968359803062E-01, + 0.1978968359803062E-01, + 0.1978968359803062E-01, + 0.1978968359803062E-01, + 0.1978968359803062E-01]) + else: + raise ValueError("Quadrature of precision this high is not implemented.") + + return QuadratureRule(xi, w) + +def create_quadrature_rule_on_quad(degree): + """Creates a Gauss-Legendre quadrature on square. + + The rule can exactly integrate 2D polynomials up to the value of + ``degree``. The domain is the square between the vertices + (1, -1)-(1, 1)-(-1, 1)-(-1, -1). + + Parameters + ---------- + degree: Highest degree polynomial to be exactly integrated by the quadrature rule + + Returns + ------- + A ``QuadratureRule`` named tuple containing the quadrature point coordinates + and the weights. + """ + + if degree <= 1: + xi = np.array([[0, 0]]) + + w = np.array([4]) + elif degree <= 3: + a = 0.577350269189626 + xi = np.array([[-a, -a], + [a, -a], + [a, a], + [-a, a]]) + + w = np.array([1, 1, 1, 1]) + elif degree <= 5: + a = 0.774596669241483 + b = 0.888888888888889 + c = 0.555555555555556 + xi = np.array([[-a, -a], + [0, -a], + [a, -a], + [-a, 0], + [0, 0], + [a, 0], + [-a, a], + [0, a], + [a, a]]) + + w = np.array([c * c, b * c, c * c, c * b, + b * b, c * b, c * c, b * c, c * c]) + + else: + raise ValueError("Quadrature of precision this high is not implemented.") + + return QuadratureRule(xi, w) + +def create_quadrature_rule_on_tetrahedron(degree): + if degree <= 1: + xi = np.array([[0.25, 0.25, 0.25]]) + + w = np.array([1.0/6.0]) + + elif degree == 2: + xi = np.array([[0.138196601125011, 0.138196601125011, 0.138196601125011], + [0.585410196624969, 0.138196601125011, 0.138196601125011], + [0.138196601125011, 0.585410196624969, 0.138196601125011], + [0.138196601125011, 0.138196601125011, 0.585410196624969]]) + + w = np.array([0.25/6.0, 0.25/6.0, 0.25/6.0, 0.25/6.0]) + + elif degree == 3: + xi = np.array([[1./4., 1./4., 1./4.], + [1./6., 1./6., 1./6.], + [1./6., 1./6., 1./2.], + [1./6., 1./2., 1./6.], + [1./2., 1./6., 1./6.]]) + + w = np.array([-2/15, 3/40, 3/40, 3/40, 3/40]) + + elif degree == 4: + xi = np.array([[0.25,0.25,0.25], + [0.785714285714286, 0.0714285714285714, 0.0714285714285714], + [0.0714285714285714, 0.0714285714285714, 0.0714285714285714], + [0.0714285714285714, 0.0714285714285714, 0.785714285714286], + [0.0714285714285714, 0.785714285714286, 0.0714285714285714], + [0.100596423833201, 0.399403576166799, 0.399403576166799], + [0.399403576166799, 0.100596423833201, 0.399403576166799], + [0.399403576166799, 0.399403576166799, 0.100596423833201], + [0.399403576166799, 0.100596423833201, 0.100596423833201], + [0.100596423833201, 0.399403576166799, 0.100596423833201], + [0.100596423833201, 0.100596423833201, 0.399403576166799]]) + + w = np.array([-0.0131555555555556, 0.00762222222222222, 0.00762222222222222, + 0.00762222222222222, 0.00762222222222222, 0.0248888888888889, + 0.0248888888888889, 0.0248888888888889, 0.0248888888888889, + 0.0248888888888889, 0.0248888888888889]) + + elif degree == 5: + xi = np.array([[0.25, 0.25, 0.25], + [0, 1./3., 1./3.], + [1./3., 1./3., 1./3.], + [1./3., 1./3., 0], + [1./3., 0, 1./3.], + [8./11., 1./11., 1./11.], + [1./11., 1./11., 1./11.], + [1./11., 1./11., 8./11.], + [1./11., 8./11., 1./11.], + [0.433449846426336, 0.0665501535736643, 0.0665501535736643], + [0.0665501535736643, 0.433449846426336, 0.0665501535736643], + [0.0665501535736643, 0.0665501535736643, 0.433449846426336], + [0.0665501535736643, 0.433449846426336, 0.433449846426336], + [0.433449846426336, 0.0665501535736643, 0.433449846426336], + [0.433449846426336, 0.433449846426336, 0.0665501535736643]]) + + w = np.array([0.0302836780970892, 0.00602678571428572, 0.00602678571428572, + 0.00602678571428572, 0.00602678571428572, 0.011645249086029, + 0.011645249086029, 0.011645249086029, 0.011645249086029, + 0.0109491415613865, 0.0109491415613865, 0.0109491415613865, + 0.0109491415613865, 0.0109491415613865, 0.0109491415613865]) + elif degree == 6: + xi = np.array([[0.356191386222545,0.214602871259152,0.214602871259152], + [0.214602871259152,0.214602871259152,0.214602871259152], + [0.214602871259152,0.214602871259152,0.356191386222545], + [0.214602871259152,0.356191386222545,0.214602871259152], + [0.877978124396166,0.0406739585346113,0.0406739585346113], + [0.0406739585346113,0.0406739585346113,0.0406739585346113], + [0.0406739585346113,0.0406739585346113,0.877978124396166], + [0.0406739585346113,0.877978124396166,0.0406739585346113], + [0.0329863295731731,0.322337890142276,0.322337890142276], + [0.322337890142276,0.322337890142276,0.322337890142276], + [0.322337890142276,0.322337890142276,0.0329863295731731], + [0.322337890142276,0.0329863295731731,0.322337890142276], + [0.269672331458316,0.0636610018750175,0.0636610018750175], + [0.0636610018750175,0.269672331458316,0.0636610018750175], + [0.0636610018750175,0.0636610018750175,0.269672331458316], + [0.603005664791649,0.0636610018750175,0.0636610018750175], + [0.0636610018750175,0.603005664791649,0.0636610018750175], + [0.0636610018750175,0.0636610018750175,0.603005664791649], + [0.0636610018750175,0.269672331458316,0.603005664791649], + [0.269672331458316,0.603005664791649,0.0636610018750175], + [0.603005664791649,0.0636610018750175,0.269672331458316], + [0.0636610018750175,0.603005664791649,0.269672331458316], + [0.269672331458316,0.0636610018750175,0.603005664791649], + [0.603005664791649,0.269672331458316,0.0636610018750175]]) + + w = np.array([0.00665379170969465, 0.00665379170969465, 0.00665379170969465, + 0.00665379170969465, 0.00167953517588678, 0.00167953517588678, + 0.00167953517588678, 0.00167953517588678, 0.0092261969239424, + 0.0092261969239424, 0.0092261969239424, 0.0092261969239424, + 0.00803571428571428, 0.00803571428571428, 0.00803571428571428, + 0.00803571428571428, 0.00803571428571428, 0.00803571428571428, + 0.00803571428571428, 0.00803571428571428, 0.00803571428571428, + 0.00803571428571428, 0.00803571428571428, 0.00803571428571428]) + else: + raise ValueError("Quadrature of precision this high is not implemented.") + + return QuadratureRule(xi, w) diff --git a/cmad/fem_utils/solve_fem.py b/cmad/fem_utils/solve_fem.py new file mode 100644 index 0000000..2e640f4 --- /dev/null +++ b/cmad/fem_utils/solve_fem.py @@ -0,0 +1,70 @@ +from cmad.fem_utils.fem_problem import fem_problem +from cmad.fem_utils.fem_functions import (initialize_equation, + assemble_module, + solve_module) +import time + +""" + +DOF_NODE: number of degrees of freedom per node +NUM_NODES: total number of nodes in the mesh +NUM_NODE_ELE: number of nodes per element +NUM_ELE: number of elements +NUM_NODES_SURF: number of nodes per surface element +NUM_FREE_DOF: number of free degrees of freedom +NUM_PRES_DOF: number of prescribed degrees of freedom +SURF_TRACTION_VECTOR: surface traction vector +E: Youngs Modulus +nu: Poisson's ratio +disp_node: (NUM_PRES_DOF x 2) array that specifies + which node and dof is prescribed +disp_val: (NUM_PRED_DOF x 1) array of values + of the prescribed displacements +eq_num: (NUM_NODES x DOF_NODE) array that specifies where + each node and its DOFs belong in the global stifness matrix +volume_conn: connectivity matrix for 3D elements +pres_surf: connectivity for surfaces that have a prescribed traction +nodal_coords: spacial coordinates for each node in mesh + +""" + +order = 3 +problem = fem_problem("patch_B", order) + +DOF_NODE, NUM_NODES, NUM_NODES_ELE, NUM_ELE, NUM_NODES_SURF, \ + nodal_coords, volume_conn = problem.get_mesh_properties() + +disp_node, disp_val, pres_surf, SURF_TRACTION_VECTOR \ + = problem.get_boundary_conditions() + +quad_rule_3D, shape_func_tetra = problem.get_3D_basis_functions() +quad_rule_2D, shape_func_triangle = problem.get_2D_basis_functions() + + +print("Number of elements:", NUM_ELE) + +E = 200 +nu = 0.3 + +eq_num, NUM_FREE_DOF, NUM_PRES_DOF \ + = initialize_equation(NUM_NODES, DOF_NODE, disp_node) + +KPP, KPF, KFF, KFP, PF, PP, UP \ + = assemble_module(NUM_PRES_DOF, NUM_FREE_DOF, NUM_ELE, NUM_NODES_ELE, + DOF_NODE, NUM_NODES_SURF, SURF_TRACTION_VECTOR, E, nu, + disp_node, disp_val, eq_num, volume_conn, nodal_coords, + pres_surf, quad_rule_3D, shape_func_tetra, quad_rule_2D, + shape_func_triangle) + +solve_start = time.time() + +UUR, UF, R = solve_module(KPP, KPF, KFF, KFP, PP, PF, UP, eq_num) + +print(UUR) + +solve_end = time.time() + +#problem.save_data("simple_shear", UUR) + +print("Solve time: ", solve_end - solve_start) + diff --git a/cmad/fem_utils/tests/test_fem.py b/cmad/fem_utils/tests/test_fem.py new file mode 100644 index 0000000..9e56bf3 --- /dev/null +++ b/cmad/fem_utils/tests/test_fem.py @@ -0,0 +1,43 @@ +import jax.numpy as np +from cmad.fem_utils.fem_problem import fem_problem +from cmad.fem_utils.fem_functions import (initialize_equation, + assemble_module, solve_module) +import unittest + +class TestFEM(unittest.TestCase): + + def test_patch_form_B(self): + + order = 3 + problem = fem_problem("patch_B", order) + + DOF_NODE, NUM_NODES, NUM_NODES_ELE, NUM_ELE, NUM_NODES_SURF, \ + nodal_coords, volume_conn = problem.get_mesh_properties() + + disp_node, disp_val, pres_surf, SURF_TRACTION_VECTOR \ + = problem.get_boundary_conditions() + + quad_rule_3D, shape_func_tetra = problem.get_3D_basis_functions() + quad_rule_2D, shape_func_triangle = problem.get_2D_basis_functions() + + E = 200 + nu = 0.3 + + eq_num, NUM_FREE_DOF, NUM_PRES_DOF\ + = initialize_equation(NUM_NODES, DOF_NODE, disp_node) + + KPP, KPF, KFF, KFP, PF, PP, UP \ + = assemble_module(NUM_PRES_DOF, NUM_FREE_DOF, NUM_ELE, NUM_NODES_ELE, + DOF_NODE, NUM_NODES_SURF,SURF_TRACTION_VECTOR, E, nu, + disp_node, disp_val, eq_num, volume_conn, nodal_coords, + pres_surf, quad_rule_3D, shape_func_tetra, quad_rule_2D, + shape_func_triangle) + + UUR, UF, R = solve_module(KPP, KPF, KFF, KFP, PP, PF, UP, eq_num) + + assert np.allclose(UUR, problem.UUR_true) + +if __name__ == "__main__": + fem_test_suite = unittest.TestLoader().loadTestsFromTestCase( + TestFEM) + unittest.TextTestRunner(verbosity=2).run(fem_test_suite) \ No newline at end of file diff --git a/cmad/fem_utils/tests/test_interpolants.py b/cmad/fem_utils/tests/test_interpolants.py new file mode 100644 index 0000000..1cd56d2 --- /dev/null +++ b/cmad/fem_utils/tests/test_interpolants.py @@ -0,0 +1,60 @@ +from cmad.fem_utils import interpolants +from cmad.fem_utils import quadrature_rule +import unittest +import jax.numpy as np + +class TestInterpolants(unittest.TestCase): + + def test_1D_shape_functions(self): + for i in range(2, 8): + test_points = quadrature_rule.create_quadrature_rule_1D(i).xigauss + shape_func_1D = interpolants.shape1d(test_points) + shape = shape_func_1D.values + dshape = shape_func_1D.gradients + + assert np.allclose(np.sum(shape, axis = 1), np.ones(len(shape))) + assert np.allclose(np.sum(dshape, axis = 1), np.zeros(len(shape))) + + def test_triangle_shape_functions(self): + for i in range(2, 7): + test_points = quadrature_rule.create_quadrature_rule_on_triangle(i).xigauss + shape_func_triangle = interpolants.shape_triangle(test_points) + shape = shape_func_triangle.values + dshape = shape_func_triangle.gradients + + assert np.allclose(np.sum(shape, axis = 1), np.ones(len(shape))) + + for j in range(len(dshape)): + gradient = dshape[j] + assert np.allclose(np.sum(gradient, axis = 1), np.zeros(len(gradient))) + + def test_quad_shape_functions(self): + for i in range(2, 6): + test_points = quadrature_rule.create_quadrature_rule_on_quad(i).xigauss + shape_func_quad = interpolants.shape_quad(test_points) + shape = shape_func_quad.values + dshape = shape_func_quad.gradients + + assert np.allclose(np.sum(shape, axis = 1), np.ones(len(shape))) + + for j in range(len(dshape)): + gradient = dshape[j] + assert np.allclose(np.sum(gradient, axis = 1), np.zeros(len(gradient))) + + def test_tetrahedron_shape_functions(self): + for i in range(2, 7): + test_points = quadrature_rule.create_quadrature_rule_on_tetrahedron(i).xigauss + shape_func_tetra = interpolants.shape_tetrahedron(test_points) + shape = shape_func_tetra.values + dshape = shape_func_tetra.gradients + + assert np.allclose(np.sum(shape, axis = 1), np.ones(len(shape))) + + for j in range(len(dshape)): + gradient = dshape[j] + assert np.allclose(np.sum(gradient, axis = 1), np.zeros(len(gradient))) + +if __name__ == "__main__": + Interpolants_test_suite = unittest.TestLoader().loadTestsFromTestCase( + TestInterpolants) + unittest.TextTestRunner(verbosity=2).run(Interpolants_test_suite) \ No newline at end of file From fb99c5d4fd1d3f50008d755b48fdc8f29f85fe7a Mon Sep 17 00:00:00 2001 From: ryany Date: Wed, 22 Jan 2025 15:04:42 -0800 Subject: [PATCH 2/5] implemented residual to calculate element stiffness matrices using AD --- cmad/fem_utils/fem_functions.py | 183 +++++++++++++++++++---------- cmad/fem_utils/fem_functions_AD.py | 62 ++++++++++ cmad/fem_utils/fem_problem.py | 76 ++++++------ cmad/fem_utils/interpolants.py | 101 ++++++++-------- cmad/fem_utils/quadrature_rule.py | 32 ++--- cmad/fem_utils/solve_fem.py | 62 ++++++---- cmad/fem_utils/tests/test_fem.py | 17 ++- 7 files changed, 328 insertions(+), 205 deletions(-) create mode 100644 cmad/fem_utils/fem_functions_AD.py diff --git a/cmad/fem_utils/fem_functions.py b/cmad/fem_utils/fem_functions.py index 4436af7..064be4c 100644 --- a/cmad/fem_utils/fem_functions.py +++ b/cmad/fem_utils/fem_functions.py @@ -1,28 +1,31 @@ import numpy as np import time +import jax -def initialize_equation(NUM_NODES, DOF_NODE, disp_node): +from cmad.fem_utils.fem_functions_AD import elem_residual - eq_num = np.zeros((NUM_NODES, DOF_NODE), dtype = int) +def initialize_equation(num_nodes, dof_node, disp_node): + + eq_num = np.zeros((num_nodes, dof_node), dtype=int) for i, node in enumerate(disp_node): node_number = node[0] direction = node[1] eq_num[node_number][direction - 1] = -(i + 1) - NUM_FREE_DOF = 0 + num_free_dof = 0 for i in range(len(eq_num)): for j in range(len(eq_num[i])): if (eq_num[i, j] == 0): - NUM_FREE_DOF += 1 - eq_num[i, j] = NUM_FREE_DOF - NUM_PRES_DOF = NUM_NODES * DOF_NODE - NUM_FREE_DOF + num_free_dof += 1 + eq_num[i, j] = num_free_dof + num_pres_dof = num_nodes * dof_node - num_free_dof - return eq_num, NUM_FREE_DOF, NUM_PRES_DOF + return eq_num, num_free_dof, num_pres_dof def assemble_prescribed_displacement( - NUM_PRES_DOF, disp_node, disp_val, eq_num): + num_pres_dof, disp_node, disp_val, eq_num): - UP = np.zeros(NUM_PRES_DOF) + UP = np.zeros(num_pres_dof) for i, row in enumerate(disp_node): node_number = row[0] dof = row[1] @@ -33,11 +36,11 @@ def assemble_prescribed_displacement( return UP def calc_element_stiffness_matrix( - elem_num, volume_conn, nodal_coords, NUM_NODES_ELE, - DOF_NODE, E, nu, quad_rule_3D, shape_func_tetra): + elem_points, num_nodes_elem, dof_node, + params, gauss_weights_3D, dshape_tetra): - gauss_weights_3D = quad_rule_3D.wgauss - dshape_tetra = shape_func_tetra.gradients + E = params[0] + nu = params[1] k = E / (3 * (1 - 2 * nu)) mu = E / (2 * (1 + nu)) @@ -51,38 +54,37 @@ def calc_element_stiffness_matrix( [0, 0, 0, 0, mu, 0], [0, 0, 0, 0, 0, mu]]) - elem_points = nodal_coords[volume_conn[elem_num], :] - - KEL = np.zeros((NUM_NODES_ELE * DOF_NODE, NUM_NODES_ELE * DOF_NODE)) + KEL = np.zeros((num_nodes_elem * dof_node, num_nodes_elem * dof_node)) for gaussPt3D in range(len(gauss_weights_3D)): - Nzeta = dshape_tetra[gaussPt3D, :, :] + w_q = gauss_weights_3D[gaussPt3D] + + dshape_tetra_q = dshape_tetra[gaussPt3D, :, :] - J = (Nzeta @ elem_points).T + J_q = (dshape_tetra_q @ elem_points).T - det_J = np.linalg.det(J) + dv_q = np.linalg.det(J_q) - gradphiXYZ = np.linalg.inv(J).T @ Nzeta + gradphiXYZ_q = np.linalg.inv(J_q).T @ dshape_tetra_q - D_phi = np.zeros((6, NUM_NODES_ELE * DOF_NODE)) - D_phi[0, 0:NUM_NODES_ELE] = gradphiXYZ[0, :] - D_phi[1, NUM_NODES_ELE:NUM_NODES_ELE * 2] = gradphiXYZ[1, :] - D_phi[2, NUM_NODES_ELE * 2:NUM_NODES_ELE * 3] = gradphiXYZ[2, :] - D_phi[3, 0:NUM_NODES_ELE * 2] \ - = np.concatenate((gradphiXYZ[1, :], gradphiXYZ[0, :])) - D_phi[4, NUM_NODES_ELE:NUM_NODES_ELE * 3] \ - = np.concatenate((gradphiXYZ[2, :], gradphiXYZ[1, :])) - D_phi[5, 0:NUM_NODES_ELE] = gradphiXYZ[2, :] - D_phi[5, NUM_NODES_ELE * 2:NUM_NODES_ELE * 3] = gradphiXYZ[0, :] + D_phi_q = np.zeros((6, num_nodes_elem * dof_node)) + D_phi_q[0, 0:num_nodes_elem] = gradphiXYZ_q[0, :] + D_phi_q[1, num_nodes_elem:num_nodes_elem * 2] = gradphiXYZ_q[1, :] + D_phi_q[2, num_nodes_elem * 2:num_nodes_elem * 3] = gradphiXYZ_q[2, :] + D_phi_q[3, 0:num_nodes_elem * 2] \ + = np.concatenate((gradphiXYZ_q[1, :], gradphiXYZ_q[0, :])) + D_phi_q[4, num_nodes_elem:num_nodes_elem * 3] \ + = np.concatenate((gradphiXYZ_q[2, :], gradphiXYZ_q[1, :])) + D_phi_q[5, 0:num_nodes_elem] = gradphiXYZ_q[2, :] + D_phi_q[5, num_nodes_elem * 2:num_nodes_elem * 3] = gradphiXYZ_q[0, :] - KEL += gauss_weights_3D[gaussPt3D] \ - * D_phi.T @ material_stiffness @ D_phi * det_J + KEL += w_q * D_phi_q.T @ material_stiffness @ D_phi_q * dv_q return KEL def calc_element_traction_vector( - surf_num, pres_surf, nodal_coords, NUM_NODES_SURF, DOF_NODE, - SURF_TRACTION_VECTOR, quad_rule_2D, shape_func_triangle): + surf_num, pres_surf, nodal_coords, num_nodes_surf, dof_node, + surf_traction_vector, quad_rule_2D, shape_func_triangle): gauss_weights_2D = quad_rule_2D.wgauss shape_triangle = shape_func_triangle.values @@ -90,26 +92,26 @@ def calc_element_traction_vector( surf_points = nodal_coords[pres_surf[surf_num], :] - PEL = np.zeros(NUM_NODES_SURF * DOF_NODE) + PEL = np.zeros(num_nodes_surf * dof_node) for gaussPt2D in range(len(gauss_weights_2D)): - N = shape_triangle[gaussPt2D, :] - Nzeta = dshape_triangle[gaussPt2D, :, :] + shape_tri_q = shape_triangle[gaussPt2D, :] + dshape_tri_q = dshape_triangle[gaussPt2D, :, :] - J = Nzeta @ surf_points + J_q = dshape_tri_q @ surf_points - Js = np.linalg.norm(np.cross(J[0, :], J[1, :])) + da_q = np.linalg.norm(np.cross(J_q[0, :], J_q[1, :])) PEL += gauss_weights_2D[gaussPt2D] \ - * (np.column_stack([N, N, N]) \ - * SURF_TRACTION_VECTOR).T.reshape(-1) * Js + * (np.column_stack([shape_tri_q, shape_tri_q, shape_tri_q]) \ + * surf_traction_vector).T.reshape(-1) * da_q return PEL def assemble_global_stiffness( - KEL, volume_conn, eq_num, ELEM_NUM, KPP, KPF, KFF, KFP): + KEL, volume_conn, eq_num, elem_num, KPP, KPF, KFF, KFP): - elem_conn = volume_conn[ELEM_NUM] + elem_conn = volume_conn[elem_num] elem_eq_num = eq_num[elem_conn, : ] global_indices = elem_eq_num.T.reshape(-1) @@ -147,31 +149,35 @@ def assemble_global_traction_vector( def assemble_module( - NUM_PRES_DOF, NUM_FREE_DOF, NUM_ELE, NUM_NODES_ELE, DOF_NODE, - NUM_NODES_SURF, SURF_TRACTION_VECTOR, E, nu, disp_node, + num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, + num_nodes_surf, surf_traction_vector, params, disp_node, disp_val, eq_num, volume_conn, nodal_coords, pres_surf, quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle): # Initialize arrays that need to be returned (KPP, KPF, KFF, KFP, PP) - KPP = np.zeros((NUM_PRES_DOF, NUM_PRES_DOF)) - KPF = np.zeros((NUM_PRES_DOF, NUM_FREE_DOF)) - KFP = np.zeros((NUM_FREE_DOF, NUM_PRES_DOF)) - KFF = np.zeros((NUM_FREE_DOF, NUM_FREE_DOF)) - PP = np.zeros(NUM_PRES_DOF) - PF = np.zeros(NUM_FREE_DOF) + KPP = np.zeros((num_pres_dof, num_pres_dof)) + KPF = np.zeros((num_pres_dof, num_free_dof)) + KFP = np.zeros((num_free_dof, num_pres_dof)) + KFF = np.zeros((num_free_dof, num_free_dof)) + PP = np.zeros(num_pres_dof) + PF = np.zeros(num_free_dof) ## Prescribe boundary conditions - UP = assemble_prescribed_displacement(NUM_PRES_DOF, disp_node, + UP = assemble_prescribed_displacement(num_pres_dof, disp_node, disp_val, eq_num) + + gauss_weights_3D = quad_rule_3D.wgauss + shape_tetra = shape_func_tetra.values + dshape_tetra = shape_func_tetra.gradients start = time.time() # assemble global stiffness and force - for elem_num in range(0, NUM_ELE): + for elem_num in range(0, num_elem): + elem_points = nodal_coords[volume_conn[elem_num], :] # get element stiffness matrix - KEL = calc_element_stiffness_matrix( - elem_num, volume_conn, nodal_coords, NUM_NODES_ELE, - DOF_NODE, E, nu, quad_rule_3D, shape_func_tetra) + KEL = calc_element_stiffness_matrix(elem_points, num_nodes_elem,dof_node, + params, gauss_weights_3D, dshape_tetra) # assemble global stiffness assemble_global_stiffness(KEL, volume_conn, eq_num, @@ -183,9 +189,68 @@ def assemble_module( for surf_num in range(len(pres_surf)): # get local traction vector - PEL = calc_element_traction_vector( - surf_num, pres_surf, nodal_coords, NUM_NODES_SURF, DOF_NODE, - SURF_TRACTION_VECTOR, quad_rule_2D, shape_func_triangle) + PEL = calc_element_traction_vector(surf_num, pres_surf, nodal_coords, num_nodes_surf, + dof_node,surf_traction_vector, quad_rule_2D, + shape_func_triangle) + + # assemble traction vector + assemble_global_traction_vector(PEL, pres_surf, + surf_num, eq_num, PF, PP) + end = time.time() + print("Time to assemble load vector: ", end - start) + + return KPP, KPF, KFF, KFP, PF, PP, UP + +def assemble_module_AD( + num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, + num_nodes_surf, surf_traction_vector, params, disp_node, + disp_val, eq_num, volume_conn, nodal_coords, pres_surf, + quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle): + + # Initialize arrays that need to be returned (KPP, KPF, KFF, KFP, PP) + KPP = np.zeros((num_pres_dof, num_pres_dof)) + KPF = np.zeros((num_pres_dof, num_free_dof)) + KFP = np.zeros((num_free_dof, num_pres_dof)) + KFF = np.zeros((num_free_dof, num_free_dof)) + PP = np.zeros(num_pres_dof) + PF = np.zeros(num_free_dof) + + ## Prescribe boundary conditions + UP = assemble_prescribed_displacement(num_pres_dof, disp_node, + disp_val, eq_num) + + start = time.time() + + grad_residual = jax.jit(jax.jacfwd(elem_residual), + static_argnames=['num_nodes_elem', 'dof_node']) + + u_guess = np.ones(num_nodes_elem * dof_node) + + gauss_weights_3D = quad_rule_3D.wgauss + shape_tetra = shape_func_tetra.values + dshape_tetra = shape_func_tetra.gradients + + # assemble global stiffness and force + for elem_num in range(0, num_elem): + + elem_points = nodal_coords[volume_conn[elem_num], :] + # get element stiffness matrix + KEL = grad_residual(u_guess, params, elem_points, num_nodes_elem, dof_node, + gauss_weights_3D, shape_tetra, dshape_tetra) + + # assemble global stiffness + assemble_global_stiffness(np.array(KEL), volume_conn, eq_num, + elem_num, KPP, KPF, KFF, KFP) + end = time.time() + print("Time to assemble stiffness matrix: ", end - start) + + start = time.time() + for surf_num in range(len(pres_surf)): + + # get local traction vector + PEL = calc_element_traction_vector(surf_num, pres_surf, nodal_coords, num_nodes_surf, + dof_node,surf_traction_vector, quad_rule_2D, + shape_func_triangle) # assemble traction vector assemble_global_traction_vector(PEL, pres_surf, diff --git a/cmad/fem_utils/fem_functions_AD.py b/cmad/fem_utils/fem_functions_AD.py new file mode 100644 index 0000000..8ce8a63 --- /dev/null +++ b/cmad/fem_utils/fem_functions_AD.py @@ -0,0 +1,62 @@ +import jax.numpy as jnp + +def compute_shape_jacobian(elem_points, dshape_tetra): + + J = (dshape_tetra @ elem_points).T + + dv = jnp.linalg.det(J) + + # derivatives of shape functions with respect to spacial coordinates + gradphiXYZ = jnp.linalg.inv(J).T @ dshape_tetra + + return dv, gradphiXYZ + +def interpolate(u, shape_tetra, gradphiXYZ, num_nodes_elem): + + ux = u[0:num_nodes_elem] + uy = u[num_nodes_elem:num_nodes_elem * 2] + uz = u[num_nodes_elem * 2:num_nodes_elem * 3] + + u = jnp.array([jnp.dot(ux, shape_tetra), + jnp.dot(uy, shape_tetra), + jnp.dot(uz, shape_tetra)]) + + grad_u = jnp.vstack([gradphiXYZ @ ux, + gradphiXYZ @ uy, + gradphiXYZ @ uz]) + + return u, grad_u + +def compute_stress(grad_u, params): + + E = params[0] + nu = params[1] + + mu = E / (2 * (1 + nu)) + lam = E * nu / ((1 + nu) * (1 - 2 * nu)) + + strain = 1 / 2 * (grad_u + grad_u.T) + + stress = lam * jnp.trace(strain) * jnp.eye(3) + 2 * mu * strain + + return stress + +def elem_residual(u, params, elem_points, num_nodes_elem, + dof_node, gauss_weights_3D, shape_tetra, dshape_tetra): + + residual = jnp.zeros((num_nodes_elem, dof_node)) + + for gaussPt3D in range(len(gauss_weights_3D)): + w_q = gauss_weights_3D[gaussPt3D] + + dshape_tetra_q = dshape_tetra[gaussPt3D, :, :] + shape_tetra_q = shape_tetra[gaussPt3D, :] + + dv_q, gradphiXYZ_q = compute_shape_jacobian(elem_points, dshape_tetra_q) + u_q, grad_u_q = interpolate(u, shape_tetra_q, gradphiXYZ_q, num_nodes_elem) + + stress = compute_stress(grad_u_q, params) + + residual += w_q * gradphiXYZ_q.T @ stress * dv_q + + return residual.reshape(-1, order='F') diff --git a/cmad/fem_utils/fem_problem.py b/cmad/fem_utils/fem_problem.py index 2e0cc6e..8663cdc 100644 --- a/cmad/fem_utils/fem_problem.py +++ b/cmad/fem_utils/fem_problem.py @@ -7,7 +7,7 @@ class fem_problem(): def __init__(self, problem_type, order): # evaluate triangle basis functions at quadrature points - self._quad_rule_2D\ + self._quad_rule_2D \ = quadrature_rule.create_quadrature_rule_on_triangle(order) gauss_pts_2D = self._quad_rule_2D.xigauss self._shape_func_triangle = interpolants.shape_triangle(gauss_pts_2D) @@ -26,30 +26,30 @@ def __init__(self, problem_type, order): self._colume_conn = self._mesh.get_volume_connectivity() self._surface_conn = self._mesh.get_surface_connectivity() - self._DOF_NODE = 3 - self._NUM_NODES = len(self._nodal_coords) - self._NUM_NODES_ELE = 4 - self._NUM_ELE = len(self._colume_conn) - self._NUM_NODES_SURF = 3 + self._dof_node = 3 + self._num_nodes = len(self._nodal_coords) + self._num_nodes_elem = 4 + self._num_elem = len(self._colume_conn) + self._num_nodes_surf = 3 # fix all nodes on plane x = 0 pres_nodes = [] - for i in range(self._NUM_NODES): + for i in range(self._num_nodes): if self._nodal_coords[i][0] == 0.0: pres_nodes.append(i) NUM_PRES_NODES = len(pres_nodes) self._disp_node = np.zeros((NUM_PRES_NODES \ - * self._DOF_NODE, 2), dtype = int) + * self._dof_node, 2), dtype = int) for i in range(NUM_PRES_NODES): - for j in range(self._DOF_NODE): - self._disp_node[i * self._DOF_NODE + j][0] \ + for j in range(self._dof_node): + self._disp_node[i * self._dof_node + j][0] \ = pres_nodes[i] - self._disp_node[i * self._DOF_NODE + j][1] = j+1 - self._disp_val = np.zeros(NUM_PRES_NODES * self._DOF_NODE) + self._disp_node[i * self._dof_node + j][1] = j + 1 + self._disp_val = np.zeros(NUM_PRES_NODES * self._dof_node) # normal traction on plane x = 1 - self._SURF_TRACTION_VECTOR = np.array([1.0, 0.0, 0.0]) + self._surf_traction_vector = np.array([1.0, 0.0, 0.0]) pres_surf = [] for surface in self._surface_conn: surface_points = self._nodal_coords[surface, :] @@ -65,15 +65,15 @@ def __init__(self, problem_type, order): self._colume_conn = self._mesh.get_volume_connectivity() self._surface_conn = self._mesh.get_surface_connectivity() - self._DOF_NODE = 3 - self._NUM_NODES = len(self._nodal_coords) - self._NUM_NODES_ELE = 4 - self._NUM_ELE = len(self._colume_conn) - self._NUM_NODES_SURF = 3 + self._dof_node = 3 + self._num_nodes = len(self._nodal_coords) + self._num_nodes_elem = 4 + self._num_elem = len(self._colume_conn) + self._num_nodes_surf = 3 # prescribe ux = 0 on x = 0, uy = 0 on y = 0, and uz = 0 on z = 0 disp_node = [] - for i in range(self._NUM_NODES): + for i in range(self._num_nodes): if self._nodal_coords[i][0] == 0.0: disp_node.append(np.array([i, 1], dtype = int)) if self._nodal_coords[i][1] == 0.0: @@ -84,7 +84,7 @@ def __init__(self, problem_type, order): self._disp_val = np.zeros(len(disp_node)) # normal traction on x = 2 - self._SURF_TRACTION_VECTOR = np.array([1.0, 0.0, 0.0]) + self._surf_traction_vector = np.array([1.0, 0.0, 0.0]) pres_surf = [] for surface in self._surface_conn: surface_points = self._nodal_coords[surface, :] @@ -101,18 +101,18 @@ def __init__(self, problem_type, order): self._colume_conn = self._mesh.get_volume_connectivity() self._surface_conn = self._mesh.get_surface_connectivity() - self._DOF_NODE = 3 - self._NUM_NODES = len(self._nodal_coords) - self._NUM_NODES_ELE = 4 - self._NUM_ELE = len(self._colume_conn) - self._NUM_NODES_SURF = 3 + self._dof_node = 3 + self._num_nodes = len(self._nodal_coords) + self._num_nodes_elem = 4 + self._num_elem = len(self._colume_conn) + self._num_nodes_surf = 3 #impose linear displacement field on the boundary disp_node = [] disp_val = [] - self.UUR_true = np.zeros((self._NUM_NODES, self._DOF_NODE)) + self.UUR_true = np.zeros((self._num_nodes, self._dof_node)) - for i in range(self._NUM_NODES): + for i in range(self._num_nodes): coord = self._nodal_coords[i] x = coord[0] y = coord[1] @@ -137,7 +137,7 @@ def __init__(self, problem_type, order): self._disp_val = np.array(disp_val) # no surface tractions - self._SURF_TRACTION_VECTOR = np.array([0.0, 0.0, 0.0]) + self._surf_traction_vector = np.array([0.0, 0.0, 0.0]) self._pres_surf = np.array([]) if problem_type == "simple_shear": @@ -148,15 +148,15 @@ def __init__(self, problem_type, order): self._colume_conn = self._mesh.get_volume_connectivity() self._surface_conn = self._mesh.get_surface_connectivity() - self._DOF_NODE = 3 - self._NUM_NODES = len(self._nodal_coords) - self._NUM_NODES_ELE = 4 - self._NUM_ELE = len(self._colume_conn) - self._NUM_NODES_SURF = 3 + self._dof_node = 3 + self._num_nodes = len(self._nodal_coords) + self._num_nodes_elem = 4 + self._num_elem = len(self._colume_conn) + self._num_nodes_surf = 3 # fix all nodes on plane z = 0, and set uz = 0 on plane z = 2 disp_node = [] - for i in range(self._NUM_NODES): + for i in range(self._num_nodes): if self._nodal_coords[i][2] == 0.0: disp_node.append(np.array([i, 1], dtype = int)) disp_node.append(np.array([i, 2], dtype = int)) @@ -168,7 +168,7 @@ def __init__(self, problem_type, order): self._disp_val = np.zeros(len(disp_node)) # shear traction in x direction on plane z = 2 - self._SURF_TRACTION_VECTOR = np.array([1.0, 0.0, 0.0]) + self._surf_traction_vector = np.array([1.0, 0.0, 0.0]) pres_surf = [] for surface in self._surface_conn: surface_points = self._nodal_coords[surface, :] @@ -183,13 +183,13 @@ def get_3D_basis_functions(self): return self._quad_rule_3D, self._shape_func_tetra def get_mesh_properties(self): - return self._DOF_NODE, self._NUM_NODES, self._NUM_NODES_ELE, \ - self._NUM_ELE, self._NUM_NODES_SURF, self._nodal_coords, \ + return self._dof_node, self._num_nodes, self._num_nodes_elem, \ + self._num_elem, self._num_nodes_surf, self._nodal_coords, \ self._colume_conn def get_boundary_conditions(self): return self._disp_node, self._disp_val, \ - self._pres_surf, self._SURF_TRACTION_VECTOR + self._pres_surf, self._surf_traction_vector def save_data(self, filename, data): self._mesh.add_point_data(data) diff --git a/cmad/fem_utils/interpolants.py b/cmad/fem_utils/interpolants.py index e1fcb45..875baa2 100644 --- a/cmad/fem_utils/interpolants.py +++ b/cmad/fem_utils/interpolants.py @@ -1,32 +1,25 @@ -from jaxtyping import Array, Float, Int -from scipy import special -import numpy as onp -import equinox as eqx -import jax.numpy as np +import numpy as np -class ShapeFunctions(eqx.Module): +class ShapeFunctions(): """ Shape functions and shape function gradients (in the parametric space). Attributes: values: Values of the shape functions at a discrete set of points. - Shape is ``(nPts, nNodes)``, where ``nPts`` is the number of - points at which the shame functinos are evaluated, and ``nNodes`` + Shape is ``(num_eval_points, num_nodes_elem)``, where ``num_eval_points`` is the number of + points at which the shape functions are evaluated, and ``num_nodes_elem`` is the number of nodes in the element (which is equal to the number of shape functions). gradients: Values of the parametric gradients of the shape functions. - Shape is ``(nPts, nDim, nNodes)``, where ``nDim`` is the number + Shape is ``(num_eval_points, dof_node, num_nodes_elem)``, where ``dof_node`` is the number of spatial dimensions. Line elements are an exception, which - have shape ``(nPts, nNdodes)``. + have shape ``(num_eval_points, num_nodes_elem)``. """ - values: Float[Array, "nq nn"] - gradients: Float[Array, "nq nn nd"] - - def __iter__(self): - yield self.values - yield self.gradients + def __init__(self, values, gradients): + self.values = values + self.gradients = gradients def shape1d(evaluationPoints): """ @@ -40,8 +33,8 @@ def shape1d(evaluationPoints): Returns: Shape function values and shape function derivatives at ``evaluationPoints``, in a tuple (``shape``, ``dshape``). - shapes: [nEvalPoints, nNodes] - dshapes: [nEvalPoints, nNodes] + shapes: [num_eval_points, num_nodes_elem] + dshapes: [num_eval_points, num_nodes_elem] """ @@ -62,23 +55,23 @@ def shape_triangle(evaluationPoints): Returns: Shape function values and shape function derivatives at ``evaluationPoints``, in a tuple (``shape``, ``dshape``). - shapes: [nEvalPoints, nNodes] - dshapes: [nEvalPoints, nDims, nNodes] + shapes: [num_eval_points, num_nodes_elem] + dshapes: [num_eval_points, dof_node, num_nodes_elem] """ - nEvalPoints = len(evaluationPoints) - nDims = 2 - nNodes = 3 + num_eval_points = len(evaluationPoints) + dof_node = 2 + num_nodes_elem = 3 shape = np.vstack((1 - evaluationPoints[:,0] - evaluationPoints[:,1], evaluationPoints[:,0], evaluationPoints[:,1])).T - dshape = onp.zeros((nEvalPoints, nDims, nNodes)) + dshape = np.zeros((num_eval_points, dof_node, num_nodes_elem)) - for i in range(nEvalPoints): - dshape[i,:,:] = onp.array([[-1, 1, 0],[-1, 0, 1]]) + for i in range(num_eval_points): + dshape[i, :, :] = np.array([[-1, 1, 0],[-1, 0, 1]]) - return ShapeFunctions(shape, np.array(dshape)) + return ShapeFunctions(shape, dshape) def shape_quad(evaluationPoints): """ @@ -92,14 +85,14 @@ def shape_quad(evaluationPoints): Returns: Shape function values and shape function derivatives at ``evaluationPoints``, in a tuple (``shape``, ``dshape``). - shapes: [nEvalPoints, nNodes] - dshapes: [nEvalPoints, nDims, nNodes] + shapes: [num_eval_points, num_nodes_elem] + dshapes: [num_eval_points, dof_node, num_nodes_elem] """ - nEvalPoints = len(evaluationPoints) - nDims = 2 - nNodes = 4 + num_eval_points = len(evaluationPoints) + dof_node = 2 + num_nodes_elem = 4 l0x = 1 - evaluationPoints[:,0] l1x = 1 + evaluationPoints[:,0] @@ -107,17 +100,17 @@ def shape_quad(evaluationPoints): l1y = 1 + evaluationPoints[:,1] shape = np.vstack((l0x * l0y / 4, l1x * l0y / 4, l1x * l1y / 4, l0x * l1y / 4)).T - dshape = onp.zeros((nEvalPoints, nDims, nNodes)) + dshape = np.zeros((num_eval_points, dof_node, num_nodes_elem)) - for i in range(nEvalPoints): + for i in range(num_eval_points): point = evaluationPoints[i] l0x = 1 - point[0] l1x = 1 + point[0] l0y = 1 - point[1] l1y = 1 + point[1] - dshape[i,:,:] = onp.array([[-l0y, l0y, l1y, -l1y],[-l0x, -l1x, l1x, l0x]]) + dshape[i, :, :] = np.array([[-l0y, l0y, l1y, -l1y],[-l0x, -l1x, l1x, l0x]]) - return ShapeFunctions(shape, np.array(dshape)) + return ShapeFunctions(shape, dshape) def shape_tetrahedron(evaluationPoints): """ @@ -131,24 +124,24 @@ def shape_tetrahedron(evaluationPoints): Returns: Shape function values and shape function derivatives at ``evaluationPoints``, in a tuple (``shape``, ``dshape``). - shapes: [nEvalPoints, nNodes] - dshapes: [nEvalPoints, nDims, nNodes] + shapes: [num_eval_points, num_nodes_elem] + dshapes: [num_eval_points, dof_node, num_nodes_elem] """ - nEvalPoints = len(evaluationPoints) - nDims = 3 - nNodes = 4 + num_eval_points = len(evaluationPoints) + dof_node = 3 + num_nodes_elem = 4 shape = np.vstack((1 - evaluationPoints[:,0] - evaluationPoints[:,1] - evaluationPoints[:,2], evaluationPoints[:,0], evaluationPoints[:,1], evaluationPoints[:,2])).T - dshape = onp.zeros((nEvalPoints, nDims, nNodes)) + dshape = np.zeros((num_eval_points, dof_node, num_nodes_elem)) - for i in range(nEvalPoints): - dshape[i,:,:] = onp.array([[-1, 1, 0, 0], [-1, 0, 1, 0], [-1, 0, 0, 1]]) + for i in range(num_eval_points): + dshape[i, :, :] = np.array([[-1, 1, 0, 0], [-1, 0, 1, 0], [-1, 0, 0, 1]]) - return ShapeFunctions(shape, np.array(dshape)) + return ShapeFunctions(shape, dshape) def shape_brick(evaluationPoints): @@ -163,14 +156,14 @@ def shape_brick(evaluationPoints): Returns: Shape function values and shape function derivatives at ``evaluationPoints``, in a tuple (``shape``, ``dshape``). - shapes: [nEvalPoints, nNodes] - dshapes: [nEvalPoints, nDims, nNodes] + shapes: [num_eval_points, num_nodes_elem] + dshapes: [num_eval_points, dof_node, num_nodes_elem] """ - nEvalPoints = len(evaluationPoints) - nDims = 3 - nNodes = 8 + num_eval_points = len(evaluationPoints) + dof_node = 3 + num_nodes_elem = 8 m1 = 1 - evaluationPoints[:,0] p1 = 1 + evaluationPoints[:,0] @@ -182,9 +175,9 @@ def shape_brick(evaluationPoints): shape = np.vstack((m1 * m2 * m3 / 8, p1 * m2 * m3 / 8, p1 * p2 * m3 / 8, m1 * p2 * m3 / 8, m1 * m2 * p3 / 8, p1 * m2 * p3 / 8, p1 * p2 * p3 / 8, m1 * p2 * p3 / 8)).T - dshape = onp.zeros((nEvalPoints, nDims, nNodes)) + dshape = np.zeros((num_eval_points, dof_node, num_nodes_elem)) - for i in range(nEvalPoints): + for i in range(num_eval_points): point = evaluationPoints[i] m1 = 1 - point[0] p1 = 1 + point[0] @@ -192,10 +185,10 @@ def shape_brick(evaluationPoints): p2 = 1 + point[1] m3 = 1 - point[2] p3 = 1 + point[2] - dshape[i,:,:] = onp.array([[-m2 * m3 / 8, m2 * m3 / 8, p2 * m3 / 8, -p2 * m3 / 8, -m2 * p3 / 8, m2 * p3 / 8, p2 * p3 / 8, -p2 * p3 / 8], + dshape[i, :, :] = np.array([[-m2 * m3 / 8, m2 * m3 / 8, p2 * m3 / 8, -p2 * m3 / 8, -m2 * p3 / 8, m2 * p3 / 8, p2 * p3 / 8, -p2 * p3 / 8], [-m1 * m3 / 8, -p1 * m3 / 8, p1 * m3 / 8, m1 * m3 / 8, -m1 * p3 / 8, -p1 * p3 / 8, p1 * p3 / 8, m1 * p3 / 8], [-m1 * m2 / 8, -p1 * m2 / 8, -p1 * p2 / 8, -m1 * p2 / 8, m1 * m2 / 8, p1 * m2 / 8, p1 * p2 / 8, m1 * p2 / 8]]) - return ShapeFunctions(shape, np.array(dshape)) + return ShapeFunctions(shape, dshape) diff --git a/cmad/fem_utils/quadrature_rule.py b/cmad/fem_utils/quadrature_rule.py index 0e7afaf..6164085 100644 --- a/cmad/fem_utils/quadrature_rule.py +++ b/cmad/fem_utils/quadrature_rule.py @@ -1,29 +1,15 @@ -from jax.lax import switch -from jaxtyping import Array, Float -import equinox as eqx -import jax.numpy as np -import math -import numpy as onp +import numpy as np import scipy.special -class QuadratureRule(eqx.Module): - """ - Quadrature rule points and weights. - An ``equinox`` ``Module`` containing ``xigauss``, a ``jax.numpy`` array of the - coordinates of the sample points in the reference domain, and - ``wgauss``, a ``jax.numpy`` array with the weights. - """ - xigauss: Float[Array, "nq 2"] - wgauss: Float[Array, "nq"] - - def __iter__(self): - yield self.xigauss - yield self.wgauss +class QuadratureRule(): + + def __init__(self, xigauss, wgauss): + self.xigauss = xigauss + self.wgauss = wgauss def __len__(self): return self.xigauss.shape[0] - def create_quadrature_rule_1D(degree): """Creates a Gauss-Legendre quadrature on the interval [-1, 1]. @@ -40,13 +26,13 @@ def create_quadrature_rule_1D(degree): and the weights. """ - n = math.ceil((degree + 1)/2) + n = np.ceil((degree + 1)/2) xi, w = scipy.special.roots_sh_legendre(n) xi = -1 + 2 * xi - xi_matrix = onp.zeros((len(xi), 1)) + xi_matrix = np.zeros((len(xi), 1)) xi_matrix[:,0]=xi - return QuadratureRule(np.array(xi_matrix), w) + return QuadratureRule(xi_matrix, w) def create_quadrature_rule_on_triangle(degree): diff --git a/cmad/fem_utils/solve_fem.py b/cmad/fem_utils/solve_fem.py index 2e640f4..e504204 100644 --- a/cmad/fem_utils/solve_fem.py +++ b/cmad/fem_utils/solve_fem.py @@ -1,57 +1,75 @@ from cmad.fem_utils.fem_problem import fem_problem from cmad.fem_utils.fem_functions import (initialize_equation, - assemble_module, + assemble_module, assemble_module_AD, solve_module) import time +import numpy as np """ -DOF_NODE: number of degrees of freedom per node -NUM_NODES: total number of nodes in the mesh -NUM_NODE_ELE: number of nodes per element -NUM_ELE: number of elements -NUM_NODES_SURF: number of nodes per surface element -NUM_FREE_DOF: number of free degrees of freedom -NUM_PRES_DOF: number of prescribed degrees of freedom -SURF_TRACTION_VECTOR: surface traction vector +dof_node: number of degrees of freedom per node +num_nodes: total number of nodes in the mesh +num_node_elem: number of nodes per element +num_elem: number of elements +num_nodes_surf: number of nodes per surface element +num_free_dof: number of free degrees of freedom +num_pres_dof: number of prescribed degrees of freedom +surf_traction_vector: surface traction vector E: Youngs Modulus nu: Poisson's ratio -disp_node: (NUM_PRES_DOF x 2) array that specifies +disp_node: (num_pres_dof x 2) array that specifies which node and dof is prescribed disp_val: (NUM_PRED_DOF x 1) array of values of the prescribed displacements -eq_num: (NUM_NODES x DOF_NODE) array that specifies where +eq_num: (num_nodes x dof_node) array that specifies where each node and its DOFs belong in the global stifness matrix volume_conn: connectivity matrix for 3D elements pres_surf: connectivity for surfaces that have a prescribed traction nodal_coords: spacial coordinates for each node in mesh - + +In the definitions below, "P" stands for prescribed displacements +and "F" stands for free displacments + +KPP: (num_pres_dof x num_pres_dof) partion of stiffness matrix with rows + and columns corresponding with prescribed DOFS +KPF: (num_pres_dof x num_free_dof) partion of stiffness matrix with rows + corresponding with prescribed DOFS and columns corresponding with free DOFS +KFF: (num_free_dof x num_free_dof) partion of stiffness matrix with rows and + columns corresponding with free DOFS +KFP: (num_free_dof x num_free_dof) partion of stiffness matrix with rows + corresponding with free DOFS and columns corresponding with prescribed DOFS +PF: (num_free_dof, ) RHS vector corresponding with free DOFS +PP: (num_pres_dof, ) RHS vector corresponding with prescribed DOFS +UP: (num_pres_dof, ) vector of prescribed displacements +UF: (num_free_dof, ) vector of free displacements (the ones that we solve for) +UUR: (num_nodes x 3) matrix of displacements at all nodes in the mesh """ order = 3 -problem = fem_problem("patch_B", order) +problem = fem_problem("simple_shear", order) -DOF_NODE, NUM_NODES, NUM_NODES_ELE, NUM_ELE, NUM_NODES_SURF, \ +dof_node, num_nodes, num_nodes_ELE, num_elem, num_nodes_surf, \ nodal_coords, volume_conn = problem.get_mesh_properties() -disp_node, disp_val, pres_surf, SURF_TRACTION_VECTOR \ +disp_node, disp_val, pres_surf, surf_traction_vector \ = problem.get_boundary_conditions() quad_rule_3D, shape_func_tetra = problem.get_3D_basis_functions() quad_rule_2D, shape_func_triangle = problem.get_2D_basis_functions() -print("Number of elements:", NUM_ELE) +print("Number of elements: ", num_elem) + +params = np.array([200, 0.3]) -E = 200 -nu = 0.3 +eq_num, num_free_dof, num_pres_dof \ + = initialize_equation(num_nodes, dof_node, disp_node) -eq_num, NUM_FREE_DOF, NUM_PRES_DOF \ - = initialize_equation(NUM_NODES, DOF_NODE, disp_node) +print("Number of free degrees of freedom: ", num_free_dof) KPP, KPF, KFF, KFP, PF, PP, UP \ - = assemble_module(NUM_PRES_DOF, NUM_FREE_DOF, NUM_ELE, NUM_NODES_ELE, - DOF_NODE, NUM_NODES_SURF, SURF_TRACTION_VECTOR, E, nu, + = assemble_module_AD(num_pres_dof, num_free_dof, num_elem, num_nodes_ELE, + dof_node, num_nodes_surf, surf_traction_vector, params, disp_node, disp_val, eq_num, volume_conn, nodal_coords, pres_surf, quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle) diff --git a/cmad/fem_utils/tests/test_fem.py b/cmad/fem_utils/tests/test_fem.py index 9e56bf3..1ce611b 100644 --- a/cmad/fem_utils/tests/test_fem.py +++ b/cmad/fem_utils/tests/test_fem.py @@ -1,4 +1,4 @@ -import jax.numpy as np +import numpy as np from cmad.fem_utils.fem_problem import fem_problem from cmad.fem_utils.fem_functions import (initialize_equation, assemble_module, solve_module) @@ -11,24 +11,23 @@ def test_patch_form_B(self): order = 3 problem = fem_problem("patch_B", order) - DOF_NODE, NUM_NODES, NUM_NODES_ELE, NUM_ELE, NUM_NODES_SURF, \ + dof_node, num_nodes, num_nodes_elem, num_elem, num_nodes_surf, \ nodal_coords, volume_conn = problem.get_mesh_properties() - disp_node, disp_val, pres_surf, SURF_TRACTION_VECTOR \ + disp_node, disp_val, pres_surf, surf_traction_vector \ = problem.get_boundary_conditions() quad_rule_3D, shape_func_tetra = problem.get_3D_basis_functions() quad_rule_2D, shape_func_triangle = problem.get_2D_basis_functions() - E = 200 - nu = 0.3 + params = np.array([200, 0.3]) - eq_num, NUM_FREE_DOF, NUM_PRES_DOF\ - = initialize_equation(NUM_NODES, DOF_NODE, disp_node) + eq_num, num_free_dof, num_pres_dof \ + = initialize_equation(num_nodes, dof_node, disp_node) KPP, KPF, KFF, KFP, PF, PP, UP \ - = assemble_module(NUM_PRES_DOF, NUM_FREE_DOF, NUM_ELE, NUM_NODES_ELE, - DOF_NODE, NUM_NODES_SURF,SURF_TRACTION_VECTOR, E, nu, + = assemble_module(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, + dof_node, num_nodes_surf,surf_traction_vector, params, disp_node, disp_val, eq_num, volume_conn, nodal_coords, pres_surf, quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle) From d825c058bb448aaa3092606381628f7c69394d08 Mon Sep 17 00:00:00 2001 From: ryany Date: Wed, 22 Jan 2025 15:18:19 -0800 Subject: [PATCH 3/5] eliminated trailing whitespaces --- cmad/fem_utils/fem_functions.py | 28 ++++++++++----------- cmad/fem_utils/fem_functions_AD.py | 6 ++--- cmad/fem_utils/fem_problem.py | 20 +++++++-------- cmad/fem_utils/interpolants.py | 14 +++++------ cmad/fem_utils/mesh.py | 8 +++--- cmad/fem_utils/quadrature_rule.py | 30 +++++++++++------------ cmad/fem_utils/solve_fem.py | 24 +++++++++--------- cmad/fem_utils/tests/test_fem.py | 6 ++--- cmad/fem_utils/tests/test_interpolants.py | 2 +- 9 files changed, 69 insertions(+), 69 deletions(-) diff --git a/cmad/fem_utils/fem_functions.py b/cmad/fem_utils/fem_functions.py index 064be4c..248b243 100644 --- a/cmad/fem_utils/fem_functions.py +++ b/cmad/fem_utils/fem_functions.py @@ -11,7 +11,7 @@ def initialize_equation(num_nodes, dof_node, disp_node): node_number = node[0] direction = node[1] eq_num[node_number][direction - 1] = -(i + 1) - + num_free_dof = 0 for i in range(len(eq_num)): for j in range(len(eq_num[i])): @@ -24,7 +24,7 @@ def initialize_equation(num_nodes, dof_node, disp_node): def assemble_prescribed_displacement( num_pres_dof, disp_node, disp_val, eq_num): - + UP = np.zeros(num_pres_dof) for i, row in enumerate(disp_node): node_number = row[0] @@ -34,7 +34,7 @@ def assemble_prescribed_displacement( UP[eqn_number - 1] = displacement return UP - + def calc_element_stiffness_matrix( elem_points, num_nodes_elem, dof_node, params, gauss_weights_3D, dshape_tetra): @@ -77,9 +77,9 @@ def calc_element_stiffness_matrix( = np.concatenate((gradphiXYZ_q[2, :], gradphiXYZ_q[1, :])) D_phi_q[5, 0:num_nodes_elem] = gradphiXYZ_q[2, :] D_phi_q[5, num_nodes_elem * 2:num_nodes_elem * 3] = gradphiXYZ_q[0, :] - + KEL += w_q * D_phi_q.T @ material_stiffness @ D_phi_q * dv_q - + return KEL def calc_element_traction_vector( @@ -105,7 +105,7 @@ def calc_element_traction_vector( PEL += gauss_weights_2D[gaussPt2D] \ * (np.column_stack([shape_tri_q, shape_tri_q, shape_tri_q]) \ * surf_traction_vector).T.reshape(-1) * da_q - + return PEL def assemble_global_stiffness( @@ -161,11 +161,11 @@ def assemble_module( KFF = np.zeros((num_free_dof, num_free_dof)) PP = np.zeros(num_pres_dof) PF = np.zeros(num_free_dof) - + ## Prescribe boundary conditions UP = assemble_prescribed_displacement(num_pres_dof, disp_node, disp_val, eq_num) - + gauss_weights_3D = quad_rule_3D.wgauss shape_tetra = shape_func_tetra.values dshape_tetra = shape_func_tetra.gradients @@ -180,14 +180,14 @@ def assemble_module( params, gauss_weights_3D, dshape_tetra) # assemble global stiffness - assemble_global_stiffness(KEL, volume_conn, eq_num, + assemble_global_stiffness(KEL, volume_conn, eq_num, elem_num, KPP, KPF, KFF, KFP) end = time.time() print("Time to assemble stiffness matrix: ", end - start) start = time.time() for surf_num in range(len(pres_surf)): - + # get local traction vector PEL = calc_element_traction_vector(surf_num, pres_surf, nodal_coords, num_nodes_surf, dof_node,surf_traction_vector, quad_rule_2D, @@ -214,7 +214,7 @@ def assemble_module_AD( KFF = np.zeros((num_free_dof, num_free_dof)) PP = np.zeros(num_pres_dof) PF = np.zeros(num_free_dof) - + ## Prescribe boundary conditions UP = assemble_prescribed_displacement(num_pres_dof, disp_node, disp_val, eq_num) @@ -223,7 +223,7 @@ def assemble_module_AD( grad_residual = jax.jit(jax.jacfwd(elem_residual), static_argnames=['num_nodes_elem', 'dof_node']) - + u_guess = np.ones(num_nodes_elem * dof_node) gauss_weights_3D = quad_rule_3D.wgauss @@ -239,14 +239,14 @@ def assemble_module_AD( gauss_weights_3D, shape_tetra, dshape_tetra) # assemble global stiffness - assemble_global_stiffness(np.array(KEL), volume_conn, eq_num, + assemble_global_stiffness(np.array(KEL), volume_conn, eq_num, elem_num, KPP, KPF, KFF, KFP) end = time.time() print("Time to assemble stiffness matrix: ", end - start) start = time.time() for surf_num in range(len(pres_surf)): - + # get local traction vector PEL = calc_element_traction_vector(surf_num, pres_surf, nodal_coords, num_nodes_surf, dof_node,surf_traction_vector, quad_rule_2D, diff --git a/cmad/fem_utils/fem_functions_AD.py b/cmad/fem_utils/fem_functions_AD.py index 8ce8a63..f49eaf5 100644 --- a/cmad/fem_utils/fem_functions_AD.py +++ b/cmad/fem_utils/fem_functions_AD.py @@ -12,7 +12,7 @@ def compute_shape_jacobian(elem_points, dshape_tetra): return dv, gradphiXYZ def interpolate(u, shape_tetra, gradphiXYZ, num_nodes_elem): - + ux = u[0:num_nodes_elem] uy = u[num_nodes_elem:num_nodes_elem * 2] uz = u[num_nodes_elem * 2:num_nodes_elem * 3] @@ -20,7 +20,7 @@ def interpolate(u, shape_tetra, gradphiXYZ, num_nodes_elem): u = jnp.array([jnp.dot(ux, shape_tetra), jnp.dot(uy, shape_tetra), jnp.dot(uz, shape_tetra)]) - + grad_u = jnp.vstack([gradphiXYZ @ ux, gradphiXYZ @ uy, gradphiXYZ @ uz]) @@ -58,5 +58,5 @@ def elem_residual(u, params, elem_points, num_nodes_elem, stress = compute_stress(grad_u_q, params) residual += w_q * gradphiXYZ_q.T @ stress * dv_q - + return residual.reshape(-1, order='F') diff --git a/cmad/fem_utils/fem_problem.py b/cmad/fem_utils/fem_problem.py index 8663cdc..62b6b6f 100644 --- a/cmad/fem_utils/fem_problem.py +++ b/cmad/fem_utils/fem_problem.py @@ -17,9 +17,9 @@ def __init__(self, problem_type, order): = quadrature_rule.create_quadrature_rule_on_tetrahedron(order) gauss_pts_3D = self._quad_rule_3D.xigauss self._shape_func_tetra = interpolants.shape_tetrahedron(gauss_pts_3D) - - if problem_type == "hole_block_example": - + + if problem_type == "hole_block_example": + self._mesh = Mesh("hole_block") self._nodal_coords = self._mesh.get_nodal_coordinates() @@ -32,7 +32,7 @@ def __init__(self, problem_type, order): self._num_elem = len(self._colume_conn) self._num_nodes_surf = 3 - # fix all nodes on plane x = 0 + # fix all nodes on plane x = 0 pres_nodes = [] for i in range(self._num_nodes): if self._nodal_coords[i][0] == 0.0: @@ -58,7 +58,7 @@ def __init__(self, problem_type, order): self._pres_surf = np.array(pres_surf) if problem_type == "uniaxial_stress": - + self._mesh = Mesh("bar") self._nodal_coords = self._mesh.get_nodal_coordinates() @@ -71,7 +71,7 @@ def __init__(self, problem_type, order): self._num_elem = len(self._colume_conn) self._num_nodes_surf = 3 - # prescribe ux = 0 on x = 0, uy = 0 on y = 0, and uz = 0 on z = 0 + # prescribe ux = 0 on x = 0, uy = 0 on y = 0, and uz = 0 on z = 0 disp_node = [] for i in range(self._num_nodes): if self._nodal_coords[i][0] == 0.0: @@ -91,10 +91,10 @@ def __init__(self, problem_type, order): if (surface_points[:, 0] == 2 * np.ones(3)).all(): pres_surf.append(surface) self._pres_surf = np.array(pres_surf) - + # setup patch test form B if problem_type == "patch_B": - + self._mesh = Mesh("cube") self._nodal_coords = self._mesh.get_nodal_coordinates() @@ -178,7 +178,7 @@ def __init__(self, problem_type, order): def get_2D_basis_functions(self): return self._quad_rule_2D, self._shape_func_triangle - + def get_3D_basis_functions(self): return self._quad_rule_3D, self._shape_func_tetra @@ -190,7 +190,7 @@ def get_mesh_properties(self): def get_boundary_conditions(self): return self._disp_node, self._disp_val, \ self._pres_surf, self._surf_traction_vector - + def save_data(self, filename, data): self._mesh.add_point_data(data) self._mesh.save_mesh(filename) \ No newline at end of file diff --git a/cmad/fem_utils/interpolants.py b/cmad/fem_utils/interpolants.py index 875baa2..a883334 100644 --- a/cmad/fem_utils/interpolants.py +++ b/cmad/fem_utils/interpolants.py @@ -65,12 +65,12 @@ def shape_triangle(evaluationPoints): shape = np.vstack((1 - evaluationPoints[:,0] - evaluationPoints[:,1], evaluationPoints[:,0], evaluationPoints[:,1])).T - + dshape = np.zeros((num_eval_points, dof_node, num_nodes_elem)) for i in range(num_eval_points): dshape[i, :, :] = np.array([[-1, 1, 0],[-1, 0, 1]]) - + return ShapeFunctions(shape, dshape) def shape_quad(evaluationPoints): @@ -109,7 +109,7 @@ def shape_quad(evaluationPoints): l0y = 1 - point[1] l1y = 1 + point[1] dshape[i, :, :] = np.array([[-l0y, l0y, l1y, -l1y],[-l0x, -l1x, l1x, l0x]]) - + return ShapeFunctions(shape, dshape) def shape_tetrahedron(evaluationPoints): @@ -135,12 +135,12 @@ def shape_tetrahedron(evaluationPoints): shape = np.vstack((1 - evaluationPoints[:,0] - evaluationPoints[:,1] - evaluationPoints[:,2], evaluationPoints[:,0], evaluationPoints[:,1], evaluationPoints[:,2])).T - + dshape = np.zeros((num_eval_points, dof_node, num_nodes_elem)) for i in range(num_eval_points): dshape[i, :, :] = np.array([[-1, 1, 0, 0], [-1, 0, 1, 0], [-1, 0, 0, 1]]) - + return ShapeFunctions(shape, dshape) @@ -158,7 +158,7 @@ def shape_brick(evaluationPoints): in a tuple (``shape``, ``dshape``). shapes: [num_eval_points, num_nodes_elem] dshapes: [num_eval_points, dof_node, num_nodes_elem] - + """ num_eval_points = len(evaluationPoints) @@ -188,7 +188,7 @@ def shape_brick(evaluationPoints): dshape[i, :, :] = np.array([[-m2 * m3 / 8, m2 * m3 / 8, p2 * m3 / 8, -p2 * m3 / 8, -m2 * p3 / 8, m2 * p3 / 8, p2 * p3 / 8, -p2 * p3 / 8], [-m1 * m3 / 8, -p1 * m3 / 8, p1 * m3 / 8, m1 * m3 / 8, -m1 * p3 / 8, -p1 * p3 / 8, p1 * p3 / 8, m1 * p3 / 8], [-m1 * m2 / 8, -p1 * m2 / 8, -p1 * p2 / 8, -m1 * p2 / 8, m1 * m2 / 8, p1 * m2 / 8, p1 * p2 / 8, m1 * p2 / 8]]) - + return ShapeFunctions(shape, dshape) diff --git a/cmad/fem_utils/mesh.py b/cmad/fem_utils/mesh.py index 1f2a18f..0dd6c18 100644 --- a/cmad/fem_utils/mesh.py +++ b/cmad/fem_utils/mesh.py @@ -55,15 +55,15 @@ def __init__(self, mesh_type): def get_nodal_coordinates(self): return self._points - + def get_volume_connectivity(self): return self._volume_conn - + def get_surface_connectivity(self): return self._surface_conn - + def add_point_data(self, data): self._mesh.point_data = {"displacement_field": data} - + def save_mesh(self, filename): self._mesh.write(filename + ".vtk") diff --git a/cmad/fem_utils/quadrature_rule.py b/cmad/fem_utils/quadrature_rule.py index 6164085..272d646 100644 --- a/cmad/fem_utils/quadrature_rule.py +++ b/cmad/fem_utils/quadrature_rule.py @@ -2,7 +2,7 @@ import scipy.special class QuadratureRule(): - + def __init__(self, xigauss, wgauss): self.xigauss = xigauss self.wgauss = wgauss @@ -13,7 +13,7 @@ def __len__(self): def create_quadrature_rule_1D(degree): """Creates a Gauss-Legendre quadrature on the interval [-1, 1]. - The rule can exactly integrate polynomials of degree up to + The rule can exactly integrate polynomials of degree up to ``degree``. Parameters @@ -38,9 +38,9 @@ def create_quadrature_rule_1D(degree): def create_quadrature_rule_on_triangle(degree): """Creates a Gauss-Legendre quadrature on the unit triangle. - The rule can exactly integrate 2D polynomials up to the value of + The rule can exactly integrate 2D polynomials up to the value of ``degree``. The domain is the triangle between the vertices - (0, 0)-(1, 0)-(0, 1). The rules here are guaranteed to be + (0, 0)-(1, 0)-(0, 1). The rules here are guaranteed to be cyclically symmetric in triangular coordinates and to have strictly positive weights. @@ -181,7 +181,7 @@ def create_quadrature_rule_on_triangle(degree): def create_quadrature_rule_on_quad(degree): """Creates a Gauss-Legendre quadrature on square. - The rule can exactly integrate 2D polynomials up to the value of + The rule can exactly integrate 2D polynomials up to the value of ``degree``. The domain is the square between the vertices (1, -1)-(1, 1)-(-1, 1)-(-1, -1). @@ -205,7 +205,7 @@ def create_quadrature_rule_on_quad(degree): [a, -a], [a, a], [-a, a]]) - + w = np.array([1, 1, 1, 1]) elif degree <= 5: a = 0.774596669241483 @@ -223,7 +223,7 @@ def create_quadrature_rule_on_quad(degree): w = np.array([c * c, b * c, c * c, c * b, b * b, c * b, c * c, b * c, c * c]) - + else: raise ValueError("Quadrature of precision this high is not implemented.") @@ -236,11 +236,11 @@ def create_quadrature_rule_on_tetrahedron(degree): w = np.array([1.0/6.0]) elif degree == 2: - xi = np.array([[0.138196601125011, 0.138196601125011, 0.138196601125011], + xi = np.array([[0.138196601125011, 0.138196601125011, 0.138196601125011], [0.585410196624969, 0.138196601125011, 0.138196601125011], [0.138196601125011, 0.585410196624969, 0.138196601125011], [0.138196601125011, 0.138196601125011, 0.585410196624969]]) - + w = np.array([0.25/6.0, 0.25/6.0, 0.25/6.0, 0.25/6.0]) elif degree == 3: @@ -249,7 +249,7 @@ def create_quadrature_rule_on_tetrahedron(degree): [1./6., 1./6., 1./2.], [1./6., 1./2., 1./6.], [1./2., 1./6., 1./6.]]) - + w = np.array([-2/15, 3/40, 3/40, 3/40, 3/40]) elif degree == 4: @@ -264,12 +264,12 @@ def create_quadrature_rule_on_tetrahedron(degree): [0.399403576166799, 0.100596423833201, 0.100596423833201], [0.100596423833201, 0.399403576166799, 0.100596423833201], [0.100596423833201, 0.100596423833201, 0.399403576166799]]) - + w = np.array([-0.0131555555555556, 0.00762222222222222, 0.00762222222222222, 0.00762222222222222, 0.00762222222222222, 0.0248888888888889, 0.0248888888888889, 0.0248888888888889, 0.0248888888888889, 0.0248888888888889, 0.0248888888888889]) - + elif degree == 5: xi = np.array([[0.25, 0.25, 0.25], [0, 1./3., 1./3.], @@ -286,7 +286,7 @@ def create_quadrature_rule_on_tetrahedron(degree): [0.0665501535736643, 0.433449846426336, 0.433449846426336], [0.433449846426336, 0.0665501535736643, 0.433449846426336], [0.433449846426336, 0.433449846426336, 0.0665501535736643]]) - + w = np.array([0.0302836780970892, 0.00602678571428572, 0.00602678571428572, 0.00602678571428572, 0.00602678571428572, 0.011645249086029, 0.011645249086029, 0.011645249086029, 0.011645249086029, @@ -317,14 +317,14 @@ def create_quadrature_rule_on_tetrahedron(degree): [0.0636610018750175,0.603005664791649,0.269672331458316], [0.269672331458316,0.0636610018750175,0.603005664791649], [0.603005664791649,0.269672331458316,0.0636610018750175]]) - + w = np.array([0.00665379170969465, 0.00665379170969465, 0.00665379170969465, 0.00665379170969465, 0.00167953517588678, 0.00167953517588678, 0.00167953517588678, 0.00167953517588678, 0.0092261969239424, 0.0092261969239424, 0.0092261969239424, 0.0092261969239424, 0.00803571428571428, 0.00803571428571428, 0.00803571428571428, 0.00803571428571428, 0.00803571428571428, 0.00803571428571428, - 0.00803571428571428, 0.00803571428571428, 0.00803571428571428, + 0.00803571428571428, 0.00803571428571428, 0.00803571428571428, 0.00803571428571428, 0.00803571428571428, 0.00803571428571428]) else: raise ValueError("Quadrature of precision this high is not implemented.") diff --git a/cmad/fem_utils/solve_fem.py b/cmad/fem_utils/solve_fem.py index e504204..2a4a777 100644 --- a/cmad/fem_utils/solve_fem.py +++ b/cmad/fem_utils/solve_fem.py @@ -1,5 +1,5 @@ from cmad.fem_utils.fem_problem import fem_problem -from cmad.fem_utils.fem_functions import (initialize_equation, +from cmad.fem_utils.fem_functions import (initialize_equation, assemble_module, assemble_module_AD, solve_module) import time @@ -17,30 +17,30 @@ surf_traction_vector: surface traction vector E: Youngs Modulus nu: Poisson's ratio -disp_node: (num_pres_dof x 2) array that specifies +disp_node: (num_pres_dof x 2) array that specifies which node and dof is prescribed -disp_val: (NUM_PRED_DOF x 1) array of values +disp_val: (NUM_PRED_DOF x 1) array of values of the prescribed displacements -eq_num: (num_nodes x dof_node) array that specifies where +eq_num: (num_nodes x dof_node) array that specifies where each node and its DOFs belong in the global stifness matrix volume_conn: connectivity matrix for 3D elements pres_surf: connectivity for surfaces that have a prescribed traction nodal_coords: spacial coordinates for each node in mesh -In the definitions below, "P" stands for prescribed displacements +In the definitions below, "P" stands for prescribed displacements and "F" stands for free displacments -KPP: (num_pres_dof x num_pres_dof) partion of stiffness matrix with rows +KPP: (num_pres_dof x num_pres_dof) partion of stiffness matrix with rows and columns corresponding with prescribed DOFS -KPF: (num_pres_dof x num_free_dof) partion of stiffness matrix with rows - corresponding with prescribed DOFS and columns corresponding with free DOFS -KFF: (num_free_dof x num_free_dof) partion of stiffness matrix with rows and +KPF: (num_pres_dof x num_free_dof) partion of stiffness matrix with rows + corresponding with prescribed DOFS and columns corresponding with free DOFS +KFF: (num_free_dof x num_free_dof) partion of stiffness matrix with rows and columns corresponding with free DOFS -KFP: (num_free_dof x num_free_dof) partion of stiffness matrix with rows +KFP: (num_free_dof x num_free_dof) partion of stiffness matrix with rows corresponding with free DOFS and columns corresponding with prescribed DOFS PF: (num_free_dof, ) RHS vector corresponding with free DOFS PP: (num_pres_dof, ) RHS vector corresponding with prescribed DOFS -UP: (num_pres_dof, ) vector of prescribed displacements +UP: (num_pres_dof, ) vector of prescribed displacements UF: (num_free_dof, ) vector of free displacements (the ones that we solve for) UUR: (num_nodes x 3) matrix of displacements at all nodes in the mesh """ @@ -70,7 +70,7 @@ KPP, KPF, KFF, KFP, PF, PP, UP \ = assemble_module_AD(num_pres_dof, num_free_dof, num_elem, num_nodes_ELE, dof_node, num_nodes_surf, surf_traction_vector, params, - disp_node, disp_val, eq_num, volume_conn, nodal_coords, + disp_node, disp_val, eq_num, volume_conn, nodal_coords, pres_surf, quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle) diff --git a/cmad/fem_utils/tests/test_fem.py b/cmad/fem_utils/tests/test_fem.py index 1ce611b..43ebf8f 100644 --- a/cmad/fem_utils/tests/test_fem.py +++ b/cmad/fem_utils/tests/test_fem.py @@ -13,10 +13,10 @@ def test_patch_form_B(self): dof_node, num_nodes, num_nodes_elem, num_elem, num_nodes_surf, \ nodal_coords, volume_conn = problem.get_mesh_properties() - + disp_node, disp_val, pres_surf, surf_traction_vector \ = problem.get_boundary_conditions() - + quad_rule_3D, shape_func_tetra = problem.get_3D_basis_functions() quad_rule_2D, shape_func_triangle = problem.get_2D_basis_functions() @@ -28,7 +28,7 @@ def test_patch_form_B(self): KPP, KPF, KFF, KFP, PF, PP, UP \ = assemble_module(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, num_nodes_surf,surf_traction_vector, params, - disp_node, disp_val, eq_num, volume_conn, nodal_coords, + disp_node, disp_val, eq_num, volume_conn, nodal_coords, pres_surf, quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle) diff --git a/cmad/fem_utils/tests/test_interpolants.py b/cmad/fem_utils/tests/test_interpolants.py index 1cd56d2..ae55fc2 100644 --- a/cmad/fem_utils/tests/test_interpolants.py +++ b/cmad/fem_utils/tests/test_interpolants.py @@ -27,7 +27,7 @@ def test_triangle_shape_functions(self): for j in range(len(dshape)): gradient = dshape[j] assert np.allclose(np.sum(gradient, axis = 1), np.zeros(len(gradient))) - + def test_quad_shape_functions(self): for i in range(2, 6): test_points = quadrature_rule.create_quadrature_rule_on_quad(i).xigauss From e156e1dcb2d74ad16be2c63d5adf0fe438086559 Mon Sep 17 00:00:00 2001 From: ryany Date: Mon, 31 Mar 2025 08:46:49 -0700 Subject: [PATCH 4/5] Dohrmann-Bochev pressure stabilization --- cmad/fem_utils/fem_functions.py | 324 +++++++++++------------ cmad/fem_utils/fem_functions_AD.py | 62 ----- cmad/fem_utils/fem_functions_bochev.py | 341 +++++++++++++++++++++++++ cmad/fem_utils/fem_problem.py | 74 ++++-- cmad/fem_utils/interpolants.py | 32 +-- cmad/fem_utils/mesh.py | 25 +- cmad/fem_utils/quadrature_rule.py | 88 +++---- cmad/fem_utils/solve_fem.py | 34 +-- cmad/fem_utils/solve_fem_DB.py | 37 +++ cmad/fem_utils/tests/test_fem.py | 2 +- 10 files changed, 670 insertions(+), 349 deletions(-) delete mode 100644 cmad/fem_utils/fem_functions_AD.py create mode 100644 cmad/fem_utils/fem_functions_bochev.py create mode 100644 cmad/fem_utils/solve_fem_DB.py diff --git a/cmad/fem_utils/fem_functions.py b/cmad/fem_utils/fem_functions.py index 248b243..d21a5ff 100644 --- a/cmad/fem_utils/fem_functions.py +++ b/cmad/fem_utils/fem_functions.py @@ -1,9 +1,7 @@ import numpy as np -import time +import jax.numpy as jnp import jax -from cmad.fem_utils.fem_functions_AD import elem_residual - def initialize_equation(num_nodes, dof_node, disp_node): eq_num = np.zeros((num_nodes, dof_node), dtype=int) @@ -35,53 +33,6 @@ def assemble_prescribed_displacement( return UP -def calc_element_stiffness_matrix( - elem_points, num_nodes_elem, dof_node, - params, gauss_weights_3D, dshape_tetra): - - E = params[0] - nu = params[1] - - k = E / (3 * (1 - 2 * nu)) - mu = E / (2 * (1 + nu)) - - c1 = k + 4/3 * mu - c2 = k - 2/3 * mu - material_stiffness = np.array([[c1, c2, c2, 0, 0, 0], - [c2, c1, c2, 0, 0, 0], - [c2, c2, c1, 0, 0, 0], - [0, 0, 0, mu, 0, 0], - [0, 0, 0, 0, mu, 0], - [0, 0, 0, 0, 0, mu]]) - - KEL = np.zeros((num_nodes_elem * dof_node, num_nodes_elem * dof_node)) - - for gaussPt3D in range(len(gauss_weights_3D)): - w_q = gauss_weights_3D[gaussPt3D] - - dshape_tetra_q = dshape_tetra[gaussPt3D, :, :] - - J_q = (dshape_tetra_q @ elem_points).T - - dv_q = np.linalg.det(J_q) - - gradphiXYZ_q = np.linalg.inv(J_q).T @ dshape_tetra_q - - D_phi_q = np.zeros((6, num_nodes_elem * dof_node)) - D_phi_q[0, 0:num_nodes_elem] = gradphiXYZ_q[0, :] - D_phi_q[1, num_nodes_elem:num_nodes_elem * 2] = gradphiXYZ_q[1, :] - D_phi_q[2, num_nodes_elem * 2:num_nodes_elem * 3] = gradphiXYZ_q[2, :] - D_phi_q[3, 0:num_nodes_elem * 2] \ - = np.concatenate((gradphiXYZ_q[1, :], gradphiXYZ_q[0, :])) - D_phi_q[4, num_nodes_elem:num_nodes_elem * 3] \ - = np.concatenate((gradphiXYZ_q[2, :], gradphiXYZ_q[1, :])) - D_phi_q[5, 0:num_nodes_elem] = gradphiXYZ_q[2, :] - D_phi_q[5, num_nodes_elem * 2:num_nodes_elem * 3] = gradphiXYZ_q[0, :] - - KEL += w_q * D_phi_q.T @ material_stiffness @ D_phi_q * dv_q - - return KEL - def calc_element_traction_vector( surf_num, pres_surf, nodal_coords, num_nodes_surf, dof_node, surf_traction_vector, quad_rule_2D, shape_func_triangle): @@ -92,7 +43,7 @@ def calc_element_traction_vector( surf_points = nodal_coords[pres_surf[surf_num], :] - PEL = np.zeros(num_nodes_surf * dof_node) + FEL = np.zeros(num_nodes_surf * dof_node) for gaussPt2D in range(len(gauss_weights_2D)): shape_tri_q = shape_triangle[gaussPt2D, :] @@ -102,18 +53,18 @@ def calc_element_traction_vector( da_q = np.linalg.norm(np.cross(J_q[0, :], J_q[1, :])) - PEL += gauss_weights_2D[gaussPt2D] \ + FEL += gauss_weights_2D[gaussPt2D] \ * (np.column_stack([shape_tri_q, shape_tri_q, shape_tri_q]) \ * surf_traction_vector).T.reshape(-1) * da_q - return PEL + return FEL -def assemble_global_stiffness( - KEL, volume_conn, eq_num, elem_num, KPP, KPF, KFF, KFP): +def assemble_global_traction_vector( + FEL, pres_surf, surf_num, eq_num, FF, FP): - elem_conn = volume_conn[elem_num] - elem_eq_num = eq_num[elem_conn, : ] - global_indices = elem_eq_num.T.reshape(-1) + surf_conn = pres_surf[surf_num] + surf_eq_num = eq_num[surf_conn, : ] + global_indices = surf_eq_num.T.reshape(-1) local_pres_indices = np.where(global_indices < 0)[0] local_free_indices = np.where(global_indices > 0)[0] @@ -121,161 +72,178 @@ def assemble_global_stiffness( global_free_indices = global_indices[local_free_indices] - 1 global_pres_indices = - global_indices[local_pres_indices] - 1 - KFF[np.ix_(global_free_indices,global_free_indices)] \ - += KEL[np.ix_(local_free_indices,local_free_indices)] - KFP[np.ix_(global_free_indices,global_pres_indices)] \ - += KEL[np.ix_(local_free_indices,local_pres_indices)] - KPF[np.ix_(global_pres_indices,global_free_indices)] \ - += KEL[np.ix_(local_pres_indices,local_free_indices)] - KPP[np.ix_(global_pres_indices,global_pres_indices)] \ - += KEL[np.ix_(local_pres_indices,local_pres_indices)] + FF[global_free_indices] += FEL[local_free_indices] + FP[global_pres_indices] += FEL[local_pres_indices] +def assemble_global_displacement_field(eq_num, UF, UP): -def assemble_global_traction_vector( - PEL, pres_surf, surf_num, eq_num, PF, PP): + UUR = np.zeros(eq_num.shape) + for i,val in enumerate(eq_num): + for j, row in enumerate(val): + if row > 0: + UUR[i, j] = UF[row - 1] + else: + UUR[i, j] = UP[- row - 1] + return UUR - surf_conn = pres_surf[surf_num] - surf_eq_num = eq_num[surf_conn, : ] - global_indices = surf_eq_num.T.reshape(-1) +def compute_shape_jacobian(elem_points, dshape_tetra): - local_pres_indices = np.where(global_indices < 0)[0] - local_free_indices = np.where(global_indices > 0)[0] + J = (dshape_tetra @ elem_points).T - global_free_indices = global_indices[local_free_indices] - 1 - global_pres_indices = - global_indices[local_pres_indices] - 1 + dv = jnp.linalg.det(J) - PF[global_free_indices] += PEL[local_free_indices] - PP[global_pres_indices] += PEL[local_pres_indices] + # derivatives of shape functions with respect to spacial coordinates + gradphiXYZ = jnp.linalg.inv(J).T @ dshape_tetra + return dv, gradphiXYZ -def assemble_module( - num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, - num_nodes_surf, surf_traction_vector, params, disp_node, - disp_val, eq_num, volume_conn, nodal_coords, pres_surf, - quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle): +def interpolate(u, shape_tetra, gradphiXYZ, num_nodes_elem): - # Initialize arrays that need to be returned (KPP, KPF, KFF, KFP, PP) - KPP = np.zeros((num_pres_dof, num_pres_dof)) - KPF = np.zeros((num_pres_dof, num_free_dof)) - KFP = np.zeros((num_free_dof, num_pres_dof)) - KFF = np.zeros((num_free_dof, num_free_dof)) - PP = np.zeros(num_pres_dof) - PF = np.zeros(num_free_dof) + ux = u[0:num_nodes_elem] + uy = u[num_nodes_elem:num_nodes_elem * 2] + uz = u[num_nodes_elem * 2:num_nodes_elem * 3] - ## Prescribe boundary conditions - UP = assemble_prescribed_displacement(num_pres_dof, disp_node, - disp_val, eq_num) + u = jnp.array([jnp.dot(ux, shape_tetra), + jnp.dot(uy, shape_tetra), + jnp.dot(uz, shape_tetra)]) - gauss_weights_3D = quad_rule_3D.wgauss - shape_tetra = shape_func_tetra.values - dshape_tetra = shape_func_tetra.gradients + grad_u = jnp.vstack([gradphiXYZ @ ux, + gradphiXYZ @ uy, + gradphiXYZ @ uz]) + + return u, grad_u + +def compute_elastic_stress(grad_u, params): - start = time.time() - # assemble global stiffness and force - for elem_num in range(0, num_elem): + E = params[0] + nu = params[1] - elem_points = nodal_coords[volume_conn[elem_num], :] - # get element stiffness matrix - KEL = calc_element_stiffness_matrix(elem_points, num_nodes_elem,dof_node, - params, gauss_weights_3D, dshape_tetra) + mu = E / (2 * (1 + nu)) + lam = E * nu / ((1 + nu) * (1 - 2 * nu)) - # assemble global stiffness - assemble_global_stiffness(KEL, volume_conn, eq_num, - elem_num, KPP, KPF, KFF, KFP) - end = time.time() - print("Time to assemble stiffness matrix: ", end - start) + strain = 1 / 2 * (grad_u + grad_u.T) - start = time.time() - for surf_num in range(len(pres_surf)): + stress = lam * jnp.trace(strain) * jnp.eye(3) + 2 * mu * strain - # get local traction vector - PEL = calc_element_traction_vector(surf_num, pres_surf, nodal_coords, num_nodes_surf, - dof_node,surf_traction_vector, quad_rule_2D, - shape_func_triangle) + return stress - # assemble traction vector - assemble_global_traction_vector(PEL, pres_surf, - surf_num, eq_num, PF, PP) - end = time.time() - print("Time to assemble load vector: ", end - start) +def compute_neo_hookean_stress(grad_u, params): - return KPP, KPF, KFF, KFP, PF, PP, UP + # computes the first Piola-Kirchoff stress tensor + E = params[0] + nu = params[1] -def assemble_module_AD( - num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, - num_nodes_surf, surf_traction_vector, params, disp_node, - disp_val, eq_num, volume_conn, nodal_coords, pres_surf, - quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle): + mu = E / (2 * (1 + nu)) + lam = E * nu / ((1 + nu) * (1 - 2 * nu)) - # Initialize arrays that need to be returned (KPP, KPF, KFF, KFP, PP) - KPP = np.zeros((num_pres_dof, num_pres_dof)) - KPF = np.zeros((num_pres_dof, num_free_dof)) - KFP = np.zeros((num_free_dof, num_pres_dof)) - KFF = np.zeros((num_free_dof, num_free_dof)) - PP = np.zeros(num_pres_dof) - PF = np.zeros(num_free_dof) + F = jnp.eye(3) + grad_u + F_inv_T = jnp.linalg.inv(F).T + J = jnp.linalg.det(F) + P = mu * (F - F_inv_T) + lam * J * (J - 1) * F_inv_T - ## Prescribe boundary conditions - UP = assemble_prescribed_displacement(num_pres_dof, disp_node, - disp_val, eq_num) + return P - start = time.time() +def compute_stress_divergence_vector(u, params, elem_points, num_nodes_elem, + dof_node, gauss_weights_3D, shape_tetra, dshape_tetra): - grad_residual = jax.jit(jax.jacfwd(elem_residual), - static_argnames=['num_nodes_elem', 'dof_node']) + SD_vec = jnp.zeros((num_nodes_elem, dof_node)) - u_guess = np.ones(num_nodes_elem * dof_node) + for gaussPt3D in range(len(gauss_weights_3D)): + w_q = gauss_weights_3D[gaussPt3D] - gauss_weights_3D = quad_rule_3D.wgauss - shape_tetra = shape_func_tetra.values - dshape_tetra = shape_func_tetra.gradients + dshape_tetra_q = dshape_tetra[gaussPt3D, :, :] + shape_tetra_q = shape_tetra[gaussPt3D, :] - # assemble global stiffness and force - for elem_num in range(0, num_elem): + dv_q, gradphiXYZ_q = compute_shape_jacobian(elem_points, dshape_tetra_q) + u_q, grad_u_q = interpolate(u, shape_tetra_q, gradphiXYZ_q, num_nodes_elem) - elem_points = nodal_coords[volume_conn[elem_num], :] - # get element stiffness matrix - KEL = grad_residual(u_guess, params, elem_points, num_nodes_elem, dof_node, - gauss_weights_3D, shape_tetra, dshape_tetra) + stress = compute_neo_hookean_stress(grad_u_q, params) - # assemble global stiffness - assemble_global_stiffness(np.array(KEL), volume_conn, eq_num, - elem_num, KPP, KPF, KFF, KFP) - end = time.time() - print("Time to assemble stiffness matrix: ", end - start) + SD_vec += w_q * gradphiXYZ_q.T @ stress.T * dv_q - start = time.time() - for surf_num in range(len(pres_surf)): + return SD_vec.reshape(-1, order='F') - # get local traction vector - PEL = calc_element_traction_vector(surf_num, pres_surf, nodal_coords, num_nodes_surf, - dof_node,surf_traction_vector, quad_rule_2D, - shape_func_triangle) +def assemble_residual(KEL, REL, volume_conn, eq_num, elem_num, KPP, KPF, KFF, KFP, RF, RP): + elem_conn = volume_conn[elem_num] + elem_eq_num = eq_num[elem_conn, :] + global_indices = elem_eq_num.T.reshape(-1) - # assemble traction vector - assemble_global_traction_vector(PEL, pres_surf, - surf_num, eq_num, PF, PP) - end = time.time() - print("Time to assemble load vector: ", end - start) + local_pres_indices = np.where(global_indices < 0)[0] + local_free_indices = np.where(global_indices > 0)[0] - return KPP, KPF, KFF, KFP, PF, PP, UP + global_free_indices = global_indices[local_free_indices] - 1 + global_pres_indices = - global_indices[local_pres_indices] - 1 -def assemble_global_displacement_field(eq_num, UF, UP): + KFF[np.ix_(global_free_indices, global_free_indices)] \ + += KEL[np.ix_(local_free_indices, local_free_indices)] + KFP[np.ix_(global_free_indices, global_pres_indices)] \ + += KEL[np.ix_(local_free_indices, local_pres_indices)] + KPF[np.ix_(global_pres_indices, global_free_indices)] \ + += KEL[np.ix_(local_pres_indices, local_free_indices)] + KPP[np.ix_(global_pres_indices, global_pres_indices)] \ + += KEL[np.ix_(local_pres_indices, local_pres_indices)] + + RF[global_free_indices] -= REL[local_free_indices] + RP[global_pres_indices] -= REL[local_pres_indices] + +def solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, + num_nodes_surf, surf_traction_vector, params, disp_node, + disp_val, eq_num, volume_conn, nodal_coords, pres_surf, + quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle, tol, max_iters): + + FP = np.zeros(num_pres_dof) + FF = np.zeros(num_free_dof) - UUR = np.zeros(eq_num.shape) - for i,val in enumerate(eq_num): - for j, row in enumerate(val): - if row > 0: - UUR[i, j] = UF[row - 1] - else: - UUR[i, j] = UP[- row - 1] - return UUR + for surf_num in range(len(pres_surf)): -def solve_module(KPP, KPF, KFF, KFP, PP, PF, UP, eq_num): + # get local traction vector + FEL = calc_element_traction_vector(surf_num, pres_surf, nodal_coords, num_nodes_surf, + dof_node, surf_traction_vector, quad_rule_2D, + shape_func_triangle) - UF = np.linalg.solve(KFF, PF - KFP @ UP) - R = KPP @ UP + KPF @ UF - PP + # assemble traction vector + assemble_global_traction_vector(FEL, pres_surf, + surf_num, eq_num, FF, FP) + + grad_SD_res = jax.jit(jax.jacfwd(compute_stress_divergence_vector), + static_argnames=['num_nodes_elem', 'dof_node']) + SD_residual_jit = jax.jit(compute_stress_divergence_vector,static_argnames=['num_nodes_elem', 'dof_node']) - UUR = assemble_global_displacement_field(eq_num, UF, UP) + gauss_weights_3D = quad_rule_3D.wgauss + shape_tetra = shape_func_tetra.values + dshape_tetra = shape_func_tetra.gradients - return UUR, UF, R \ No newline at end of file + ## Prescribed displacements + UP = assemble_prescribed_displacement(num_pres_dof, disp_node, + disp_val, eq_num) + # free displacements + UF = np.zeros(num_free_dof) + + for i in range(max_iters): + UUR = assemble_global_displacement_field(eq_num, UF, UP) + RF = np.zeros(num_free_dof) + RP = np.zeros(num_pres_dof) + KPP = np.zeros((num_pres_dof, num_pres_dof)) + KPF = np.zeros((num_pres_dof, num_free_dof)) + KFP = np.zeros((num_free_dof, num_pres_dof)) + KFF = np.zeros((num_free_dof, num_free_dof)) + # assemble global stiffness and force + for elem_num in range(0, num_elem): + u_elem = UUR[volume_conn[elem_num], :].reshape(-1, order='F') + elem_points = nodal_coords[volume_conn[elem_num], :] + # get element tangent stiffness matrix + KEL = grad_SD_res(u_elem, params, elem_points, num_nodes_elem, dof_node, + gauss_weights_3D, shape_tetra, dshape_tetra) + + REL = SD_residual_jit(u_elem, params, elem_points, num_nodes_elem, dof_node, + gauss_weights_3D, shape_tetra, dshape_tetra) + + assemble_residual(np.array(KEL), np.array(REL), volume_conn, + eq_num, elem_num, KPP, KPF, KFF, KFP, RF, RP) + + print("||R||: ", np.linalg.norm(RF + FF)) + + if (np.linalg.norm(RF + FF) < tol): + return UUR + + UF += np.linalg.solve(KFF, RF + FF) \ No newline at end of file diff --git a/cmad/fem_utils/fem_functions_AD.py b/cmad/fem_utils/fem_functions_AD.py deleted file mode 100644 index f49eaf5..0000000 --- a/cmad/fem_utils/fem_functions_AD.py +++ /dev/null @@ -1,62 +0,0 @@ -import jax.numpy as jnp - -def compute_shape_jacobian(elem_points, dshape_tetra): - - J = (dshape_tetra @ elem_points).T - - dv = jnp.linalg.det(J) - - # derivatives of shape functions with respect to spacial coordinates - gradphiXYZ = jnp.linalg.inv(J).T @ dshape_tetra - - return dv, gradphiXYZ - -def interpolate(u, shape_tetra, gradphiXYZ, num_nodes_elem): - - ux = u[0:num_nodes_elem] - uy = u[num_nodes_elem:num_nodes_elem * 2] - uz = u[num_nodes_elem * 2:num_nodes_elem * 3] - - u = jnp.array([jnp.dot(ux, shape_tetra), - jnp.dot(uy, shape_tetra), - jnp.dot(uz, shape_tetra)]) - - grad_u = jnp.vstack([gradphiXYZ @ ux, - gradphiXYZ @ uy, - gradphiXYZ @ uz]) - - return u, grad_u - -def compute_stress(grad_u, params): - - E = params[0] - nu = params[1] - - mu = E / (2 * (1 + nu)) - lam = E * nu / ((1 + nu) * (1 - 2 * nu)) - - strain = 1 / 2 * (grad_u + grad_u.T) - - stress = lam * jnp.trace(strain) * jnp.eye(3) + 2 * mu * strain - - return stress - -def elem_residual(u, params, elem_points, num_nodes_elem, - dof_node, gauss_weights_3D, shape_tetra, dshape_tetra): - - residual = jnp.zeros((num_nodes_elem, dof_node)) - - for gaussPt3D in range(len(gauss_weights_3D)): - w_q = gauss_weights_3D[gaussPt3D] - - dshape_tetra_q = dshape_tetra[gaussPt3D, :, :] - shape_tetra_q = shape_tetra[gaussPt3D, :] - - dv_q, gradphiXYZ_q = compute_shape_jacobian(elem_points, dshape_tetra_q) - u_q, grad_u_q = interpolate(u, shape_tetra_q, gradphiXYZ_q, num_nodes_elem) - - stress = compute_stress(grad_u_q, params) - - residual += w_q * gradphiXYZ_q.T @ stress * dv_q - - return residual.reshape(-1, order='F') diff --git a/cmad/fem_utils/fem_functions_bochev.py b/cmad/fem_utils/fem_functions_bochev.py new file mode 100644 index 0000000..d146a6e --- /dev/null +++ b/cmad/fem_utils/fem_functions_bochev.py @@ -0,0 +1,341 @@ +import jax +import jax.numpy as jnp +import numpy as np + +def initialize_equation(num_nodes, dof_node, disp_node): + + eq_num = np.zeros((num_nodes, dof_node), dtype=int) + for i, node in enumerate(disp_node): + node_number = node[0] + direction = node[1] + eq_num[node_number][direction - 1] = -(i + 1) + + num_free_dof = 0 + for i in range(len(eq_num)): + for j in range(len(eq_num[i])): + if (eq_num[i, j] == 0): + num_free_dof += 1 + eq_num[i, j] = num_free_dof + num_pres_dof = num_nodes * dof_node - num_free_dof + + return eq_num, num_free_dof, num_pres_dof + +def assemble_prescribed_displacement( + num_pres_dof, disp_node, disp_val, eq_num): + + UP = np.zeros(num_pres_dof) + for i, row in enumerate(disp_node): + node_number = row[0] + dof = row[1] + displacement = disp_val[i] + eqn_number = -eq_num[node_number][dof - 1] + UP[eqn_number - 1] = displacement + + return UP + +def calc_element_traction_vector( + surf_num, pres_surf, nodal_coords, num_nodes_surf, dof_node, + surf_traction_vector, quad_rule_2D, shape_func_triangle): + + gauss_weights_2D = quad_rule_2D.wgauss + shape_triangle = shape_func_triangle.values + dshape_triangle = shape_func_triangle.gradients + + surf_points = nodal_coords[pres_surf[surf_num], :] + + FEL = np.zeros(num_nodes_surf * dof_node) + + for gaussPt2D in range(len(gauss_weights_2D)): + shape_tri_q = shape_triangle[gaussPt2D, :] + dshape_tri_q = dshape_triangle[gaussPt2D, :, :] + + J_q = dshape_tri_q @ surf_points + + da_q = np.linalg.norm(np.cross(J_q[0, :], J_q[1, :])) + + FEL += gauss_weights_2D[gaussPt2D] \ + * (np.column_stack([shape_tri_q, shape_tri_q, shape_tri_q]) \ + * surf_traction_vector).T.reshape(-1) * da_q + + return FEL + +def assemble_global_traction_vector( + FEL, pres_surf, surf_num, eq_num, FF, FP): + + surf_conn = pres_surf[surf_num] + surf_eq_num = eq_num[surf_conn, : ] + global_indices = surf_eq_num.T.reshape(-1) + + local_pres_indices = np.where(global_indices < 0)[0] + local_free_indices = np.where(global_indices > 0)[0] + + global_free_indices = global_indices[local_free_indices] - 1 + global_pres_indices = - global_indices[local_pres_indices] - 1 + + FF[global_free_indices] += FEL[local_free_indices] + FP[global_pres_indices] += FEL[local_pres_indices] + +def assemble_global_displacement_field(eq_num, UF, UP): + + UUR = np.zeros(eq_num.shape) + for i,val in enumerate(eq_num): + for j, row in enumerate(val): + if row > 0: + UUR[i, j] = UF[row - 1] + else: + UUR[i, j] = UP[- row - 1] + return UUR + +def compute_shape_jacobian(elem_points, dshape_tetra): + + J = (dshape_tetra @ elem_points).T + + dv = jnp.linalg.det(J) + + # derivatives of shape functions with respect to spacial coordinates + gradphiXYZ = jnp.linalg.inv(J).T @ dshape_tetra + + return dv, gradphiXYZ + +def interpolate_u(u, shape_tetra, gradphiXYZ, num_nodes_elem): + + ux = u[0:num_nodes_elem] + uy = u[num_nodes_elem:num_nodes_elem * 2] + uz = u[num_nodes_elem * 2:num_nodes_elem * 3] + + u = jnp.array([jnp.dot(ux, shape_tetra), + jnp.dot(uy, shape_tetra), + jnp.dot(uz, shape_tetra)]) + + grad_u = jnp.vstack([gradphiXYZ @ ux, + gradphiXYZ @ uy, + gradphiXYZ @ uz]) + + return u, grad_u + +def interpolate_p(p, shape_tetra): + p = jnp.dot(p, shape_tetra) + return p + +def compute_piola_stress(grad_u, p, params): + + # computes the first Piola-Kirchoff stress tensor (set J = 1) + E = params[0] + nu = params[1] + + mu = E / (2 * (1 + nu)) + lam = E * nu / ((1 + nu) * (1 - 2 * nu)) + + F = jnp.eye(3) + grad_u + F_inv_T = jnp.linalg.inv(F).T + T = mu * F @ F.T - 1 / 3 * jnp.trace(mu * F @ F.T) * jnp.eye(3) + p * jnp.eye(3) + P = T @ F_inv_T + + return P + +def compute_stress_divergence_vector(u, p, params, elem_points, num_nodes_elem, + dof_node, gauss_weights_3D, shape_tetra, dshape_tetra): + + S_D_vec = jnp.zeros((num_nodes_elem, dof_node)) + + for gaussPt3D in range(len(gauss_weights_3D)): + w_q = gauss_weights_3D[gaussPt3D] + + dshape_tetra_q = dshape_tetra[gaussPt3D, :, :] + shape_tetra_q = shape_tetra[gaussPt3D, :] + + dv_q, gradphiXYZ_q = compute_shape_jacobian(elem_points, dshape_tetra_q) + u_q, grad_u_q = interpolate_u(u, shape_tetra_q, gradphiXYZ_q, num_nodes_elem) + p_q = interpolate_p(p, shape_tetra_q) + + stress = compute_piola_stress(grad_u_q, p_q, params) + + S_D_vec += w_q * gradphiXYZ_q.T @ stress.T * dv_q + + return S_D_vec.reshape(-1, order='F') + +def compute_incompressibility_residual(u, p, params, elem_points, num_nodes_elem, + gauss_weights_3D, shape_tetra, dshape_tetra): + + E = params[0] + nu = params[1] + G_param = E/(2 * (1 + nu)) + alpha = 1.0 + + H = 0 + G = jnp.zeros(num_nodes_elem) + residual = jnp.zeros(num_nodes_elem) + + for gaussPt3D in range(len(gauss_weights_3D)): + w_q = gauss_weights_3D[gaussPt3D] + + dshape_tetra_q = dshape_tetra[gaussPt3D, :, :] + shape_tetra_q = shape_tetra[gaussPt3D, :] + + dv_q, gradphiXYZ_q = compute_shape_jacobian(elem_points, dshape_tetra_q) + u_q, grad_u_q = interpolate_u(u, shape_tetra_q, gradphiXYZ_q, num_nodes_elem) + + F = jnp.eye(3) + grad_u_q + + residual += w_q * shape_tetra_q * (jnp.linalg.det(F) - 1) * dv_q + + # DB contibution (projection onto constant polynomial space) + H += w_q * 1.0 * dv_q + G += w_q * shape_tetra_q * dv_q + + # (N.T)(alpha / G)(N)(p) + residual -= alpha / G_param * w_q * shape_tetra_q * jnp.dot(shape_tetra_q, p) * dv_q + + # alpha / G * (G.T)(H^-1)(G)(p) + residual += alpha / G_param * G * (1 / H) * jnp.dot(G, p) + + return residual + +def assemble_module(KEL, C1_EL, C2_EL, VEL, SD_EL_res, J_EL_res, volume_conn, + eq_num_u, eq_num_p, elem_num, KFF, C1FF, C2FF, VFF, SD_F, J_F): + + elem_conn = volume_conn[elem_num] + + elem_eq_num_u = eq_num_u[elem_conn, :] + global_indices_u = elem_eq_num_u.T.reshape(-1) + local_free_indices_u = np.where(global_indices_u > 0)[0] + global_free_indices_u = global_indices_u[local_free_indices_u] - 1 + + global_indices_p = eq_num_p[elem_conn] + local_free_indices_p = np.where(global_indices_p > 0)[0] + global_free_indices_p = global_indices_p[local_free_indices_p] - 1 + + KFF[np.ix_(global_free_indices_u, global_free_indices_u)] \ + += KEL[np.ix_(local_free_indices_u, local_free_indices_u)] + + C1FF[np.ix_(global_free_indices_u, global_free_indices_p)] \ + += C1_EL[np.ix_(local_free_indices_u, local_free_indices_p)] + + C2FF[np.ix_(global_free_indices_p, global_free_indices_u)] \ + += C2_EL[np.ix_(local_free_indices_p, local_free_indices_u)] + + VFF[np.ix_(global_free_indices_p, global_free_indices_p)] \ + += VEL[np.ix_(local_free_indices_p, local_free_indices_p)] + + SD_F[global_free_indices_u] += SD_EL_res[local_free_indices_u] + + J_F[global_free_indices_p] += J_EL_res[local_free_indices_p] + +def solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, + num_nodes_surf, surf_traction_vector, params, disp_node, + disp_val, eq_num_u, eq_num_p, volume_conn, nodal_coords, pres_surf, + quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle, tol, max_iters): + + num_nodes = len(nodal_coords) + + FP = np.zeros(num_pres_dof) + FF = np.zeros(num_free_dof) + + for surf_num in range(len(pres_surf)): + + # get local traction vector + FEL = calc_element_traction_vector(surf_num, pres_surf, nodal_coords, num_nodes_surf, + dof_node, surf_traction_vector, quad_rule_2D, + shape_func_triangle) + + # assemble traction vector + assemble_global_traction_vector(FEL, pres_surf, + surf_num, eq_num_u, FF, FP) + + #derivative of stress-divergence residual w.r.t u + grad_SD_res_u = jax.jit(jax.jacfwd(compute_stress_divergence_vector), + static_argnames=['num_nodes_elem', 'dof_node']) + #derivative of stress-divergence residual w.r.t p + grad_SD_res_p = jax.jit(jax.jacfwd(compute_stress_divergence_vector, argnums=1), + static_argnames=['num_nodes_elem', 'dof_node']) + #evaluate stress-divergence residual + SD_residual_jit = jax.jit(compute_stress_divergence_vector, + static_argnames=['num_nodes_elem', 'dof_node']) + + #derivative of incompressibility residual w.r.t u + grad_J_res_u = jax.jit(jax.jacfwd(compute_incompressibility_residual), + static_argnames=['num_nodes_elem']) + #derivative of incompressibility residual w.r.t p + grad_J_res_p = jax.jit(jax.jacfwd(compute_incompressibility_residual, argnums=1), + static_argnames=['num_nodes_elem']) + #evaluate incompressibility residual + J_residual_jit = jax.jit(compute_incompressibility_residual, + static_argnames=['num_nodes_elem']) + + gauss_weights_3D = quad_rule_3D.wgauss + shape_tetra = shape_func_tetra.values + dshape_tetra = shape_func_tetra.gradients + + ## Prescribed displacements + UP = assemble_prescribed_displacement(num_pres_dof, disp_node, + disp_val, eq_num_u) + # free displacements + UF = np.zeros(num_free_dof) + + # free pressure nodes (fix pressure at one node) + pF = np.zeros(num_nodes - 1) + + M = np.zeros((num_free_dof + num_nodes - 1, num_free_dof + num_nodes - 1)) + f = np.zeros(num_free_dof + num_nodes - 1) + + for i in range(max_iters): + # global displacement and pressure fields + UUR = assemble_global_displacement_field(eq_num_u, UF, UP) + pUR = np.append(pF.copy(), 0.0) + + # initialize + KFF = np.zeros((num_free_dof, num_free_dof)) + C1FF = np.zeros((num_free_dof, num_nodes - 1)) + C2FF = np.zeros((num_nodes - 1, num_free_dof)) + VFF = np.zeros((num_nodes - 1, num_nodes - 1)) + SD_F = np.zeros(num_free_dof) + J_F = np.zeros(num_nodes - 1) + + # assemble global stiffness and force + for elem_num in range(0, num_elem): + u_elem = UUR[volume_conn[elem_num], :].reshape(-1, order='F') + p_elem = pUR[volume_conn[elem_num]] + elem_points = nodal_coords[volume_conn[elem_num], :] + + # get element tangent matrices + KEL = grad_SD_res_u(u_elem, p_elem, params, elem_points, num_nodes_elem, + dof_node, gauss_weights_3D, shape_tetra, dshape_tetra) + + C1_EL = grad_SD_res_p(u_elem, p_elem, params, elem_points, num_nodes_elem, + dof_node, gauss_weights_3D, shape_tetra, dshape_tetra) + + C2_EL = grad_J_res_u(u_elem, p_elem, params, elem_points, num_nodes_elem, + gauss_weights_3D, shape_tetra, dshape_tetra) + + VEL = grad_J_res_p(u_elem, p_elem, params, elem_points, num_nodes_elem, + gauss_weights_3D, shape_tetra, dshape_tetra) + + SD_EL_res = SD_residual_jit(u_elem, p_elem, params, elem_points, num_nodes_elem, + dof_node, gauss_weights_3D, shape_tetra, dshape_tetra) + + J_EL_res = J_residual_jit(u_elem, p_elem, params, elem_points, num_nodes_elem, + gauss_weights_3D, shape_tetra, dshape_tetra) + + assemble_module(np.array(KEL), np.array(C1_EL), np.array(C2_EL), np.array(VEL), + np.array(SD_EL_res), np.array(J_EL_res), volume_conn, eq_num_u, + eq_num_p, elem_num, KFF, C1FF, C2FF, VFF, SD_F, J_F) + + f[0:num_free_dof] = SD_F - FF + f[num_free_dof:] = J_F + + print("||R||: ", np.linalg.norm(f)) + + if (np.linalg.norm(f) < tol): + return UUR, pUR + + M[0:num_free_dof, 0:num_free_dof] = KFF + M[0:num_free_dof, num_free_dof:] = C1FF + M[num_free_dof:, 0:num_free_dof] = C2FF + M[num_free_dof:, num_free_dof:] = VFF + + # print('Cond(M): ', np.linalg.cond(M)) + + delta = np.linalg.solve(M, -f) + + UF += delta[0:num_free_dof] + pF += delta[num_free_dof:] diff --git a/cmad/fem_utils/fem_problem.py b/cmad/fem_utils/fem_problem.py index 62b6b6f..a5f9b31 100644 --- a/cmad/fem_utils/fem_problem.py +++ b/cmad/fem_utils/fem_problem.py @@ -40,7 +40,7 @@ def __init__(self, problem_type, order): NUM_PRES_NODES = len(pres_nodes) self._disp_node = np.zeros((NUM_PRES_NODES \ - * self._dof_node, 2), dtype = int) + * self._dof_node, 2), dtype=int) for i in range(NUM_PRES_NODES): for j in range(self._dof_node): self._disp_node[i * self._dof_node + j][0] \ @@ -49,7 +49,7 @@ def __init__(self, problem_type, order): self._disp_val = np.zeros(NUM_PRES_NODES * self._dof_node) # normal traction on plane x = 1 - self._surf_traction_vector = np.array([1.0, 0.0, 0.0]) + self._surf_traction_vector = np.array([0.1, 0.0, 0.0]) pres_surf = [] for surface in self._surface_conn: surface_points = self._nodal_coords[surface, :] @@ -75,16 +75,16 @@ def __init__(self, problem_type, order): disp_node = [] for i in range(self._num_nodes): if self._nodal_coords[i][0] == 0.0: - disp_node.append(np.array([i, 1], dtype = int)) + disp_node.append(np.array([i, 1], dtype=int)) if self._nodal_coords[i][1] == 0.0: - disp_node.append(np.array([i, 2], dtype = int)) + disp_node.append(np.array([i, 2], dtype=int)) if self._nodal_coords[i][2] == 0.0: - disp_node.append(np.array([i, 3], dtype = int)) - self._disp_node = np.array(disp_node, dtype = int) + disp_node.append(np.array([i, 3], dtype=int)) + self._disp_node = np.array(disp_node, dtype=int) self._disp_val = np.zeros(len(disp_node)) # normal traction on x = 2 - self._surf_traction_vector = np.array([1.0, 0.0, 0.0]) + self._surf_traction_vector = np.array([0.1, 0.0, 0.0]) pres_surf = [] for surface in self._surface_conn: surface_points = self._nodal_coords[surface, :] @@ -118,22 +118,23 @@ def __init__(self, problem_type, order): y = coord[1] z = coord[2] - ux = (2 * x + y + z - 4) / 2 - uy = (x + 2 * y + z - 4) / 2 - uz = (x + y + 2 * z - 4) / 2 + scaling = 1e-4 + ux = scaling * (2 * x + y + z - 4) / 2 + uy = scaling * (x + 2 * y + z - 4) / 2 + uz = scaling * (x + y + 2 * z - 4) / 2 self.UUR_true[i, :] = np.array([ux, uy, uz]) if (x == 0.0 or x == 2.0 or y == 0.0 or y == 2.0 or z == 0 or z == 2.0): - disp_node.append(np.array([i, 1], dtype = int)) - disp_node.append(np.array([i, 2], dtype = int)) - disp_node.append(np.array([i, 3], dtype = int)) + disp_node.append(np.array([i, 1], dtype=int)) + disp_node.append(np.array([i, 2], dtype=int)) + disp_node.append(np.array([i, 3], dtype=int)) disp_val.append(ux) disp_val.append(uy) disp_val.append(uz) - self._disp_node = np.array(disp_node, dtype = int) + self._disp_node = np.array(disp_node, dtype=int) self._disp_val = np.array(disp_val) # no surface tractions @@ -158,23 +159,56 @@ def __init__(self, problem_type, order): disp_node = [] for i in range(self._num_nodes): if self._nodal_coords[i][2] == 0.0: - disp_node.append(np.array([i, 1], dtype = int)) - disp_node.append(np.array([i, 2], dtype = int)) - disp_node.append(np.array([i, 3], dtype = int)) + disp_node.append(np.array([i, 1], dtype=int)) + disp_node.append(np.array([i, 2], dtype=int)) + disp_node.append(np.array([i, 3], dtype=int)) if self._nodal_coords[i][2] == 2.0: - disp_node.append(np.array([i, 3], dtype = int)) + disp_node.append(np.array([i, 3], dtype=int)) - self._disp_node = np.array(disp_node, dtype = int) + self._disp_node = np.array(disp_node, dtype=int) self._disp_val = np.zeros(len(disp_node)) # shear traction in x direction on plane z = 2 - self._surf_traction_vector = np.array([1.0, 0.0, 0.0]) + self._surf_traction_vector = np.array([0.1, 0.0, 0.0]) pres_surf = [] for surface in self._surface_conn: surface_points = self._nodal_coords[surface, :] if (surface_points[:, 2] == 2 * np.ones(3)).all(): pres_surf.append(surface) self._pres_surf = np.array(pres_surf) + + if problem_type == "cook_membrane": + + self._mesh = Mesh("cook") + + self._nodal_coords = self._mesh.get_nodal_coordinates() + self._colume_conn = self._mesh.get_volume_connectivity() + self._surface_conn = self._mesh.get_surface_connectivity() + + self._dof_node = 3 + self._num_nodes = len(self._nodal_coords) + self._num_nodes_elem = 4 + self._num_elem = len(self._colume_conn) + self._num_nodes_surf = 3 + + # fix all nodes on plane x = 0 + disp_node = [] + for i in range(self._num_nodes): + if self._nodal_coords[i][0] == 0.0: + disp_node.append(np.array([i, 1], dtype=int)) + disp_node.append(np.array([i, 2], dtype=int)) + disp_node.append(np.array([i, 3], dtype=int)) + self._disp_node = np.array(disp_node, dtype=int) + self._disp_val = np.zeros(len(disp_node)) + + # traction on x = 48 + self._surf_traction_vector = np.array([0.0, 1.0, 0.0]) + pres_surf = [] + for surface in self._surface_conn: + surface_points = self._nodal_coords[surface, :] + if (surface_points[:, 0] == 48 * np.ones(3)).all(): + pres_surf.append(surface) + self._pres_surf = np.array(pres_surf) def get_2D_basis_functions(self): return self._quad_rule_2D, self._shape_func_triangle diff --git a/cmad/fem_utils/interpolants.py b/cmad/fem_utils/interpolants.py index a883334..ffb35be 100644 --- a/cmad/fem_utils/interpolants.py +++ b/cmad/fem_utils/interpolants.py @@ -38,7 +38,7 @@ def shape1d(evaluationPoints): """ - shape = np.vstack(((1 - evaluationPoints[:,0])/2.0, (1 + evaluationPoints[:,0])/2.0)).T + shape = np.vstack(((1 - evaluationPoints[:, 0]) / 2.0, (1 + evaluationPoints[:, 0]) / 2.0)).T dshape = np.vstack((-0.5 * np.ones(len(evaluationPoints)), 0.5 * np.ones(len(evaluationPoints)))).T return ShapeFunctions(shape, dshape) @@ -63,8 +63,8 @@ def shape_triangle(evaluationPoints): dof_node = 2 num_nodes_elem = 3 - shape = np.vstack((1 - evaluationPoints[:,0] - evaluationPoints[:,1], - evaluationPoints[:,0], evaluationPoints[:,1])).T + shape = np.vstack((1 - evaluationPoints[:, 0] - evaluationPoints[:, 1], + evaluationPoints[:, 0], evaluationPoints[:, 1])).T dshape = np.zeros((num_eval_points, dof_node, num_nodes_elem)) @@ -94,10 +94,10 @@ def shape_quad(evaluationPoints): dof_node = 2 num_nodes_elem = 4 - l0x = 1 - evaluationPoints[:,0] - l1x = 1 + evaluationPoints[:,0] - l0y = 1 - evaluationPoints[:,1] - l1y = 1 + evaluationPoints[:,1] + l0x = 1 - evaluationPoints[:, 0] + l1x = 1 + evaluationPoints[:, 0] + l0y = 1 - evaluationPoints[:, 1] + l1y = 1 + evaluationPoints[:, 1] shape = np.vstack((l0x * l0y / 4, l1x * l0y / 4, l1x * l1y / 4, l0x * l1y / 4)).T dshape = np.zeros((num_eval_points, dof_node, num_nodes_elem)) @@ -108,7 +108,7 @@ def shape_quad(evaluationPoints): l1x = 1 + point[0] l0y = 1 - point[1] l1y = 1 + point[1] - dshape[i, :, :] = np.array([[-l0y, l0y, l1y, -l1y],[-l0x, -l1x, l1x, l0x]]) + dshape[i, :, :] = np.array([[-l0y, l0y, l1y, -l1y],[-l0x, -l1x, l1x, l0x]]) / 4 return ShapeFunctions(shape, dshape) @@ -133,8 +133,8 @@ def shape_tetrahedron(evaluationPoints): dof_node = 3 num_nodes_elem = 4 - shape = np.vstack((1 - evaluationPoints[:,0] - evaluationPoints[:,1] - evaluationPoints[:,2], - evaluationPoints[:,0], evaluationPoints[:,1], evaluationPoints[:,2])).T + shape = np.vstack((1 - evaluationPoints[:, 0] - evaluationPoints[:, 1] - evaluationPoints[:, 2], + evaluationPoints[:, 0], evaluationPoints[:, 1], evaluationPoints[:, 2])).T dshape = np.zeros((num_eval_points, dof_node, num_nodes_elem)) @@ -165,12 +165,12 @@ def shape_brick(evaluationPoints): dof_node = 3 num_nodes_elem = 8 - m1 = 1 - evaluationPoints[:,0] - p1 = 1 + evaluationPoints[:,0] - m2 = 1 - evaluationPoints[:,1] - p2 = 1 + evaluationPoints[:,1] - m3 = 1 - evaluationPoints[:,2] - p3 = 1 + evaluationPoints[:,3] + m1 = 1 - evaluationPoints[:, 0] + p1 = 1 + evaluationPoints[:, 0] + m2 = 1 - evaluationPoints[:, 1] + p2 = 1 + evaluationPoints[:, 1] + m3 = 1 - evaluationPoints[:, 2] + p3 = 1 + evaluationPoints[:, 3] shape = np.vstack((m1 * m2 * m3 / 8, p1 * m2 * m3 / 8, p1 * p2 * m3 / 8, m1 * p2 * m3 / 8, m1 * m2 * p3 / 8, p1 * m2 * p3 / 8, p1 * p2 * p3 / 8, m1 * p2 * p3 / 8)).T diff --git a/cmad/fem_utils/mesh.py b/cmad/fem_utils/mesh.py index 0dd6c18..8111331 100644 --- a/cmad/fem_utils/mesh.py +++ b/cmad/fem_utils/mesh.py @@ -11,8 +11,8 @@ def __init__(self, mesh_type): # square prism with hole through center if self._mesh_type == "hole_block": - geom.characteristic_length_min = 0.1 - geom.characteristic_length_max = 0.1 + geom.characteristic_length_min = 0.05 + geom.characteristic_length_max = 0.05 rectangle = geom.add_rectangle([0.0, 0.0, 0.0], 1.0, 1.0) @@ -22,7 +22,7 @@ def __init__(self, mesh_type): rectangle, disk ) - geom.extrude(flat, [0.0, 0.0, 0.25], num_layers = 5) + geom.extrude(flat, [0.0, 0.0, 0.25], num_layers=10) self._mesh = geom.generate_mesh() # thin beam @@ -33,7 +33,7 @@ def __init__(self, mesh_type): rectangle = geom.add_rectangle([0.0, 0.0, 0.0], 2, 0.5) - geom.extrude(rectangle, [0.0, 0.0, 0.5], num_layers = 5) + geom.extrude(rectangle, [0.0, 0.0, 0.5], num_layers=5) self._mesh = geom.generate_mesh() # 2 x 2 x 2 cube @@ -44,7 +44,20 @@ def __init__(self, mesh_type): rectangle = geom.add_rectangle([0.0, 0.0, 0.0], 2, 2) - geom.extrude(rectangle, [0.0, 0.0, 2], num_layers = 10) + geom.extrude(rectangle, [0.0, 0.0, 2], num_layers=10) + self._mesh = geom.generate_mesh() + + if self._mesh_type == "cook": + + geom.characteristic_length_min = 1.5 + geom.characteristic_length_max = 1.5 + + beam = geom.add_polygon([[0.0, 0.0], + [48., 44.], + [48., 60.], + [0.0, 44.]]) + + geom.extrude(beam, [0.0, 0.0, 1.5], num_layers=1) self._mesh = geom.generate_mesh() self._points = self._mesh.points @@ -63,7 +76,7 @@ def get_surface_connectivity(self): return self._surface_conn def add_point_data(self, data): - self._mesh.point_data = {"displacement_field": data} + self._mesh.point_data = data def save_mesh(self, filename): self._mesh.write(filename + ".vtk") diff --git a/cmad/fem_utils/quadrature_rule.py b/cmad/fem_utils/quadrature_rule.py index 272d646..4e4538c 100644 --- a/cmad/fem_utils/quadrature_rule.py +++ b/cmad/fem_utils/quadrature_rule.py @@ -26,13 +26,13 @@ def create_quadrature_rule_1D(degree): and the weights. """ - n = np.ceil((degree + 1)/2) + n = np.ceil((degree + 1) / 2) xi, w = scipy.special.roots_sh_legendre(n) xi = -1 + 2 * xi xi_matrix = np.zeros((len(xi), 1)) - xi_matrix[:,0]=xi - return QuadratureRule(xi_matrix, w) + xi_matrix[:, 0]=xi + return QuadratureRule(xi_matrix, w * 2) def create_quadrature_rule_on_triangle(degree): @@ -222,7 +222,7 @@ def create_quadrature_rule_on_quad(degree): [a, a]]) w = np.array([c * c, b * c, c * c, c * b, - b * b, c * b, c * c, b * c, c * c]) + b * b, c * b, c * c, b * c, c * c]) / 4 else: raise ValueError("Quadrature of precision this high is not implemented.") @@ -241,19 +241,19 @@ def create_quadrature_rule_on_tetrahedron(degree): [0.138196601125011, 0.585410196624969, 0.138196601125011], [0.138196601125011, 0.138196601125011, 0.585410196624969]]) - w = np.array([0.25/6.0, 0.25/6.0, 0.25/6.0, 0.25/6.0]) + w = np.array([0.25 / 6.0, 0.25 / 6.0, 0.25 / 6.0, 0.25 / 6.0]) elif degree == 3: - xi = np.array([[1./4., 1./4., 1./4.], - [1./6., 1./6., 1./6.], - [1./6., 1./6., 1./2.], - [1./6., 1./2., 1./6.], - [1./2., 1./6., 1./6.]]) + xi = np.array([[1. / 4., 1. / 4., 1. / 4.], + [1. / 6., 1. / 6., 1. / 6.], + [1. / 6., 1. / 6., 1. / 2.], + [1./6., 1./2., 1. / 6.], + [1. / 2., 1. / 6., 1. / 6.]]) - w = np.array([-2/15, 3/40, 3/40, 3/40, 3/40]) + w = np.array([-2 / 15, 3 / 40, 3 / 40, 3 / 40, 3 / 40]) elif degree == 4: - xi = np.array([[0.25,0.25,0.25], + xi = np.array([[0.25, 0.25, 0.25], [0.785714285714286, 0.0714285714285714, 0.0714285714285714], [0.0714285714285714, 0.0714285714285714, 0.0714285714285714], [0.0714285714285714, 0.0714285714285714, 0.785714285714286], @@ -272,14 +272,14 @@ def create_quadrature_rule_on_tetrahedron(degree): elif degree == 5: xi = np.array([[0.25, 0.25, 0.25], - [0, 1./3., 1./3.], - [1./3., 1./3., 1./3.], - [1./3., 1./3., 0], - [1./3., 0, 1./3.], - [8./11., 1./11., 1./11.], - [1./11., 1./11., 1./11.], - [1./11., 1./11., 8./11.], - [1./11., 8./11., 1./11.], + [0, 1. / 3., 1. / 3.], + [1. / 3., 1. / 3., 1. / 3.], + [1. / 3., 1. / 3., 0], + [1. / 3., 0, 1. / 3.], + [8. / 11., 1. / 11., 1. / 11.], + [1. / 11., 1. / 11., 1. / 11.], + [1. / 11., 1. / 11., 8. / 11.], + [1. / 11., 8. / 11., 1. / 11.], [0.433449846426336, 0.0665501535736643, 0.0665501535736643], [0.0665501535736643, 0.433449846426336, 0.0665501535736643], [0.0665501535736643, 0.0665501535736643, 0.433449846426336], @@ -293,30 +293,30 @@ def create_quadrature_rule_on_tetrahedron(degree): 0.0109491415613865, 0.0109491415613865, 0.0109491415613865, 0.0109491415613865, 0.0109491415613865, 0.0109491415613865]) elif degree == 6: - xi = np.array([[0.356191386222545,0.214602871259152,0.214602871259152], - [0.214602871259152,0.214602871259152,0.214602871259152], - [0.214602871259152,0.214602871259152,0.356191386222545], - [0.214602871259152,0.356191386222545,0.214602871259152], - [0.877978124396166,0.0406739585346113,0.0406739585346113], - [0.0406739585346113,0.0406739585346113,0.0406739585346113], - [0.0406739585346113,0.0406739585346113,0.877978124396166], - [0.0406739585346113,0.877978124396166,0.0406739585346113], - [0.0329863295731731,0.322337890142276,0.322337890142276], - [0.322337890142276,0.322337890142276,0.322337890142276], - [0.322337890142276,0.322337890142276,0.0329863295731731], - [0.322337890142276,0.0329863295731731,0.322337890142276], - [0.269672331458316,0.0636610018750175,0.0636610018750175], - [0.0636610018750175,0.269672331458316,0.0636610018750175], - [0.0636610018750175,0.0636610018750175,0.269672331458316], - [0.603005664791649,0.0636610018750175,0.0636610018750175], - [0.0636610018750175,0.603005664791649,0.0636610018750175], - [0.0636610018750175,0.0636610018750175,0.603005664791649], - [0.0636610018750175,0.269672331458316,0.603005664791649], - [0.269672331458316,0.603005664791649,0.0636610018750175], - [0.603005664791649,0.0636610018750175,0.269672331458316], - [0.0636610018750175,0.603005664791649,0.269672331458316], - [0.269672331458316,0.0636610018750175,0.603005664791649], - [0.603005664791649,0.269672331458316,0.0636610018750175]]) + xi = np.array([[0.356191386222545, 0.214602871259152, 0.214602871259152], + [0.214602871259152, 0.214602871259152, 0.214602871259152], + [0.214602871259152, 0.214602871259152, 0.356191386222545], + [0.214602871259152, 0.356191386222545, 0.214602871259152], + [0.877978124396166, 0.0406739585346113, 0.0406739585346113], + [0.0406739585346113, 0.0406739585346113, 0.0406739585346113], + [0.0406739585346113, 0.0406739585346113, 0.877978124396166], + [0.0406739585346113, 0.877978124396166, 0.0406739585346113], + [0.0329863295731731, 0.322337890142276, 0.322337890142276], + [0.322337890142276, 0.322337890142276, 0.322337890142276], + [0.322337890142276, 0.322337890142276, 0.0329863295731731], + [0.322337890142276, 0.0329863295731731, 0.322337890142276], + [0.269672331458316, 0.0636610018750175, 0.0636610018750175], + [0.0636610018750175, 0.269672331458316, 0.0636610018750175], + [0.0636610018750175, 0.0636610018750175, 0.269672331458316], + [0.603005664791649, 0.0636610018750175, 0.0636610018750175], + [0.0636610018750175, 0.603005664791649, 0.0636610018750175], + [0.0636610018750175, 0.0636610018750175, 0.603005664791649], + [0.0636610018750175, 0.269672331458316, 0.603005664791649], + [0.269672331458316, 0.603005664791649, 0.0636610018750175], + [0.603005664791649, 0.0636610018750175, 0.269672331458316], + [0.0636610018750175, 0.603005664791649, 0.269672331458316], + [0.269672331458316, 0.0636610018750175, 0.603005664791649], + [0.603005664791649, 0.269672331458316, 0.0636610018750175]]) w = np.array([0.00665379170969465, 0.00665379170969465, 0.00665379170969465, 0.00665379170969465, 0.00167953517588678, 0.00167953517588678, diff --git a/cmad/fem_utils/solve_fem.py b/cmad/fem_utils/solve_fem.py index 2a4a777..d49d122 100644 --- a/cmad/fem_utils/solve_fem.py +++ b/cmad/fem_utils/solve_fem.py @@ -1,8 +1,6 @@ from cmad.fem_utils.fem_problem import fem_problem from cmad.fem_utils.fem_functions import (initialize_equation, - assemble_module, assemble_module_AD, - solve_module) -import time + solve_fem_newton) import numpy as np """ @@ -38,17 +36,17 @@ columns corresponding with free DOFS KFP: (num_free_dof x num_free_dof) partion of stiffness matrix with rows corresponding with free DOFS and columns corresponding with prescribed DOFS -PF: (num_free_dof, ) RHS vector corresponding with free DOFS -PP: (num_pres_dof, ) RHS vector corresponding with prescribed DOFS +FF: (num_free_dof, ) RHS vector corresponding with free DOFS +FP: (num_pres_dof, ) RHS vector corresponding with prescribed DOFS UP: (num_pres_dof, ) vector of prescribed displacements UF: (num_free_dof, ) vector of free displacements (the ones that we solve for) UUR: (num_nodes x 3) matrix of displacements at all nodes in the mesh """ order = 3 -problem = fem_problem("simple_shear", order) +problem = fem_problem("cook_membrane", order) -dof_node, num_nodes, num_nodes_ELE, num_elem, num_nodes_surf, \ +dof_node, num_nodes, num_nodes_elem, num_elem, num_nodes_surf, \ nodal_coords, volume_conn = problem.get_mesh_properties() disp_node, disp_val, pres_surf, surf_traction_vector \ @@ -67,22 +65,14 @@ print("Number of free degrees of freedom: ", num_free_dof) -KPP, KPF, KFF, KFP, PF, PP, UP \ - = assemble_module_AD(num_pres_dof, num_free_dof, num_elem, num_nodes_ELE, - dof_node, num_nodes_surf, surf_traction_vector, params, - disp_node, disp_val, eq_num, volume_conn, nodal_coords, - pres_surf, quad_rule_3D, shape_func_tetra, quad_rule_2D, - shape_func_triangle) +tol = 5e-12 +max_iters = 10 -solve_start = time.time() - -UUR, UF, R = solve_module(KPP, KPF, KFF, KFP, PP, PF, UP, eq_num) +UUR = solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, + num_nodes_surf, surf_traction_vector, params, disp_node, + disp_val, eq_num, volume_conn, nodal_coords, pres_surf, + quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle, tol, max_iters) print(UUR) -solve_end = time.time() - -#problem.save_data("simple_shear", UUR) - -print("Solve time: ", solve_end - solve_start) - +#problem.save_data("cook_membrane", {"displacement_field": UUR}) diff --git a/cmad/fem_utils/solve_fem_DB.py b/cmad/fem_utils/solve_fem_DB.py new file mode 100644 index 0000000..0d30672 --- /dev/null +++ b/cmad/fem_utils/solve_fem_DB.py @@ -0,0 +1,37 @@ +from cmad.fem_utils.fem_problem import fem_problem +from cmad.fem_utils.fem_functions_bochev import (initialize_equation, + solve_fem_newton) +import numpy as np + +order = 3 +problem = fem_problem("cook_membrane", order) + +dof_node, num_nodes, num_nodes_elem, num_elem, num_nodes_surf, \ + nodal_coords, volume_conn = problem.get_mesh_properties() + +disp_node, disp_val, pres_surf, surf_traction_vector \ + = problem.get_boundary_conditions() + +quad_rule_3D, shape_func_tetra = problem.get_3D_basis_functions() +quad_rule_2D, shape_func_triangle = problem.get_2D_basis_functions() + +print("Number of elements: ", num_elem) + +params = np.array([200, 0.5]) + +eq_num_u, num_free_dof, num_pres_dof \ + = initialize_equation(num_nodes, dof_node, disp_node) + +print("Number of free degrees of freedom: ", num_free_dof) + +eq_num_p = np.append(np.linspace(1, num_nodes - 1, num_nodes - 1, dtype = int), -1) + +tol = 5e-12 +max_iters = 10 + +UUR, pUR = solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, + num_nodes_surf, surf_traction_vector, params, disp_node, + disp_val, eq_num_u, eq_num_p, volume_conn, nodal_coords, pres_surf, + quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle, tol, max_iters) + +problem.save_data("cook_membrane_incompressible", {'displacement_field': UUR, 'pressure_field': pUR}) \ No newline at end of file diff --git a/cmad/fem_utils/tests/test_fem.py b/cmad/fem_utils/tests/test_fem.py index 43ebf8f..11fc6a5 100644 --- a/cmad/fem_utils/tests/test_fem.py +++ b/cmad/fem_utils/tests/test_fem.py @@ -33,7 +33,7 @@ def test_patch_form_B(self): shape_func_triangle) UUR, UF, R = solve_module(KPP, KPF, KFF, KFP, PP, PF, UP, eq_num) - + print(UUR) assert np.allclose(UUR, problem.UUR_true) if __name__ == "__main__": From 7f4722a006ea2f1aad0eb1b3adf4110655b71282 Mon Sep 17 00:00:00 2001 From: ryany Date: Wed, 2 Apr 2025 08:11:18 -0700 Subject: [PATCH 5/5] formatting and added extra test --- cmad/fem_utils/fem_functions.py | 36 ++++++----- cmad/fem_utils/fem_functions_bochev.py | 88 ++++++++++++++------------ cmad/fem_utils/fem_problem.py | 36 +++++++++++ cmad/fem_utils/mesh.py | 24 +++++-- cmad/fem_utils/solve_fem.py | 9 +-- cmad/fem_utils/solve_fem_DB.py | 6 +- 6 files changed, 128 insertions(+), 71 deletions(-) diff --git a/cmad/fem_utils/fem_functions.py b/cmad/fem_utils/fem_functions.py index d21a5ff..d645c4d 100644 --- a/cmad/fem_utils/fem_functions.py +++ b/cmad/fem_utils/fem_functions.py @@ -104,12 +104,12 @@ def interpolate(u, shape_tetra, gradphiXYZ, num_nodes_elem): uz = u[num_nodes_elem * 2:num_nodes_elem * 3] u = jnp.array([jnp.dot(ux, shape_tetra), - jnp.dot(uy, shape_tetra), - jnp.dot(uz, shape_tetra)]) + jnp.dot(uy, shape_tetra), + jnp.dot(uz, shape_tetra)]) grad_u = jnp.vstack([gradphiXYZ @ ux, - gradphiXYZ @ uy, - gradphiXYZ @ uz]) + gradphiXYZ @ uy, + gradphiXYZ @ uz]) return u, grad_u @@ -143,8 +143,9 @@ def compute_neo_hookean_stress(grad_u, params): return P -def compute_stress_divergence_vector(u, params, elem_points, num_nodes_elem, - dof_node, gauss_weights_3D, shape_tetra, dshape_tetra): +def compute_stress_divergence_vector( + u, params, elem_points, num_nodes_elem, dof_node, + gauss_weights_3D, shape_tetra, dshape_tetra): SD_vec = jnp.zeros((num_nodes_elem, dof_node)) @@ -182,15 +183,16 @@ def assemble_residual(KEL, REL, volume_conn, eq_num, elem_num, KPP, KPF, KFF, KF += KEL[np.ix_(local_pres_indices, local_free_indices)] KPP[np.ix_(global_pres_indices, global_pres_indices)] \ += KEL[np.ix_(local_pres_indices, local_pres_indices)] - + RF[global_free_indices] -= REL[local_free_indices] RP[global_pres_indices] -= REL[local_pres_indices] -def solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, - num_nodes_surf, surf_traction_vector, params, disp_node, - disp_val, eq_num, volume_conn, nodal_coords, pres_surf, - quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle, tol, max_iters): - +def solve_fem_newton( + num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, + num_nodes_surf, surf_traction_vector, params, disp_node, disp_val, + eq_num, volume_conn, nodal_coords, pres_surf, quad_rule_3D, + shape_func_tetra, quad_rule_2D, shape_func_triangle, tol, max_iters): + FP = np.zeros(num_pres_dof) FF = np.zeros(num_free_dof) @@ -204,7 +206,7 @@ def solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_n # assemble traction vector assemble_global_traction_vector(FEL, pres_surf, surf_num, eq_num, FF, FP) - + grad_SD_res = jax.jit(jax.jacfwd(compute_stress_divergence_vector), static_argnames=['num_nodes_elem', 'dof_node']) SD_residual_jit = jax.jit(compute_stress_divergence_vector,static_argnames=['num_nodes_elem', 'dof_node']) @@ -233,14 +235,14 @@ def solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_n elem_points = nodal_coords[volume_conn[elem_num], :] # get element tangent stiffness matrix KEL = grad_SD_res(u_elem, params, elem_points, num_nodes_elem, dof_node, - gauss_weights_3D, shape_tetra, dshape_tetra) - + gauss_weights_3D, shape_tetra, dshape_tetra) + REL = SD_residual_jit(u_elem, params, elem_points, num_nodes_elem, dof_node, - gauss_weights_3D, shape_tetra, dshape_tetra) + gauss_weights_3D, shape_tetra, dshape_tetra) assemble_residual(np.array(KEL), np.array(REL), volume_conn, eq_num, elem_num, KPP, KPF, KFF, KFP, RF, RP) - + print("||R||: ", np.linalg.norm(RF + FF)) if (np.linalg.norm(RF + FF) < tol): diff --git a/cmad/fem_utils/fem_functions_bochev.py b/cmad/fem_utils/fem_functions_bochev.py index d146a6e..8f035f3 100644 --- a/cmad/fem_utils/fem_functions_bochev.py +++ b/cmad/fem_utils/fem_functions_bochev.py @@ -83,7 +83,7 @@ def assemble_global_displacement_field(eq_num, UF, UP): if row > 0: UUR[i, j] = UF[row - 1] else: - UUR[i, j] = UP[- row - 1] + UUR[i, j] = UP[-row - 1] return UUR def compute_shape_jacobian(elem_points, dshape_tetra): @@ -104,12 +104,12 @@ def interpolate_u(u, shape_tetra, gradphiXYZ, num_nodes_elem): uz = u[num_nodes_elem * 2:num_nodes_elem * 3] u = jnp.array([jnp.dot(ux, shape_tetra), - jnp.dot(uy, shape_tetra), - jnp.dot(uz, shape_tetra)]) + jnp.dot(uy, shape_tetra), + jnp.dot(uz, shape_tetra)]) grad_u = jnp.vstack([gradphiXYZ @ ux, - gradphiXYZ @ uy, - gradphiXYZ @ uz]) + gradphiXYZ @ uy, + gradphiXYZ @ uz]) return u, grad_u @@ -133,8 +133,9 @@ def compute_piola_stress(grad_u, p, params): return P -def compute_stress_divergence_vector(u, p, params, elem_points, num_nodes_elem, - dof_node, gauss_weights_3D, shape_tetra, dshape_tetra): +def compute_stress_divergence_vector( + u, p, params, elem_points, num_nodes_elem, dof_node, + gauss_weights_3D, shape_tetra, dshape_tetra): S_D_vec = jnp.zeros((num_nodes_elem, dof_node)) @@ -154,12 +155,13 @@ def compute_stress_divergence_vector(u, p, params, elem_points, num_nodes_elem, return S_D_vec.reshape(-1, order='F') -def compute_incompressibility_residual(u, p, params, elem_points, num_nodes_elem, +def compute_incompressibility_residual( + u, p, params, elem_points, num_nodes_elem, gauss_weights_3D, shape_tetra, dshape_tetra): E = params[0] nu = params[1] - G_param = E/(2 * (1 + nu)) + G_param = E / (2 * (1 + nu)) alpha = 1.0 H = 0 @@ -175,10 +177,10 @@ def compute_incompressibility_residual(u, p, params, elem_points, num_nodes_elem dv_q, gradphiXYZ_q = compute_shape_jacobian(elem_points, dshape_tetra_q) u_q, grad_u_q = interpolate_u(u, shape_tetra_q, gradphiXYZ_q, num_nodes_elem) - F = jnp.eye(3) + grad_u_q + F = jnp.eye(3) + grad_u_q residual += w_q * shape_tetra_q * (jnp.linalg.det(F) - 1) * dv_q - + # DB contibution (projection onto constant polynomial space) H += w_q * 1.0 * dv_q G += w_q * shape_tetra_q * dv_q @@ -187,13 +189,14 @@ def compute_incompressibility_residual(u, p, params, elem_points, num_nodes_elem residual -= alpha / G_param * w_q * shape_tetra_q * jnp.dot(shape_tetra_q, p) * dv_q # alpha / G * (G.T)(H^-1)(G)(p) - residual += alpha / G_param * G * (1 / H) * jnp.dot(G, p) + residual += alpha / G_param * G * (1 / H) * jnp.dot(G, p) + + return residual - return residual +def assemble_module( + KEL, C1_EL, C2_EL, VEL, SD_EL_res, J_EL_res, volume_conn, eq_num_u, + eq_num_p, elem_num, KFF, C1FF, C2FF, VFF, SD_F, J_F): -def assemble_module(KEL, C1_EL, C2_EL, VEL, SD_EL_res, J_EL_res, volume_conn, - eq_num_u, eq_num_p, elem_num, KFF, C1FF, C2FF, VFF, SD_F, J_F): - elem_conn = volume_conn[elem_num] elem_eq_num_u = eq_num_u[elem_conn, :] @@ -207,12 +210,12 @@ def assemble_module(KEL, C1_EL, C2_EL, VEL, SD_EL_res, J_EL_res, volume_conn, KFF[np.ix_(global_free_indices_u, global_free_indices_u)] \ += KEL[np.ix_(local_free_indices_u, local_free_indices_u)] - + C1FF[np.ix_(global_free_indices_u, global_free_indices_p)] \ - += C1_EL[np.ix_(local_free_indices_u, local_free_indices_p)] + += C1_EL[np.ix_(local_free_indices_u, local_free_indices_p)] C2FF[np.ix_(global_free_indices_p, global_free_indices_u)] \ - += C2_EL[np.ix_(local_free_indices_p, local_free_indices_u)] + += C2_EL[np.ix_(local_free_indices_p, local_free_indices_u)] VFF[np.ix_(global_free_indices_p, global_free_indices_p)] \ += VEL[np.ix_(local_free_indices_p, local_free_indices_p)] @@ -221,13 +224,14 @@ def assemble_module(KEL, C1_EL, C2_EL, VEL, SD_EL_res, J_EL_res, volume_conn, J_F[global_free_indices_p] += J_EL_res[local_free_indices_p] -def solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, +def solve_fem_newton( + num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, num_nodes_surf, surf_traction_vector, params, disp_node, disp_val, eq_num_u, eq_num_p, volume_conn, nodal_coords, pres_surf, quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle, tol, max_iters): - + num_nodes = len(nodal_coords) - + FP = np.zeros(num_pres_dof) FF = np.zeros(num_free_dof) @@ -251,17 +255,17 @@ def solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_n #evaluate stress-divergence residual SD_residual_jit = jax.jit(compute_stress_divergence_vector, static_argnames=['num_nodes_elem', 'dof_node']) - + #derivative of incompressibility residual w.r.t u grad_J_res_u = jax.jit(jax.jacfwd(compute_incompressibility_residual), - static_argnames=['num_nodes_elem']) + static_argnames=['num_nodes_elem']) #derivative of incompressibility residual w.r.t p - grad_J_res_p = jax.jit(jax.jacfwd(compute_incompressibility_residual, argnums=1), - static_argnames=['num_nodes_elem']) + grad_J_res_p = jax.jit(jax.jacfwd(compute_incompressibility_residual, argnums=1), + static_argnames=['num_nodes_elem']) #evaluate incompressibility residual J_residual_jit = jax.jit(compute_incompressibility_residual, static_argnames=['num_nodes_elem']) - + gauss_weights_3D = quad_rule_3D.wgauss shape_tetra = shape_func_tetra.values dshape_tetra = shape_func_tetra.gradients @@ -296,45 +300,45 @@ def solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_n u_elem = UUR[volume_conn[elem_num], :].reshape(-1, order='F') p_elem = pUR[volume_conn[elem_num]] elem_points = nodal_coords[volume_conn[elem_num], :] - + # get element tangent matrices KEL = grad_SD_res_u(u_elem, p_elem, params, elem_points, num_nodes_elem, dof_node, gauss_weights_3D, shape_tetra, dshape_tetra) - + C1_EL = grad_SD_res_p(u_elem, p_elem, params, elem_points, num_nodes_elem, - dof_node, gauss_weights_3D, shape_tetra, dshape_tetra) - + dof_node, gauss_weights_3D, shape_tetra, dshape_tetra) + C2_EL = grad_J_res_u(u_elem, p_elem, params, elem_points, num_nodes_elem, gauss_weights_3D, shape_tetra, dshape_tetra) - + VEL = grad_J_res_p(u_elem, p_elem, params, elem_points, num_nodes_elem, - gauss_weights_3D, shape_tetra, dshape_tetra) - + gauss_weights_3D, shape_tetra, dshape_tetra) + SD_EL_res = SD_residual_jit(u_elem, p_elem, params, elem_points, num_nodes_elem, - dof_node, gauss_weights_3D, shape_tetra, dshape_tetra) - + dof_node, gauss_weights_3D, shape_tetra, dshape_tetra) + J_EL_res = J_residual_jit(u_elem, p_elem, params, elem_points, num_nodes_elem, - gauss_weights_3D, shape_tetra, dshape_tetra) + gauss_weights_3D, shape_tetra, dshape_tetra) assemble_module(np.array(KEL), np.array(C1_EL), np.array(C2_EL), np.array(VEL), np.array(SD_EL_res), np.array(J_EL_res), volume_conn, eq_num_u, eq_num_p, elem_num, KFF, C1FF, C2FF, VFF, SD_F, J_F) - + f[0:num_free_dof] = SD_F - FF f[num_free_dof:] = J_F - + print("||R||: ", np.linalg.norm(f)) if (np.linalg.norm(f) < tol): return UUR, pUR - + M[0:num_free_dof, 0:num_free_dof] = KFF M[0:num_free_dof, num_free_dof:] = C1FF M[num_free_dof:, 0:num_free_dof] = C2FF M[num_free_dof:, num_free_dof:] = VFF - + # print('Cond(M): ', np.linalg.cond(M)) - + delta = np.linalg.solve(M, -f) UF += delta[0:num_free_dof] diff --git a/cmad/fem_utils/fem_problem.py b/cmad/fem_utils/fem_problem.py index a5f9b31..ed20bdf 100644 --- a/cmad/fem_utils/fem_problem.py +++ b/cmad/fem_utils/fem_problem.py @@ -192,12 +192,15 @@ def __init__(self, problem_type, order): self._num_nodes_surf = 3 # fix all nodes on plane x = 0 + # fix z displacments on plane z = 0 disp_node = [] for i in range(self._num_nodes): if self._nodal_coords[i][0] == 0.0: disp_node.append(np.array([i, 1], dtype=int)) disp_node.append(np.array([i, 2], dtype=int)) disp_node.append(np.array([i, 3], dtype=int)) + elif self._nodal_coords[i][2] == 0.0: + disp_node.append(np.array([i, 3], dtype=int)) self._disp_node = np.array(disp_node, dtype=int) self._disp_val = np.zeros(len(disp_node)) @@ -209,6 +212,39 @@ def __init__(self, problem_type, order): if (surface_points[:, 0] == 48 * np.ones(3)).all(): pres_surf.append(surface) self._pres_surf = np.array(pres_surf) + + if problem_type == "vert_beam": + + self._mesh = Mesh("vert_beam") + + self._nodal_coords = self._mesh.get_nodal_coordinates() + self._colume_conn = self._mesh.get_volume_connectivity() + self._surface_conn = self._mesh.get_surface_connectivity() + + self._dof_node = 3 + self._num_nodes = len(self._nodal_coords) + self._num_nodes_elem = 4 + self._num_elem = len(self._colume_conn) + self._num_nodes_surf = 3 + + # fix all nodes on plane z = 0 + disp_node = [] + for i in range(self._num_nodes): + if self._nodal_coords[i][2] == 0.0: + disp_node.append(np.array([i, 1], dtype=int)) + disp_node.append(np.array([i, 2], dtype=int)) + disp_node.append(np.array([i, 3], dtype=int)) + self._disp_node = np.array(disp_node, dtype=int) + self._disp_val = np.zeros(len(disp_node)) + + # diagonal traction on z = 5 + self._surf_traction_vector = np.array([0.2, 0.2, 0.0]) + pres_surf = [] + for surface in self._surface_conn: + surface_points = self._nodal_coords[surface, :] + if (surface_points[:, 2] == 5 * np.ones(3)).all(): + pres_surf.append(surface) + self._pres_surf = np.array(pres_surf) def get_2D_basis_functions(self): return self._quad_rule_2D, self._shape_func_triangle diff --git a/cmad/fem_utils/mesh.py b/cmad/fem_utils/mesh.py index 8111331..764af0e 100644 --- a/cmad/fem_utils/mesh.py +++ b/cmad/fem_utils/mesh.py @@ -46,18 +46,32 @@ def __init__(self, mesh_type): geom.extrude(rectangle, [0.0, 0.0, 2], num_layers=10) self._mesh = geom.generate_mesh() - + if self._mesh_type == "cook": - - geom.characteristic_length_min = 1.5 - geom.characteristic_length_max = 1.5 + + geom.characteristic_length_min = 1.0 + geom.characteristic_length_max = 1.0 beam = geom.add_polygon([[0.0, 0.0], [48., 44.], [48., 60.], [0.0, 44.]]) + + geom.extrude(beam, [0.0, 0.0, 1.0], num_layers=1) + self._mesh = geom.generate_mesh() + + if self._mesh_type == "vert_beam": + + geom.characteristic_length_min = 0.125 + geom.characteristic_length_max = 0.125 + + beam = geom.add_polygon([[0.0, 0.0], + [1.0, 0.0], + [1.0, 1.0], + [0.0, 1.0]]) - geom.extrude(beam, [0.0, 0.0, 1.5], num_layers=1) + geom.extrude(beam, [0.0, 0.0, 5.0], num_layers=40) + self._mesh = geom.generate_mesh() self._points = self._mesh.points diff --git a/cmad/fem_utils/solve_fem.py b/cmad/fem_utils/solve_fem.py index d49d122..4524f2b 100644 --- a/cmad/fem_utils/solve_fem.py +++ b/cmad/fem_utils/solve_fem.py @@ -68,10 +68,11 @@ tol = 5e-12 max_iters = 10 -UUR = solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, - num_nodes_surf, surf_traction_vector, params, disp_node, - disp_val, eq_num, volume_conn, nodal_coords, pres_surf, - quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle, tol, max_iters) +UUR = solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, + dof_node, num_nodes_surf, surf_traction_vector, params, + disp_node, disp_val, eq_num, volume_conn, nodal_coords, + pres_surf, quad_rule_3D, shape_func_tetra, quad_rule_2D, + shape_func_triangle, tol, max_iters) print(UUR) diff --git a/cmad/fem_utils/solve_fem_DB.py b/cmad/fem_utils/solve_fem_DB.py index 0d30672..4c24a72 100644 --- a/cmad/fem_utils/solve_fem_DB.py +++ b/cmad/fem_utils/solve_fem_DB.py @@ -30,8 +30,8 @@ max_iters = 10 UUR, pUR = solve_fem_newton(num_pres_dof, num_free_dof, num_elem, num_nodes_elem, dof_node, - num_nodes_surf, surf_traction_vector, params, disp_node, - disp_val, eq_num_u, eq_num_p, volume_conn, nodal_coords, pres_surf, - quad_rule_3D, shape_func_tetra, quad_rule_2D, shape_func_triangle, tol, max_iters) + num_nodes_surf, surf_traction_vector, params, disp_node, disp_val, + eq_num_u, eq_num_p, volume_conn, nodal_coords, pres_surf, quad_rule_3D, + shape_func_tetra, quad_rule_2D, shape_func_triangle, tol, max_iters) problem.save_data("cook_membrane_incompressible", {'displacement_field': UUR, 'pressure_field': pUR}) \ No newline at end of file