1"""The discrete shell: element data, lumped mass, rigid-body modes, internal force and the
2stable time-step bound.
4Given a mesh and a CymbalProfile, this module samples the thickness and the surface slope on the
5elements, builds the DKT bending stiffness and the membrane stiffness of every element, lumps the
6mass, and exposes the internal nodal force f_int(q) = -dU/dq obtained by reverse-mode automatic
7differentiation of the summed element strain energies. It also provides the two operators the
8integrator needs besides f_int: the sigma1 damping force on w and the projection that removes
9the six rigid-body modes of the free shell. See docs/model.md, sections "Spatial discretization"
10and "Time integration".
11"""
12from __future__ import annotations
14import logging
15import math
16from dataclasses import dataclass
18import jax
19import jax.numpy as jnp
20import numpy as np
22from .config import Material, Numerics
23from .elements import dkt_stiffness, element_energy, element_energy_parts
24from .geometry import CymbalProfile
25from .mesh import Mesh
27N_DOF = 5 # (u, v, w, theta_x, theta_y) per node
29log = logging.getLogger(__name__)
32def global_omega_max(K_e: np.ndarray, M: np.ndarray, tri: np.ndarray, tol: float = 1e-7,
33 max_iter: int = 20000) -> tuple[float, int, bool]:
34 """sqrt(lambda_max(M^-1 K)) of the stiffness assembled from the element matrices ``K_e``
35 (E, 15, 15), node-major DOFs, by power iteration on M^-1/2 K M^-1/2 (lumped M, (N, 5)).
37 The start vector x_i = sin(i + 1) and the stopping rule (the Rayleigh quotient changes by at
38 most ``tol`` relative over 10 iterations) are fixed, so that cymsimweb, which implements the
39 same iteration, arrives at the same value. The Rayleigh quotient approaches lambda_max from
40 below. Returns (omega, iterations, converged)."""
41 E, L = len(tri), M.size
42 dof = (N_DOF * tri[:, :, None] + np.arange(N_DOF)).reshape(E, 15)
43 isq = 1.0 / np.sqrt(M.ravel())
44 x = np.sin(np.arange(L) + 1.0)
45 x /= np.linalg.norm(x)
46 hist = []
47 for it in range(1, max_iter + 1):
48 ye = np.einsum("eab,eb->ea", K_e, (x * isq)[dof])
49 y = np.bincount(dof.ravel(), weights=ye.ravel(), minlength=L) * isq
50 lam = float(x @ y)
51 hist.append(lam)
52 x = y / np.linalg.norm(y)
53 if it >= 20 and it % 10 == 0 and abs(lam - hist[-11]) <= tol * lam:
54 return math.sqrt(lam), it, True
55 return math.sqrt(hist[-1]), max_iter, False
59class ElementChecks:
60 """Residuals of the element self-tests; all should be small (1e-6 or better)."""
61 rigid_tilt_residual: float # |K_e U_tilt| / |K_e| for a rigid tilt w = y, thx = 1
62 constant_curvature_error: float # relative error of the energy of w = x^2 / 2
63 zero_energy_modes_min: int # zero eigenvalues of M_e^-1 K_e per element (expect 6)
64 zero_energy_modes_max: int
65 rigid_mode_ratio: float # max |K r| / |K random| over the 6 rigid modes of the shell
68class ShellModel:
69 """Assembled shell on a mesh. All device arrays are in the simulation precision ``dtype``."""
71 def __init__(self, mesh: Mesh, profile: CymbalProfile, material: Material, numerics: Numerics):
72 self.mesh = mesh
73 self.material = material
74 self.numerics = numerics
75 self.dtype = jnp.float32 if numerics.dtype == "float32" else jnp.float64
76 self.nl = 1.0 if numerics.nonlinear else 0.0
77 xy, tri = mesh.xy, mesh.tri
78 N, E = mesh.n_nodes, mesh.n_elements
79 x, y = xy[tri, 0], xy[tri, 1]
80 area = mesh.area
82 # ---- profile sampled on the mesh: thickness per element and node, height per node.
83 # The slope that enters the membrane strain is the CST gradient of the nodal
84 # heights, i.e. the initial surface is represented in the same piecewise-linear
85 # space as w. This keeps a rigid tilt of the shell exactly force-free; the analytic
86 # slope at the centroid would not (see docs/model.md).
87 self.h_e = profile.thickness_at(mesh.centroid)
88 self.h_node = profile.thickness_at(xy)
89 self.z0_node = profile.z0_at(xy)
90 z0_e = self.z0_node[tri]
91 self.z0x = (mesh.Gx * z0_e).sum(1)
92 self.z0y = (mesh.Gy * z0_e).sum(1)
93 self.max_slope = float(np.hypot(self.z0x, self.z0y).max())
94 # ---- unit normal of the rest surface at the nodes, from the area-weighted mean of the
95 # CST slopes of the adjacent elements; +z side is the front (bell up). The normal
96 # acceleration n . (u_dd, v_dd, w_dd) is what is recorded and what radiates.
97 gx, gy = np.zeros(N), np.zeros(N)
98 for k in range(3):
99 np.add.at(gx, tri[:, k], area * self.z0x)
100 np.add.at(gy, tri[:, k], area * self.z0y)
101 a_node = np.zeros(N)
102 for k in range(3):
103 np.add.at(a_node, tri[:, k], area)
104 gx, gy = gx / a_node, gy / a_node
105 nrm = np.sqrt(1 + gx ** 2 + gy ** 2)
106 self.normal_node = np.column_stack([-gx / nrm, -gy / nrm, 1 / nrm])
108 # ---- element stiffness: bending rigidity D, membrane stiffness C, DKT matrix
109 E_, nu, rho = material.E, material.nu, material.rho
110 self.D_e = E_ * self.h_e ** 3 / (12 * (1 - nu ** 2))
111 self.Cm_e = (E_ * self.h_e / (1 - nu ** 2))[:, None] * np.array([1.0, nu, 0.5 * (1 - nu)])
112 self.Ke = np.asarray(jax.vmap(dkt_stiffness, in_axes=(0, 0, 0, None))(
113 jnp.asarray(x), jnp.asarray(y), jnp.asarray(self.D_e), nu))
115 # ---- lumped mass per DOF: a third of the element mass to each node; rotational inertia
116 # I = m (h^2 / 12 + f A_e), the second term being the mass-scaling device discussed
117 # in the documentation
118 m_e = rho * self.h_e * area / 3
119 I_e = m_e * (self.h_e ** 2 / 12 + numerics.rot_inertia_factor * area)
120 M = np.zeros((N, N_DOF))
121 for k in range(3):
122 np.add.at(M[:, 0], tri[:, k], m_e * numerics.inplane_mass_scale)
123 np.add.at(M[:, 1], tri[:, k], m_e * numerics.inplane_mass_scale)
124 np.add.at(M[:, 2], tri[:, k], m_e)
125 np.add.at(M[:, 3], tri[:, k], I_e)
126 np.add.at(M[:, 4], tri[:, k], I_e)
127 self.M = M
128 self.m_e, self.I_e = m_e, I_e
129 self.mass = float(M[:, 2].sum())
131 # ---- rigid-body modes of the curved shell, M-orthonormalised (see docs). A tilt drags
132 # the in-plane displacements along: rotation about x by an angle a gives v = -a z0,
133 # w = a y; rotation about y gives u = +a z0, w = -a x (first order, shallow shell).
134 X, Y = xy[:, 0] - xy[:, 0].mean(), xy[:, 1] - xy[:, 1].mean()
135 Rm = np.zeros((6, N, N_DOF))
136 Rm[0, :, 0] = 1
137 Rm[1, :, 1] = 1
138 Rm[2, :, 0], Rm[2, :, 1] = -Y, X
139 Rm[3, :, 2] = 1
140 Rm[4, :, 2], Rm[4, :, 3], Rm[4, :, 1] = Y, 1, -self.z0_node
141 Rm[5, :, 2], Rm[5, :, 4], Rm[5, :, 0] = -X, 1, self.z0_node
142 for k in range(6):
143 for l in range(k):
144 Rm[k] -= (Rm[l] * M * Rm[k]).sum() * Rm[l]
145 Rm[k] /= math.sqrt((Rm[k] * M * Rm[k]).sum())
146 self.Rm = Rm
148 # ---- stability bound (Irons): lambda_max(M^-1 K) <= max_e lambda_max(M_e^-1 K_e), with
149 # K_e the tangent stiffness at rest (Hessian of the element energy at q = 0)
150 Khess = jax.vmap(jax.hessian(element_energy), in_axes=(0, 0, 0, 0, 0, 0, 0, 0, None))(
151 jnp.zeros((E, 3, N_DOF)), jnp.asarray(mesh.Gx), jnp.asarray(mesh.Gy), jnp.asarray(area),
152 jnp.asarray(self.z0x), jnp.asarray(self.z0y), jnp.asarray(self.Cm_e), jnp.asarray(self.Ke),
153 self.nl).reshape(E, 15, 15)
154 Me = np.tile(np.stack([m_e * numerics.inplane_mass_scale, m_e * numerics.inplane_mass_scale,
155 m_e, I_e, I_e], 1), (1, 3)) # (E, 15), node-major
156 S = np.asarray(Khess) / np.sqrt(Me[:, :, None] * Me[:, None, :])
157 self._lam = np.linalg.eigvalsh(S) # (E, 15) ascending
158 self.limiting_element = int(np.argmax(self._lam[:, -1]))
159 self.omega_max = math.sqrt(self._lam[:, -1].max())
160 # ---- the global lambda_max(M^-1 K) of the assembled tangent, which the element bound
161 # overestimates (by about 1.45 at 6 mm elements); see docs/model.md
162 self.K_rest = np.asarray(Khess) # (E, 15, 15) tangent at rest
163 self.omega_global, self.omega_global_iterations, converged = global_omega_max(self.K_rest, M, tri)
164 if numerics.stability_bound == "global" and not converged:
165 log.warning("the power iteration for the global stability bound did not converge in %d "
166 "iterations; using the element bound", self.omega_global_iterations)
167 use_global = numerics.stability_bound == "global" and converged
168 self.omega_step = min(self.omega_global, self.omega_max) if use_global else self.omega_max
169 self.dt_stable = 2.0 / self.omega_step
171 # ---- device arrays in simulation precision
172 f = lambda a: jnp.asarray(a, dtype=self.dtype)
173 self.tri_j = jnp.asarray(tri, dtype=jnp.int32)
174 self.Gx_j, self.Gy_j, self.area_j = f(mesh.Gx), f(mesh.Gy), f(area)
175 self.z0x_j, self.z0y_j, self.Cm_j, self.Ke_j = f(self.z0x), f(self.z0y), f(self.Cm_e), f(self.Ke)
176 self.rhoh_area_j = f(rho * self.h_e * area)
177 self.h_node_j = f(self.h_node)
178 self.normal_j = f(self.normal_node)
179 self.M_j, self.Minv_j, self.Rm_j = f(M), f(1.0 / M), f(Rm)
180 self._sigma1 = float(material.sigma1)
181 self._N = N
182 self._f_int = jax.grad(lambda q: -self.total_energy(q))
184 # ---- energies and forces (plain JAX functions; the integrator jits the whole step)
185 def energy_parts_per_element(self, q):
186 """(membrane, bending) energy of every element for nodal state q (N, 5)."""
187 return jax.vmap(element_energy_parts, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, None))(
188 q[self.tri_j], self.Gx_j, self.Gy_j, self.area_j, self.z0x_j, self.z0y_j,
189 self.Cm_j, self.Ke_j, self.nl)
191 def linear_energy_parts(self, q):
192 """(membrane, bending) strain energy of q without the von Karman terms; both quadratic forms."""
193 Em, Eb = jax.vmap(element_energy_parts, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, None))(
194 q[self.tri_j], self.Gx_j, self.Gy_j, self.area_j, self.z0x_j, self.z0y_j, self.Cm_j, self.Ke_j, 0.0)
195 return Em.sum(), Eb.sum()
197 def linear_energy(self, q):
198 """Strain energy of q with the quadratic (von Karman) strain terms dropped: a quadratic
199 form, so that sqrt(kinetic_energy(v) + linear_energy(q)) is a norm on the state."""
200 Em, Eb = jax.vmap(element_energy_parts, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, None))(
201 q[self.tri_j], self.Gx_j, self.Gy_j, self.area_j, self.z0x_j, self.z0y_j, self.Cm_j, self.Ke_j, 0.0)
202 return Em.sum() + Eb.sum()
204 def total_energy(self, q):
205 Em, Eb = self.energy_parts_per_element(q)
206 return Em.sum() + Eb.sum()
208 def internal_force(self, q):
209 """f_int = -dU/dq, shape (N, 5)."""
210 return self._f_int(q)
212 def kinetic_energy(self, v):
213 return 0.5 * (self.M_j * v * v).sum()
215 def damping_force_w(self, vw):
216 """The sigma1 term as a nodal force on w: -2 sigma1 rho h Laplacian(w_dot) integrated
217 with the CST gradients, i.e. 2 sigma1 sum_e (rho h A)_e G^T (G w_dot_e)."""
218 ve = vw[self.tri_j]
219 vx, vy = (self.Gx_j * ve).sum(1), (self.Gy_j * ve).sum(1)
220 c = self.rhoh_area_j[:, None] * (self.Gx_j * vx[:, None] + self.Gy_j * vy[:, None])
221 return 2 * self._sigma1 * jax.ops.segment_sum(c.ravel(), self.tri_j.ravel(), num_segments=self._N)
223 def project_rigid(self, v):
224 """Remove the M-orthogonal projection onto the six rigid-body modes."""
225 coef = jnp.einsum("knd,nd->k", self.Rm_j, self.M_j * v)
226 return v - jnp.einsum("k,knd->nd", coef, self.Rm_j)
228 # ---- diagnostics
229 def element_checks(self) -> ElementChecks:
230 """Element self-tests (float64) and the rigid-mode test of the assembled tangent."""
231 mesh = self.mesh
232 x, y = mesh.xy[mesh.tri, 0], mesh.xy[mesh.tri, 1]
233 E = mesh.n_elements
234 U_tilt = np.stack([y, np.ones_like(y), np.zeros_like(y)], 2).reshape(E, 9)
235 res = np.abs(np.einsum("eij,ej->ei", self.Ke, U_tilt)).max() / np.abs(self.Ke).max()
236 U_cur = np.stack([0.5 * x ** 2, np.zeros_like(x), -x], 2).reshape(E, 9)
237 Eb = 0.5 * np.einsum("ei,eij,ej->e", U_cur, self.Ke, U_cur)
238 err = np.abs(Eb / (0.5 * self.D_e * mesh.area) - 1).max()
239 n_zero = (self._lam < 1e-9 * self._lam[:, -1:]).sum(1)
241 q0 = jnp.zeros((mesh.n_nodes, N_DOF), self.dtype)
242 K_apply = lambda d: jax.jvp(self._f_int, (q0,), (d,))[1]
243 rng = np.random.default_rng(0)
244 worst = 0.0
245 for k in range(6):
246 rnd = rng.standard_normal((mesh.n_nodes, N_DOF))
247 rnd *= np.linalg.norm(self.Rm[k]) / np.linalg.norm(rnd)
248 worst = max(worst, float(jnp.abs(K_apply(self.Rm_j[k])).max()
249 / jnp.abs(K_apply(jnp.asarray(rnd, self.dtype))).max()))
250 return ElementChecks(float(res), float(err), int(n_zero.min()), int(n_zero.max()), worst)