simulate_blasts.py

#!/usr/bin/env python
# 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/>.
"""
Entrypoint for the Warp blast-pressure simulation.

Massively parallel explosion sampling on the GPU: for each blast point in a
configurable cube, one Warp thread per ray traces through five armor
configurations (helmet + both plates / helmet + front plate / helmet only /
front plate only / none) and the solid-angle-weighted total impulse at the
sensor is written to CSV.

    uv run simulate_blasts.py --config config.yaml
    uv run simulate_blasts.py --config config.yaml --set blast.cube_segments=12
"""

import argparse
import csv
import logging
from datetime import datetime

import numpy as np

import config as config_mod
import diffraction
import geometry
import sampling
import warp_core

# Armor configurations, in CSV column order. "front_plate_helmet" (helmet +
# front plate) is what "full_armor" meant before the back plate existed.
CONFIGS = ["full_armor", "front_plate_helmet", "helmet_only", "vest_only",
           "no_armor"]
COLUMNS = ["bp_x", "bp_y", "bp_z", *CONFIGS]


def parse_arguments():
    p = argparse.ArgumentParser(description="Simulate blast pressure (Warp backend).")
    p.add_argument("--config", type=str, default="config.yaml",
                   help="Path to the YAML configuration file. Default: config.yaml")
    p.add_argument("--set", dest="overrides", action="append", default=[],
                   metavar="KEY.PATH=VALUE",
                   help="Override a single config key (repeatable), "
                        "e.g. --set blast.cube_segments=12")
    p.add_argument("--output-file", type=str, default=None,
                   help="Override output.file from the config.")
    p.add_argument("--verbose", action="store_true",
                   help="Enable DEBUG-level logging.")
    return p.parse_args()


def run(cfg):
    device = warp_core.configure(cfg.compute.device)
    logging.info("Warp device: %s", device)

    helmet = geometry.load_helmet(cfg.geometry.helmet)
    vest = geometry.load_vest(cfg.geometry.vest)
    back = geometry.load_back_plate(cfg.geometry.vest)
    sensor = sampling.Sensor(origin=np.zeros(3), radius=cfg.sensor.radius_m)

    vertices_all = None
    if cfg.rays.fan_mode == "geometry":
        vertices_all = np.concatenate(
            [np.asarray(helmet.vertices, dtype=np.float64),
             np.asarray(vest.vertices, dtype=np.float64),
             np.asarray(back.vertices, dtype=np.float64)], axis=0)

    grid = sampling.build_grid(cfg, sensor, vertices_all)
    e_total = grid.blast_os.shape[0]
    logging.info("%d blast points survive the min-distance filter; "
                 "grid=%dx%d rays/explosion (device=%s).",
                 e_total, grid.az_off.shape[1], grid.el_off.shape[1], device)

    # Warn if the sensor cone is under-resolved (its direct hit can fall between
    # rays and read 0). Real runs want a finer rays.resolution_arcmin.
    rays_axis = grid.valid_az.sum(axis=1)
    under = int((rays_axis < 3).sum())
    if under:
        logging.warning("%d/%d explosions have < 3 rays/axis; the sensor cone is "
                        "under-resolved (no_armor may read 0). Consider lowering "
                        "rays.resolution_arcmin.", under, e_total)

    k_const = sampling.precompute_waveform_constant(
        cfg.physics.positive_phase_duration_s, cfg.physics.decay_alpha)

    sensor_coupling, edge_sets = None, None
    if cfg.diffraction.enabled:
        sensor_coupling = diffraction.build_sensor_coupling(grid)
        edge_sets = diffraction.build_edge_sets(cfg, helmet, vest, back)
        n_seg = {k: (0 if v is None else v.apex.shape[0]) for k, v in edge_sets.items()}
        logging.info("Edge diffraction on (mode=%s): segments/config %s",
                     cfg.diffraction.edge_mode, n_seg)

    results = warp_core.simulate(cfg, helmet, vest, back, sensor, grid, k_const,
                                 device, sensor_coupling, edge_sets)

    rows = [
        (float(grid.blast_os[e, 0]), float(grid.blast_os[e, 1]), float(grid.blast_os[e, 2]),
         *(float(results[c][e]) for c in CONFIGS))
        for e in range(e_total)
    ]
    rows.sort(key=lambda r: (r[0], r[1], r[2]))
    return rows


def write_csv(path, rows):
    with open(path, mode="w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(COLUMNS)
        for row in rows:
            w.writerow([f"{row[0]:.2f}", f"{row[1]:.2f}", f"{row[2]:.2f}",
                        *(f"{v:.6f}" for v in row[3:])])


def main():
    args = parse_arguments()
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s - %(levelname)s - %(message)s")

    cfg = config_mod.load_config(args.config, args.overrides)
    out_path = args.output_file or cfg.output.file or \
        f"./blast-results-{datetime.now():%Y-%m-%d-%H%M}.csv"

    logging.info("Config: %s  ->  output: %s", args.config, out_path)
    logging.info("Cube: center=%s extents=%sm segments=%s | rays: %s arcmin, "
                 "max_bounces=%s, fan=%s",
                 cfg.blast.cube_center_m, cfg.blast.cube_extents_m,
                 cfg.blast.cube_segments, cfg.rays.resolution_arcmin,
                 cfg.rays.max_bounces, cfg.rays.fan_mode)

    rows = run(cfg)
    write_csv(out_path, rows)
    logging.info("Wrote %d rows to %s", len(rows), out_path)


if __name__ == "__main__":
    main()