"""Kirchhoff-Love plate on a Cartesian grid masked to a disk, with Bogner-Fox-Schmit elements. The disk of radius R is approximated by the union of the square cells of a uniform Cartesian grid of spacing a whose centers lie inside the disk (a staircase boundary; the grid is centered on the disk so the approximation has the symmetries of the square lattice). Each cell is a C1-conforming bicubic Hermite rectangle with four degrees of freedom per node (w, w_x, w_y, w_xy). Free edges, including the staircase corners, are natural boundary conditions of the Kirchhoff energy, so nothing is imposed. Element matrices are integrated exactly with 4 x 4 Gauss-Legendre points and assembled into sparse matrices; the lowest modes come from a shift-invert Lanczos eigensolver (ARPACK through scipy) and are mass-normalized. Three of them are rigid-body modes at zero frequency. The class also provides the closed-form linear modal response to the challenge's Hann strike (a particular solution during the pulse and free decay afterwards) and an exact evaluation of the retarded Rayleigh integral, both taken from the square-plate solver of challenges/plate-solutions, so that the linear limit of the nonlinear solver can be checked on the same mesh. Coordinates are those of the challenge: the disk is centered at the origin, w is positive in +z, the strike enters as -p, and w is reported directly as u_z. """ from __future__ import annotations import json from pathlib import Path import numpy as np import scipy.sparse as sp from numpy.polynomial import legendre as leg from scipy.sparse.linalg import eigsh from scipy.spatial import cKDTree RIGID_OMEGA = 1.0 # rad/s; eigenvalues below this are the three rigid-body modes def load_challenge(path: Path) -> dict: with open(path) as f: return json.load(f) # ---------------------------------------------------------------- Hermite cubics on [0, 1] def hermite(s: np.ndarray): """Values and first two derivatives (w.r.t. s) of the four cubic Hermite functions. Shapes (len(s), 4).""" s = np.asarray(s, dtype=float) H = np.stack([1 - 3 * s**2 + 2 * s**3, s - 2 * s**2 + s**3, 3 * s**2 - 2 * s**3, -(s**2) + s**3], axis=-1) dH = np.stack([-6 * s + 6 * s**2, 1 - 4 * s + 3 * s**2, 6 * s - 6 * s**2, -2 * s + 3 * s**2], axis=-1) d2H = np.stack([-6 + 12 * s, -4 + 6 * s, 6 - 12 * s, -2 + 6 * s], axis=-1) return H, dH, d2H CORNERS = [(0, 0), (1, 0), (0, 1), (1, 1)] def shape_functions(xi: np.ndarray, eta: np.ndarray, a: float): """BFS shape functions and second derivatives at local coordinates (xi, eta) for square a x a elements. Returns N, N_xx, N_yy, N_xy, each of shape (npts, 16). Local dof order: for corner c in [(0,0), (1,0), (0,1), (1,1)]: w, w_x, w_y, w_xy. """ Hx, dHx, d2Hx = hermite(xi) Hy, dHy, d2Hy = hermite(eta) N = np.zeros((len(xi), 16)) Nxx = np.zeros_like(N) Nyy = np.zeros_like(N) Nxy = np.zeros_like(N) for c, (cx, cy) in enumerate(CORNERS): ix0, ix1 = (0, 1) if cx == 0 else (2, 3) iy0, iy1 = (0, 1) if cy == 0 else (2, 3) combos = [(ix0, iy0, 1.0), (ix1, iy0, a), (ix0, iy1, a), (ix1, iy1, a * a)] for d, (fx, fy, sc) in enumerate(combos): k = 4 * c + d N[:, k] = sc * Hx[:, fx] * Hy[:, fy] Nxx[:, k] = sc * d2Hx[:, fx] * Hy[:, fy] / a**2 Nyy[:, k] = sc * Hx[:, fx] * d2Hy[:, fy] / a**2 Nxy[:, k] = sc * dHx[:, fx] * dHy[:, fy] / (a * a) return N, Nxx, Nyy, Nxy class BFSDisk: def __init__(self, ch: dict, spacing: float, n_modes: int, symmetric_y: bool = False, max_hz: float | None = None, thickness=None): """symmetric_y restricts the basis to functions even in y (w(x, -y) = w(x, y)), which is exact for a strike centered on the x axis; the grid is symmetric about that axis by construction. max_hz drops modes above it. thickness, if given, is a function h(x, y) (variable thickness); otherwise the uniform parameters["h"].""" p = ch["parameters"] self.p = p R, nu = p["disk_radius"], p["nu"] self.R = R a = float(spacing) self.a = a self.h_at = thickness if thickness is not None else (lambda x, y: np.full(np.shape(x), float(p["h"]))) h = float(p["h"]) if "h" in p else float(np.mean(self.h_at(np.zeros(1), np.zeros(1)))) D = p["E"] * h**3 / (12.0 * (1.0 - nu**2)) m = p["rho"] * h self.D, self.m = D, m # the uniform values (used only when the thickness is uniform) # ---- the masked grid half = int(np.ceil(R / a)) + 1 ncell = 2 * half self.ncell = ncell self.half = half self.x0 = -half * a # coordinate of grid line 0 centers = self.x0 + (np.arange(ncell) + 0.5) * a CX, CY = np.meshgrid(centers, centers, indexing="ij") active = CX**2 + CY**2 <= R**2 cells = np.argwhere(active) # (nc, 2): ix, iy self.cells = cells self.n_cells = len(cells) self.cell_id = -np.ones((ncell, ncell), dtype=int) self.cell_id[cells[:, 0], cells[:, 1]] = np.arange(self.n_cells) self.cell_center = np.stack([centers[cells[:, 0]], centers[cells[:, 1]]], axis=1) self._tree = cKDTree(self.cell_center) nn = ncell + 1 gnode = np.zeros((self.n_cells, 4), dtype=int) for c, (cx, cy) in enumerate(CORNERS): gnode[:, c] = (cells[:, 0] + cx) * nn + (cells[:, 1] + cy) touch = np.bincount(gnode.ravel(), minlength=nn * nn) used = touch > 0 self.node_map = -np.ones(nn * nn, dtype=int) self.node_map[used] = np.arange(int(used.sum())) self.n_nodes = int(used.sum()) self.ndof = 4 * self.n_nodes cnode = self.node_map[gnode] # (nc, 4) compact node ids gidx = np.flatnonzero(used) self.node_i, self.node_j = gidx // nn, gidx % nn # grid indices of each compact node self.cell_dofs = (4 * cnode[:, :, None] + np.arange(4)[None, None, :]).reshape(self.n_cells, 16) bnodes = self.node_map[used & (touch < 4)] self.boundary_dofs = (4 * bnodes[:, None] + np.arange(4)[None, :]).ravel() self.interior_dofs = np.setdiff1d(np.arange(self.ndof), self.boundary_dofs) self.area = self.n_cells * a * a # ---- element matrices g, wg = leg.leggauss(4) s = (g + 1) / 2 ws = wg / 2 XI, ETA = np.meshgrid(s, s, indexing="ij") WQ = (ws[:, None] * ws[None, :]).ravel() * a * a N, Nxx, Nyy, Nxy = shape_functions(XI.ravel(), ETA.ravel(), a) self._shape = (N, Nxx, Nyy, Nxy, WQ) C = np.array([[1.0, nu, 0.0], [nu, 1.0, 0.0], [0.0, 0.0, 2.0 * (1.0 - nu)]]) # Gauss-point coordinates of every cell, and the thickness there gx = (self.x0 + (self.cells[:, 0][:, None] + XI.ravel()[None, :]) * a) gy = (self.x0 + (self.cells[:, 1][:, None] + ETA.ravel()[None, :]) * a) self._gauss_xy = (gx, gy) hq = self.h_at(gx, gy) # (n_cells, nq) self.K = self.assemble_coeff(C, p["E"] * hq**3 / (12.0 * (1.0 - nu**2))) self.M1 = self.assemble_mass(np.ones_like(hq)) # unit-density mass matrix self.M = self.assemble_mass(p["rho"] * hq) # ---- modes: the n_modes lowest, mass-normalized (shift-invert Lanczos around a negative shift) self.symmetric_y = symmetric_y evals, evecs = self.eigenmodes(self.K, self.M, n_modes, sigma=-1.0e3) omega = np.sqrt(np.maximum(evals, 0.0)) omega[omega < RIGID_OMEGA] = 0.0 if max_hz is not None: keep = omega <= 2 * np.pi * max_hz omega, evecs = omega[keep], evecs[:, keep] self.omega = omega self.V = evecs self.n_modes = len(omega) self.alpha = p["sigma0"] + 0.5 * p["eta"] * omega**2 self.rigid = omega == 0.0 self.n_rigid = int(self.rigid.sum()) # ---- modal forces: f_n = -int p_hat Phi_n dA over the strike disk (polar Gauss quadrature) Rs, xs, ys = p["strike_radius"], p["strike_x"], p["strike_y"] c_phi = 0.5 - 2.0 / np.pi**2 gr, gw = leg.leggauss(24) r = Rs * (gr + 1.0) / 2.0 wr = gw * Rs / 2.0 nth = 64 th = 2 * np.pi * np.arange(nth) / nth rr, tt = np.meshgrid(r, th, indexing="ij") ww = (wr[:, None] * rr * (2 * np.pi / nth)).ravel() px = (xs + rr * np.cos(tt)).ravel() py = (ys + rr * np.sin(tt)).ravel() phat = p["strike_peak_force"] / (c_phi * np.pi * Rs**2) * 0.5 * (1.0 + np.cos(np.pi * rr.ravel() / Rs)) self.f = -(ww * phat) @ self.modes_at(px, py) self._listener_q = None # optional override of the listener quadrature (used by the shell solver) self._setup_linear_response() def _setup_linear_response(self): """Closed-form constants of the damped modal response to the Hann strike (see the square-plate solver). Depends on omega, alpha, f and rigid; call again after changing them.""" p = self.p # ---- closed-form linear response constants (see the square-plate solver) T = p["strike_duration"] Om = 2 * np.pi / T self.T, self.Om = T, Om al, w, f = self.alpha, self.omega, self.f with np.errstate(divide="ignore", invalid="ignore"): disc = np.sqrt(al**2 - w**2 + 0j) self.s1 = -al + disc self.s2 = -al - disc H = 1.0 / (w**2 - Om**2 + 2j * al * Om) self.qconst = np.where(self.rigid, 0.0, f / (2.0 * w**2)) self.Hc = np.where(self.rigid, 0.0, (f / 2.0) * H) rhs1 = -(self.qconst - self.Hc.real) rhs2 = (1j * Om * self.Hc).real self.B = (rhs2 - self.s1 * rhs1) / (self.s2 - self.s1) self.A = rhs1 - self.B self.Gc = (f / 2.0) / (2.0 * al + 1j * Om) self.vconst = f / (4.0 * al) self.Ar = -(self.vconst - self.Gc.real) qT, vT, _ = self._pulse_phase(np.array([T])) self.qT, self.vT = qT[:, 0], vT[:, 0] with np.errstate(divide="ignore", invalid="ignore"): self.B2 = (self.vT - self.s1 * self.qT) / (self.s2 - self.s1) self.A2 = self.qT - self.B2 # ------------------------------------------------------------ symmetry and eigensolves def symmetric_prolongation(self, dofs: np.ndarray | None = None) -> sp.csc_matrix: """Sparse matrix whose columns span the functions even in y: each column pairs a node with its mirror image (w and w_x equal, w_y and w_xy opposite), and nodes on the axis keep only w and w_x. If dofs is given, the matrix is restricted to those rows (in that order) and to the columns that touch them.""" j_axis = self.half gi, gj = self.node_i, self.node_j nn = self.ncell + 1 mirror = self.node_map[gi * nn + (2 * j_axis - gj)] rows, cols, vals = [], [], [] col = 0 signs = (1.0, 1.0, -1.0, -1.0) for n in range(self.n_nodes): if gj[n] > j_axis: nm = mirror[n] for d in range(4): rows += [4 * n + d, 4 * nm + d] cols += [col, col] vals += [1.0, signs[d]] col += 1 elif gj[n] == j_axis: for d in range(2): rows.append(4 * n + d) cols.append(col) vals.append(1.0) col += 1 P = sp.csc_matrix((vals, (rows, cols)), shape=(self.ndof, col)) if dofs is not None: P = P[dofs] P = P[:, np.unique(P.nonzero()[1])] return P.tocsc() def eigenmodes(self, K, M, k: int, sigma: float, dofs: np.ndarray | None = None): """The k lowest eigenpairs of K v = lambda M v, restricted to the dofs given (default all) and to the y-symmetric subspace if the mesh was built with symmetric_y. Sorted ascending, M-normalized, and returned on the full dof set (zero on excluded dofs).""" Kr, Mr = (K, M) if dofs is None else (K[dofs][:, dofs], M[dofs][:, dofs]) P = self.symmetric_prolongation(dofs) if self.symmetric_y else None if P is not None: Kr, Mr = P.T @ Kr @ P, P.T @ Mr @ P Kr, Mr = sp.csc_matrix(Kr), sp.csc_matrix(Mr) evals, evecs = eigsh(Kr, k=min(k, Kr.shape[0] - 2), M=Mr, sigma=sigma, which="LM") order = np.argsort(evals) evals, evecs = evals[order], evecs[:, order] if P is not None: evecs = P @ evecs Md = M if dofs is None else M[dofs][:, dofs] evecs = evecs / np.sqrt(np.einsum("ij,ij->j", evecs, Md @ evecs)) if dofs is not None: full = np.zeros((self.ndof, evecs.shape[1])) full[dofs] = evecs evecs = full return evals, evecs # ------------------------------------------------------------ assembly def element_matrix(self, Cmat: np.ndarray) -> np.ndarray: """Element matrix of int [w_xx, w_yy, w_xy] Cmat [v_xx, v_yy, v_xy]^T dA (no stiffness factor).""" _, Nxx, Nyy, Nxy, WQ = self._shape Ae = np.zeros((16, 16)) for q in range(len(WQ)): B2 = np.stack([Nxx[q], Nyy[q], Nxy[q]]) Ae += WQ[q] * (B2.T @ Cmat @ B2) return Ae def assemble_coeff(self, Cmat: np.ndarray, coef: np.ndarray) -> sp.csc_matrix: """Assemble int coef(x) [w_xx, w_yy, w_xy] Cmat [v_xx, v_yy, v_xy]^T dA with a coefficient given at the Gauss points of every cell (coef of shape (n_cells, nq)).""" _, Nxx, Nyy, Nxy, WQ = self._shape B = np.stack([Nxx, Nyy, Nxy], axis=1) # (nq, 3, 16) BCB = np.einsum("qai,ab,qbj->qij", B, Cmat, B) # (nq, 16, 16) Ke = np.einsum("eq,qij->eij", coef * WQ[None, :], BCB) return self._assemble_all(Ke) def assemble_mass(self, coef: np.ndarray) -> sp.csc_matrix: """Assemble int coef(x) w v dA with a Gauss-point coefficient of shape (n_cells, nq).""" N, _, _, _, WQ = self._shape NN = np.einsum("qi,qj->qij", N, N) Me = np.einsum("eq,qij->eij", coef * WQ[None, :], NN) return self._assemble_all(Me) def _assemble_all(self, Ae_all: np.ndarray) -> sp.csc_matrix: rows = np.repeat(self.cell_dofs, 16, axis=1).ravel() cols = np.tile(self.cell_dofs, (1, 16)).ravel() A = sp.coo_matrix((Ae_all.ravel(), (rows, cols)), shape=(self.ndof, self.ndof)).tocsc() A.sum_duplicates() return A def assemble(self, Ae: np.ndarray) -> sp.csc_matrix: rows = np.repeat(self.cell_dofs, 16, axis=1).ravel() cols = np.tile(self.cell_dofs, (1, 16)).ravel() vals = np.tile(Ae.ravel(), self.n_cells) A = sp.coo_matrix((vals, (rows, cols)), shape=(self.ndof, self.ndof)).tocsc() A.sum_duplicates() return A # ------------------------------------------------------------ quadrature over the masked domain def quadrature(self, order: int = 3): """Gauss-Legendre points and weights over all active cells: xq, yq, wq.""" g, wg = leg.leggauss(order) s = (g + 1) / 2 ws = wg / 2 XI, ETA = np.meshgrid(s, s, indexing="ij") a = self.a xq = (self.x0 + (self.cells[:, 0][:, None] + XI.ravel()[None, :]) * a).ravel() yq = (self.x0 + (self.cells[:, 1][:, None] + ETA.ravel()[None, :]) * a).ravel() wq = np.tile((ws[:, None] * ws[None, :]).ravel() * a * a, self.n_cells) return xq, yq, wq # ------------------------------------------------------------ evaluation at points def locate(self, x: np.ndarray, y: np.ndarray, tol: float = 1e-9): """Cell id and local coordinates of points. Points inside the disk that fall in an inactive cell of the staircase are assigned the nearest active cell (the element polynomial is extrapolated slightly). Returns cid (-1 for points outside the disk), xi, eta.""" x = np.asarray(x, dtype=float) y = np.asarray(y, dtype=float) ix = np.clip(np.floor((x - self.x0) / self.a).astype(int), 0, self.ncell - 1) iy = np.clip(np.floor((y - self.x0) / self.a).astype(int), 0, self.ncell - 1) cid = self.cell_id[ix, iy].copy() inside = x**2 + y**2 <= self.R**2 * (1 + tol) fix = inside & (cid < 0) if fix.any(): _, near = self._tree.query(np.stack([x[fix], y[fix]], axis=1)) cid[fix] = near cid[~inside] = -1 c = np.maximum(cid, 0) cx = self.x0 + self.cells[c, 0] * self.a cy = self.x0 + self.cells[c, 1] * self.a return cid, (x - cx) / self.a, (y - cy) / self.a def _sparse_basis(self, vals: np.ndarray, cid: np.ndarray) -> sp.csr_matrix: npts = len(cid) ok = cid >= 0 rows = np.repeat(np.arange(npts)[ok], 16) cols = self.cell_dofs[cid[ok]].ravel() data = vals[ok].ravel() return sp.csr_matrix((data, (rows, cols)), shape=(npts, self.ndof)) def basis_at(self, x, y) -> sp.csr_matrix: """Sparse (npts x ndof) matrix of shape-function values; rows of points outside the disk are zero.""" cid, xi, eta = self.locate(x, y) N, _, _, _ = shape_functions(xi, eta, self.a) return self._sparse_basis(N, cid) def basis_derivs_at(self, x, y): cid, xi, eta = self.locate(x, y) _, Nxx, Nyy, Nxy = shape_functions(xi, eta, self.a) return self._sparse_basis(Nxx, cid), self._sparse_basis(Nyy, cid), self._sparse_basis(Nxy, cid) def modes_at(self, x, y) -> np.ndarray: return self.basis_at(x, y) @ self.V def inside(self, x, y) -> np.ndarray: return np.asarray(x) ** 2 + np.asarray(y) ** 2 <= self.R**2 * (1 + 1e-9) # ------------------------------------------------------------ modal time response (linear) def _pulse_phase(self, t: np.ndarray): t = t[None, :] Om = self.Om e = np.exp(1j * Om * t) s1, s2, A, B = self.s1[:, None], self.s2[:, None], self.A[:, None], self.B[:, None] E1, E2 = np.exp(s1 * t), np.exp(s2 * t) Hc = self.Hc[:, None] q = (A * E1 + B * E2).real + self.qconst[:, None] - (Hc * e).real v = (A * s1 * E1 + B * s2 * E2).real - (1j * Om * Hc * e).real acc = (A * s1**2 * E1 + B * s2**2 * E2).real + (Om**2 * Hc * e).real rg = self.rigid if rg.any(): al = self.alpha[rg][:, None] Ar, Gc, vc = self.Ar[rg][:, None], self.Gc[rg][:, None], self.vconst[rg][:, None] Ed = np.exp(-2.0 * al * t) q[rg] = Ar * (1.0 - Ed) / (2.0 * al) + vc * t - ((e - 1.0) * Gc / (1j * Om)).real v[rg] = Ar * Ed + vc - (Gc * e).real acc[rg] = -2.0 * al * Ar * Ed - (1j * Om * Gc * e).real return q, v, acc def _free_phase(self, tau: np.ndarray): tau = tau[None, :] s1, s2, A, B = self.s1[:, None], self.s2[:, None], self.A2[:, None], self.B2[:, None] E1, E2 = np.exp(s1 * tau), np.exp(s2 * tau) q = (A * E1 + B * E2).real v = (A * s1 * E1 + B * s2 * E2).real acc = (A * s1**2 * E1 + B * s2**2 * E2).real rg = self.rigid if rg.any(): al = self.alpha[rg][:, None] qT, vT = self.qT[rg][:, None], self.vT[rg][:, None] Ed = np.exp(-2.0 * al * tau) q[rg] = qT + vT * (1.0 - Ed) / (2.0 * al) v[rg] = vT * Ed acc[rg] = -2.0 * al * vT * Ed return q, v, acc def response(self, t: np.ndarray): q = np.zeros((self.n_modes, len(t))) v = np.zeros_like(q) acc = np.zeros_like(q) m1 = (t >= 0) & (t <= self.T) m2 = t > self.T if m1.any(): q[:, m1], v[:, m1], acc[:, m1] = self._pulse_phase(t[m1]) if m2.any(): q[:, m2], v[:, m2], acc[:, m2] = self._free_phase(t[m2] - self.T) return q, v, acc # ------------------------------------------------------------ Rayleigh integral (linear, exact retardation) def listener_quadrature(self, order: int = 3): if self._listener_q is not None: return self._listener_q p = self.p xq, yq, wq = self.quadrature(order) r = np.sqrt((p["listener_x"] - xq) ** 2 + (p["listener_y"] - yq) ** 2 + p["listener_height"] ** 2) tau = r / p["c_air"] cq = p["rho_air"] / (2 * np.pi) * wq / r return xq, yq, tau, cq def pressure(self, fs: int, nsamp: int) -> np.ndarray: xq, yq, tau, cq = self.listener_quadrature() Phi = self.modes_at(xq, yq) # (nq, n_modes) t = np.arange(nsamp) / fs out = np.zeros(nsamp) t_split = self.T + tau.max() early = (t < t_split) & (t >= tau.min()) late = t >= t_split ne = int(early.sum()) if ne: chunk = max(1, 4_000_000 // max(1, self.n_modes * ne)) for i0 in range(0, len(xq), chunk): sl = slice(i0, i0 + chunk) tt = (t[early][None, :] - tau[sl][:, None]).ravel() _, _, acc = self.response(tt) acc = acc.reshape(self.n_modes, -1, ne) out[early] += np.einsum("q,qn,nqk->k", cq[sl], Phi[sl], acc, optimize=True) if late.any(): tl = t[late] - self.T rg = self.rigid G1 = (cq[:, None] * Phi * np.exp(-np.outer(tau, self.s1))).sum(axis=0) G2 = (cq[:, None] * Phi * np.exp(-np.outer(tau, self.s2))).sum(axis=0) c1 = self.A2 * self.s1**2 * G1 c2 = self.B2 * self.s2**2 * G2 s1, s2 = self.s1.copy(), self.s2.copy() if rg.any(): al = self.alpha[rg] Gr = (cq[:, None] * Phi[:, rg] * np.exp(np.outer(tau, 2.0 * al))).sum(axis=0) c1[rg] = -2.0 * al * self.vT[rg] * Gr c2[rg] = 0.0 s1[rg] = -2.0 * al s2[rg] = 0.0 coef = np.concatenate([c1, c2]) rates = np.concatenate([s1, s2]) res = np.zeros(len(tl)) step = max(1, 20_000_000 // len(coef)) for k0 in range(0, len(tl), step): tk = tl[k0:k0 + step] res[k0:k0 + step] = (np.exp(np.outer(tk, rates)) @ coef).real out[late] = res return out