Skip to content

Lazy compilation for numba gufuncs #103

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.

Sign up for GitHub

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

Merged
merged 4 commits into from
Aug 4, 2025
Merged

Lazy compilation for numba gufuncs #103

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
abcd201
lazy crps gufuncs
frazane Jul 26, 2025
9232433
add lazy gufuncs wrappers to utils
frazane Jul 27, 2025
f750b21
adopt lazy gufuncs wrappers
frazane Jul 27, 2025
657c5e2
eagerly compile qd crps estimator
frazane Jul 31, 2025
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
86 changes: 19 additions & 67 deletions scoringrules/core/crps/_gufuncs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,39 +3,28 @@
import numpy as np
from numba import guvectorize, njit, vectorize

from scoringrules.core.utils import lazy_gufunc_wrapper_uv

INV_SQRT_PI = 1 / np.sqrt(np.pi)
EPSILON = 1e-6


@guvectorize(
[
"void(float32[:], float32[:], float32[:], float32[:])",
"void(float64[:], float64[:], float64[:], float64[:])",
],
"(),(a),(a)->()",
)
@lazy_gufunc_wrapper_uv
@guvectorize("(),(a),(a)->()")
def quantile_pinball_gufunc(
obs: np.ndarray, fct: np.ndarray, alpha: np.ndarray, out: np.ndarray
):
"""Pinball loss based estimator for the CRPS from quantiles."""
obs = obs[0]
Q = fct.shape[0]
pb = 0
for fc, level in zip(fct, alpha): # noqa: B905
pb += (obs > fc) * level * (obs - fc) + (obs <= fc) * (1 - level) * (fc - obs)
out[0] = 2 * (pb / Q)


@guvectorize(
[
"void(float32[:], float32[:], float32[:])",
"void(float64[:], float64[:], float64[:])",
],
"(),(n)->()",
)
@guvectorize("(),(n)->()")
def _crps_ensemble_int_gufunc(obs: np.ndarray, fct: np.ndarray, out: np.ndarray):
"""CRPS estimator based on the integral form."""
obs = obs[0]
M = fct.shape[0]

