# SPDX-License-Identifier: AGPL-3.0-or-later
# Copyright (C) 2025 SWGY, Inc
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""
NVIDIA Warp device core for the blast ray-tracer.
One kernel thread per (explosion, ray). Each thread generates its ray from the
compact per-explosion frame + az/el offsets, marches up to `max_bounces`
reflections with an early-exit `break`, queries the armor meshes with Warp's
BVH `wp.mesh_query_ray` (nearest hit, double-sided), tests the sensor sphere,
and atomic-adds its solid-angle-weighted impulse to the per-explosion
accumulator. Mesh queries are float32; the impulse sum is float64.
"""
import logging
import numpy as np
import warp as wp
_T_MIN = 1e-6 # self-intersection offset for reflected rays (metres)
_initialized = False
def configure(device: str = "cuda") -> str:
"""Initialise Warp and resolve the compute device (falls back to CPU)."""
global _initialized
if not _initialized:
wp.init()
_initialized = True
if device == "cuda" and not wp.get_cuda_devices():
logging.warning("No CUDA device found; falling back to Warp CPU backend.")
return "cpu"
return device
def build_mesh(vertices, faces, device):
"""Build a Warp BVH mesh from trimesh-style vertices (V,3) + faces (F,3)."""
pts = wp.array(np.ascontiguousarray(vertices, dtype=np.float32),
dtype=wp.vec3, device=device)
idx = wp.array(np.ascontiguousarray(faces.reshape(-1), dtype=np.int32),
dtype=wp.int32, device=device)
return wp.Mesh(points=pts, indices=idx)
@wp.kernel
def _trace(
frames: wp.array(dtype=wp.mat33),
blast_os: wp.array(dtype=wp.vec3),
az_off: wp.array2d(dtype=wp.float32),
el_off: wp.array2d(dtype=wp.float32),
valid_az: wp.array2d(dtype=wp.int32),
valid_el: wp.array2d(dtype=wp.int32),
n_az: wp.int32,
mesh_a: wp.uint64,
mesh_b: wp.uint64,
n_mesh: wp.int32,
sensor_o: wp.vec3,
sensor_r: wp.float32,
p0: wp.float32,
k_const: wp.float32,
refl_loss: wp.float32,
res_rad: wp.float32,
max_len: wp.float32,
t_min: wp.float32,
max_bounces: wp.int32,
out: wp.array(dtype=wp.float64),
):
e, rr = wp.tid()
ie = rr // n_az
iaz = rr % n_az
if valid_el[e, ie] == 0 or valid_az[e, iaz] == 0:
return # padded / out-of-cone ray -> skip
f = frames[e]
fwd = wp.vec3(f[0, 0], f[0, 1], f[0, 2])
right = wp.vec3(f[1, 0], f[1, 1], f[1, 2])
up = wp.vec3(f[2, 0], f[2, 1], f[2, 2])
ela = el_off[e, ie]
aza = az_off[e, iaz]
ce = wp.cos(ela)
d = wp.normalize(fwd * (ce * wp.cos(aza)) + right * (ce * wp.sin(aza)) + up * wp.sin(ela))
o = blast_os[e]
dist = wp.float32(0.0)
refl = wp.float32(1.0)
impulse = wp.float32(0.0)
hit_sensor = wp.int32(0)
b = wp.int32(0)
while b <= max_bounces:
remaining = max_len - dist
if remaining <= 0.0:
break
# nearest geometry hit across this config's mesh(es); best_t bounds the search
best_t = remaining
got_geom = wp.int32(0)
gn = wp.vec3(0.0, 0.0, 0.0)
if n_mesh >= 1:
qa = wp.mesh_query_ray(mesh_a, o, d, best_t)
if qa.result:
best_t = qa.t
got_geom = wp.int32(1)
gn = qa.normal
if n_mesh == 2:
qb = wp.mesh_query_ray(mesh_b, o, d, best_t)
if qb.result:
best_t = qb.t
got_geom = wp.int32(1)
gn = qb.normal
# sensor sphere along the segment [o, o + d*remaining]; nearest in-range root
oc = o - sensor_o
bb = wp.dot(oc, d)
cc = wp.dot(oc, oc) - sensor_r * sensor_r
disc = bb * bb - cc
ds = wp.float32(0.0)
got_sensor = wp.int32(0)
if disc >= 0.0:
sq = wp.sqrt(disc)
t0 = -bb - sq
t1 = -bb + sq
if t0 >= 0.0 and t0 <= remaining:
ds = t0
got_sensor = wp.int32(1)
elif t1 >= 0.0 and t1 <= remaining:
ds = t1
got_sensor = wp.int32(1)
# sensor hit before geometry (ties -> geometry, strict <)
if got_sensor == 1 and (got_geom == 0 or ds < best_t):
eff = wp.max(dist + ds, wp.float32(1.0e-6))
impulse = refl * p0 / (eff * eff) * k_const
hit_sensor = wp.int32(1)
break
if got_geom == 0:
break
# advance to the surface, offset off it (self-intersection), reflect
o = o + d * best_t + d * t_min
dist = dist + best_t
nrm = wp.normalize(gn) # reflection is sign-invariant in the normal
d = wp.normalize(d - 2.0 * wp.dot(d, nrm) * nrm)
refl = refl * refl_loss
b = b + 1
if hit_sensor == 1:
d_omega = ce * res_rad * res_rad
wp.atomic_add(out, e, wp.float64(impulse) * wp.float64(d_omega))
_INV_4PI = 0.07957747154594767 # 1/(4 pi)
@wp.func
def _occluded(mesh_a: wp.uint64, mesh_b: wp.uint64, n_mesh: wp.int32,
P: wp.vec3, Q: wp.vec3, t_min: wp.float32) -> wp.int32:
"""1 if geometry blocks the open segment (P, Q); endpoints excluded by t_min."""
d = Q - P
dist = wp.length(d)
if dist <= 2.0 * t_min:
return wp.int32(0)
dir = d / dist
o = P + dir * t_min
lim = dist - 2.0 * t_min
if n_mesh >= 1:
qa = wp.mesh_query_ray(mesh_a, o, dir, lim)
if qa.result:
return wp.int32(1)
if n_mesh == 2:
qb = wp.mesh_query_ray(mesh_b, o, dir, lim)
if qb.result:
return wp.int32(1)
return wp.int32(0)
@wp.kernel
def _diffract(
blast_os: wp.array(dtype=wp.vec3),
apex: wp.array(dtype=wp.vec3),
tangent: wp.array(dtype=wp.vec3),
n_ref: wp.array(dtype=wp.vec3),
seg_nu: wp.array(dtype=wp.float32),
seg_len: wp.array(dtype=wp.float32),
coupling: wp.array(dtype=wp.float64), # S_fan(e), the direct-path fan weight
mesh_a: wp.uint64,
mesh_b: wp.uint64,
n_mesh: wp.int32,
sensor_o: wp.vec3,
p0: wp.float32,
k_const: wp.float32,
t_min: wp.float32,
out: wp.array(dtype=wp.float64),
):
# One thread per (explosion, edge segment). Adds the segment's diffracted
# impulse in the engine's units, matched to the direct path so GO+diffraction
# is continuous across the shadow boundary:
# dI = (-nu/4pi) * S_fan(e) * (p0 k / R_SR) * (beta/(m l)) * seg_len
e, i = wp.tid()
S = blast_os[e]
A = apex[i]
t = tangent[i]
e0 = n_ref[i]
e1 = wp.cross(t, e0)
dS = S - A
zS = wp.dot(dS, t)
perpS = dS - zS * t
rS = wp.length(perpS)
dR = sensor_o - A
zR = wp.dot(dR, t)
perpR = dR - zR * t
rR = wp.length(perpR)
if rS < t_min or rR < t_min:
return # source or receiver on the edge line (measure zero)
# azimuth about the edge from the reference face, wrapped to [0, 2pi) so the
# wedge coordinate is consistent (a 2pi shift flips beta's sign for nu=1/2)
thS = wp.atan2(wp.dot(perpS, e1), wp.dot(perpS, e0))
thR = wp.atan2(wp.dot(perpR, e1), wp.dot(perpR, e0))
if thS < 0.0:
thS = thS + 6.28318530717959
if thR < 0.0:
thR = thR + 6.28318530717959
m = wp.length(dS)
l = wp.length(dR)
r_sr = wp.length(S - sensor_o)
cosh_eta = (m * l + zS * zR) / (rS * rR)
if cosh_eta < 1.0:
cosh_eta = 1.0
eta = wp.log(cosh_eta + wp.sqrt(cosh_eta * cosh_eta - 1.0))
nu = seg_nu[i]
en = wp.exp(nu * eta)
cnu = 0.5 * (en + 1.0 / en) # cosh(nu * eta)
beta = wp.float32(0.0)
phi1 = 3.14159265358979 + thS + thR
phi2 = 3.14159265358979 + thS - thR
phi3 = 3.14159265358979 - thS + thR
phi4 = 3.14159265358979 - thS - thR
beta += wp.sin(nu * phi1) / (cnu - wp.cos(nu * phi1))
beta += wp.sin(nu * phi2) / (cnu - wp.cos(nu * phi2))
beta += wp.sin(nu * phi3) / (cnu - wp.cos(nu * phi3))
beta += wp.sin(nu * phi4) / (cnu - wp.cos(nu * phi4))
val = (-nu * _INV_4PI) * (p0 * k_const / r_sr) * (beta / (m * l)) * seg_len[i]
# diffracted path is blast -> apex -> sensor; both legs must be unblocked
if _occluded(mesh_a, mesh_b, n_mesh, S, A, t_min) == 1:
return
if _occluded(mesh_a, mesh_b, n_mesh, A, sensor_o, t_min) == 1:
return
wp.atomic_add(out, e, wp.float64(val) * coupling[e])
def simulate(cfg, helmet, vest, back, sensor, grid, k_const, device,
sensor_coupling=None, edge_sets=None):
"""Run all five armor configs. Returns {config_name: np.ndarray(E) float64}.
When ``cfg.diffraction.enabled`` and ``sensor_coupling``/``edge_sets`` are
provided, a second (``_diffract``) launch adds first-order edge diffraction
into the same per-explosion accumulator.
"""
e_total = grid.blast_os.shape[0]
gel = grid.el_off.shape[1]
gaz = grid.az_off.shape[1]
def dev(a, dt):
return wp.array(np.ascontiguousarray(a), dtype=dt, device=device)
frames = dev(grid.frames.astype(np.float32), wp.mat33)
blast_os = dev(grid.blast_os.astype(np.float32), wp.vec3)
az = dev(grid.az_off.astype(np.float32), wp.float32)
el = dev(grid.el_off.astype(np.float32), wp.float32)
vaz = dev(grid.valid_az.astype(np.int32), wp.int32)
vel = dev(grid.valid_el.astype(np.int32), wp.int32)
helmet_mesh = build_mesh(helmet.vertices, helmet.faces, device)
vest_mesh = build_mesh(vest.vertices, vest.faces, device)
# Front + back plates merged into one BVH, so full_armor (helmet + both
# plates) still fits the kernel's two-mesh limit.
plates_mesh = build_mesh(
np.concatenate([np.asarray(vest.vertices), np.asarray(back.vertices)]),
np.concatenate([np.asarray(vest.faces),
np.asarray(back.faces) + len(vest.vertices)]),
device)
sensor_o = wp.vec3(float(sensor.origin[0]), float(sensor.origin[1]), float(sensor.origin[2]))
# (mesh_a, mesh_b, n_mesh) per config; no_armor queries nothing (direct only)
plan = {
"full_armor": (helmet_mesh.id, plates_mesh.id, 2),
"front_plate_helmet": (helmet_mesh.id, vest_mesh.id, 2),
"helmet_only": (helmet_mesh.id, helmet_mesh.id, 1),
"vest_only": (vest_mesh.id, vest_mesh.id, 1),
"no_armor": (helmet_mesh.id, helmet_mesh.id, 0),
}
diffract_on = (cfg.diffraction.enabled
and sensor_coupling is not None and edge_sets is not None)
if diffract_on:
coupling_d = dev(np.asarray(sensor_coupling, np.float64), wp.float64)
edev = {}
for name, ea in edge_sets.items():
if ea is None or ea.apex.shape[0] == 0:
edev[name] = None
continue
edev[name] = (
dev(ea.apex.astype(np.float32), wp.vec3),
dev(ea.tangent.astype(np.float32), wp.vec3),
dev(ea.n_ref.astype(np.float32), wp.vec3),
dev(ea.nu.astype(np.float32), wp.float32),
dev(ea.seg_len.astype(np.float32), wp.float32),
int(ea.apex.shape[0]),
)
results = {}
for name, (ma, mb, nm) in plan.items():
out = wp.zeros(e_total, dtype=wp.float64, device=device)
wp.launch(_trace, dim=(e_total, gel * gaz), device=device, inputs=[
frames, blast_os, az, el, vaz, vel, gaz,
wp.uint64(ma), wp.uint64(mb), nm,
sensor_o, float(sensor.radius),
float(cfg.physics.initial_peak_pressure_kpa), float(k_const),
float(cfg.physics.reflection_loss), float(grid.res_rad),
float(cfg.rays.max_length_m), float(_T_MIN), int(cfg.rays.max_bounces),
out,
])
if diffract_on and edev.get(name) is not None:
ap, tg, nr, nu_a, sl, nseg = edev[name]
wp.launch(_diffract, dim=(e_total, nseg), device=device, inputs=[
blast_os, ap, tg, nr, nu_a, sl, coupling_d,
wp.uint64(ma), wp.uint64(mb), nm, sensor_o,
float(cfg.physics.initial_peak_pressure_kpa), float(k_const),
float(_T_MIN), out,
])
wp.synchronize()
combined = out.numpy()
if diffract_on and edev.get(name) is not None:
# First-order diffraction is signed; where the discrete ray-traced GO
# under-resolves a near-zero-field transition cell, the sum can dip
# slightly negative. Impulse is nonnegative -> clamp, and report it
# (the effect shrinks with finer rays.resolution_arcmin).
n_clamped = int((combined < 0.0).sum())
if n_clamped:
logging.info("diffraction: clamped %d/%d negative transition "
"cells for %s", n_clamped, combined.size, name)
combined = np.maximum(combined, 0.0)
results[name] = combined
logging.info("Finished config: %s", name)
return results