Im calc rust

pyproject.toml

@@ -7,11 +7,11 @@ name = "workflow"
 authors = [{ name = "QuakeCoRE" }]
 description = "Next generation workflow powered by cylc."
 readme = "README.md"
-requires-python = ">=3.12,<3.14"
+requires-python = ">=3.12"
 dynamic = ["version"]
 dependencies = [
   # UCGMSim Dependencies
-  "im-calculation",
+  "im-calculation>=2025.12.5",
   "velocity-modelling>=2025.12.1",
   "nshmdb>=2025.12.1",
   "oq_wrapper>=2025.12.3",
@@ -24,7 +24,6 @@ dependencies = [
   "xarray[io]",
 
   # Numerics
-  "numexpr",
   "numpy",
   "scipy",
   "shapely",

workflow/scripts/bb_sim.py

@@ -162,7 +162,7 @@ def combine_hf_and_lf(
         station_vs30_ffp,
         sep=r"\s+",
         header=None,
-        names=["station", "vsite"],  # type: ignore[invalid-argument-type]
+        names=["station", "vsite"],
     ).set_index("station")
     vs30_df["vsite"] = vs30_df["vsite"].astype(np.float32)
     vs30_df = vs30_df.loc[common_stations]

workflow/scripts/generate_station_coordinates.py

@@ -79,7 +79,7 @@ def generate_fd_files(
         stat_file,
         delimiter=r"\s+",
         comment="#",
-        names=["lon", "lat", "name"],  # type: ignore[invalid-argument-type]
+        names=["lon", "lat", "name"],
     )
 
     x, y = proj(

workflow/scripts/hf_sim.py

@@ -301,7 +301,7 @@ def run_hf(
         station_file,
         delimiter=r"\s+",
         header=None,
-        names=["longitude", "latitude", "name"],  # type: ignore[invalid-argument-type]
+        names=["longitude", "latitude", "name"],
     ).set_index("name")
     station_hashes = np.array(
         [stable_hash(name) for name in stations.index], dtype=np.int32
@@ -332,8 +332,8 @@ def run_hf(
                 hf_sim_path,
                 hf_input_template,
                 station["latitude"],
-                station["longitude"],  # type: ignore[invalid-argument-type]
-                name,
+                station["longitude"],
+                str(name),
                 int(station["seed"]),
             )
             for name, station in stations.iterrows()

workflow/scripts/im_calc.py

@@ -29,9 +29,8 @@
 
 import functools
 from pathlib import Path
-from typing import Annotated, Optional
+from typing import Annotated
 
-import numexpr as ne
 import numpy as np
 import pandas as pd
 import tqdm
@@ -50,6 +49,8 @@
     SourceConfig,
 )
 
+PSA_STEP = 10000
+
 app = typer.Typer()
 
 
@@ -63,13 +64,12 @@ def calculate_instensity_measures(
     ],
     output_path: Annotated[Path, typer.Argument(dir_okay=False, writable=True)],
     simulated_stations: Annotated[bool, typer.Option()] = True,
-    psa_rotd_maximum_memory_allocation: Annotated[
-        Optional[float], typer.Option(min=0)
-    ] = None,
+    psa_step: Annotated[int, typer.Option()] = PSA_STEP,
     ko_directory: Annotated[
         Path | None, typer.Option(exists=True, file_okay=False)
     ] = None,
     override_ims: Annotated[list[IM] | None, typer.Option("-i", "--im")] = None,
+    cores: Annotated[int | None, typer.Option(min=1)] = None,
 ) -> None:
     """Calculate intensity measures for simulation data.
 
@@ -83,14 +83,18 @@ def calculate_instensity_measures(
         Output directory for IM calc summary statistics.
     simulated_stations : bool, default True
         If passed, calculate for simulated stations.
-    psa_rotd_maximum_memory_allocation : Optional[float]
-        Maximum amount of memory allocated for rotated PSA calculation station buffer, in gigabytes.
+    psa_step : int
+        Maximum number of stations to read from disk at once for pSA calculation
     ko_directory : Path
         Directory containing the KO matrix files for FAS calculation. Not required for other IMs.
     override_ims : list of str
         Intensity measures to calculate. If not set, reads from the realisation file.
+    cores : int or None
+        Set the number of cores for parallel processing of IMs. If set
+        to `None`, will default to the available cores from
+        `utils.get_available_cores`.
     """
-    ne.set_num_threads(utils.get_available_cores())
+    cores = cores or utils.get_available_cores()
 
     metadata = RealisationMetadata.read_from_realisation(realisation_ffp)
     resolution = Resolution.read_from_realisation_or_defaults(
@@ -121,28 +125,27 @@ def calculate_instensity_measures(
     nyquist_frequency = 1 / (2 * resolution.dt)
 
     im_function_map = {
-        IM.PGA: ims.peak_ground_acceleration,
-        IM.PGV: functools.partial(ims.peak_ground_velocity, dt=resolution.dt),
-        IM.CAV: functools.partial(ims.cumulative_absolute_velocity, dt=resolution.dt),
-        IM.AI: functools.partial(ims.arias_intensity, dt=resolution.dt),
-        IM.Ds575: functools.partial(
-            ims.ds575,
-            dt=resolution.dt,
+        IM.PGA: functools.partial(ims.peak_ground_acceleration, cores=cores),
+        IM.PGV: functools.partial(
+            ims.peak_ground_velocity, dt=resolution.dt, cores=cores
         ),
-        IM.Ds595: functools.partial(
-            ims.ds595,
-            dt=resolution.dt,
+        IM.PGD: functools.partial(
+            ims.peak_ground_displacement, dt=resolution.dt, cores=cores
         ),
+        IM.CAV: functools.partial(
+            ims.cumulative_absolute_velocity, dt=resolution.dt, cores=cores
+        ),
+        IM.AI: functools.partial(ims.arias_intensity, dt=resolution.dt, cores=cores),
+        IM.Ds575: functools.partial(ims.ds575, dt=resolution.dt, cores=cores),
+        IM.Ds595: functools.partial(ims.ds595, dt=resolution.dt, cores=cores),
         IM.pSA: functools.partial(
             ims.pseudo_spectral_acceleration,
             periods=np.array(
-                intensity_measure_parameters.valid_periods, dtype=np.float32
+                intensity_measure_parameters.valid_periods, dtype=np.float64
             ),
             dt=resolution.dt,
-            psa_rotd_maximum_memory_allocation=psa_rotd_maximum_memory_allocation * 1e9
-            if psa_rotd_maximum_memory_allocation
-            else None,
-            cores=utils.get_available_cores(),
+            step=psa_step,
+            cores=cores,
         ),
         IM.FAS: functools.partial(
             ims.fourier_amplitude_spectra,
@@ -151,7 +154,7 @@ def calculate_instensity_measures(
                 intensity_measure_parameters.fas_frequencies <= nyquist_frequency
             ],
             ko_directory=ko_directory,
-            cores=utils.get_available_cores(),
+            cores=cores,
         ),
     }
     latitude = broadband.latitude.values
@@ -215,13 +218,12 @@ def calculate_instensity_measures(
         attrs={"hypo_lat": hypocentre[0], "hypo_lon": hypocentre[1]},
     )
 
-    # Add each column of the DataFrame as a coordinate
-    # TODO: Refactor IM Calculation to use waveforms in (component, station, time) format
-    # Convert (component, station, time) to (station, time, component).
-    waveform = np.transpose(broadband.waveform.values, (1, 2, 0))
+    waveform = broadband.waveform.values.astype(np.float64)
+
     for im_name in (pbar := tqdm.tqdm(intensity_measures)):
         pbar.set_description(im_name)
         im_fn = im_function_map[im_name]
+
         result = im_fn(waveform)
 
         if isinstance(result, pd.DataFrame):

workflow/scripts/realisation_to_srf.py

@@ -782,12 +782,18 @@ def generate_srf(
         single_threaded=single_threaded,
     )
     srf_name = normalise_name(metadata.name)
-    stitch_srf_files(
-        source_config.source_geometries,
-        rupture_propagation,
-        work_directory,
-        srf_name,
-    )
+    if len(source_config.source_geometries) > 1:
+        stitch_srf_files(
+            source_config.source_geometries,
+            rupture_propagation,
+            work_directory,
+            srf_name,
+        )
+    else:
+        source_name = list(source_config.source_geometries)[0]
+        input_srf_path = work_directory / "srf" / f"{normalise_name(source_name)}.srf"
+        output_srf_path = work_directory / f"{srf_name}.srf"
+        shutil.move(input_srf_path, output_srf_path)
     srf_config.write_to_realisation(realisation_ffp)
 
     shutil.copyfile(work_directory / (srf_name + ".srf"), output_srf_filepath)