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 total_energy(self, q):
192 Em, Eb = self.energy_parts_per_element(q)
193 return Em.sum() + Eb.sum()
195 def internal_force(self, q):
196 """f_int = -dU/dq, shape (N, 5)."""
197 return self._f_int(q)
199 def kinetic_energy(self, v):
200 return 0.5 * (self.M_j * v * v).sum()
202 def damping_force_w(self, vw):
203 """The sigma1 term as a nodal force on w: -2 sigma1 rho h Laplacian(w_dot) integrated
204 with the CST gradients, i.e. 2 sigma1 sum_e (rho h A)_e G^T (G w_dot_e)."""
205 ve = vw[self.tri_j]
206 vx, vy = (self.Gx_j * ve).sum(1), (self.Gy_j * ve).sum(1)
207 c = self.rhoh_area_j[:, None] * (self.Gx_j * vx[:, None] + self.Gy_j * vy[:, None])
208 return 2 * self._sigma1 * jax.ops.segment_sum(c.ravel(), self.tri_j.ravel(), num_segments=self._N)
210 def project_rigid(self, v):
211 """Remove the M-orthogonal projection onto the six rigid-body modes."""
212 coef = jnp.einsum("knd,nd->k", self.Rm_j, self.M_j * v)
213 return v - jnp.einsum("k,knd->nd", coef, self.Rm_j)
215 # ---- diagnostics
216 def element_checks(self) -> ElementChecks:
217 """Element self-tests (float64) and the rigid-mode test of the assembled tangent."""
218 mesh = self.mesh
219 x, y = mesh.xy[mesh.tri, 0], mesh.xy[mesh.tri, 1]
220 E = mesh.n_elements
221 U_tilt = np.stack([y, np.ones_like(y), np.zeros_like(y)], 2).reshape(E, 9)
222 res = np.abs(np.einsum("eij,ej->ei", self.Ke, U_tilt)).max() / np.abs(self.Ke).max()
223 U_cur = np.stack([0.5 * x ** 2, np.zeros_like(x), -x], 2).reshape(E, 9)
224 Eb = 0.5 * np.einsum("ei,eij,ej->e", U_cur, self.Ke, U_cur)
225 err = np.abs(Eb / (0.5 * self.D_e * mesh.area) - 1).max()
226 n_zero = (self._lam < 1e-9 * self._lam[:, -1:]).sum(1)
228 q0 = jnp.zeros((mesh.n_nodes, N_DOF), self.dtype)
229 K_apply = lambda d: jax.jvp(self._f_int, (q0,), (d,))[1]
230 rng = np.random.default_rng(0)
231 worst = 0.0
232 for k in range(6):
233 rnd = rng.standard_normal((mesh.n_nodes, N_DOF))
234 rnd *= np.linalg.norm(self.Rm[k]) / np.linalg.norm(rnd)
235 worst = max(worst, float(jnp.abs(K_apply(self.Rm_j[k])).max()
236 / jnp.abs(K_apply(jnp.asarray(rnd, self.dtype))).max()))
237 return ElementChecks(float(res), float(err), int(n_zero.min()), int(n_zero.max()), worst)