diff --git a/pyrpod/plume/PlumeStrikeCalculator.py b/pyrpod/plume/PlumeStrikeCalculator.py index 241ed6f..bbadf98 100644 --- a/pyrpod/plume/PlumeStrikeCalculator.py +++ b/pyrpod/plume/PlumeStrikeCalculator.py @@ -7,20 +7,187 @@ 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, Tuple +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. @@ -30,10 +197,42 @@ def compute_plume_strikes( - 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) @@ -117,6 +316,87 @@ def compute_plume_strikes( 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 32679c5..f270e86 100644 --- a/pyrpod/rpod/PlumeStrikeEstimationStudy.py +++ b/pyrpod/rpod/PlumeStrikeEstimationStudy.py @@ -27,7 +27,12 @@ ) 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, + extract_plume_params, + run_parallel_plume_strikes, +) logger = get_logger("pyrpod.rpod.PlumeStrikeEstimationStudy") @@ -754,24 +759,102 @@ 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() 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' @@ -782,22 +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, - ) + # 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 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() 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()