3688ceacymsim: Python port of cymbalsim with a modal linear model and a PyTorch von Karman modelJeremy Magland 1"""One explicit time step (symplectic Euler) of the von Karman plate with a
2thickness field, a stress-free profile, the losses, and the drumstick, as a
3pure function of PyTorch tensors over the ``(ny, nx)`` grid.
5The bending forces are the gradient of the discrete Kirchhoff energy with the
6local bending stiffness D(h) at the nodes and at the cells (natural free-edge
7conditions); the membrane forces come from bilinear cells with 2 x 2 Gauss
8integration, the cell's mean thickness, and the Marguerre form of the von
9Karman strains. Nodes outside the plate (weight 0) stay at zero: a second
10difference exists only where the node and both neighbors belong to the plate,
11and a cell contributes only when its four corners do. That gives the free edge
12of the rectangle and a staircase free edge for any other outline, hole included.
14Three places would otherwise turn single-precision rounding into audible
15noise, and are written as in the reference implementation: the membrane
16strains carry the deviation from the stress-free shape in every term (the
17direct form 0.5 (wx^2 - w0x^2) leaves a rounding error proportional to w0x^2,
18which does not decay); the update of w is compensated, with the rounding error
19carried in a separate array; and the listener sum (in ``simulation.py``) uses
20the first difference of the delayed velocity rather than the instantaneous
21acceleration.
22"""
24from __future__ import annotations
26import math
27from dataclasses import dataclass
29import numpy as np
30import torch
31import torch.nn.functional as F
33from ..constants import STICK_RELEASE
34from ..linear.assemble import shape_derivatives
35from ..params import Strike
36from ..plate import Plate
39def pad(t: torch.Tensor, left: int, right: int, top: int, bottom: int) -> torch.Tensor:
40 return F.pad(t, (left, right, top, bottom))
44class Static:
45 """Per-node and per-cell data of a plate on the device, fixed for a run."""
47 ins: torch.Tensor
48 """Node in the plate, as 0/1."""
49 wgt: torch.Tensor
50 inv_mass: torch.Tensor
51 """1 / lumped mass, zero outside."""
52 Dn: torch.Tensor
53 """Bending stiffness at the nodes."""
54 Dc: torch.Tensor
55 """Bending stiffness at the cells (zero outside), ``(ny - 1, nx - 1)``."""
56 Cm: torch.Tensor
57 """Membrane stiffness at the cells times the Gauss weight hx hy / 4 (zero outside)."""
58 Gs: torch.Tensor
59 has_a: torch.Tensor
60 has_b: torch.Tensor
61 nb: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
62 """Left, right, lower, upper neighbor in the plate (0/1)."""
63 w0x: torch.Tensor
64 """Slopes of the stress-free shape at the Gauss points, ``(4, ny - 1, nx - 1)``."""
65 w0y: torch.Tensor
66 dNdx: torch.Tensor
67 """Shape function derivatives, ``(4 gp, 4 corners)``."""
68 dNdy: torch.Tensor
69 relief: torch.Tensor
70 """Inertia-relief acceleration field per unit stick force."""
71 patch_w: torch.Tensor
72 """m_k / M_patch on the contact patch, zero elsewhere: the contact coordinate is its dot with w."""
73 patch_f: torch.Tensor
74 """1 / M_patch on the contact patch: the force distribution per unit mass."""
75 hx: float
76 hy: float
77 nu: float
78 dt: float
79 sigma0: float
80 sigma1: float
81 sigma_in: float
82 K: float
83 alpha: float
84 stick_mass: float
85 strike: bool
88def static_data(pl: Plate, strike: Strike | None, dt: float, sigma0: float, sigma1: float, sigma_in: float, device, dtype) -> Static:
89 g, m = pl.grid, pl.material
90 t = lambda a: torch.as_tensor(np.asarray(a, dtype=np.float64), dtype=dtype, device=device) # noqa: E731
91 ins = g.mask.astype(float)
92 mass = pl.node_mass
93 inv_mass = np.where(mass > 0, 1 / np.where(mass > 0, mass, 1), 0.0)
94 cm = g.cell_mask.astype(float)
95 hc = pl.cell_thickness
96 left, right, down, up = g.neighbor_masks()
97 w0c = [np.asarray(c) for c in (pl.w0[:-1, :-1], pl.w0[:-1, 1:], pl.w0[1:, 1:], pl.w0[1:, :-1])]
98 dNdx = np.zeros((4, 4))
99 dNdy = np.zeros((4, 4))
100 w0x = np.zeros((4,) + cm.shape)
101 w0y = np.zeros((4,) + cm.shape)
102 for gp in range(4):
103 dNdx[gp], dNdy[gp] = shape_derivatives(gp, g.hx, g.hy)
104 w0x[gp] = sum(dNdx[gp, b] * w0c[b] for b in range(4))
105 w0y[gp] = sum(dNdy[gp, b] * w0c[b] for b in range(4))
106 if strike is not None:
107 patch = pl.strike_patch(strike)
108 cx, cy = pl.center
109 Ix, Iy = pl.inertia
110 relief = ins * (1 / pl.mass + (g.Y - cy) * (patch.cy - cy) / Ix + (g.X - cx) * (patch.cx - cx) / Iy)
111 patch_w = np.where(patch.mask, mass / patch.mass, 0.0)
112 patch_f = np.where(patch.mask, 1 / patch.mass, 0.0)
113 K, alpha, ms = strike.stiffness, strike.exponent, strike.mass
114 else:
115 relief = patch_w = patch_f = np.zeros(g.shape)
116 K, alpha, ms = 0.0, 1.0, 1.0
117 return Static(
118 ins=t(ins),
119 wgt=t(g.node_weight),
120 inv_mass=t(inv_mass),
121 Dn=t(m.bending_stiffness(pl.h)),
122 Dc=t(m.bending_stiffness(hc) * cm),
123 Cm=t(g.hx * g.hy / 4 * m.membrane_stiffness(hc) * cm),
124 Gs=t(g.hx * g.hy / 4 * m.shear_stiffness(hc) * cm),
125 has_a=t(left & right),
126 has_b=t(down & up),
127 nb=(t(left), t(right), t(down), t(up)),
128 w0x=t(w0x),
129 w0y=t(w0y),
130 dNdx=t(dNdx),
131 dNdy=t(dNdy),
132 relief=t(relief),
133 patch_w=t(patch_w),
134 patch_f=t(patch_f),
135 hx=g.hx,
136 hy=g.hy,
137 nu=m.nu,
138 dt=dt,
139 sigma0=sigma0,
140 sigma1=sigma1,
141 sigma_in=sigma_in,
142 K=K,
143 alpha=alpha,
144 stick_mass=ms,
145 strike=strike is not None,
146 )
149def accelerations(S: Static, u: torch.Tensor, v: torch.Tensor, w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
150 """Undamped accelerations of every node for the displacements ``(u, v, w)`` from the stress-free shape."""
151 hx, hy, nu = S.hx, S.hy, S.nu
152 ihx2, ihy2, ihxy = 1 / hx**2, 1 / hy**2, 1 / (hx * hy)
154 # Bending: weighted moments at the nodes, mixed differences at the cells.
155 wp = pad(w, 1, 1, 1, 1)
156 a = S.has_a * (wp[1:-1, 2:] - 2 * w + wp[1:-1, :-2]) * ihx2
157 b = S.has_b * (wp[2:, 1:-1] - 2 * w + wp[:-2, 1:-1]) * ihy2
158 wD = S.wgt * S.Dn
159 Ma = S.has_a * wD * (a + nu * b)
160 Mb = S.has_b * wD * (b + nu * a)
161 Cc = S.Dc * (w[1:, 1:] - w[1:, :-1] - w[:-1, 1:] + w[:-1, :-1]) * ihxy
162 Map = pad(Ma, 1, 1, 0, 0)
163 Mbp = pad(Mb, 0, 0, 1, 1)
164 Cp = pad(Cc, 1, 1, 1, 1)
165 bend = (Map[:, 2:] - 2 * Ma + Map[:, :-2]) * ihx2 + (Mbp[2:, :] - 2 * Mb + Mbp[:-2, :]) * ihy2 + 2 * (1 - nu) * (Cp[1:, 1:] - Cp[1:, :-1] - Cp[:-1, 1:] + Cp[:-1, :-1]) * ihxy
166 Fw = -hx * hy * bend
168 # Membrane: bilinear cells, 2 x 2 Gauss points, Marguerre strains with the deviation in every term.
169 U = (u[:-1, :-1], u[:-1, 1:], u[1:, 1:], u[1:, :-1])
170 V = (v[:-1, :-1], v[:-1, 1:], v[1:, 1:], v[1:, :-1])
171 W = (w[:-1, :-1], w[:-1, 1:], w[1:, 1:], w[1:, :-1])
172 Fu_c = [None] * 4
173 Fv_c = [None] * 4
174 Fw_c = [None] * 4
175 for gp in range(4):
176 dNdx, dNdy = S.dNdx[gp], S.dNdy[gp]
177 ux = dNdx[0] * U[0] + dNdx[1] * U[1] + dNdx[2] * U[2] + dNdx[3] * U[3]
178 uy = dNdy[0] * U[0] + dNdy[1] * U[1] + dNdy[2] * U[2] + dNdy[3] * U[3]
179 vx = dNdx[0] * V[0] + dNdx[1] * V[1] + dNdx[2] * V[2] + dNdx[3] * V[3]
180 vy = dNdy[0] * V[0] + dNdy[1] * V[1] + dNdy[2] * V[2] + dNdy[3] * V[3]
181 dwx = dNdx[0] * W[0] + dNdx[1] * W[1] + dNdx[2] * W[2] + dNdx[3] * W[3]
182 dwy = dNdy[0] * W[0] + dNdy[1] * W[1] + dNdy[2] * W[2] + dNdy[3] * W[3]
183 w0x, w0y = S.w0x[gp], S.w0y[gp]
184 wx, wy = w0x + dwx, w0y + dwy
185 exx = ux + dwx * (w0x + 0.5 * dwx)
186 eyy = vy + dwy * (w0y + 0.5 * dwy)
187 gxy = uy + vx + w0x * dwy + w0y * dwx + dwx * dwy
188 Nxx = S.Cm * (exx + nu * eyy)
189 Nyy = S.Cm * (eyy + nu * exx)
190 Nxy = S.Gs * gxy
191 fx = Nxx * wx + Nxy * wy
192 fy = Nxy * wx + Nyy * wy
193 for k in range(4):
194 fu = -(Nxx * dNdx[k] + Nxy * dNdy[k])
195 fv = -(Nyy * dNdy[k] + Nxy * dNdx[k])
196 fw = -(fx * dNdx[k] + fy * dNdy[k])
197 Fu_c[k] = fu if Fu_c[k] is None else Fu_c[k] + fu
198 Fv_c[k] = fv if Fv_c[k] is None else Fv_c[k] + fv
199 Fw_c[k] = fw if Fw_c[k] is None else Fw_c[k] + fw
200 # Corner k of every cell goes to node offsets (0,0), (0,1), (1,1), (1,0) in (j, i).
201 pads = ((0, 1, 0, 1), (1, 0, 0, 1), (1, 0, 1, 0), (0, 1, 1, 0))
202 Fu = sum(pad(Fu_c[k], *pads[k]) for k in range(4))
203 Fv = sum(pad(Fv_c[k], *pads[k]) for k in range(4))
204 Fw = Fw + sum(pad(Fw_c[k], *pads[k]) for k in range(4))
205 return Fu * S.inv_mass, Fv * S.inv_mass, Fw * S.inv_mass
208def damping(S: Static, du: torch.Tensor, dv: torch.Tensor, dw: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
209 """The loss terms: -2 sigma0 v on every component, +2 sigma1 (graph Laplacian
210 of w_t) on the transverse component, -2 sigma_in v on the in-plane components."""
211 left, right, down, up = S.nb
212 p = pad(dw, 1, 1, 1, 1)
213 lap = (left * (p[1:-1, :-2] - dw) + right * (p[1:-1, 2:] - dw)) / S.hx**2 + (down * (p[:-2, 1:-1] - dw) + up * (p[2:, 1:-1] - dw)) / S.hy**2
214 r = 2 * (S.sigma0 + S.sigma_in)
215 return -r * du, -r * dv, -2 * S.sigma0 * dw + 2 * S.sigma1 * lap
218def step(
219 S: Static,
220 u: torch.Tensor,
221 v: torch.Tensor,
222 w: torch.Tensor,
223 wlo: torch.Tensor,
224 du: torch.Tensor,
225 dv: torch.Tensor,
226 dw: torch.Tensor,
227 stick: torch.Tensor,
228) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
229 """One step. ``stick`` holds (position, velocity, force, max force, time since release, withdrawn)."""
230 dt = S.dt
231 au, av, aw = accelerations(S, u, v, w)
232 lu, lv, lw = damping(S, du, dv, dw)
233 au = au + lu
234 av = av + lv
235 aw = aw + lw
236 if S.strike:
237 # Hertz contact between the free stick mass and the contact patch (the
238 # plate nodes within the tip radius, plus the nearest one). The contact
239 # coordinate is the mass-weighted mean displacement of the patch and the
240 # force is spread over the patch in proportion to the node masses, so
241 # the pair is conjugate. Once the contact has been released for
242 # STICK_RELEASE the stick is withdrawn. The force of the previous step
243 # is applied to the whole plate as the rigid acceleration field that
244 # balances the patch load and its moments (inertia relief of the free
245 # plate), so the plate gains no net momentum.
246 sx, sv, f_prev, f_max, released, withdrawn = stick.unbind()
247 aw = aw + f_prev * S.relief
248 wc = (S.patch_w * w).sum()
249 pen = wc - sx
250 active = withdrawn == 0
251 contact = active & (pen > 0)
252 fc = torch.where(contact, S.K * torch.clamp(pen, min=0) ** S.alpha, torch.zeros_like(pen))
253 released = torch.where(contact, torch.zeros_like(released), torch.where(active & (f_max > 0), released + dt, released))
254 withdrawn = torch.where(active & (released > STICK_RELEASE), torch.ones_like(withdrawn), withdrawn)
255 aw = aw - fc * S.patch_f
256 svn = sv + dt * fc / S.stick_mass
257 stick = torch.stack([sx + dt * svn, svn, fc, torch.maximum(f_max, fc), released, withdrawn])
258 du = S.ins * (du + dt * au)
259 dv = S.ins * (dv + dt * av)
260 dw = S.ins * (dw + dt * aw)
261 u = S.ins * (u + dt * du)
262 v = S.ins * (v + dt * dv)
263 # Compensated update of w: the increment dt v can be far below the float32
264 # resolution of w, so the rounding error is carried separately.
265 y = dt * dw + wlo
266 t = w + y
267 lo = (w - t) + y
268 w = S.ins * t
269 wlo = S.ins * lo
270 return u, v, w, wlo, du, dv, dw, stick