Im calc rust

workflow/scripts/im_calc.py

@@ -31,7 +31,6 @@
 from pathlib import Path
 from typing import Annotated, Optional
 
-import numexpr as ne
 import numpy as np
 import pandas as pd
 import tqdm
@@ -70,6 +69,7 @@ def calculate_instensity_measures(
         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.
 
@@ -89,8 +89,12 @@ def calculate_instensity_measures(
         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.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=10000,
+            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)