Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fixe: #2757 #2914

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion .mailmap
Original file line number Diff line number Diff line change
Expand Up @@ -287,4 +287,6 @@ Israel Roldan <[email protected]> airv_zxf <[email protected]

Michael Zingale <[email protected]>

Akash Kaushik <[email protected]>
Akash Kaushik <[email protected]>

Bhavin umatiya <[email protected]>
81 changes: 44 additions & 37 deletions tardis/plasma/properties/continuum_processes/fast_array_util.py
Original file line number Diff line number Diff line change
@@ -1,65 +1,72 @@
# 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
# Define the configuration dictionary
njit_dict = {
"parallel": True, # Enables parallel execution
"cache": True, # Caches the compiled function for reuse
"fastmath": True # Enables aggressive optimizations for floating-point operations
}


@njit(**njit_dict)
@njit(**njit_dict) # Use only the dictionary
def numba_cumulative_trapezoid(f, x):
"""
Cumulatively integrate f(x) using the composite trapezoidal rule.

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.")

dx = np.diff(x)
cumulative = (dx * (f[1:] + f[:-1]) / 2.0).cumsum()

return np.concatenate(([0], cumulative / cumulative[-1]))

@njit(**njit_dict)
@njit(**njit_dict) # Use only the dictionary
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

integrated[start:stop, col] = numba_cumulative_trapezoid(
f[start:stop, col], x[start:stop]
)

return integrated
Loading