GPU path for the tensor assembly and time stepping (PyTorch); record the hardware per solution
3 changed files+134−29
solvers/vonkarman-modal-disk/__pycache__/bfs_disk.cpython-313.pycadded+0−0View file
Binary file not shown.
solvers/vonkarman-modal-disk/make_submission.pymodified+1−1View file
@@ -25,7 +25,7 @@ def main():
2525 "parameters": m["parameters"],
2626 "timing_seconds": m["timing_seconds"],
2727 "more_accurate_than": [prev] if prev else [],
28- "hardware": HARDWARE,
28+ "hardware": m.get("hardware", HARDWARE),
2929 }
3030 if "fields" in m:
3131 entry["fields"] = m["fields"]
solvers/vonkarman-modal-disk/vk_modal_disk.pymodified+133−28View file
@@ -27,12 +27,16 @@ Because the strike is centered on the x axis, the solution is even in y, and by
2727 modes needed for a given frequency reach (the coupling cost scales with the cube of the count). --full
2828 uses the unrestricted bases instead.
2929
30-Usage: python vk_modal_disk.py <challenge.json> <out-dir> [--fields tag,...] [--force F] [--seconds S] [--full] [tag ...]
30+With --device cuda the coupling tensor is assembled and the time stepping runs on a GPU through PyTorch
31+(the modes and quadrature stay in numpy/scipy); the CPU path is the reference and the two agree to round-off.
32+
33+Usage: python vk_modal_disk.py <challenge.json> <out-dir> [--fields tag,...] [--force F] [--seconds S] [--full] [--device cuda] [tag ...]
3134 where a tag is g<grid spacing in mm>-w<transverse modes>-a<airy modes>-m<oversampling>, e.g. g8-w200-a200-m1
3235 """
3336 from __future__ import annotations
3437
3538 import json
39+import os
3640 import sys
3741 import time
3842 from pathlib import Path
@@ -44,8 +48,14 @@ from bfs_disk import BFSDisk, load_challenge # noqa: E402
4448
4549
4650 class VonKarmanDisk:
47- def __init__(self, ch: dict, spacing: float, n_w: int, n_airy: int, force_override: float | None = None, log=None, symmetric_y: bool | None = None):
51+ def __init__(self, ch: dict, spacing: float, n_w: int, n_airy: int, force_override: float | None = None, log=None, symmetric_y: bool | None = None, device: str = "cpu"):
4852 self.ch = ch
53+ self.device = device
54+ self.torch = None
55+ if device != "cpu":
56+ import torch
57+
58+ self.torch = torch
4959 p = dict(ch["parameters"])
5060 if force_override is not None:
5161 p["strike_peak_force"] = force_override
@@ -80,15 +90,33 @@ class VonKarmanDisk:
8090 Pxx, Pyy, Pxy = Bxx @ V, Byy @ V, Bxy @ V # (Q, n_w)
8191 Bn = pl.basis_at(xq, yq)
8292 PsiQ = Bn @ Psi # (Q, n_airy)
83- H = np.zeros((self.n_airy, self.n_w, self.n_w))
84- for pi in range(self.n_airy):
85- wp = (WQ * PsiQ[:, pi])[:, None]
86- A = (Pxx * wp).T @ Pyy
87- H[pi] = A + A.T - 2.0 * (Pxy * wp).T @ Pxy
88- self.H = H
89- self.Hflat = H.reshape(self.n_airy * self.n_w, self.n_w)
9093 self.coef = -(p["E"] * p["h"]) / (2.0 * self.zeta4)
91- self.log(f" coupling tensor {H.shape} in {time.perf_counter() - t0:.1f}s")
94+ if self.torch is None:
95+ H = np.zeros((self.n_airy, self.n_w, self.n_w))
96+ for pi in range(self.n_airy):
97+ wp = (WQ * PsiQ[:, pi])[:, None]
98+ A = (Pxx * wp).T @ Pyy
99+ H[pi] = A + A.T - 2.0 * (Pxy * wp).T @ Pxy
100+ self.H = H
101+ self.Hflat = H.reshape(self.n_airy * self.n_w, self.n_w)
102+ shape = H.shape
103+ else:
104+ # the same assembly on the device; the tensor stays there for the time stepping
105+ t = self.torch
106+ dev = t.device(self.device)
107+ Pxx_t, Pyy_t, Pxy_t = (t.as_tensor(a, device=dev) for a in (Pxx, Pyy, Pxy))
108+ WPsi = t.as_tensor(WQ[:, None] * PsiQ, device=dev) # (Q, n_airy)
109+ H = t.empty((self.n_airy, self.n_w, self.n_w), dtype=t.float64, device=dev)
110+ for pi in range(self.n_airy):
111+ wp = WPsi[:, pi : pi + 1]
112+ A = (Pxx_t * wp).T @ Pyy_t
113+ H[pi] = A + A.T - 2.0 * (Pxy_t * wp).T @ Pxy_t
114+ t.cuda.synchronize() if self.device.startswith("cuda") else None
115+ self.H = None
116+ self.Hflat_t = H.reshape(self.n_airy * self.n_w, self.n_w)
117+ shape = tuple(H.shape)
118+ del Pxx_t, Pyy_t, Pxy_t, WPsi
119+ self.log(f" coupling tensor {shape} on {self.device} in {time.perf_counter() - t0:.1f}s")
92120 # Rayleigh integral: quadrature weights, delays, modal shapes at the points
93121 _, _, self.tau, self.cq = pl.listener_quadrature(3)
94122 self.PhiQ = Bn @ V # (Q, n_w)
@@ -98,12 +126,10 @@ class VonKarmanDisk:
98126 eta = self.coef * (v @ q) # eta_p
99127 return v.T @ eta # F_k = sum_p v[p,k] eta_p
100128
101- def run(self, fs: int, nsamp: int, m: int, field_grid=None, field_rate: float | None = None):
129+ def _prepare(self, fs: int, nsamp: int, m: int, field_grid, field_rate):
102130 p = self.p
103131 dt = 1.0 / (fs * m)
104- T, Om = p["strike_duration"], 2 * np.pi / p["strike_duration"]
105- nsteps = nsamp * m
106- w2, al, f = self.omega**2, self.alpha, self.f
132+ w2, al = self.omega**2, self.alpha
107133 # exact two-step recurrence of the damped linear oscillator: q+ = a1 q + a2 q- + b F
108134 rigid = self.omega == 0.0
109135 disc = np.sqrt(al**2 - w2 + 0j)
@@ -134,18 +160,26 @@ class VonKarmanDisk:
134160 for u in np.unique(dd):
135161 sel = dd == u
136162 G[u] += (self.cq[sel] * lag[sel, j]) @ self.PhiQ[sel]
137- ring = np.zeros((dmax + 1, self.n_w))
138- out = np.zeros(nsamp)
139- q_prev = np.zeros(self.n_w)
140- q = np.zeros(self.n_w)
141- frames = None
163+ fld = None
142164 if field_grid is not None:
143165 fx, fy, ft = field_grid
144166 XX, YY = np.meshgrid(fx, fy, indexing="xy")
145167 inside = self.plate.inside(XX.ravel(), YY.ravel())
146168 PhiG = self.plate.basis_at(XX.ravel(), YY.ravel()) @ self.plate.V
147- frame_every = int(round(fs * m / field_rate))
148- frames = []
169+ fld = dict(PhiG=PhiG, inside=inside, every=int(round(fs * m / field_rate)), nt=len(ft), ny=len(fy), nx=len(fx))
170+ return dict(dt=dt, T=p["strike_duration"], Om=2 * np.pi / p["strike_duration"], nsteps=nsamp * m, a1=a1, a2=a2, b=b, w2=w2, al=al, G=G, dmax=dmax, fld=fld)
171+
172+ def run(self, fs: int, nsamp: int, m: int, field_grid=None, field_rate: float | None = None):
173+ pr = self._prepare(fs, nsamp, m, field_grid, field_rate)
174+ if self.torch is not None:
175+ return self._run_device(pr, m, nsamp)
176+ dt, T, Om, nsteps, a1, a2, b, w2, al, G, dmax, fld = (pr[k] for k in ("dt", "T", "Om", "nsteps", "a1", "a2", "b", "w2", "al", "G", "dmax", "fld"))
177+ f = self.f
178+ ring = np.zeros((dmax + 1, self.n_w))
179+ out = np.zeros(nsamp)
180+ q_prev = np.zeros(self.n_w)
181+ q = np.zeros(self.n_w)
182+ frames = [] if fld else None
149183 t_start = time.perf_counter()
150184 report = max(1, nsteps // 10)
151185 for nstep in range(nsteps + 1):
@@ -160,9 +194,9 @@ class VonKarmanDisk:
160194 if k < nsamp:
161195 idx = (nstep - np.arange(dmax + 1)) % (dmax + 1)
162196 out[k] = np.sum(G * ring[idx])
163- if frames is not None and nstep % frame_every == 0 and len(frames) < len(ft):
164- fr_ = (PhiG @ q).astype(np.float32)
165- fr_[~inside] = np.nan
197+ if frames is not None and nstep % fld["every"] == 0 and len(frames) < fld["nt"]:
198+ fr_ = (fld["PhiG"] @ q).astype(np.float32)
199+ fr_[~fld["inside"]] = np.nan
166200 frames.append(fr_)
167201 q_prev, q = q, q_next
168202 if nstep % report == 0 and nstep:
@@ -172,11 +206,80 @@ class VonKarmanDisk:
172206 self.log(f" step {nstep}/{nsteps} t={t:.2f}s elapsed={el:.0f}s eta={el * (nsteps - nstep) / nstep:.0f}s max|q|={np.abs(q).max():.3e}")
173207 if not np.all(np.isfinite(out)):
174208 raise RuntimeError("output is not finite; increase the oversampling m")
175- fields = None
176- if frames is not None:
177- fields = np.stack(frames).reshape(len(frames), len(field_grid[1]), len(field_grid[0]))
209+ fields = np.stack(frames).reshape(len(frames), fld["ny"], fld["nx"]) if frames is not None else None
178210 return out, fields
179211
212+ def _run_device(self, pr, m: int, nsamp: int):
213+ """The same loop with the tensor contraction and the recurrence on the device (PyTorch)."""
214+ t = self.torch
215+ dev = t.device(self.device)
216+ T64 = lambda a: t.as_tensor(np.asarray(a, dtype=np.float64), device=dev)
217+ dt, T, Om, nsteps, dmax, fld = pr["dt"], pr["T"], pr["Om"], pr["nsteps"], pr["dmax"], pr["fld"]
218+ a1, a2, b, w2, al, G, f, coef = (T64(x) for x in (pr["a1"], pr["a2"], pr["b"], pr["w2"], pr["al"], pr["G"], self.f, self.coef))
219+ Hf = self.Hflat_t
220+ L = dmax + 1
221+ ar = t.arange(L, device=dev)
222+ ring = t.zeros((L, self.n_w), dtype=t.float64, device=dev)
223+ out = t.zeros(nsamp, dtype=t.float64, device=dev)
224+ q = t.zeros(self.n_w, dtype=t.float64, device=dev)
225+ q_prev = t.zeros_like(q)
226+ frames = [] if fld else None
227+ PhiG = T64(fld["PhiG"]) if fld else None
228+ inv2dt = 1.0 / (2.0 * dt)
229+ t_start = time.perf_counter()
230+ report = max(1, nsteps // 10)
231+ for nstep in range(nsteps + 1):
232+ tt = nstep * dt
233+ psi = 0.5 * (1.0 - np.cos(Om * tt)) if 0.0 <= tt <= T else 0.0
234+ v = (Hf @ q).reshape(self.n_airy, self.n_w)
235+ eta = coef * (v @ q)
236+ Fext = v.T @ eta
237+ if psi:
238+ Fext = Fext + f * psi
239+ q_next = a1 * q + a2 * q_prev + b * Fext
240+ qdd = Fext - w2 * q - al * ((q_next - q_prev) * inv2dt)
241+ ring[nstep % L] = qdd
242+ if nstep % m == 0:
243+ k = nstep // m
244+ if k < nsamp:
245+ idx = (nstep - ar) % L
246+ out[k] = t.sum(G * ring[idx])
247+ if frames is not None and nstep % fld["every"] == 0 and len(frames) < fld["nt"]:
248+ fr_ = (PhiG @ q).cpu().numpy().astype(np.float32)
249+ fr_[~fld["inside"]] = np.nan
250+ frames.append(fr_)
251+ q_prev, q = q, q_next
252+ if nstep % report == 0 and nstep:
253+ qc = q.cpu().numpy()
254+ if not np.all(np.isfinite(qc)):
255+ raise RuntimeError(f"solution diverged at step {nstep} (t = {tt:.4f} s); increase the oversampling m")
256+ el = time.perf_counter() - t_start
257+ self.log(f" step {nstep}/{nsteps} t={tt:.2f}s elapsed={el:.0f}s eta={el * (nsteps - nstep) / nstep:.0f}s max|q|={np.abs(qc).max():.3e}")
258+ out = out.cpu().numpy()
259+ if not np.all(np.isfinite(out)):
260+ raise RuntimeError("output is not finite; increase the oversampling m")
261+ fields = np.stack(frames).reshape(len(frames), fld["ny"], fld["nx"]) if frames is not None else None
262+ return out, fields
263+
264+
265+def hardware_string(device: str) -> str:
266+ cpu = "unknown CPU"
267+ try:
268+ for line in open("/proc/cpuinfo"):
269+ if line.startswith("model name"):
270+ cpu = " ".join(line.split(":", 1)[1].split())
271+ break
272+ except OSError:
273+ pass
274+ s = f"{cpu} ({os.cpu_count()} threads), numpy/scipy for the setup"
275+ if device != "cpu":
276+ import torch
277+
278+ s += f"; coupling tensor and time stepping on {torch.cuda.get_device_name(0)} (PyTorch {torch.__version__})"
279+ else:
280+ s += " and the time stepping"
281+ return s
282+
180283
181284 def parse_tag(tag: str):
182285 parts = dict((kv[0], int(kv[1:])) for kv in tag.split("-"))
@@ -204,6 +307,7 @@ def main():
204307 fields_for = v or set()
205308 force = take("--force", float)
206309 seconds = take("--seconds", float)
310+ device = take("--device", str) or "cpu"
207311 full = "--full" in args
208312 if full:
209313 args.remove("--full")
@@ -217,7 +321,7 @@ def main():
217321 spacing, n_w, n_a, m = parse_tag(tag)
218322 log(f"{tag}: force={force if force is not None else ch['parameters']['strike_peak_force']} N")
219323 t0 = time.perf_counter()
220- model = VonKarmanDisk(ch, spacing, n_w, n_a, force, log=log, symmetric_y=False if full else None)
324+ model = VonKarmanDisk(ch, spacing, n_w, n_a, force, log=log, symmetric_y=False if full else None, device=device)
221325 t1 = time.perf_counter()
222326 grid = None
223327 if tag in fields_for:
@@ -246,6 +350,7 @@ def main():
246350 "highest_transverse_mode_hz": round(float(model.omega.max() / (2 * np.pi)), 1),
247351 },
248352 "timing_seconds": round(dt_total, 3),
353+ "hardware": hardware_string(device),
249354 "peak_si": float(np.max(np.abs(s))),
250355 "setup_seconds": round(t1 - t0, 3),
251356 }