From 5f4416384bbfa1b527e14d4e4f89e918feb78a8b Mon Sep 17 00:00:00 2001 From: andytorrestb Date: Sat, 18 Jul 2026 16:33:45 -0500 Subject: [PATCH 1/7] vectorize face-level plume strike detection Replace the per-face Python loop in compute_plume_strikes() with NumPy operations over all faces per active thruster. The original loop is preserved verbatim as _compute_plume_strikes_scalar() for regression comparison, debugging, and benchmarking. The vectorized core operates on plain serializable inputs (centroids, normals, thruster dicts, config values) so it can later run in process workers. Gas kinetics stays scalar, evaluated only for struck faces. Zero-distance faces and the legacy 3.14 theta constant are preserved so strike counts and struck-face IDs match the scalar path bit-for-bit (verified on all case/rpod/* cases, kinetics arrays exactly equal). Co-Authored-By: Claude Fable 5 --- pyrpod/plume/PlumeStrikeCalculator.py | 472 ++++++++++++++++++-------- 1 file changed, 335 insertions(+), 137 deletions(-) diff --git a/pyrpod/plume/PlumeStrikeCalculator.py b/pyrpod/plume/PlumeStrikeCalculator.py index 241ed6f..eeb12f8 100644 --- a/pyrpod/plume/PlumeStrikeCalculator.py +++ b/pyrpod/plume/PlumeStrikeCalculator.py @@ -1,137 +1,335 @@ -""" -Plume impingement computations for RPOD. - -Responsibilities: -- Given target mesh, VV pose, and active thrusters, compute per-face strike metrics -- Return numpy arrays/dicts; do not write files - -This consolidates logic currently in RPOD.jfh_plume_strikes into -reusable, testable functions. -""" -from __future__ import annotations - -from typing import Any, Dict, Tuple -import numpy as np -from pyrpod.plume.RarefiedPlumeGasKinetics import SimplifiedGasKinetics - - -def compute_plume_strikes( - target_mesh: Any, - target_unit_normals: np.ndarray, - vv: Any, - jfh_step: Dict[str, Any], - environment: Any, -) -> Dict[str, np.ndarray]: - """Compute plume strike arrays for a single JFH step. - - Inputs - - target_mesh: numpy-stl Mesh-like, exposes .vectors (N x 3 x 3) - - target_unit_normals: (N x 3) array of per-face unit normals - - vv: Visiting vehicle with thruster_data and thruster_metrics - - jfh_step: dict with keys 'thrusters' (list[int]), 'xyz' (pos), 'dcm' (3x3) - - environment: provides config for plume and kinetics - - Returns - - dict with per-face arrays for current step: strikes and optionally pressures, shear_stress, heat_flux_rate, heat_flux_load - """ - num_faces = len(target_mesh.vectors) - strikes = np.zeros(num_faces) - - use_kinetics = environment.config['pm']['kinetics'] != 'None' - if use_kinetics: - pressures = np.zeros(num_faces) - shear_stresses = np.zeros(num_faces) - heat_flux = np.zeros(num_faces) - heat_flux_load = np.zeros(num_faces) - - vv_pos = np.array(jfh_step['xyz']) - vv_orientation = np.array(jfh_step['dcm']).transpose() - thrusters = jfh_step['thrusters'] - firing_time = float(jfh_step['t']) if 't' in jfh_step else 0.0 - - # Build mapping from numeric JFH indices to thruster ids consistent with legacy - link = {} - i = 1 - for thruster in vv.thruster_data: - link[str(i)] = vv.thruster_data[thruster]['name'] - i += 1 - - plume_radius = float(environment.config['plume']['radius']) - wedge_theta = float(environment.config['plume']['wedge_theta']) - - for thr in thrusters: - thruster_id = link[str(thr)][0] - - thruster_orientation = np.array(vv.thruster_data[thruster_id]['dcm']).transpose() - thruster_orientation = thruster_orientation.dot(vv_orientation) - plume_normal = np.array(thruster_orientation[0]) - norm_plume_normal = np.linalg.norm(plume_normal) - unit_plume_normal = plume_normal / norm_plume_normal - - thr_exit = np.array(vv.thruster_data[thruster_id]['exit']) - thruster_pos = vv_pos + thr_exit - thruster_pos = thruster_pos[0] - - for idx, face in enumerate(target_mesh.vectors): - face = np.array(face).transpose() - centroid = np.array([face[0].mean(), face[1].mean(), face[2].mean()]) - distance = thruster_pos - centroid - norm_distance = np.linalg.norm(distance) - if norm_distance == 0: - continue - unit_distance = distance / norm_distance - - theta = 3.14 - np.arccos(np.dot(np.squeeze(unit_distance), np.squeeze(unit_plume_normal))) - - n = np.squeeze(target_unit_normals[idx]) - unit_plume = np.squeeze(plume_normal / norm_plume_normal) - surface_dot_plume = np.dot(n, unit_plume) - - within_distance = float(norm_distance) < plume_radius - within_theta = float(theta) < wedge_theta - facing_thruster = surface_dot_plume < 0 - - if within_distance and within_theta and facing_thruster: - strikes[idx] += 1 - if use_kinetics: - T_w = float(environment.config['tv']['surface_temp']) - sigma = float(environment.config['tv']['sigma']) - t_type = vv.thruster_data[thruster_id]['type'][0] - thruster_metrics = vv.thruster_metrics[t_type] - simple_plume = SimplifiedGasKinetics(norm_distance, theta, thruster_metrics, T_w, sigma) - pressures[idx] += simple_plume.get_pressure() - shear = simple_plume.get_shear_pressure() - shear_stresses[idx] += abs(shear) - hf = simple_plume.get_heat_flux() - heat_flux[idx] += hf - heat_flux_load[idx] += hf * firing_time - - result = {"strikes": strikes} - if use_kinetics: - result.update({ - "pressures": pressures, - "shear_stress": shear_stresses, - "heat_flux_rate": heat_flux, - "heat_flux_load": heat_flux_load, - }) - return result - - -def accumulate_cumulative( - cumulative: Dict[str, np.ndarray], - current: Dict[str, np.ndarray], -) -> Dict[str, np.ndarray]: - """Accumulate per-step arrays into cumulative tallies (e.g., cum_strikes, max_pressures).""" - if "cum_strikes" in cumulative and "strikes" in current: - cumulative["cum_strikes"] = cumulative["cum_strikes"] + current["strikes"] - - # Max trackers if available - if "max_pressures" in cumulative and "pressures" in current: - cumulative["max_pressures"] = np.maximum(cumulative["max_pressures"], current["pressures"]) - if "max_shears" in cumulative and "shear_stress" in current: - cumulative["max_shears"] = np.maximum(cumulative["max_shears"], current["shear_stress"]) - - if "cum_heat_flux_load" in cumulative and "heat_flux_load" in current: - cumulative["cum_heat_flux_load"] = cumulative["cum_heat_flux_load"] + current["heat_flux_load"] - - return cumulative +""" +Plume impingement computations for RPOD. + +Responsibilities: +- Given target mesh, VV pose, and active thrusters, compute per-face strike metrics +- Return numpy arrays/dicts; do not write files + +This consolidates logic currently in RPOD.jfh_plume_strikes into +reusable, testable functions. + +Implementation notes: +- compute_plume_strikes() runs a NumPy-vectorized strike-detection path by + default. _compute_plume_strikes_scalar() preserves the original per-face + loop verbatim as a reference implementation for tests and benchmarking. +- The vectorized core operates on plain serializable inputs (arrays, dicts, + floats) so it can also run inside process-based workers. + +Future work (no new dependencies planned): +- Vectorize the SimplifiedGasKinetics evaluations for struck faces. +- Shared-memory arrays (multiprocessing.shared_memory) for very large meshes. +- Chunking strategy to batch many small firings per worker task. +""" +from __future__ import annotations + +from typing import Any, Dict, Optional + +import numpy as np +from pyrpod.plume.RarefiedPlumeGasKinetics import SimplifiedGasKinetics + + +def compute_face_centroids(vectors: np.ndarray) -> np.ndarray: + """Compute per-face centroids for an (N x 3 x 3) array of face vertices. + + Averages the three vertices of each face, matching the scalar reference + (mean over each coordinate in the face's native dtype). The target is + stationary during a run, so callers should compute this once and pass it + to compute_plume_strikes() via face_centroids. + """ + return np.asarray(vectors).mean(axis=1) + + +def _build_thruster_link(thruster_data: Dict[str, Any]) -> Dict[str, Any]: + """Map numeric JFH thruster indices ('1', '2', ...) to thruster names, + consistent with legacy ordering of the thruster configuration.""" + link = {} + i = 1 + for thruster in thruster_data: + link[str(i)] = thruster_data[thruster]['name'] + i += 1 + return link + + +def extract_plume_params(environment: Any) -> Dict[str, Any]: + """Extract the plain config values needed for strike computation. + + Returns a picklable dict (radius, wedge_theta, use_kinetics, and — only + when kinetics is enabled — surface_temp and sigma) so workers never need + the full environment object. + """ + config = environment.config + use_kinetics = config['pm']['kinetics'] != 'None' + params: Dict[str, Any] = { + 'radius': float(config['plume']['radius']), + 'wedge_theta': float(config['plume']['wedge_theta']), + 'use_kinetics': use_kinetics, + 'surface_temp': None, + 'sigma': None, + } + if use_kinetics: + params['surface_temp'] = float(config['tv']['surface_temp']) + params['sigma'] = float(config['tv']['sigma']) + return params + + +def _compute_plume_strikes_core( + face_centroids: np.ndarray, + target_unit_normals: np.ndarray, + thruster_data: Dict[str, Any], + thruster_metrics: Optional[Dict[str, Any]], + jfh_step: Dict[str, Any], + plume_params: Dict[str, Any], +) -> Dict[str, np.ndarray]: + """Vectorized strike computation on plain serializable inputs. + + Geometry is evaluated with NumPy over all faces per active thruster. + Gas-kinetics quantities remain scalar: SimplifiedGasKinetics is + instantiated only for struck face indices, exactly as in the scalar + reference. Memory scales with the number of faces (a few (N,) and (N,3) + temporaries), independent of the number of firings. + """ + num_faces = len(face_centroids) + strikes = np.zeros(num_faces) + + use_kinetics = plume_params['use_kinetics'] + if use_kinetics: + pressures = np.zeros(num_faces) + shear_stresses = np.zeros(num_faces) + heat_flux = np.zeros(num_faces) + heat_flux_load = np.zeros(num_faces) + + vv_pos = np.array(jfh_step['xyz']) + vv_orientation = np.array(jfh_step['dcm']).transpose() + thrusters = jfh_step['thrusters'] + firing_time = float(jfh_step['t']) if 't' in jfh_step else 0.0 + + link = _build_thruster_link(thruster_data) + + plume_radius = float(plume_params['radius']) + wedge_theta = float(plume_params['wedge_theta']) + + normals = np.asarray(target_unit_normals) + + for thr in thrusters: + thruster_id = link[str(thr)][0] + + thruster_orientation = np.array(thruster_data[thruster_id]['dcm']).transpose() + thruster_orientation = thruster_orientation.dot(vv_orientation) + plume_normal = np.array(thruster_orientation[0]) + norm_plume_normal = np.linalg.norm(plume_normal) + unit_plume_normal = plume_normal / norm_plume_normal + + thr_exit = np.array(thruster_data[thruster_id]['exit']) + thruster_pos = vv_pos + thr_exit + thruster_pos = thruster_pos[0] + + distance = thruster_pos - face_centroids + norm_distance = np.linalg.norm(distance, axis=1) + + # Faces whose centroid coincides with the thruster exit are skipped, + # matching the scalar reference's `norm_distance == 0` guard. + valid = norm_distance != 0.0 + unit_distance = np.zeros_like(distance) + np.divide( + distance, + norm_distance[:, np.newaxis], + out=unit_distance, + where=valid[:, np.newaxis], + ) + + # NOTE: 3.14 (not np.pi) is kept deliberately to reproduce the legacy + # scalar reference bit-for-bit; changing it shifts theta by ~1.6e-3 rad + # and can alter struck-face IDs near the wedge boundary. + theta = 3.14 - np.arccos((unit_distance * unit_plume_normal).sum(axis=1)) + + surface_dot_plume = (normals * unit_plume_normal).sum(axis=1) + + hit = ( + valid + & (norm_distance < plume_radius) + & (theta < wedge_theta) + & (surface_dot_plume < 0) + ) + + strikes[hit] += 1 + + if use_kinetics: + T_w = plume_params['surface_temp'] + sigma = plume_params['sigma'] + t_type = thruster_data[thruster_id]['type'][0] + metrics = thruster_metrics[t_type] + for idx in np.nonzero(hit)[0]: + simple_plume = SimplifiedGasKinetics( + norm_distance[idx], theta[idx], metrics, T_w, sigma + ) + pressures[idx] += simple_plume.get_pressure() + shear = simple_plume.get_shear_pressure() + shear_stresses[idx] += abs(shear) + hf = simple_plume.get_heat_flux() + heat_flux[idx] += hf + heat_flux_load[idx] += hf * firing_time + + result = {"strikes": strikes} + if use_kinetics: + result.update({ + "pressures": pressures, + "shear_stress": shear_stresses, + "heat_flux_rate": heat_flux, + "heat_flux_load": heat_flux_load, + }) + return result + + +def compute_plume_strikes( + target_mesh: Any, + target_unit_normals: np.ndarray, + vv: Any, + jfh_step: Dict[str, Any], + environment: Any, + face_centroids: Optional[np.ndarray] = None, +) -> Dict[str, np.ndarray]: + """Compute plume strike arrays for a single JFH step. + + Inputs + - target_mesh: numpy-stl Mesh-like, exposes .vectors (N x 3 x 3) + - target_unit_normals: (N x 3) array of per-face unit normals + - vv: Visiting vehicle with thruster_data and thruster_metrics + - jfh_step: dict with keys 'thrusters' (list[int]), 'xyz' (pos), 'dcm' (3x3) + - environment: provides config for plume and kinetics + - face_centroids: optional (N x 3) precomputed face centroids + (see compute_face_centroids). When the target is stationary, callers + should compute centroids once per run and pass them here; if omitted, + they are computed from target_mesh for this step. + + Returns + - dict with per-face arrays for current step: strikes and optionally pressures, shear_stress, heat_flux_rate, heat_flux_load + """ + if face_centroids is None: + face_centroids = compute_face_centroids(target_mesh.vectors) + plume_params = extract_plume_params(environment) + return _compute_plume_strikes_core( + face_centroids=face_centroids, + target_unit_normals=target_unit_normals, + thruster_data=vv.thruster_data, + # Only defined/needed when kinetics is enabled; the core only reads it + # for struck faces, matching the scalar reference. + thruster_metrics=getattr(vv, 'thruster_metrics', None), + jfh_step=jfh_step, + plume_params=plume_params, + ) + + +def _compute_plume_strikes_scalar( + target_mesh: Any, + target_unit_normals: np.ndarray, + vv: Any, + jfh_step: Dict[str, Any], + environment: Any, +) -> Dict[str, np.ndarray]: + """Scalar reference implementation of compute_plume_strikes(). + + Preserved verbatim from the original per-face loop. Kept for regression + tests, debugging, and benchmarking against the vectorized path; the two + must produce identical strike arrays and struck-face IDs. + """ + num_faces = len(target_mesh.vectors) + strikes = np.zeros(num_faces) + + use_kinetics = environment.config['pm']['kinetics'] != 'None' + if use_kinetics: + pressures = np.zeros(num_faces) + shear_stresses = np.zeros(num_faces) + heat_flux = np.zeros(num_faces) + heat_flux_load = np.zeros(num_faces) + + vv_pos = np.array(jfh_step['xyz']) + vv_orientation = np.array(jfh_step['dcm']).transpose() + thrusters = jfh_step['thrusters'] + firing_time = float(jfh_step['t']) if 't' in jfh_step else 0.0 + + # Build mapping from numeric JFH indices to thruster ids consistent with legacy + link = {} + i = 1 + for thruster in vv.thruster_data: + link[str(i)] = vv.thruster_data[thruster]['name'] + i += 1 + + plume_radius = float(environment.config['plume']['radius']) + wedge_theta = float(environment.config['plume']['wedge_theta']) + + for thr in thrusters: + thruster_id = link[str(thr)][0] + + thruster_orientation = np.array(vv.thruster_data[thruster_id]['dcm']).transpose() + thruster_orientation = thruster_orientation.dot(vv_orientation) + plume_normal = np.array(thruster_orientation[0]) + norm_plume_normal = np.linalg.norm(plume_normal) + unit_plume_normal = plume_normal / norm_plume_normal + + thr_exit = np.array(vv.thruster_data[thruster_id]['exit']) + thruster_pos = vv_pos + thr_exit + thruster_pos = thruster_pos[0] + + for idx, face in enumerate(target_mesh.vectors): + face = np.array(face).transpose() + centroid = np.array([face[0].mean(), face[1].mean(), face[2].mean()]) + distance = thruster_pos - centroid + norm_distance = np.linalg.norm(distance) + if norm_distance == 0: + continue + unit_distance = distance / norm_distance + + theta = 3.14 - np.arccos(np.dot(np.squeeze(unit_distance), np.squeeze(unit_plume_normal))) + + n = np.squeeze(target_unit_normals[idx]) + unit_plume = np.squeeze(plume_normal / norm_plume_normal) + surface_dot_plume = np.dot(n, unit_plume) + + within_distance = float(norm_distance) < plume_radius + within_theta = float(theta) < wedge_theta + facing_thruster = surface_dot_plume < 0 + + if within_distance and within_theta and facing_thruster: + strikes[idx] += 1 + if use_kinetics: + T_w = float(environment.config['tv']['surface_temp']) + sigma = float(environment.config['tv']['sigma']) + t_type = vv.thruster_data[thruster_id]['type'][0] + thruster_metrics = vv.thruster_metrics[t_type] + simple_plume = SimplifiedGasKinetics(norm_distance, theta, thruster_metrics, T_w, sigma) + pressures[idx] += simple_plume.get_pressure() + shear = simple_plume.get_shear_pressure() + shear_stresses[idx] += abs(shear) + hf = simple_plume.get_heat_flux() + heat_flux[idx] += hf + heat_flux_load[idx] += hf * firing_time + + result = {"strikes": strikes} + if use_kinetics: + result.update({ + "pressures": pressures, + "shear_stress": shear_stresses, + "heat_flux_rate": heat_flux, + "heat_flux_load": heat_flux_load, + }) + return result + + +def accumulate_cumulative( + cumulative: Dict[str, np.ndarray], + current: Dict[str, np.ndarray], +) -> Dict[str, np.ndarray]: + """Accumulate per-step arrays into cumulative tallies (e.g., cum_strikes, max_pressures).""" + if "cum_strikes" in cumulative and "strikes" in current: + cumulative["cum_strikes"] = cumulative["cum_strikes"] + current["strikes"] + + # Max trackers if available + if "max_pressures" in cumulative and "pressures" in current: + cumulative["max_pressures"] = np.maximum(cumulative["max_pressures"], current["pressures"]) + if "max_shears" in cumulative and "shear_stress" in current: + cumulative["max_shears"] = np.maximum(cumulative["max_shears"], current["shear_stress"]) + + if "cum_heat_flux_load" in cumulative and "heat_flux_load" in current: + cumulative["cum_heat_flux_load"] = cumulative["cum_heat_flux_load"] + current["heat_flux_load"] + + return cumulative From 7d567e8640c202f27c4ffa17e2bb4d50c847efff Mon Sep 17 00:00:00 2001 From: andytorrestb Date: Sat, 18 Jul 2026 16:34:13 -0500 Subject: [PATCH 2/7] precompute target face centroids once per plume strike run The target mesh is stationary during a JFH run, so jfh_plume_strikes() now computes face centroids once and passes them to compute_plume_strikes() via the optional face_centroids parameter instead of recomputing them for every firing. Co-Authored-By: Claude Fable 5 --- pyrpod/rpod/PlumeStrikeEstimationStudy.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/pyrpod/rpod/PlumeStrikeEstimationStudy.py b/pyrpod/rpod/PlumeStrikeEstimationStudy.py index 32679c5..66b47db 100644 --- a/pyrpod/rpod/PlumeStrikeEstimationStudy.py +++ b/pyrpod/rpod/PlumeStrikeEstimationStudy.py @@ -27,7 +27,10 @@ ) from pyrpod.rpod.io import ensure_results_dirs, write_jfh from pyrpod.rpod.PlumeStudyExport import PlumeStudyExport -from pyrpod.plume.PlumeStrikeCalculator import compute_plume_strikes +from pyrpod.plume.PlumeStrikeCalculator import ( + compute_face_centroids, + compute_plume_strikes, +) logger = get_logger("pyrpod.rpod.PlumeStrikeEstimationStudy") @@ -772,6 +775,9 @@ def jfh_plume_strikes(self): self.create_results_dir() target = self.target.mesh target_normals = target.get_unit_normals() + # The target is stationary for the whole run, so face centroids are + # computed once here and reused for every firing. + target_centroids = compute_face_centroids(target.vectors) # Initialize cumulative arrays kinetics_on = self.environment.config['pm']['kinetics'] != 'None' @@ -797,6 +803,7 @@ def jfh_plume_strikes(self): vv=self.vv, jfh_step=step, environment=self.environment, + face_centroids=target_centroids, ) strikes = result["strikes"] From 7dd75b24c282f99958df394afe892a9736bb3ac7 Mon Sep 17 00:00:00 2001 From: andytorrestb Date: Sat, 18 Jul 2026 16:36:16 -0500 Subject: [PATCH 3/7] add optional process-based parallelization across JFH firings jfh_plume_strikes() gains optional parallel/workers arguments layered over new optional config keys ([exec] parallel, [exec] workers); default behavior remains serial and byte-identical to before. When enabled, independent per-firing strike arrays are computed in a ProcessPoolExecutor (one firing per task). Workers receive only plain serializable inputs -- precomputed centroids, unit normals, thruster dicts, JFH step dicts, and the few needed config scalars -- shipped once per worker via a pool initializer, so memory scales with workers x faces rather than firings x faces. Cumulative arrays (cum_strikes, max_pressures, max_shears, cum_heat_flux_load) are still accumulated serially in firing order and VTK files are written serially by the parent, preserving the firing_data return structure exactly. Invalid parallel/workers values raise clear ValueErrors; runtime parallel failures log a warning and fall back to serial execution. Co-Authored-By: Claude Fable 5 --- pyrpod/plume/PlumeStrikeCalculator.py | 84 ++++++++++++- pyrpod/rpod/PlumeStrikeEstimationStudy.py | 144 +++++++++++++++++++--- 2 files changed, 211 insertions(+), 17 deletions(-) diff --git a/pyrpod/plume/PlumeStrikeCalculator.py b/pyrpod/plume/PlumeStrikeCalculator.py index eeb12f8..755f8c9 100644 --- a/pyrpod/plume/PlumeStrikeCalculator.py +++ b/pyrpod/plume/PlumeStrikeCalculator.py @@ -22,7 +22,8 @@ """ from __future__ import annotations -from typing import Any, Dict, Optional +from concurrent.futures import ProcessPoolExecutor +from typing import Any, Dict, List, Optional, Sequence import numpy as np from pyrpod.plume.RarefiedPlumeGasKinetics import SimplifiedGasKinetics @@ -315,6 +316,87 @@ def _compute_plume_strikes_scalar( return result +# Per-process state for parallel workers. Populated once per worker by +# _parallel_worker_init so the (N,3) target arrays are shipped to each worker +# a single time instead of once per submitted firing. Memory therefore scales +# with workers x faces, never firings x faces. +_WORKER_STATE: Dict[str, Any] = {} + + +def _parallel_worker_init( + face_centroids: np.ndarray, + target_unit_normals: np.ndarray, + thruster_data: Dict[str, Any], + thruster_metrics: Optional[Dict[str, Any]], + plume_params: Dict[str, Any], +) -> None: + """ProcessPoolExecutor initializer: cache shared per-run inputs.""" + _WORKER_STATE['face_centroids'] = face_centroids + _WORKER_STATE['target_unit_normals'] = target_unit_normals + _WORKER_STATE['thruster_data'] = thruster_data + _WORKER_STATE['thruster_metrics'] = thruster_metrics + _WORKER_STATE['plume_params'] = plume_params + + +def _parallel_worker_compute(task) -> Any: + """Compute strikes for one firing inside a worker process. + + task is (firing_index, jfh_step); returns (firing_index, result dict). + """ + firing_index, jfh_step = task + result = _compute_plume_strikes_core( + face_centroids=_WORKER_STATE['face_centroids'], + target_unit_normals=_WORKER_STATE['target_unit_normals'], + thruster_data=_WORKER_STATE['thruster_data'], + thruster_metrics=_WORKER_STATE['thruster_metrics'], + jfh_step=jfh_step, + plume_params=_WORKER_STATE['plume_params'], + ) + return firing_index, result + + +def run_parallel_plume_strikes( + jfh_steps: Sequence[Dict[str, Any]], + face_centroids: np.ndarray, + target_unit_normals: np.ndarray, + thruster_data: Dict[str, Any], + thruster_metrics: Optional[Dict[str, Any]], + plume_params: Dict[str, Any], + workers: int, +) -> List[Dict[str, np.ndarray]]: + """Compute per-firing strike results across processes, one firing per task. + + All inputs must be plain serializable data (NumPy arrays, dicts, + primitives) — full study/vehicle/environment objects are never pickled. + Results are returned as a list indexed by firing, preserving JFH order + regardless of completion order; cumulative accumulation and VTK output + remain the caller's responsibility (serial, in the parent process). + + Raises whatever the executor or workers raise; callers are expected to + fall back to the serial path with a clear message. + """ + results: List[Optional[Dict[str, np.ndarray]]] = [None] * len(jfh_steps) + with ProcessPoolExecutor( + max_workers=workers, + initializer=_parallel_worker_init, + initargs=( + face_centroids, + target_unit_normals, + thruster_data, + thruster_metrics, + plume_params, + ), + ) as executor: + futures = [ + executor.submit(_parallel_worker_compute, (i, step)) + for i, step in enumerate(jfh_steps) + ] + for future in futures: + firing_index, result = future.result() + results[firing_index] = result + return results + + def accumulate_cumulative( cumulative: Dict[str, np.ndarray], current: Dict[str, np.ndarray], diff --git a/pyrpod/rpod/PlumeStrikeEstimationStudy.py b/pyrpod/rpod/PlumeStrikeEstimationStudy.py index 66b47db..f270e86 100644 --- a/pyrpod/rpod/PlumeStrikeEstimationStudy.py +++ b/pyrpod/rpod/PlumeStrikeEstimationStudy.py @@ -30,6 +30,8 @@ from pyrpod.plume.PlumeStrikeCalculator import ( compute_face_centroids, compute_plume_strikes, + extract_plume_params, + run_parallel_plume_strikes, ) logger = get_logger("pyrpod.rpod.PlumeStrikeEstimationStudy") @@ -757,19 +759,94 @@ def set_face_distance(self, thruster_pos, centroid): return distance, norm_distance, unit_distance - def jfh_plume_strikes(self): + def _resolve_parallel_options(self, parallel, workers, n_firings): + """ + Resolves parallel execution settings for jfh_plume_strikes(). + + Precedence: explicit method arguments override the optional + [exec] config section, which defaults to serial execution. + + Config keys (both optional): + - [exec] parallel : bool — enable process-based parallelization + across firings (default false). + - [exec] workers : int — number of worker processes. Defaults to + min(os.cpu_count(), n_firings) when parallel is enabled. + + Returns + ------- + (bool, int) + (parallel_enabled, workers) — workers is capped at n_firings; + workers <= 1 resolves to serial execution. + """ + config = self.environment.config + if parallel is None: + try: + parallel = config.getboolean('exec', 'parallel', fallback=False) + except ValueError as exc: + raise ValueError( + "Invalid config value for [exec] parallel: expected a " + "boolean (true/false)." + ) from exc + if workers is None: + try: + workers = config.getint('exec', 'workers', fallback=None) + except ValueError as exc: + raise ValueError( + "Invalid config value for [exec] workers: expected a " + "positive integer." + ) from exc + if workers is not None and workers < 1: + raise ValueError( + f"workers must be a positive integer, got {workers}." + ) + + if not parallel: + return False, 1 + + if workers is None: + workers = min(os.cpu_count() or 1, n_firings) + # Never spawn more workers than there are firings to compute. + workers = min(workers, n_firings) + if workers <= 1: + return False, 1 + return True, workers + + def jfh_plume_strikes(self, parallel=None, workers=None): """ Calculates number of plume strikes according to data provided for RPOD analysis. - Method does not take any parameters but assumes that study assets are correctly configured. + Method assumes that study assets are correctly configured. These assets include one JetFiringHistory, one TargetVehicle, and one VisitingVehicle. A Simple plume model is used. It does not calculate plume physics, only strikes. which are determined with a user defined "plume cone" geometry. Simple vector mathematics is used to determine if an VTK surface elements is struck by the "plume cone". + Parameters + ---------- + parallel : bool, optional + Enable process-based parallelization across firings. Defaults + to None, meaning "use the optional [exec] parallel config key", + which itself defaults to false (serial, legacy behavior). + workers : int, optional + Number of worker processes when parallel execution is enabled. + Defaults to None, meaning "use the optional [exec] workers + config key", falling back to min(os.cpu_count(), n_firings). + + Notes + ----- + Only independent per-firing strike arrays are computed in worker + processes; cumulative arrays (cum_strikes, max_pressures, + max_shears, cum_heat_flux_load) are accumulated serially in + original firing order, and VTK output is written serially by the + parent process. If the parallel path fails to initialize or run, + the method logs a warning and falls back to serial execution. + Returns ------- - Method doesn't currently return anything. Simply produces data as needed. - Does the method need to return a status message? or pass similar data? + dict + firing_data keyed by firing number ('1'..'N'), each holding + per-face arrays: strikes, cum_strikes, and, when kinetics is + enabled, pressures, max_pressures, shear_stress, max_shears, + heat_flux_rate, heat_flux_load, cum_heat_flux_load. """ # Prepare results directories and target data self.create_results_dir() @@ -788,23 +865,58 @@ def jfh_plume_strikes(self): firing_data = {} - # Loop through each firing in the JFH and delegate to impingement module - for firing in range(len(self.jfh.JFH)): - step = { + n_firings = len(self.jfh.JFH) + + # Build serializable step dicts once; shared by serial and parallel paths. + steps = [] + for firing in range(n_firings): + steps.append({ 'thrusters': self.jfh.JFH[firing]['thrusters'], 'xyz': np.array(self.jfh.JFH[firing]['xyz']), 'dcm': np.array(self.jfh.JFH[firing]['dcm']), 't': float(self.jfh.JFH[firing]['t']) - } + }) + + parallel_enabled, n_workers = self._resolve_parallel_options(parallel, workers, n_firings) + + # Optionally compute independent per-firing results in worker + # processes. Workers receive only plain serializable inputs (arrays, + # dicts, config scalars) — never the study/vehicle/environment objects. + per_firing_results = None + if parallel_enabled: + try: + per_firing_results = run_parallel_plume_strikes( + jfh_steps=steps, + face_centroids=target_centroids, + target_unit_normals=target_normals, + thruster_data=self.vv.thruster_data, + thruster_metrics=getattr(self.vv, 'thruster_metrics', None), + plume_params=extract_plume_params(self.environment), + workers=n_workers, + ) + except Exception as exc: + logger.warning( + "Parallel plume strike execution failed (%s: %s); " + "falling back to serial execution.", + type(exc).__name__, exc, + ) + per_firing_results = None - result = compute_plume_strikes( - target_mesh=target, - target_unit_normals=target_normals, - vv=self.vv, - jfh_step=step, - environment=self.environment, - face_centroids=target_centroids, - ) + # Loop through each firing in the JFH and delegate to impingement module + for firing in range(n_firings): + step = steps[firing] + + if per_firing_results is not None: + result = per_firing_results[firing] + else: + result = compute_plume_strikes( + target_mesh=target, + target_unit_normals=target_normals, + vv=self.vv, + jfh_step=step, + environment=self.environment, + face_centroids=target_centroids, + ) strikes = result["strikes"] cum_strikes = cum_strikes + strikes From a943ef10fab778aa1730ac46d46207a4161351ed Mon Sep 17 00:00:00 2001 From: andytorrestb Date: Sat, 18 Jul 2026 16:37:58 -0500 Subject: [PATCH 4/7] add regression tests for vectorized and parallel plume strikes plume_unit_test_04 asserts the vectorized compute_plume_strikes() path (with and without precomputed centroids) reproduces the scalar reference exactly -- strike arrays and struck-face IDs bit-for-bit -- for every firing of the 1d_approach (geometry-only) and multi_thrusters_square (Simplified kinetics) cases, plus explicit zero-distance face skipping. plume_integration_test_02 asserts jfh_plume_strikes() default behavior stays serial and return-compatible, that the parallel path (workers=2) matches serial per-firing and cumulative outputs, that workers=1 resolves to serial, and that invalid worker counts raise ValueError. Kinetics arrays use np.allclose(rtol=1e-12, atol=1e-12): observed bit-identical on this platform, tolerance only guards platform BLAS reduction-order differences (rationale documented in each file). Co-Authored-By: Claude Fable 5 --- tests/plume/plume_integration_test_02.py | 133 ++++++++++++++++++++++ tests/plume/plume_unit_test_04.py | 138 +++++++++++++++++++++++ 2 files changed, 271 insertions(+) create mode 100644 tests/plume/plume_integration_test_02.py create mode 100644 tests/plume/plume_unit_test_04.py diff --git a/tests/plume/plume_integration_test_02.py b/tests/plume/plume_integration_test_02.py new file mode 100644 index 0000000..4ce4678 --- /dev/null +++ b/tests/plume/plume_integration_test_02.py @@ -0,0 +1,133 @@ +# ======================== +# PyRPOD: tests/plume/plume_integration_test_02.py +# ======================== +# Asserts that PlumeStrikeEstimationStudy.jfh_plume_strikes(): +# 1. keeps its default behavior serial and return-compatible (dict keyed +# '1'..'N' with the documented per-firing arrays), and +# 2. produces identical per-firing and cumulative outputs when the optional +# process-based parallel path is enabled, and +# 3. rejects invalid worker counts with a clear error. +# +# Two representative existing cases are used: +# - case/rpod/1d_approach (kinetics disabled) +# - case/rpod/multi_thrusters_square (Simplified kinetics) +# +# Strike arrays (strikes, cum_strikes) and struck-face IDs must match +# exactly. Kinetics arrays are compared with np.allclose(rtol=1e-12, +# atol=1e-12): serial and parallel paths execute the same vectorized code on +# the same inputs (observed bit-identical here), but the tolerance guards +# against BLAS/SIMD reduction-order differences across platforms/processes. +# +# VTK output goes to each case's existing results/ directory, the same safe +# pattern the existing rpod integration tests rely on. + +import unittest + +import numpy as np + +from pyrpod.mission import MissionEnvironment +from pyrpod.rpod import JetFiringHistory, PlumeStrikeEstimationStudy +from pyrpod.vehicle import TargetVehicle, VisitingVehicle + +KINETICS_RTOL = 1e-12 +KINETICS_ATOL = 1e-12 +EXACT_KEYS = ("strikes", "cum_strikes") + +GEOMETRY_KEYS = {"strikes", "cum_strikes"} +KINETICS_KEYS = GEOMETRY_KEYS | { + "pressures", "max_pressures", + "shear_stress", "max_shears", + "heat_flux_rate", "heat_flux_load", "cum_heat_flux_load", +} + + +def make_study(case_dir): + jfh = JetFiringHistory.JetFiringHistory(case_dir) + jfh.read_jfh() + + tv = TargetVehicle.TargetVehicle(case_dir) + tv.set_stl() + + vv = VisitingVehicle.VisitingVehicle(case_dir) + vv.set_stl() + vv.set_thruster_config() + vv.set_thruster_metrics() + + me = MissionEnvironment.MissionEnvironment(case_dir) + study = PlumeStrikeEstimationStudy.PlumeStrikeEstimationStudy(me) + study.study_init(jfh, tv, vv) + return study + + +class ParallelMatchesSerialChecks(unittest.TestCase): + + def assert_firing_data_equal(self, serial, parallel, case_dir): + self.assertEqual(set(serial.keys()), set(parallel.keys())) + for firing in serial: + self.assertEqual( + set(serial[firing].keys()), set(parallel[firing].keys()), + msg=f"{case_dir} firing {firing}: cellData keys differ", + ) + for key in serial[firing]: + if key in EXACT_KEYS: + np.testing.assert_array_equal( + parallel[firing][key], serial[firing][key], + err_msg=f"{case_dir} firing {firing}: {key} differs", + ) + else: + self.assertTrue( + np.allclose( + parallel[firing][key], serial[firing][key], + rtol=KINETICS_RTOL, atol=KINETICS_ATOL, + ), + msg=f"{case_dir} firing {firing}: {key} differs", + ) + # Struck-face IDs must be identical. + np.testing.assert_array_equal( + np.nonzero(parallel[firing]["strikes"])[0], + np.nonzero(serial[firing]["strikes"])[0], + err_msg=f"{case_dir} firing {firing}: struck-face IDs differ", + ) + + def assert_return_structure(self, firing_data, n_firings, expected_keys): + self.assertEqual( + sorted(firing_data.keys(), key=int), + [str(i + 1) for i in range(n_firings)], + ) + for firing in firing_data: + self.assertEqual(set(firing_data[firing].keys()), expected_keys) + + def run_case(self, case_dir, expected_keys): + study = make_study(case_dir) + n_firings = len(study.jfh.JFH) + + # Default call: no arguments, serial legacy behavior. + serial = study.jfh_plume_strikes() + self.assert_return_structure(serial, n_firings, expected_keys) + + # Parallel path with an explicit worker count. + parallel = make_study(case_dir).jfh_plume_strikes(parallel=True, workers=2) + self.assert_return_structure(parallel, n_firings, expected_keys) + + self.assert_firing_data_equal(serial, parallel, case_dir) + + # workers=1 must resolve to the serial path and identical output. + one_worker = make_study(case_dir).jfh_plume_strikes(parallel=True, workers=1) + self.assert_firing_data_equal(serial, one_worker, case_dir) + + def test_geometry_only_case(self): + self.run_case('../case/rpod/1d_approach/', GEOMETRY_KEYS) + + def test_kinetics_multi_thruster_case(self): + self.run_case('../case/rpod/multi_thrusters_square/', KINETICS_KEYS) + + def test_invalid_workers_rejected(self): + study = make_study('../case/rpod/1d_approach/') + with self.assertRaises(ValueError): + study.jfh_plume_strikes(parallel=True, workers=0) + with self.assertRaises(ValueError): + study.jfh_plume_strikes(parallel=True, workers=-2) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/plume/plume_unit_test_04.py b/tests/plume/plume_unit_test_04.py new file mode 100644 index 0000000..97ea335 --- /dev/null +++ b/tests/plume/plume_unit_test_04.py @@ -0,0 +1,138 @@ +# ======================== +# PyRPOD: tests/plume/plume_unit_test_04.py +# ======================== +# Asserts that the NumPy-vectorized plume strike detection in +# compute_plume_strikes() reproduces the scalar reference implementation +# (_compute_plume_strikes_scalar) exactly, for every firing of two +# representative existing cases: +# - case/rpod/1d_approach (kinetics disabled, geometry only) +# - case/rpod/multi_thrusters_square (Simplified kinetics, multiple thrusters) +# +# Strike arrays and struck-face IDs must match exactly (integer counts from +# strict geometric comparisons). Kinetics arrays are compared with +# np.allclose(rtol=1e-12, atol=1e-12): on the tested platform they are +# bit-for-bit identical (both paths feed identical scalar inputs to +# SimplifiedGasKinetics), but the tolerance guards against BLAS/SIMD +# reduction-order differences on other platforms without weakening the test +# in any practically meaningful way. +# +# These functions compute arrays only; no files are written. + +import unittest + +import numpy as np + +from pyrpod.mission import MissionEnvironment +from pyrpod.plume.PlumeStrikeCalculator import ( + _compute_plume_strikes_scalar, + compute_face_centroids, + compute_plume_strikes, +) +from pyrpod.rpod import JetFiringHistory +from pyrpod.vehicle import TargetVehicle, VisitingVehicle + +KINETICS_RTOL = 1e-12 +KINETICS_ATOL = 1e-12 +KINETICS_KEYS = ("pressures", "shear_stress", "heat_flux_rate", "heat_flux_load") + + +def load_case(case_dir): + jfh = JetFiringHistory.JetFiringHistory(case_dir) + jfh.read_jfh() + + tv = TargetVehicle.TargetVehicle(case_dir) + tv.set_stl() + + vv = VisitingVehicle.VisitingVehicle(case_dir) + vv.set_thruster_config() + vv.set_thruster_metrics() + + me = MissionEnvironment.MissionEnvironment(case_dir) + return jfh, tv, vv, me + + +def build_step(jfh, firing): + return { + 'thrusters': jfh.JFH[firing]['thrusters'], + 'xyz': np.array(jfh.JFH[firing]['xyz']), + 'dcm': np.array(jfh.JFH[firing]['dcm']), + 't': float(jfh.JFH[firing]['t']), + } + + +class VectorizedMatchesScalarChecks(unittest.TestCase): + + def assert_case_equivalence(self, case_dir): + jfh, tv, vv, me = load_case(case_dir) + target = tv.mesh + normals = target.get_unit_normals() + centroids = compute_face_centroids(target.vectors) + + for firing in range(len(jfh.JFH)): + step = build_step(jfh, firing) + + scalar = _compute_plume_strikes_scalar(target, normals, vv, step, me) + vectorized = compute_plume_strikes( + target, normals, vv, step, me, face_centroids=centroids + ) + # Backward-compatible call without precomputed centroids. + vectorized_no_centroids = compute_plume_strikes( + target, normals, vv, step, me + ) + + for vec in (vectorized, vectorized_no_centroids): + self.assertEqual(set(vec.keys()), set(scalar.keys())) + + # Exact strike counts and exact struck-face IDs are required. + np.testing.assert_array_equal( + vec["strikes"], scalar["strikes"], + err_msg=f"{case_dir} firing {firing}: strikes differ", + ) + np.testing.assert_array_equal( + np.nonzero(vec["strikes"])[0], + np.nonzero(scalar["strikes"])[0], + err_msg=f"{case_dir} firing {firing}: struck-face IDs differ", + ) + + # Kinetics arrays: tolerance documented in the module header. + for key in KINETICS_KEYS: + if key in scalar: + self.assertTrue( + np.allclose( + vec[key], scalar[key], + rtol=KINETICS_RTOL, atol=KINETICS_ATOL, + ), + msg=f"{case_dir} firing {firing}: {key} differs", + ) + + def test_geometry_only_case(self): + self.assert_case_equivalence('../case/rpod/1d_approach/') + + def test_kinetics_multi_thruster_case(self): + self.assert_case_equivalence('../case/rpod/multi_thrusters_square/') + + def test_zero_distance_face_skipped(self): + # A face whose centroid coincides with the thruster exit must be + # skipped by both implementations (legacy `norm_distance == 0` guard). + case_dir = '../case/rpod/1d_approach/' + jfh, tv, vv, me = load_case(case_dir) + target = tv.mesh + normals = target.get_unit_normals() + + step = build_step(jfh, 0) + # Place the vehicle so the first thruster exit lands exactly on the + # centroid of face 0. + centroid0 = compute_face_centroids(target.vectors)[0].astype(np.float64) + thruster_id = vv.thruster_data[next(iter(vv.thruster_data))]['name'][0] + exit_pos = np.array(vv.thruster_data[thruster_id]['exit'])[0] + step['xyz'] = centroid0 - exit_pos + + scalar = _compute_plume_strikes_scalar(target, normals, vv, step, me) + vectorized = compute_plume_strikes(target, normals, vv, step, me) + + self.assertEqual(scalar["strikes"][0], 0.0) + np.testing.assert_array_equal(vectorized["strikes"], scalar["strikes"]) + + +if __name__ == '__main__': + unittest.main() From efd26cdd0bc5467e13dfd8c03ed2b155f419c3fe Mon Sep 17 00:00:00 2001 From: andytorrestb Date: Sat, 18 Jul 2026 16:39:08 -0500 Subject: [PATCH 5/7] add plume strike performance benchmark script python scripts/benchmark_plume_strikes.py times the scalar reference vs the vectorized geometry path, and serial vs process-parallel execution, on existing repo cases (default case/rpod/1d_approach; --case, --workers, --repeat, --skip-scalar options). It reports elapsed times and verifies strike equivalence but asserts no speedup thresholds, keeping it stable on shared machines. Measured on this machine: vectorized geometry is ~245x faster than the scalar reference on 1d_approach (10 firings x 3584 faces) and ~24x on multi_thrusters_square with Simplified kinetics. Process-parallel is slower than serial on these small cases because Windows process spawn dominates; it is intended for large meshes/many firings. Co-Authored-By: Claude Fable 5 --- scripts/benchmark_plume_strikes.py | 164 +++++++++++++++++++++++++++++ 1 file changed, 164 insertions(+) create mode 100644 scripts/benchmark_plume_strikes.py diff --git a/scripts/benchmark_plume_strikes.py b/scripts/benchmark_plume_strikes.py new file mode 100644 index 0000000..3e48de2 --- /dev/null +++ b/scripts/benchmark_plume_strikes.py @@ -0,0 +1,164 @@ +""" +Benchmark plume strike computation: scalar reference vs vectorized geometry, +and serial vs optional process-parallel execution across JFH firings. + +Usage (from the repository root): + python scripts/benchmark_plume_strikes.py + python scripts/benchmark_plume_strikes.py --case case/rpod/multi_thrusters_square/ + python scripts/benchmark_plume_strikes.py --workers 4 --repeat 3 + +Reports elapsed wall-clock times only; it asserts result equivalence but no +speedup thresholds, so it stays robust on shared/CI machines. Correctness is +covered by tests/plume/plume_unit_test_04.py and plume_integration_test_02.py. + +Uses only cases and data already present in the repository; no new +dependencies (NumPy + standard library). +""" +import argparse +import os +import sys +import time + +import numpy as np + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO_ROOT) + +from pyrpod.mission import MissionEnvironment # noqa: E402 +from pyrpod.plume.PlumeStrikeCalculator import ( # noqa: E402 + _compute_plume_strikes_scalar, + compute_face_centroids, + compute_plume_strikes, + extract_plume_params, + run_parallel_plume_strikes, +) +from pyrpod.rpod import JetFiringHistory # noqa: E402 +from pyrpod.vehicle import TargetVehicle, VisitingVehicle # noqa: E402 + + +def load_case(case_dir): + jfh = JetFiringHistory.JetFiringHistory(case_dir) + jfh.read_jfh() + tv = TargetVehicle.TargetVehicle(case_dir) + tv.set_stl() + vv = VisitingVehicle.VisitingVehicle(case_dir) + vv.set_thruster_config() + vv.set_thruster_metrics() + me = MissionEnvironment.MissionEnvironment(case_dir) + return jfh, tv, vv, me + + +def build_steps(jfh): + return [ + { + 'thrusters': jfh.JFH[firing]['thrusters'], + 'xyz': np.array(jfh.JFH[firing]['xyz']), + 'dcm': np.array(jfh.JFH[firing]['dcm']), + 't': float(jfh.JFH[firing]['t']), + } + for firing in range(len(jfh.JFH)) + ] + + +def time_best_of(fn, repeat): + best = float('inf') + result = None + for _ in range(repeat): + start = time.perf_counter() + result = fn() + best = min(best, time.perf_counter() - start) + return best, result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + '--case', default='case/rpod/1d_approach/', + help='Case directory (relative to the repo root or absolute).', + ) + parser.add_argument( + '--workers', type=int, default=min(os.cpu_count() or 1, 4), + help='Worker processes for the parallel benchmark (default: up to 4).', + ) + parser.add_argument( + '--repeat', type=int, default=1, + help='Repetitions per measurement; best time is reported.', + ) + parser.add_argument( + '--skip-scalar', action='store_true', + help='Skip the (slow) scalar reference benchmark.', + ) + args = parser.parse_args() + + case_dir = args.case + if not os.path.isabs(case_dir): + case_dir = os.path.join(REPO_ROOT, case_dir) + case_dir = case_dir.replace('\\', '/') + if not case_dir.endswith('/'): + case_dir += '/' + + jfh, tv, vv, me = load_case(case_dir) + target = tv.mesh + normals = target.get_unit_normals() + centroids = compute_face_centroids(target.vectors) + steps = build_steps(jfh) + n_firings = len(steps) + n_faces = len(target.vectors) + + print(f"case: {case_dir}") + print(f"firings: {n_firings}, faces: {n_faces}, " + f"kinetics: {me.config['pm']['kinetics']}") + print(f"repeat: {args.repeat} (best time reported)\n") + + # --- Scalar reference vs vectorized geometry (per-firing compute only) --- + def run_vectorized(): + return [ + compute_plume_strikes(target, normals, vv, step, me, + face_centroids=centroids) + for step in steps + ] + + t_vec, vec_results = time_best_of(run_vectorized, args.repeat) + print(f"vectorized (serial): {t_vec:8.3f} s") + + if not args.skip_scalar: + def run_scalar(): + return [ + _compute_plume_strikes_scalar(target, normals, vv, step, me) + for step in steps + ] + + t_scalar, scalar_results = time_best_of(run_scalar, args.repeat) + print(f"scalar reference: {t_scalar:8.3f} s" + f" ({t_scalar / t_vec:5.1f}x slower than vectorized)") + for vec, ref in zip(vec_results, scalar_results): + assert np.array_equal(vec['strikes'], ref['strikes']), \ + "vectorized/scalar strike mismatch" + + # --- Serial vs parallel across firings (vectorized in both) --- + workers = min(args.workers, n_firings) + if workers > 1: + def run_parallel(): + return run_parallel_plume_strikes( + jfh_steps=steps, + face_centroids=centroids, + target_unit_normals=normals, + thruster_data=vv.thruster_data, + thruster_metrics=getattr(vv, 'thruster_metrics', None), + plume_params=extract_plume_params(me), + workers=workers, + ) + + t_par, par_results = time_best_of(run_parallel, args.repeat) + print(f"parallel ({workers} workers): {t_par:8.3f} s" + f" ({t_vec / t_par:5.2f}x vs serial vectorized; includes " + f"process startup)") + for vec, par in zip(vec_results, par_results): + assert np.array_equal(vec['strikes'], par['strikes']), \ + "serial/parallel strike mismatch" + else: + print(f"parallel benchmark skipped (workers={workers})") + + +if __name__ == '__main__': + main() From eb9bfcbcce8f9ccb7840a038fc86aecd54d32795 Mon Sep 17 00:00:00 2001 From: andytorrestb Date: Sat, 18 Jul 2026 17:59:45 -0500 Subject: [PATCH 6/7] move plume strike benchmark into tests/rpod as integration test 06 Relocate scripts/benchmark_plume_strikes.py to tests/rpod/rpod_integration_test_06.py, following the established rpod test format (unittest.TestCase, header description block, ../case/rpod paths, conftest-derived rpod/integration markers). The benchmark prints elapsed times (visible with pytest -s) for scalar-vs-vectorized geometry on 1d_approach and multi_thrusters_square, and serial-vs-parallel execution with 2 workers; it asserts only exact strike equality, never timing thresholds. No existing test files are modified. Co-Authored-By: Claude Fable 5 --- scripts/benchmark_plume_strikes.py | 164 ------------------------- tests/rpod/rpod_integration_test_06.py | 142 +++++++++++++++++++++ 2 files changed, 142 insertions(+), 164 deletions(-) delete mode 100644 scripts/benchmark_plume_strikes.py create mode 100644 tests/rpod/rpod_integration_test_06.py diff --git a/scripts/benchmark_plume_strikes.py b/scripts/benchmark_plume_strikes.py deleted file mode 100644 index 3e48de2..0000000 --- a/scripts/benchmark_plume_strikes.py +++ /dev/null @@ -1,164 +0,0 @@ -""" -Benchmark plume strike computation: scalar reference vs vectorized geometry, -and serial vs optional process-parallel execution across JFH firings. - -Usage (from the repository root): - python scripts/benchmark_plume_strikes.py - python scripts/benchmark_plume_strikes.py --case case/rpod/multi_thrusters_square/ - python scripts/benchmark_plume_strikes.py --workers 4 --repeat 3 - -Reports elapsed wall-clock times only; it asserts result equivalence but no -speedup thresholds, so it stays robust on shared/CI machines. Correctness is -covered by tests/plume/plume_unit_test_04.py and plume_integration_test_02.py. - -Uses only cases and data already present in the repository; no new -dependencies (NumPy + standard library). -""" -import argparse -import os -import sys -import time - -import numpy as np - -REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) -sys.path.insert(0, REPO_ROOT) - -from pyrpod.mission import MissionEnvironment # noqa: E402 -from pyrpod.plume.PlumeStrikeCalculator import ( # noqa: E402 - _compute_plume_strikes_scalar, - compute_face_centroids, - compute_plume_strikes, - extract_plume_params, - run_parallel_plume_strikes, -) -from pyrpod.rpod import JetFiringHistory # noqa: E402 -from pyrpod.vehicle import TargetVehicle, VisitingVehicle # noqa: E402 - - -def load_case(case_dir): - jfh = JetFiringHistory.JetFiringHistory(case_dir) - jfh.read_jfh() - tv = TargetVehicle.TargetVehicle(case_dir) - tv.set_stl() - vv = VisitingVehicle.VisitingVehicle(case_dir) - vv.set_thruster_config() - vv.set_thruster_metrics() - me = MissionEnvironment.MissionEnvironment(case_dir) - return jfh, tv, vv, me - - -def build_steps(jfh): - return [ - { - 'thrusters': jfh.JFH[firing]['thrusters'], - 'xyz': np.array(jfh.JFH[firing]['xyz']), - 'dcm': np.array(jfh.JFH[firing]['dcm']), - 't': float(jfh.JFH[firing]['t']), - } - for firing in range(len(jfh.JFH)) - ] - - -def time_best_of(fn, repeat): - best = float('inf') - result = None - for _ in range(repeat): - start = time.perf_counter() - result = fn() - best = min(best, time.perf_counter() - start) - return best, result - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - '--case', default='case/rpod/1d_approach/', - help='Case directory (relative to the repo root or absolute).', - ) - parser.add_argument( - '--workers', type=int, default=min(os.cpu_count() or 1, 4), - help='Worker processes for the parallel benchmark (default: up to 4).', - ) - parser.add_argument( - '--repeat', type=int, default=1, - help='Repetitions per measurement; best time is reported.', - ) - parser.add_argument( - '--skip-scalar', action='store_true', - help='Skip the (slow) scalar reference benchmark.', - ) - args = parser.parse_args() - - case_dir = args.case - if not os.path.isabs(case_dir): - case_dir = os.path.join(REPO_ROOT, case_dir) - case_dir = case_dir.replace('\\', '/') - if not case_dir.endswith('/'): - case_dir += '/' - - jfh, tv, vv, me = load_case(case_dir) - target = tv.mesh - normals = target.get_unit_normals() - centroids = compute_face_centroids(target.vectors) - steps = build_steps(jfh) - n_firings = len(steps) - n_faces = len(target.vectors) - - print(f"case: {case_dir}") - print(f"firings: {n_firings}, faces: {n_faces}, " - f"kinetics: {me.config['pm']['kinetics']}") - print(f"repeat: {args.repeat} (best time reported)\n") - - # --- Scalar reference vs vectorized geometry (per-firing compute only) --- - def run_vectorized(): - return [ - compute_plume_strikes(target, normals, vv, step, me, - face_centroids=centroids) - for step in steps - ] - - t_vec, vec_results = time_best_of(run_vectorized, args.repeat) - print(f"vectorized (serial): {t_vec:8.3f} s") - - if not args.skip_scalar: - def run_scalar(): - return [ - _compute_plume_strikes_scalar(target, normals, vv, step, me) - for step in steps - ] - - t_scalar, scalar_results = time_best_of(run_scalar, args.repeat) - print(f"scalar reference: {t_scalar:8.3f} s" - f" ({t_scalar / t_vec:5.1f}x slower than vectorized)") - for vec, ref in zip(vec_results, scalar_results): - assert np.array_equal(vec['strikes'], ref['strikes']), \ - "vectorized/scalar strike mismatch" - - # --- Serial vs parallel across firings (vectorized in both) --- - workers = min(args.workers, n_firings) - if workers > 1: - def run_parallel(): - return run_parallel_plume_strikes( - jfh_steps=steps, - face_centroids=centroids, - target_unit_normals=normals, - thruster_data=vv.thruster_data, - thruster_metrics=getattr(vv, 'thruster_metrics', None), - plume_params=extract_plume_params(me), - workers=workers, - ) - - t_par, par_results = time_best_of(run_parallel, args.repeat) - print(f"parallel ({workers} workers): {t_par:8.3f} s" - f" ({t_vec / t_par:5.2f}x vs serial vectorized; includes " - f"process startup)") - for vec, par in zip(vec_results, par_results): - assert np.array_equal(vec['strikes'], par['strikes']), \ - "serial/parallel strike mismatch" - else: - print(f"parallel benchmark skipped (workers={workers})") - - -if __name__ == '__main__': - main() diff --git a/tests/rpod/rpod_integration_test_06.py b/tests/rpod/rpod_integration_test_06.py new file mode 100644 index 0000000..69b7f3e --- /dev/null +++ b/tests/rpod/rpod_integration_test_06.py @@ -0,0 +1,142 @@ +# ======================== +# PyRPOD: tests/rpod/rpod_integration_test_06.py +# ======================== +# Performance benchmark for plume strike computation on existing RPOD cases: +# - scalar reference vs NumPy-vectorized geometry +# (1d_approach: geometry only; multi_thrusters_square: Simplified kinetics) +# - serial vs process-parallel execution across JFH firings +# +# Elapsed wall-clock times are printed (run pytest with -s to see them); no +# speedup thresholds are asserted, so timing noise on shared/CI machines can +# never fail the build. Each benchmark only asserts that the compared paths +# produce identical strike arrays, i.e. that the benchmark code ran +# successfully. Correctness itself is covered by +# tests/plume/plume_unit_test_04.py and tests/plume/plume_integration_test_02.py. +# +# Note: the parallel benchmark includes worker process startup (on Windows, +# spawn re-imports pyrpod/scipy per worker), which dominates on the small +# repository cases; parallel execution is intended for large meshes and/or +# many firings. + +import time +import unittest + +import numpy as np + +from pyrpod.mission import MissionEnvironment +from pyrpod.plume.PlumeStrikeCalculator import ( + _compute_plume_strikes_scalar, + compute_face_centroids, + compute_plume_strikes, + extract_plume_params, + run_parallel_plume_strikes, +) +from pyrpod.rpod import JetFiringHistory +from pyrpod.vehicle import TargetVehicle, VisitingVehicle + + +def load_case(case_dir): + jfh = JetFiringHistory.JetFiringHistory(case_dir) + jfh.read_jfh() + + tv = TargetVehicle.TargetVehicle(case_dir) + tv.set_stl() + + vv = VisitingVehicle.VisitingVehicle(case_dir) + vv.set_thruster_config() + vv.set_thruster_metrics() + + me = MissionEnvironment.MissionEnvironment(case_dir) + return jfh, tv, vv, me + + +def build_steps(jfh): + return [ + { + 'thrusters': jfh.JFH[firing]['thrusters'], + 'xyz': np.array(jfh.JFH[firing]['xyz']), + 'dcm': np.array(jfh.JFH[firing]['dcm']), + 't': float(jfh.JFH[firing]['t']), + } + for firing in range(len(jfh.JFH)) + ] + + +def timed(fn): + start = time.perf_counter() + result = fn() + return time.perf_counter() - start, result + + +class PlumeStrikeBenchmarkChecks(unittest.TestCase): + + def benchmark_scalar_vs_vectorized(self, case_dir): + jfh, tv, vv, me = load_case(case_dir) + target = tv.mesh + normals = target.get_unit_normals() + centroids = compute_face_centroids(target.vectors) + steps = build_steps(jfh) + + t_vec, vec_results = timed(lambda: [ + compute_plume_strikes(target, normals, vv, step, me, + face_centroids=centroids) + for step in steps + ]) + t_scalar, scalar_results = timed(lambda: [ + _compute_plume_strikes_scalar(target, normals, vv, step, me) + for step in steps + ]) + + print(f"\n[benchmark] {case_dir} " + f"({len(steps)} firings x {len(target.vectors)} faces, " + f"kinetics: {me.config['pm']['kinetics']})") + print(f"[benchmark] vectorized (serial): {t_vec:8.3f} s") + print(f"[benchmark] scalar reference: {t_scalar:8.3f} s " + f"({t_scalar / t_vec:5.1f}x slower than vectorized)") + + # Success criterion: both paths ran and agree exactly. + for vec, ref in zip(vec_results, scalar_results): + np.testing.assert_array_equal(vec['strikes'], ref['strikes']) + + def test_scalar_vs_vectorized_geometry(self): + self.benchmark_scalar_vs_vectorized('../case/rpod/1d_approach/') + + def test_scalar_vs_vectorized_kinetics(self): + self.benchmark_scalar_vs_vectorized('../case/rpod/multi_thrusters_square/') + + def test_serial_vs_parallel(self): + case_dir = '../case/rpod/1d_approach/' + jfh, tv, vv, me = load_case(case_dir) + target = tv.mesh + normals = target.get_unit_normals() + centroids = compute_face_centroids(target.vectors) + steps = build_steps(jfh) + workers = min(2, len(steps)) + + t_serial, serial_results = timed(lambda: [ + compute_plume_strikes(target, normals, vv, step, me, + face_centroids=centroids) + for step in steps + ]) + t_par, par_results = timed(lambda: run_parallel_plume_strikes( + jfh_steps=steps, + face_centroids=centroids, + target_unit_normals=normals, + thruster_data=vv.thruster_data, + thruster_metrics=getattr(vv, 'thruster_metrics', None), + plume_params=extract_plume_params(me), + workers=workers, + )) + + print(f"\n[benchmark] {case_dir} ({len(steps)} firings)") + print(f"[benchmark] serial (vectorized): {t_serial:8.3f} s") + print(f"[benchmark] parallel ({workers} workers): {t_par:8.3f} s " + f"(includes process startup)") + + # Success criterion: both paths ran and agree exactly. + for ser, par in zip(serial_results, par_results): + np.testing.assert_array_equal(ser['strikes'], par['strikes']) + + +if __name__ == '__main__': + unittest.main() From 732373ef1b86893e9e5ed45cd0d62c301083caa5 Mon Sep 17 00:00:00 2001 From: andytorrestb Date: Sat, 18 Jul 2026 19:21:41 -0500 Subject: [PATCH 7/7] normalize PlumeStrikeCalculator.py to CRLF line endings The full-file rewrite in 5f44163 wrote the file with LF endings, while the rest of the repository (including this file on master) uses CRLF. Because core.autocrlf is true with no .gitattributes, the EOL flip made git treat all 137 original lines as deleted-and-re-added, inflating the diff to +417/-137 when the real change is +281/-1. Restoring CRLF (content byte-identical to the previous commit) collapses the diff to the actual change. No functional change. Co-Authored-By: Claude Fable 5 --- pyrpod/plume/PlumeStrikeCalculator.py | 834 +++++++++++++------------- 1 file changed, 417 insertions(+), 417 deletions(-) diff --git a/pyrpod/plume/PlumeStrikeCalculator.py b/pyrpod/plume/PlumeStrikeCalculator.py index 755f8c9..bbadf98 100644 --- a/pyrpod/plume/PlumeStrikeCalculator.py +++ b/pyrpod/plume/PlumeStrikeCalculator.py @@ -1,417 +1,417 @@ -""" -Plume impingement computations for RPOD. - -Responsibilities: -- Given target mesh, VV pose, and active thrusters, compute per-face strike metrics -- Return numpy arrays/dicts; do not write files - -This consolidates logic currently in RPOD.jfh_plume_strikes into -reusable, testable functions. - -Implementation notes: -- compute_plume_strikes() runs a NumPy-vectorized strike-detection path by - default. _compute_plume_strikes_scalar() preserves the original per-face - loop verbatim as a reference implementation for tests and benchmarking. -- The vectorized core operates on plain serializable inputs (arrays, dicts, - floats) so it can also run inside process-based workers. - -Future work (no new dependencies planned): -- Vectorize the SimplifiedGasKinetics evaluations for struck faces. -- Shared-memory arrays (multiprocessing.shared_memory) for very large meshes. -- Chunking strategy to batch many small firings per worker task. -""" -from __future__ import annotations - -from concurrent.futures import ProcessPoolExecutor -from typing import Any, Dict, List, Optional, Sequence - -import numpy as np -from pyrpod.plume.RarefiedPlumeGasKinetics import SimplifiedGasKinetics - - -def compute_face_centroids(vectors: np.ndarray) -> np.ndarray: - """Compute per-face centroids for an (N x 3 x 3) array of face vertices. - - Averages the three vertices of each face, matching the scalar reference - (mean over each coordinate in the face's native dtype). The target is - stationary during a run, so callers should compute this once and pass it - to compute_plume_strikes() via face_centroids. - """ - return np.asarray(vectors).mean(axis=1) - - -def _build_thruster_link(thruster_data: Dict[str, Any]) -> Dict[str, Any]: - """Map numeric JFH thruster indices ('1', '2', ...) to thruster names, - consistent with legacy ordering of the thruster configuration.""" - link = {} - i = 1 - for thruster in thruster_data: - link[str(i)] = thruster_data[thruster]['name'] - i += 1 - return link - - -def extract_plume_params(environment: Any) -> Dict[str, Any]: - """Extract the plain config values needed for strike computation. - - Returns a picklable dict (radius, wedge_theta, use_kinetics, and — only - when kinetics is enabled — surface_temp and sigma) so workers never need - the full environment object. - """ - config = environment.config - use_kinetics = config['pm']['kinetics'] != 'None' - params: Dict[str, Any] = { - 'radius': float(config['plume']['radius']), - 'wedge_theta': float(config['plume']['wedge_theta']), - 'use_kinetics': use_kinetics, - 'surface_temp': None, - 'sigma': None, - } - if use_kinetics: - params['surface_temp'] = float(config['tv']['surface_temp']) - params['sigma'] = float(config['tv']['sigma']) - return params - - -def _compute_plume_strikes_core( - face_centroids: np.ndarray, - target_unit_normals: np.ndarray, - thruster_data: Dict[str, Any], - thruster_metrics: Optional[Dict[str, Any]], - jfh_step: Dict[str, Any], - plume_params: Dict[str, Any], -) -> Dict[str, np.ndarray]: - """Vectorized strike computation on plain serializable inputs. - - Geometry is evaluated with NumPy over all faces per active thruster. - Gas-kinetics quantities remain scalar: SimplifiedGasKinetics is - instantiated only for struck face indices, exactly as in the scalar - reference. Memory scales with the number of faces (a few (N,) and (N,3) - temporaries), independent of the number of firings. - """ - num_faces = len(face_centroids) - strikes = np.zeros(num_faces) - - use_kinetics = plume_params['use_kinetics'] - if use_kinetics: - pressures = np.zeros(num_faces) - shear_stresses = np.zeros(num_faces) - heat_flux = np.zeros(num_faces) - heat_flux_load = np.zeros(num_faces) - - vv_pos = np.array(jfh_step['xyz']) - vv_orientation = np.array(jfh_step['dcm']).transpose() - thrusters = jfh_step['thrusters'] - firing_time = float(jfh_step['t']) if 't' in jfh_step else 0.0 - - link = _build_thruster_link(thruster_data) - - plume_radius = float(plume_params['radius']) - wedge_theta = float(plume_params['wedge_theta']) - - normals = np.asarray(target_unit_normals) - - for thr in thrusters: - thruster_id = link[str(thr)][0] - - thruster_orientation = np.array(thruster_data[thruster_id]['dcm']).transpose() - thruster_orientation = thruster_orientation.dot(vv_orientation) - plume_normal = np.array(thruster_orientation[0]) - norm_plume_normal = np.linalg.norm(plume_normal) - unit_plume_normal = plume_normal / norm_plume_normal - - thr_exit = np.array(thruster_data[thruster_id]['exit']) - thruster_pos = vv_pos + thr_exit - thruster_pos = thruster_pos[0] - - distance = thruster_pos - face_centroids - norm_distance = np.linalg.norm(distance, axis=1) - - # Faces whose centroid coincides with the thruster exit are skipped, - # matching the scalar reference's `norm_distance == 0` guard. - valid = norm_distance != 0.0 - unit_distance = np.zeros_like(distance) - np.divide( - distance, - norm_distance[:, np.newaxis], - out=unit_distance, - where=valid[:, np.newaxis], - ) - - # NOTE: 3.14 (not np.pi) is kept deliberately to reproduce the legacy - # scalar reference bit-for-bit; changing it shifts theta by ~1.6e-3 rad - # and can alter struck-face IDs near the wedge boundary. - theta = 3.14 - np.arccos((unit_distance * unit_plume_normal).sum(axis=1)) - - surface_dot_plume = (normals * unit_plume_normal).sum(axis=1) - - hit = ( - valid - & (norm_distance < plume_radius) - & (theta < wedge_theta) - & (surface_dot_plume < 0) - ) - - strikes[hit] += 1 - - if use_kinetics: - T_w = plume_params['surface_temp'] - sigma = plume_params['sigma'] - t_type = thruster_data[thruster_id]['type'][0] - metrics = thruster_metrics[t_type] - for idx in np.nonzero(hit)[0]: - simple_plume = SimplifiedGasKinetics( - norm_distance[idx], theta[idx], metrics, T_w, sigma - ) - pressures[idx] += simple_plume.get_pressure() - shear = simple_plume.get_shear_pressure() - shear_stresses[idx] += abs(shear) - hf = simple_plume.get_heat_flux() - heat_flux[idx] += hf - heat_flux_load[idx] += hf * firing_time - - result = {"strikes": strikes} - if use_kinetics: - result.update({ - "pressures": pressures, - "shear_stress": shear_stresses, - "heat_flux_rate": heat_flux, - "heat_flux_load": heat_flux_load, - }) - return result - - -def compute_plume_strikes( - target_mesh: Any, - target_unit_normals: np.ndarray, - vv: Any, - jfh_step: Dict[str, Any], - environment: Any, - face_centroids: Optional[np.ndarray] = None, -) -> Dict[str, np.ndarray]: - """Compute plume strike arrays for a single JFH step. - - Inputs - - target_mesh: numpy-stl Mesh-like, exposes .vectors (N x 3 x 3) - - target_unit_normals: (N x 3) array of per-face unit normals - - vv: Visiting vehicle with thruster_data and thruster_metrics - - jfh_step: dict with keys 'thrusters' (list[int]), 'xyz' (pos), 'dcm' (3x3) - - environment: provides config for plume and kinetics - - face_centroids: optional (N x 3) precomputed face centroids - (see compute_face_centroids). When the target is stationary, callers - should compute centroids once per run and pass them here; if omitted, - they are computed from target_mesh for this step. - - Returns - - dict with per-face arrays for current step: strikes and optionally pressures, shear_stress, heat_flux_rate, heat_flux_load - """ - if face_centroids is None: - face_centroids = compute_face_centroids(target_mesh.vectors) - plume_params = extract_plume_params(environment) - return _compute_plume_strikes_core( - face_centroids=face_centroids, - target_unit_normals=target_unit_normals, - thruster_data=vv.thruster_data, - # Only defined/needed when kinetics is enabled; the core only reads it - # for struck faces, matching the scalar reference. - thruster_metrics=getattr(vv, 'thruster_metrics', None), - jfh_step=jfh_step, - plume_params=plume_params, - ) - - -def _compute_plume_strikes_scalar( - target_mesh: Any, - target_unit_normals: np.ndarray, - vv: Any, - jfh_step: Dict[str, Any], - environment: Any, -) -> Dict[str, np.ndarray]: - """Scalar reference implementation of compute_plume_strikes(). - - Preserved verbatim from the original per-face loop. Kept for regression - tests, debugging, and benchmarking against the vectorized path; the two - must produce identical strike arrays and struck-face IDs. - """ - num_faces = len(target_mesh.vectors) - strikes = np.zeros(num_faces) - - use_kinetics = environment.config['pm']['kinetics'] != 'None' - if use_kinetics: - pressures = np.zeros(num_faces) - shear_stresses = np.zeros(num_faces) - heat_flux = np.zeros(num_faces) - heat_flux_load = np.zeros(num_faces) - - vv_pos = np.array(jfh_step['xyz']) - vv_orientation = np.array(jfh_step['dcm']).transpose() - thrusters = jfh_step['thrusters'] - firing_time = float(jfh_step['t']) if 't' in jfh_step else 0.0 - - # Build mapping from numeric JFH indices to thruster ids consistent with legacy - link = {} - i = 1 - for thruster in vv.thruster_data: - link[str(i)] = vv.thruster_data[thruster]['name'] - i += 1 - - plume_radius = float(environment.config['plume']['radius']) - wedge_theta = float(environment.config['plume']['wedge_theta']) - - for thr in thrusters: - thruster_id = link[str(thr)][0] - - thruster_orientation = np.array(vv.thruster_data[thruster_id]['dcm']).transpose() - thruster_orientation = thruster_orientation.dot(vv_orientation) - plume_normal = np.array(thruster_orientation[0]) - norm_plume_normal = np.linalg.norm(plume_normal) - unit_plume_normal = plume_normal / norm_plume_normal - - thr_exit = np.array(vv.thruster_data[thruster_id]['exit']) - thruster_pos = vv_pos + thr_exit - thruster_pos = thruster_pos[0] - - for idx, face in enumerate(target_mesh.vectors): - face = np.array(face).transpose() - centroid = np.array([face[0].mean(), face[1].mean(), face[2].mean()]) - distance = thruster_pos - centroid - norm_distance = np.linalg.norm(distance) - if norm_distance == 0: - continue - unit_distance = distance / norm_distance - - theta = 3.14 - np.arccos(np.dot(np.squeeze(unit_distance), np.squeeze(unit_plume_normal))) - - n = np.squeeze(target_unit_normals[idx]) - unit_plume = np.squeeze(plume_normal / norm_plume_normal) - surface_dot_plume = np.dot(n, unit_plume) - - within_distance = float(norm_distance) < plume_radius - within_theta = float(theta) < wedge_theta - facing_thruster = surface_dot_plume < 0 - - if within_distance and within_theta and facing_thruster: - strikes[idx] += 1 - if use_kinetics: - T_w = float(environment.config['tv']['surface_temp']) - sigma = float(environment.config['tv']['sigma']) - t_type = vv.thruster_data[thruster_id]['type'][0] - thruster_metrics = vv.thruster_metrics[t_type] - simple_plume = SimplifiedGasKinetics(norm_distance, theta, thruster_metrics, T_w, sigma) - pressures[idx] += simple_plume.get_pressure() - shear = simple_plume.get_shear_pressure() - shear_stresses[idx] += abs(shear) - hf = simple_plume.get_heat_flux() - heat_flux[idx] += hf - heat_flux_load[idx] += hf * firing_time - - result = {"strikes": strikes} - if use_kinetics: - result.update({ - "pressures": pressures, - "shear_stress": shear_stresses, - "heat_flux_rate": heat_flux, - "heat_flux_load": heat_flux_load, - }) - return result - - -# Per-process state for parallel workers. Populated once per worker by -# _parallel_worker_init so the (N,3) target arrays are shipped to each worker -# a single time instead of once per submitted firing. Memory therefore scales -# with workers x faces, never firings x faces. -_WORKER_STATE: Dict[str, Any] = {} - - -def _parallel_worker_init( - face_centroids: np.ndarray, - target_unit_normals: np.ndarray, - thruster_data: Dict[str, Any], - thruster_metrics: Optional[Dict[str, Any]], - plume_params: Dict[str, Any], -) -> None: - """ProcessPoolExecutor initializer: cache shared per-run inputs.""" - _WORKER_STATE['face_centroids'] = face_centroids - _WORKER_STATE['target_unit_normals'] = target_unit_normals - _WORKER_STATE['thruster_data'] = thruster_data - _WORKER_STATE['thruster_metrics'] = thruster_metrics - _WORKER_STATE['plume_params'] = plume_params - - -def _parallel_worker_compute(task) -> Any: - """Compute strikes for one firing inside a worker process. - - task is (firing_index, jfh_step); returns (firing_index, result dict). - """ - firing_index, jfh_step = task - result = _compute_plume_strikes_core( - face_centroids=_WORKER_STATE['face_centroids'], - target_unit_normals=_WORKER_STATE['target_unit_normals'], - thruster_data=_WORKER_STATE['thruster_data'], - thruster_metrics=_WORKER_STATE['thruster_metrics'], - jfh_step=jfh_step, - plume_params=_WORKER_STATE['plume_params'], - ) - return firing_index, result - - -def run_parallel_plume_strikes( - jfh_steps: Sequence[Dict[str, Any]], - face_centroids: np.ndarray, - target_unit_normals: np.ndarray, - thruster_data: Dict[str, Any], - thruster_metrics: Optional[Dict[str, Any]], - plume_params: Dict[str, Any], - workers: int, -) -> List[Dict[str, np.ndarray]]: - """Compute per-firing strike results across processes, one firing per task. - - All inputs must be plain serializable data (NumPy arrays, dicts, - primitives) — full study/vehicle/environment objects are never pickled. - Results are returned as a list indexed by firing, preserving JFH order - regardless of completion order; cumulative accumulation and VTK output - remain the caller's responsibility (serial, in the parent process). - - Raises whatever the executor or workers raise; callers are expected to - fall back to the serial path with a clear message. - """ - results: List[Optional[Dict[str, np.ndarray]]] = [None] * len(jfh_steps) - with ProcessPoolExecutor( - max_workers=workers, - initializer=_parallel_worker_init, - initargs=( - face_centroids, - target_unit_normals, - thruster_data, - thruster_metrics, - plume_params, - ), - ) as executor: - futures = [ - executor.submit(_parallel_worker_compute, (i, step)) - for i, step in enumerate(jfh_steps) - ] - for future in futures: - firing_index, result = future.result() - results[firing_index] = result - return results - - -def accumulate_cumulative( - cumulative: Dict[str, np.ndarray], - current: Dict[str, np.ndarray], -) -> Dict[str, np.ndarray]: - """Accumulate per-step arrays into cumulative tallies (e.g., cum_strikes, max_pressures).""" - if "cum_strikes" in cumulative and "strikes" in current: - cumulative["cum_strikes"] = cumulative["cum_strikes"] + current["strikes"] - - # Max trackers if available - if "max_pressures" in cumulative and "pressures" in current: - cumulative["max_pressures"] = np.maximum(cumulative["max_pressures"], current["pressures"]) - if "max_shears" in cumulative and "shear_stress" in current: - cumulative["max_shears"] = np.maximum(cumulative["max_shears"], current["shear_stress"]) - - if "cum_heat_flux_load" in cumulative and "heat_flux_load" in current: - cumulative["cum_heat_flux_load"] = cumulative["cum_heat_flux_load"] + current["heat_flux_load"] - - return cumulative +""" +Plume impingement computations for RPOD. + +Responsibilities: +- Given target mesh, VV pose, and active thrusters, compute per-face strike metrics +- Return numpy arrays/dicts; do not write files + +This consolidates logic currently in RPOD.jfh_plume_strikes into +reusable, testable functions. + +Implementation notes: +- compute_plume_strikes() runs a NumPy-vectorized strike-detection path by + default. _compute_plume_strikes_scalar() preserves the original per-face + loop verbatim as a reference implementation for tests and benchmarking. +- The vectorized core operates on plain serializable inputs (arrays, dicts, + floats) so it can also run inside process-based workers. + +Future work (no new dependencies planned): +- Vectorize the SimplifiedGasKinetics evaluations for struck faces. +- Shared-memory arrays (multiprocessing.shared_memory) for very large meshes. +- Chunking strategy to batch many small firings per worker task. +""" +from __future__ import annotations + +from concurrent.futures import ProcessPoolExecutor +from typing import Any, Dict, List, Optional, Sequence + +import numpy as np +from pyrpod.plume.RarefiedPlumeGasKinetics import SimplifiedGasKinetics + + +def compute_face_centroids(vectors: np.ndarray) -> np.ndarray: + """Compute per-face centroids for an (N x 3 x 3) array of face vertices. + + Averages the three vertices of each face, matching the scalar reference + (mean over each coordinate in the face's native dtype). The target is + stationary during a run, so callers should compute this once and pass it + to compute_plume_strikes() via face_centroids. + """ + return np.asarray(vectors).mean(axis=1) + + +def _build_thruster_link(thruster_data: Dict[str, Any]) -> Dict[str, Any]: + """Map numeric JFH thruster indices ('1', '2', ...) to thruster names, + consistent with legacy ordering of the thruster configuration.""" + link = {} + i = 1 + for thruster in thruster_data: + link[str(i)] = thruster_data[thruster]['name'] + i += 1 + return link + + +def extract_plume_params(environment: Any) -> Dict[str, Any]: + """Extract the plain config values needed for strike computation. + + Returns a picklable dict (radius, wedge_theta, use_kinetics, and — only + when kinetics is enabled — surface_temp and sigma) so workers never need + the full environment object. + """ + config = environment.config + use_kinetics = config['pm']['kinetics'] != 'None' + params: Dict[str, Any] = { + 'radius': float(config['plume']['radius']), + 'wedge_theta': float(config['plume']['wedge_theta']), + 'use_kinetics': use_kinetics, + 'surface_temp': None, + 'sigma': None, + } + if use_kinetics: + params['surface_temp'] = float(config['tv']['surface_temp']) + params['sigma'] = float(config['tv']['sigma']) + return params + + +def _compute_plume_strikes_core( + face_centroids: np.ndarray, + target_unit_normals: np.ndarray, + thruster_data: Dict[str, Any], + thruster_metrics: Optional[Dict[str, Any]], + jfh_step: Dict[str, Any], + plume_params: Dict[str, Any], +) -> Dict[str, np.ndarray]: + """Vectorized strike computation on plain serializable inputs. + + Geometry is evaluated with NumPy over all faces per active thruster. + Gas-kinetics quantities remain scalar: SimplifiedGasKinetics is + instantiated only for struck face indices, exactly as in the scalar + reference. Memory scales with the number of faces (a few (N,) and (N,3) + temporaries), independent of the number of firings. + """ + num_faces = len(face_centroids) + strikes = np.zeros(num_faces) + + use_kinetics = plume_params['use_kinetics'] + if use_kinetics: + pressures = np.zeros(num_faces) + shear_stresses = np.zeros(num_faces) + heat_flux = np.zeros(num_faces) + heat_flux_load = np.zeros(num_faces) + + vv_pos = np.array(jfh_step['xyz']) + vv_orientation = np.array(jfh_step['dcm']).transpose() + thrusters = jfh_step['thrusters'] + firing_time = float(jfh_step['t']) if 't' in jfh_step else 0.0 + + link = _build_thruster_link(thruster_data) + + plume_radius = float(plume_params['radius']) + wedge_theta = float(plume_params['wedge_theta']) + + normals = np.asarray(target_unit_normals) + + for thr in thrusters: + thruster_id = link[str(thr)][0] + + thruster_orientation = np.array(thruster_data[thruster_id]['dcm']).transpose() + thruster_orientation = thruster_orientation.dot(vv_orientation) + plume_normal = np.array(thruster_orientation[0]) + norm_plume_normal = np.linalg.norm(plume_normal) + unit_plume_normal = plume_normal / norm_plume_normal + + thr_exit = np.array(thruster_data[thruster_id]['exit']) + thruster_pos = vv_pos + thr_exit + thruster_pos = thruster_pos[0] + + distance = thruster_pos - face_centroids + norm_distance = np.linalg.norm(distance, axis=1) + + # Faces whose centroid coincides with the thruster exit are skipped, + # matching the scalar reference's `norm_distance == 0` guard. + valid = norm_distance != 0.0 + unit_distance = np.zeros_like(distance) + np.divide( + distance, + norm_distance[:, np.newaxis], + out=unit_distance, + where=valid[:, np.newaxis], + ) + + # NOTE: 3.14 (not np.pi) is kept deliberately to reproduce the legacy + # scalar reference bit-for-bit; changing it shifts theta by ~1.6e-3 rad + # and can alter struck-face IDs near the wedge boundary. + theta = 3.14 - np.arccos((unit_distance * unit_plume_normal).sum(axis=1)) + + surface_dot_plume = (normals * unit_plume_normal).sum(axis=1) + + hit = ( + valid + & (norm_distance < plume_radius) + & (theta < wedge_theta) + & (surface_dot_plume < 0) + ) + + strikes[hit] += 1 + + if use_kinetics: + T_w = plume_params['surface_temp'] + sigma = plume_params['sigma'] + t_type = thruster_data[thruster_id]['type'][0] + metrics = thruster_metrics[t_type] + for idx in np.nonzero(hit)[0]: + simple_plume = SimplifiedGasKinetics( + norm_distance[idx], theta[idx], metrics, T_w, sigma + ) + pressures[idx] += simple_plume.get_pressure() + shear = simple_plume.get_shear_pressure() + shear_stresses[idx] += abs(shear) + hf = simple_plume.get_heat_flux() + heat_flux[idx] += hf + heat_flux_load[idx] += hf * firing_time + + result = {"strikes": strikes} + if use_kinetics: + result.update({ + "pressures": pressures, + "shear_stress": shear_stresses, + "heat_flux_rate": heat_flux, + "heat_flux_load": heat_flux_load, + }) + return result + + +def compute_plume_strikes( + target_mesh: Any, + target_unit_normals: np.ndarray, + vv: Any, + jfh_step: Dict[str, Any], + environment: Any, + face_centroids: Optional[np.ndarray] = None, +) -> Dict[str, np.ndarray]: + """Compute plume strike arrays for a single JFH step. + + Inputs + - target_mesh: numpy-stl Mesh-like, exposes .vectors (N x 3 x 3) + - target_unit_normals: (N x 3) array of per-face unit normals + - vv: Visiting vehicle with thruster_data and thruster_metrics + - jfh_step: dict with keys 'thrusters' (list[int]), 'xyz' (pos), 'dcm' (3x3) + - environment: provides config for plume and kinetics + - face_centroids: optional (N x 3) precomputed face centroids + (see compute_face_centroids). When the target is stationary, callers + should compute centroids once per run and pass them here; if omitted, + they are computed from target_mesh for this step. + + Returns + - dict with per-face arrays for current step: strikes and optionally pressures, shear_stress, heat_flux_rate, heat_flux_load + """ + if face_centroids is None: + face_centroids = compute_face_centroids(target_mesh.vectors) + plume_params = extract_plume_params(environment) + return _compute_plume_strikes_core( + face_centroids=face_centroids, + target_unit_normals=target_unit_normals, + thruster_data=vv.thruster_data, + # Only defined/needed when kinetics is enabled; the core only reads it + # for struck faces, matching the scalar reference. + thruster_metrics=getattr(vv, 'thruster_metrics', None), + jfh_step=jfh_step, + plume_params=plume_params, + ) + + +def _compute_plume_strikes_scalar( + target_mesh: Any, + target_unit_normals: np.ndarray, + vv: Any, + jfh_step: Dict[str, Any], + environment: Any, +) -> Dict[str, np.ndarray]: + """Scalar reference implementation of compute_plume_strikes(). + + Preserved verbatim from the original per-face loop. Kept for regression + tests, debugging, and benchmarking against the vectorized path; the two + must produce identical strike arrays and struck-face IDs. + """ + num_faces = len(target_mesh.vectors) + strikes = np.zeros(num_faces) + + use_kinetics = environment.config['pm']['kinetics'] != 'None' + if use_kinetics: + pressures = np.zeros(num_faces) + shear_stresses = np.zeros(num_faces) + heat_flux = np.zeros(num_faces) + heat_flux_load = np.zeros(num_faces) + + vv_pos = np.array(jfh_step['xyz']) + vv_orientation = np.array(jfh_step['dcm']).transpose() + thrusters = jfh_step['thrusters'] + firing_time = float(jfh_step['t']) if 't' in jfh_step else 0.0 + + # Build mapping from numeric JFH indices to thruster ids consistent with legacy + link = {} + i = 1 + for thruster in vv.thruster_data: + link[str(i)] = vv.thruster_data[thruster]['name'] + i += 1 + + plume_radius = float(environment.config['plume']['radius']) + wedge_theta = float(environment.config['plume']['wedge_theta']) + + for thr in thrusters: + thruster_id = link[str(thr)][0] + + thruster_orientation = np.array(vv.thruster_data[thruster_id]['dcm']).transpose() + thruster_orientation = thruster_orientation.dot(vv_orientation) + plume_normal = np.array(thruster_orientation[0]) + norm_plume_normal = np.linalg.norm(plume_normal) + unit_plume_normal = plume_normal / norm_plume_normal + + thr_exit = np.array(vv.thruster_data[thruster_id]['exit']) + thruster_pos = vv_pos + thr_exit + thruster_pos = thruster_pos[0] + + for idx, face in enumerate(target_mesh.vectors): + face = np.array(face).transpose() + centroid = np.array([face[0].mean(), face[1].mean(), face[2].mean()]) + distance = thruster_pos - centroid + norm_distance = np.linalg.norm(distance) + if norm_distance == 0: + continue + unit_distance = distance / norm_distance + + theta = 3.14 - np.arccos(np.dot(np.squeeze(unit_distance), np.squeeze(unit_plume_normal))) + + n = np.squeeze(target_unit_normals[idx]) + unit_plume = np.squeeze(plume_normal / norm_plume_normal) + surface_dot_plume = np.dot(n, unit_plume) + + within_distance = float(norm_distance) < plume_radius + within_theta = float(theta) < wedge_theta + facing_thruster = surface_dot_plume < 0 + + if within_distance and within_theta and facing_thruster: + strikes[idx] += 1 + if use_kinetics: + T_w = float(environment.config['tv']['surface_temp']) + sigma = float(environment.config['tv']['sigma']) + t_type = vv.thruster_data[thruster_id]['type'][0] + thruster_metrics = vv.thruster_metrics[t_type] + simple_plume = SimplifiedGasKinetics(norm_distance, theta, thruster_metrics, T_w, sigma) + pressures[idx] += simple_plume.get_pressure() + shear = simple_plume.get_shear_pressure() + shear_stresses[idx] += abs(shear) + hf = simple_plume.get_heat_flux() + heat_flux[idx] += hf + heat_flux_load[idx] += hf * firing_time + + result = {"strikes": strikes} + if use_kinetics: + result.update({ + "pressures": pressures, + "shear_stress": shear_stresses, + "heat_flux_rate": heat_flux, + "heat_flux_load": heat_flux_load, + }) + return result + + +# Per-process state for parallel workers. Populated once per worker by +# _parallel_worker_init so the (N,3) target arrays are shipped to each worker +# a single time instead of once per submitted firing. Memory therefore scales +# with workers x faces, never firings x faces. +_WORKER_STATE: Dict[str, Any] = {} + + +def _parallel_worker_init( + face_centroids: np.ndarray, + target_unit_normals: np.ndarray, + thruster_data: Dict[str, Any], + thruster_metrics: Optional[Dict[str, Any]], + plume_params: Dict[str, Any], +) -> None: + """ProcessPoolExecutor initializer: cache shared per-run inputs.""" + _WORKER_STATE['face_centroids'] = face_centroids + _WORKER_STATE['target_unit_normals'] = target_unit_normals + _WORKER_STATE['thruster_data'] = thruster_data + _WORKER_STATE['thruster_metrics'] = thruster_metrics + _WORKER_STATE['plume_params'] = plume_params + + +def _parallel_worker_compute(task) -> Any: + """Compute strikes for one firing inside a worker process. + + task is (firing_index, jfh_step); returns (firing_index, result dict). + """ + firing_index, jfh_step = task + result = _compute_plume_strikes_core( + face_centroids=_WORKER_STATE['face_centroids'], + target_unit_normals=_WORKER_STATE['target_unit_normals'], + thruster_data=_WORKER_STATE['thruster_data'], + thruster_metrics=_WORKER_STATE['thruster_metrics'], + jfh_step=jfh_step, + plume_params=_WORKER_STATE['plume_params'], + ) + return firing_index, result + + +def run_parallel_plume_strikes( + jfh_steps: Sequence[Dict[str, Any]], + face_centroids: np.ndarray, + target_unit_normals: np.ndarray, + thruster_data: Dict[str, Any], + thruster_metrics: Optional[Dict[str, Any]], + plume_params: Dict[str, Any], + workers: int, +) -> List[Dict[str, np.ndarray]]: + """Compute per-firing strike results across processes, one firing per task. + + All inputs must be plain serializable data (NumPy arrays, dicts, + primitives) — full study/vehicle/environment objects are never pickled. + Results are returned as a list indexed by firing, preserving JFH order + regardless of completion order; cumulative accumulation and VTK output + remain the caller's responsibility (serial, in the parent process). + + Raises whatever the executor or workers raise; callers are expected to + fall back to the serial path with a clear message. + """ + results: List[Optional[Dict[str, np.ndarray]]] = [None] * len(jfh_steps) + with ProcessPoolExecutor( + max_workers=workers, + initializer=_parallel_worker_init, + initargs=( + face_centroids, + target_unit_normals, + thruster_data, + thruster_metrics, + plume_params, + ), + ) as executor: + futures = [ + executor.submit(_parallel_worker_compute, (i, step)) + for i, step in enumerate(jfh_steps) + ] + for future in futures: + firing_index, result = future.result() + results[firing_index] = result + return results + + +def accumulate_cumulative( + cumulative: Dict[str, np.ndarray], + current: Dict[str, np.ndarray], +) -> Dict[str, np.ndarray]: + """Accumulate per-step arrays into cumulative tallies (e.g., cum_strikes, max_pressures).""" + if "cum_strikes" in cumulative and "strikes" in current: + cumulative["cum_strikes"] = cumulative["cum_strikes"] + current["strikes"] + + # Max trackers if available + if "max_pressures" in cumulative and "pressures" in current: + cumulative["max_pressures"] = np.maximum(cumulative["max_pressures"], current["pressures"]) + if "max_shears" in cumulative and "shear_stress" in current: + cumulative["max_shears"] = np.maximum(cumulative["max_shears"], current["shear_stress"]) + + if "cum_heat_flux_load" in cumulative and "heat_flux_load" in current: + cumulative["cum_heat_flux_load"] = cumulative["cum_heat_flux_load"] + current["heat_flux_load"] + + return cumulative