Exact pSA solving

IM/ims.py

@@ -581,6 +581,123 @@ def ds595(waveform: Waveform, dt: float) -> xr.Dataset:
     return significant_duration(waveform, dt, 5, 95, IM.Ds595.value)
 
 
+OSCILLATOR_SAMPLES_PER_PERIOD = 10
+"""Minimum samples per oscillator period before a record is upsampled for pSA.
+
+This is the conventional rule (Nigam and Jennings 1969; Boore and Goulet
+2014) for resolving the peak oscillator response.
+"""
+
+_UPSAMPLE_BUDGET = 2**24
+"""Samples of upsampling FFT (across all components) to hold in memory at once."""
+
+
+def _nigam_jennings_coefficients(
+    periods: npt.NDArray[np.float64], dt: float, damping: float = DAMPING
+) -> npt.NDArray[np.float64]:
+    """Exact recurrence coefficients for unit-mass SDOF oscillators.
+
+    The oscillator equation `u'' + 2 ξ ω u' + ω² u = -a_g(t)` is integrated
+    exactly over each step assuming `a_g` varies linearly between samples
+    (Nigam and Jennings 1969). This is the first-order-hold discretisation of
+    the state `[u, u']`, computed as the matrix exponential of the state
+    matrix augmented with the load and its slope, which avoids the
+    cancellation the textbook closed-form expressions suffer when `ω dt` is
+    small.
+
+    Parameters
+    ----------
+    periods : ndarray of float64
+        Natural periods of the oscillators (s).
+    dt : float
+        Timestep of the ground acceleration (s).
+    damping : float, optional
+        Fraction of critical damping.
+
+    Returns
+    -------
+    ndarray of float64
+        Shape `(len(periods), 8)`, one oscillator per row, laid out as
+        documented on `N_COEFFICIENTS` in `src-rust/psa.rs`. The state is
+        rescaled to `[ω² u, u']` so the recurrence yields pseudo-acceleration.
+    """
+    w = 2 * np.pi / periods
+    # Augmented state [u, u', p, p'] with load p = -a_g and constant slope p'.
+    state = np.zeros((len(w), 4, 4))
+    state[:, 0, 1] = 1.0
+    state[:, 1, 0] = -(w**2)
+    state[:, 1, 1] = -2 * damping * w
+    state[:, 1, 2] = 1.0
+    state[:, 2, 3] = 1.0 / dt
+    transition = sp.linalg.expm(state * dt)
+    # This transition matrix predicts p(n + 1) and p'(n + 1), but we only need
+    # the part of the transition matrix that corresponds to the u, u'. The
+    # matrix E = exp(state * dt) will have shape
+    # [Φ Γ0 Γ1; 0 0 1 1; 0 0 0 1]
+    # Where Φ, Γ0, Γ1 are 2 x 2 matrices. Matrix exponentiation shows [u(n + 1),
+    # u'(n + 1), p(n + 1), p'(n + 1)] = E [u(n), u'(n), p(n), p'(n)].
+    # Noting that p'(n) = p(n + 1) - p(n),
+    # u(n+1) = Φ u(n) + Γ0 p(n) + Γ1 (p(n+1) - p(n)), with p = -a_g.
+    phi = transition[:, :2, :2]
+    gamma_0 = transition[:, :2, 2]
+    gamma_1 = transition[:, :2, 3]
+    coefficients = np.concatenate(
+        [phi.reshape(-1, 4), gamma_1 - gamma_0, -gamma_1], axis=1
+    )
+    # Rescale the displacement row and column to pseudo-acceleration.
+    coefficients[:, [1, 4, 6]] *= (w**2)[:, np.newaxis]
+    coefficients[:, 2] /= w**2
+    return coefficients
+
+
+def _upsampling_factors(
+    periods: npt.NDArray[np.float64], dt: float
+) -> npt.NDArray[np.int64]:
+    """Power-of-two upsampling factor giving each period enough samples.
+
+    Parameters
+    ----------
+    periods : ndarray of float64
+        Natural periods of the oscillators (s).
+    dt : float
+        Timestep of the record (s).
+
+    Returns
+    -------
+    ndarray of int64
+        The smallest power of two `f` with `dt / f <= period / 10`.
+    """
+    ratio = OSCILLATOR_SAMPLES_PER_PERIOD * dt / periods
+    exponent = np.ceil(np.log2(np.maximum(ratio, 1.0)))
+    return (2**exponent).astype(np.int64)
+
+
+def _upsample_record(waveforms: np.ndarray, factor: int) -> np.ndarray:
+    """Band-limited (sinc) interpolation of waveforms by an integer factor.
+
+    The record is zero padded to at least twice its length so the periodic
+    interpolation does not wrap the end of the record onto its start.
+
+    Parameters
+    ----------
+    waveforms : ndarray
+        Waveforms of shape `(..., nt)`.
+    factor : int
+        Upsampling factor.
+
+    Returns
+    -------
+    ndarray
+        Shape `(..., (nt - 1) * factor + 1)`, spanning the same time interval
+        and preserving the original samples.
+    """
+    nt = waveforms.shape[-1]
+    padded = np.zeros(waveforms.shape[:-1] + (sp.fft.next_fast_len(2 * nt, real=True),))
+    padded[..., :nt] = waveforms
+    upsampled = sp.signal.resample(padded, padded.shape[-1] * factor, axis=-1)
+    return upsampled[..., : (nt - 1) * factor + 1]
+
+
 def _psa_kernel(
     block: np.ndarray,
     *,
@@ -603,8 +720,31 @@ def _psa_kernel(
     np.ndarray
         A chunk of solved pSA values.
     """
-    (comp_0, comp_90, comp_ver), lead = _components(block)
-    psa = _core._psa(comp_0, comp_90, comp_ver, periods, dt, DAMPING)
+    components, lead = _components(block)
+    n_stations, nt = components.shape[1:]
+    factors = _upsampling_factors(periods, dt)
+    groups = {
+        factor: _nigam_jennings_coefficients(periods[factors == factor], dt / factor)
+        for factor in np.unique(factors)
+    }
+    max_factor = int(factors.max(initial=1))
+    batch = (
+        n_stations
+        if max_factor == 1
+        else max(1, _UPSAMPLE_BUDGET // (len(components) * 2 * nt * max_factor))
+    )
+    psa = np.empty((n_stations, len(periods), len(ROTD_COMPONENTS)))
+    for start in range(0, n_stations, batch):
+        stations = slice(start, start + batch)
+        original = components[:, stations]
+        upsampled = (
+            _upsample_record(original, max_factor) if max_factor > 1 else original
+        )
+        for factor, coefficients in groups.items():
+            # Keep the original samples exactly rather than their
+            # interpolated round trip.
+            chunk = original if factor == 1 else upsampled[..., :: max_factor // factor]
+            psa[stations, factors == factor] = _core._psa(*chunk, coefficients)
     return psa.reshape(lead + (len(periods), len(ROTD_COMPONENTS)))
 
 
@@ -616,7 +756,7 @@ def pseudo_spectral_acceleration(
     """Compute pseudo-spectral acceleration (PSA) statistics.
 
     Calculates PSA for single-degree-of-freedom oscillators across various
-    periods using the Newmark-beta method and computes rotated (RotD) statistics.
+    periods.
 
     Parameters
     ----------

benches/intensity_benchmarks.rs

@@ -17,8 +17,19 @@ const SAMPLE_LENGTHS: &[usize] = &[
     20_000, // 100 seconds
     40_000, // 200 seconds
 ];
-const PSA_PERIODS: &[f64] = &[1.0]; // Should be independent of pSA period, but added here for safety in the future
-const DAMPING: f64 = 0.05;
+// Oscillator coefficients for T = 1 s, 5% damping at SAMPLING_RATE, from
+// `IM.ims._nigam_jennings_coefficients(np.array([1.0]), 0.005)`. The cost of
+// the recurrence does not depend on their values.
+const PSA_COEFFICIENTS: [f64; psa::N_COEFFICIENTS] = [
+    0.9995070766804318,
+    0.19704993255509468,
+    -0.004991333100781255,
+    0.996370929600225,
+    -0.00032856713331003905,
+    -0.0024941543916031455,
+    -0.00016435618625817934,
+    -0.00249717870917811,
+];
 
 /// Benchmark CAV (Cumulative Absolute Velocity) calculations
 fn bench_cav(c: &mut Criterion) {
@@ -131,30 +142,26 @@ fn bench_psa(c: &mut Criterion) {
     // PSA is expensive, so we might want to use a smaller sample size
     group.sample_size(10);
 
-    for &period in PSA_PERIODS {
-        for &stations in STATION_COUNTS {
-            // For PSA, use smaller sample sets to keep benchmark times reasonable
-            for &samples in SAMPLE_LENGTHS {
-                let waveforms = generate_waveforms(stations, samples);
-                let view = waveforms.view();
-                let periods = array![period];
-                let param = format!("T{:.1}s_{}stn_{}smp", period, stations, samples);
-
-                group.throughput(Throughput::Bytes((stations * samples * 8) as u64));
-
-                group.bench_with_input(BenchmarkId::new("Sequential", &param), &view, |b, &v| {
-                    b.iter(|| {
-                        psa::psa(
-                            black_box(&v),
-                            black_box(&v),
-                            black_box(&v),
-                            black_box(&periods.view()),
-                            black_box(SAMPLING_RATE),
-                            black_box(DAMPING),
-                        )
-                    })
-                });
-            }
+    let coefficients = array![PSA_COEFFICIENTS];
+    for &stations in STATION_COUNTS {
+        // For PSA, use smaller sample sets to keep benchmark times reasonable
+        for &samples in SAMPLE_LENGTHS {
+            let waveforms = generate_waveforms(stations, samples);
+            let view = waveforms.view();
+            let param = format!("T1.0s_{}stn_{}smp", stations, samples);
+
+            group.throughput(Throughput::Bytes((stations * samples * 8) as u64));
+
+            group.bench_with_input(BenchmarkId::new("Sequential", &param), &view, |b, &v| {
+                b.iter(|| {
+                    psa::psa(
+                        black_box(&v),
+                        black_box(&v),
+                        black_box(&v),
+                        black_box(&coefficients.view()),
+                    )
+                })
+            });
         }
     }
 

src-rust/lib.rs

@@ -13,7 +13,7 @@ use pyo3::prelude::*;
 /// import the module.
 #[pymodule]
 mod _core {
-    use numpy::{IntoPyArray, PyArray1, PyArray2, PyArray3, PyReadonlyArray1, PyReadonlyArray2};
+    use numpy::{IntoPyArray, PyArray1, PyArray2, PyArray3, PyReadonlyArray2};
     use pyo3::prelude::*;
 
     use crate::arias_intensity;
@@ -75,8 +75,9 @@ mod _core {
 
     /// Pseudo-spectral acceleration statistics for every station and period.
     ///
-    /// The three components have shape `(ns, nt)`. Returns an `(ns,
-    /// n_periods, 9)` array whose last axis is laid out as
+    /// The three components have shape `(ns, nt)` and `coefficients` has
+    /// shape `(n_periods, 8)`, one oscillator per row as documented on
+    /// [`psa::N_COEFFICIENTS`]. Returns an `(ns, n_periods, 9)` array whose last axis is laid out as
     /// `ROTD_COMPONENTS`: the 000, 090, vertical and geometric mean peaks,
     /// then the five columns [`_rotd`] returns.
     #[pyfunction]
@@ -85,15 +86,13 @@ mod _core {
         comp_0_py: PyReadonlyArray2<f64>,
         comp_90_py: PyReadonlyArray2<f64>,
         comp_ver_py: PyReadonlyArray2<f64>,
-        periods_py: PyReadonlyArray1<f64>,
-        dt: f64,
-        xi: f64,
+        coefficients_py: PyReadonlyArray2<f64>,
     ) -> Bound<'py, PyArray3<f64>> {
         let comp_0 = comp_0_py.as_array();
         let comp_90 = comp_90_py.as_array();
         let comp_ver = comp_ver_py.as_array();
-        let periods = periods_py.as_array();
-        let psa = py.detach(|| psa::psa(&comp_0, &comp_90, &comp_ver, &periods, dt, xi));
+        let coefficients = coefficients_py.as_array();
+        let psa = py.detach(|| psa::psa(&comp_0, &comp_90, &comp_ver, &coefficients));
         psa.into_pyarray(py)
     }