if np.isnan(obs):
Expand Down Expand Up @@ -153,16 +142,9 @@ def _crps_ensemble_fair_gufunc(obs: np.ndarray, fct: np.ndarray, out: np.ndarray
out[0] = e_1 / M - 0.5 * e_2 / (M * (M - 1))


@guvectorize(
[
"void(float32[:], float32[:], float32[:])",
"void(float64[:], float64[:], float64[:])",
],
"(),(n)->()",
)
@guvectorize("(),(n)->()")
def _crps_ensemble_pwm_gufunc(obs: np.ndarray, fct: np.ndarray, out: np.ndarray):
"""CRPS estimator based on the probability weighted moment (PWM) form."""
obs = obs[0]
M = fct.shape[-1]

if np.isnan(obs):
Expand All @@ -181,19 +163,12 @@ def _crps_ensemble_pwm_gufunc(obs: np.ndarray, fct: np.ndarray, out: np.ndarray)
out[0] = expected_diff / M + β_0 / M - 2 * β_1 / (M * (M - 1))


@guvectorize(
[
"void(float32[:], float32[:], float32[:])",
"void(float64[:], float64[:], float64[:])",
],
"(),(n)->()",
)
@guvectorize("(),(n)->()")
def _crps_ensemble_akr_gufunc(obs: np.ndarray, fct: np.ndarray, out: np.ndarray):
"""CRPS estimaton based on the approximate kernel representation."""
M = fct.shape[-1]
obs = obs[0]
e_1 = 0
e_2 = 0
e_1 = 0.0
e_2 = 0.0
for i, forecast in enumerate(fct):
if i == 0:
i = M - 1
Expand All @@ -202,19 +177,12 @@ def _crps_ensemble_akr_gufunc(obs: np.ndarray, fct: np.ndarray, out: np.ndarray)
out[0] = e_1 / M - 0.5 * 1 / M * e_2


@guvectorize(
[
"void(float32[:], float32[:], float32[:])",
"void(float64[:], float64[:], float64[:])",
],
"(),(n)->()",
)
@guvectorize("(),(n)->()")
def _crps_ensemble_akr_circperm_gufunc(
obs: np.ndarray, fct: np.ndarray, out: np.ndarray
):
"""CRPS estimaton based on the AKR with cyclic permutation."""
M = fct.shape[-1]
obs = obs[0]
e_1 = 0.0
e_2 = 0.0
for i, forecast in enumerate(fct):
Expand All @@ -224,13 +192,7 @@ def _crps_ensemble_akr_circperm_gufunc(
out[0] = e_1 / M - 0.5 * 1 / M * e_2


@guvectorize(
[
"void(float32[:], float32[:], float32[:], float32[:], float32[:])",
"void(float64[:], float64[:], float64[:], float64[:], float64[:])",
],
"(),(n),(),(n)->()",
)
@guvectorize("(),(n),(),(n)->()")
def _owcrps_ensemble_nrg_gufunc(
obs: np.ndarray,
fct: np.ndarray,
Expand All @@ -239,8 +201,6 @@ def _owcrps_ensemble_nrg_gufunc(
out: np.ndarray,
):
"""Outcome-weighted CRPS estimator based on the energy form."""
obs = obs[0]
ow = ow[0]
M = fct.shape[-1]

if np.isnan(obs):
Expand All @@ -260,13 +220,7 @@ def _owcrps_ensemble_nrg_gufunc(
out[0] = e_1 / (M * wbar) - 0.5 * e_2 / ((M * wbar) ** 2)


@guvectorize(
[
"void(float32[:], float32[:], float32[:], float32[:], float32[:])",
"void(float64[:], float64[:], float64[:], float64[:], float64[:])",
],
"(),(n),(),(n)->()",
)
@guvectorize("(),(n),(),(n)->()")
def _vrcrps_ensemble_nrg_gufunc(
obs: np.ndarray,
fct: np.ndarray,
Expand All @@ -275,8 +229,6 @@ def _vrcrps_ensemble_nrg_gufunc(
out: np.ndarray,
):
"""Vertically re-scaled CRPS estimator based on the energy form."""
obs = obs[0]
ow = ow[0]
M = fct.shape[-1]

if np.isnan(obs):
Expand Down Expand Up @@ -344,15 +296,15 @@ def _crps_logistic_ufunc(obs: float, mu: float, sigma: float) -> float:


estimator_gufuncs = {
"akr_circperm": _crps_ensemble_akr_circperm_gufunc,
"akr": _crps_ensemble_akr_gufunc,
"akr_circperm": lazy_gufunc_wrapper_uv(_crps_ensemble_akr_circperm_gufunc),
"akr": lazy_gufunc_wrapper_uv(_crps_ensemble_akr_gufunc),
"fair": _crps_ensemble_fair_gufunc,
"int": _crps_ensemble_int_gufunc,
"nrg": _crps_ensemble_nrg_gufunc,
"pwm": _crps_ensemble_pwm_gufunc,
"int": lazy_gufunc_wrapper_uv(_crps_ensemble_int_gufunc),
"nrg": lazy_gufunc_wrapper_uv(_crps_ensemble_nrg_gufunc),
Copy link
Collaborator

@sallen12 sallen12 Aug 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nrg estimator could also be lazily loaded - currently it also has the longer decorator, as if it is not lazily loaded

"pwm": lazy_gufunc_wrapper_uv(_crps_ensemble_pwm_gufunc),
"qd": _crps_ensemble_qd_gufunc,
"ownrg": _owcrps_ensemble_nrg_gufunc,
"vrnrg": _vrcrps_ensemble_nrg_gufunc,
"ownrg": lazy_gufunc_wrapper_uv(_owcrps_ensemble_nrg_gufunc),
"vrnrg": lazy_gufunc_wrapper_uv(_vrcrps_ensemble_nrg_gufunc),
}

__all__ = [
Expand Down
23 changes: 7 additions & 16 deletions scoringrules/core/energy/_gufuncs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,16 @@
import numpy as np
from numba import guvectorize

from scoringrules.core.utils import lazy_gufunc_wrapper_mv


@guvectorize(
[
"void(float32[:], float32[:,:], float32[:])",
"void(float64[:], float64[:,:], float64[:])",
],
"(d),(m,d)->()",
cache=True,
Copy link
Collaborator

@sallen12 sallen12 Aug 4, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't fully understand the motivation for caching the energy score gufunc but not the other gufuncs imported at import time, e.g. (fair and nrg) crps and variogram score?

The weighted interval score is also cached but lazily loaded. Again, why this score and not the other scores?

)
def _energy_score_gufunc(
obs: np.ndarray,
Expand All @@ -27,13 +30,8 @@ def _energy_score_gufunc(
out[0] = e_1 / M - 0.5 / (M**2) * e_2


@guvectorize(
[
"void(float32[:], float32[:,:], float32[:], float32[:], float32[:])",
"void(float64[:], float64[:,:], float64[:], float64[:], float64[:])",
],
"(d),(m,d),(),(m)->()",
)
@lazy_gufunc_wrapper_mv
@guvectorize("(d),(m,d),(),(m)->()")
def _owenergy_score_gufunc(
obs: np.ndarray,
fct: np.ndarray,
Expand All @@ -43,7 +41,6 @@ def _owenergy_score_gufunc(
):
"""Compute the Outcome-Weighted Energy Score for a finite ensemble."""
M = fct.shape[0]
ow = ow[0]

e_1 = 0.0
e_2 = 0.0
Expand All @@ -57,13 +54,8 @@ def _owenergy_score_gufunc(
out[0] = e_1 / (M * wbar) - 0.5 * e_2 / (M**2 * wbar**2)


@guvectorize(
[
"void(float32[:], float32[:,:], float32[:], float32[:], float32[:])",
"void(float64[:], float64[:,:], float64[:], float64[:], float64[:])",
],
"(d),(m,d),(),(m)->()",
)
@lazy_gufunc_wrapper_mv
@guvectorize("(d),(m,d),(),(m)->()")
def _vrenergy_score_gufunc(
obs: np.ndarray,
fct: np.ndarray,
Expand All @@ -73,7 +65,6 @@ def _vrenergy_score_gufunc(
):
"""Compute the Vertically Re-scaled Energy Score for a finite ensemble."""
M = fct.shape[0]
ow = ow[0]

e_1 = 0.0
e_2 = 0.0
Expand Down
14 changes: 4 additions & 10 deletions scoringrules/core/error_spread/_gufunc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,17 +1,11 @@
import numpy as np
from numba import guvectorize

INV_SQRT_PI = 1 / np.sqrt(np.pi)
EPSILON = 1e-6
from scoringrules.core.utils import lazy_gufunc_wrapper_uv


@guvectorize(
[
"void(float32[:], float32[:], float32[:])",
"void(float64[:], float64[:], float64[:])",
],
"(),(n)->()",
)
@lazy_gufunc_wrapper_uv
@guvectorize("(),(n)->()")
def _error_spread_score_gufunc(obs: np.ndarray, fct: np.ndarray, out: np.ndarray):
"""Error spread score."""
M = fct.shape[0]
Expand All @@ -34,7 +28,7 @@ def _error_spread_score_gufunc(obs: np.ndarray, fct: np.ndarray, out: np.ndarray
g /= (M - 1) * (M - 2)

for i in range(M):
e += fct[i] - obs[i]
e += fct[i] - obs
e /= M

out[0] = (s**2 - e**2 - e * s * g) ** 2
16 changes: 5 additions & 11 deletions scoringrules/core/interval/_gufunc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,11 @@
import numpy as np
from numba import guvectorize

from scoringrules.core.utils import lazy_gufunc_wrapper_uv

@guvectorize(
[
"void(float32[:], float32[:], float32[:], float32[:], float32[:], float32[:], float32[:], float32[:])",
"void(float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:])",
],
"(),(),(a),(a),(a),(),(a) ->()",
)

@lazy_gufunc_wrapper_uv
@guvectorize("(),(),(a),(a),(a),(),(a) ->()", cache=True)
def _weighted_interval_score_gufunc(
obs: np.ndarray,
median: np.ndarray,
Expand All @@ -20,9 +17,6 @@ def _weighted_interval_score_gufunc(
out: np.ndarray,
):
"""Weighted Interval score (WIS)."""
obs = obs[0]
med = median[0]
w_med = weight_median[0]
K = alpha.shape[0]
wis = 0
for low, upp, a, w_a in zip(lower, upper, alpha, weight_alpha): # noqa: B905
Expand All @@ -32,6 +26,6 @@ def _weighted_interval_score_gufunc(
+ (obs > upp) * (2 / a) * (obs - upp)
)
wis += w_a * ws_a
wis += w_med * np.abs(obs - med)
wis += weight_median * np.abs(obs - median)
wis /= K + 1 / 2
out[0] = wis
Loading
Loading