Update fast_array_util.py

tardis/plasma/properties/continuum_processes/fast_array_util.py

@@ -1,8 +1,5 @@
-# It is currently not possible to use scipy.integrate.cumulative_trapezoid in
-# numba. So here is my own implementation.
 import numpy as np
 from numba import njit, prange
-
 from tardis.transport.montecarlo import njit_dict
 
 
@@ -13,53 +10,61 @@ def numba_cumulative_trapezoid(f, x):
 
     Parameters
     ----------
-    f : numpy.ndarray, dtype float
-        Input array to integrate.
-    x : numpy.ndarray, dtype float
-        The coordinate to integrate along.
+    f : numpy.ndarray
+        Input array to integrate. Shape: (N,)
+    x : numpy.ndarray
+        The coordinate array. Shape: (N,)
 
     Returns
     -------
-    numpy.ndarray, dtype float
-        The result of cumulative integration of f along x
+    numpy.ndarray
+        Cumulative integral of f along x, normalized by the final value. Shape: (N,)
     """
-    integ = (np.diff(x) * (f[1:] + f[:-1]) / 2.0).cumsum()
-    return integ / integ[-1]
+    if len(f) != len(x):
+        raise ValueError("Input arrays f and x must have the same length.")
+    if len(f) < 2:
+        raise ValueError("Input arrays must have at least two elements for integration.")
 
+    # Compute the cumulative trapezoidal integral
+    dx = np.diff(x)
+    cumulative = (dx * (f[1:] + f[:-1]) / 2.0).cumsum()
 
-@njit(**njit_dict)
+    # Normalize by the final value
+    return np.concatenate(([0], cumulative / cumulative[-1]))
+
+
+@njit(**njit_dict, parallel=True)
 def cumulative_integrate_array_by_blocks(f, x, block_references):
     """
-    Cumulatively integrate a function over blocks.
-
-    This function cumulatively integrates a function `f` defined at
-    locations `x` over blocks given in `block_references`.
+    Cumulatively integrate a 2D array over blocks defined by block references.
 
     Parameters
     ----------
-    f : numpy.ndarray, dtype float
-        Input array to integrate. Shape is (N_freq, N_shells), where
-        N_freq is the number of frequency values and N_shells is the number
-        of computational shells.
-    x : numpy.ndarray, dtype float
-        The sample points corresponding to the `f` values. Shape is (N_freq,).
-    block_references : numpy.ndarray, dtype int
-        The start indices of the blocks to be integrated. Shape is (N_blocks,).
+    f : numpy.ndarray
+        2D input array to integrate. Shape: (N_freq, N_shells)
+    x : numpy.ndarray
+        The coordinate array. Shape: (N_freq,)
+    block_references : numpy.ndarray
+        Start indices of the blocks to be integrated. Shape: (N_blocks,)
 
     Returns
     -------
-    numpy.ndarray, dtype float
-        Array with cumulatively integrated values. Shape is (N_freq, N_shells)
-        same as f.
+    numpy.ndarray
+        2D array with cumulative integrals for each block. Shape: (N_freq, N_shells)
     """
-    n_rows = len(block_references) - 1
+    n_blocks = len(block_references) - 1
     integrated = np.zeros_like(f)
-    for i in prange(f.shape[1]):  # columns
-        # TODO: Avoid this loop through vectorization of cumulative_trapezoid
-        for j in prange(n_rows):  # rows
-            start = block_references[j]
-            stop = block_references[j + 1]
-            integrated[start + 1 : stop, i] = numba_cumulative_trapezoid(
-                f[start:stop, i], x[start:stop]
+
+    for col in prange(f.shape[1]):  # Iterate over columns (N_shells)
+        for block_idx in range(n_blocks):  # Iterate over blocks
+            start = block_references[block_idx]
+            stop = block_references[block_idx + 1]
+
+            if stop - start < 2:
+                continue  # Skip blocks that are too small to integrate
+
+            integrated[start:stop, col] = numba_cumulative_trapezoid(
+                f[start:stop, col], x[start:stop]
             )
+
     return integrated