FIX: benchmark, caching, fasthmath, opt. gpu/cuda

Author: rahulranjansah

README.md

@@ -46,6 +46,29 @@ uv sync --all-extras
 pip install -e .
 ```
 
+### GPU Acceleration (optional)
+
+If you have an NVIDIA GPU, install the GPU extra to enable CUDA acceleration.
+This installs `cupy-cuda12x` (GPU array ops) and `numba-cuda` (Numba `@cuda.jit`
+kernels). No separate CUDA toolkit download needed — only an NVIDIA driver
+(version 450+). The code automatically detects GPU availability and falls back
+to CPU if CUDA is not present.
+
+**With uv (recommended):**
+```bash
+uv pip install -e ".[gpu]"
+```
+
+**With pip:**
+```bash
+pip install -e ".[gpu]"
+```
+
+**Manual install (if you only want the packages directly):**
+```bash
+pip install cupy-cuda12x numba-cuda
+```
+
 ### Development
 
 ```bash
@@ -85,9 +108,10 @@ pulsesuite/
 - NumPy ≥1.26.4
 - SciPy ≥1.15.2
 - Matplotlib ≥3.10.0
-- pyFFTW ≥0.15.0 (recommended for performance)
-- Numba ≥0.61.2 (optional, for JIT acceleration)
-- Numba-CUDA==0.23.0 (optional, for CUDA acceleration)
+- pyFFTW ≥0.15.0
+- Numba ≥0.61.2
+- CuPy-CUDA12x ≥13.5.0 (optional, GPU array ops — install via `pip install -e ".[gpu]"`)
+- numba-cuda ≥0.6.0 (optional, Numba CUDA kernels — included in `[gpu]` extra)
 
 ## Running a simulation
 

pyproject.toml

@@ -72,6 +72,10 @@ test = [
     "pytest-cov>=4.0.0",
     "pytest-xdist>=3.0.0",
 ]
+gpu = [
+    "cupy-cuda12x>=13.5.0",    # GPU array ops (requires CUDA 12 driver)
+    "numba-cuda>=0.6.0",       # Numba @cuda.jit kernels
+]
 jupyter = ["notebook", "ipywidgets>=7.6"]
 doc = [
     "httpx",

run_pulsesuite_sbetestpropnov30.sh

@@ -10,8 +10,8 @@ cd $PBS_O_WORKDIR/src/pulsesuite/PSTD3D
 export PATH="$HOME/.local/bin:$PATH"
 
 # Match Numba thread count to PBS allocation
-export NUMBA_NUM_THREADS=8
-export OMP_NUM_THREADS=8
-export MKL_NUM_THREADS=8
+export NUMBA_NUM_THREADS=48
+export OMP_NUM_THREADS=48
+export MKL_NUM_THREADS=48
 
 uv run python -u -m pulsesuite.PSTD3D.sbetestpropnov30

src/pulsesuite/PSTD3D/SBEs.py

@@ -805,7 +805,7 @@ def CalcH(Meh, Wee, Whh, C, D, p, VC, Heh, Hee, Hhh):
         )
 
 
-@jit(nopython=True, cache=True, parallel=True)
+@jit(nopython=True, cache=True, parallel=True, fastmath=True)
 def _CalcH_jit(
     Meh, Wee, Whh, C, D, p, VC, Heh, Hee, Hhh, Nk, Ee, Eh, gap, excitons, lf, freepot
 ):
@@ -1938,7 +1938,7 @@ def dpdt(C, D, p, Heh, Hee, Hhh, GamE, GamH, OffP):
             return _dpdt_fallback(C, D, p, Heh, Hee, Hhh, GamE, GamH, OffP)
 
 
-@jit(nopython=True, cache=True, parallel=True)
+@jit(nopython=True, cache=True, parallel=True, fastmath=True)
 def _dpdt_jit(C, D, p, Heh, Hee, Hhh, GamE, GamH, OffP, ii_val, hbar_val):
     """JIT-compiled dpdt calculation."""
     Nk = p.shape[0]
@@ -2305,7 +2305,7 @@ def dCdt(Cee, Dhh, Phe, Heh, Hee, Hhh, GamE, GamH, OffE):
             return _dCdt_fallback(Cee, Dhh, Phe, Heh, Hee, Hhh, GamE, GamH, OffE)
 
 
-@jit(nopython=True, cache=True, parallel=True)
+@jit(nopython=True, cache=True, parallel=True, fastmath=True)
 def _dCdt_jit(Cee, Dhh, Phe, Heh, Hee, Hhh, GamE, GamH, OffE, ii_val, hbar_val):
     """JIT-compiled dCdt calculation."""
     Nk = Cee.shape[0]
@@ -2603,7 +2603,7 @@ def dDdt(Cee, Dhh, Phe, Heh, Hee, Hhh, GamE, GamH, OffH):
             return _dDdt_fallback(Cee, Dhh, Phe, Heh, Hee, Hhh, GamE, GamH, OffH)
 
 
-@jit(nopython=True, cache=True, parallel=True)
+@jit(nopython=True, cache=True, parallel=True, fastmath=True)
 def _dDdt_jit(Cee, Dhh, Phe, Heh, Hee, Hhh, GamE, GamH, OffH, ii_val, hbar_val):
     """JIT-compiled dDdt calculation."""
     Nk = Dhh.shape[0]
@@ -3161,7 +3161,7 @@ def CalcXqw(iq, w, kr, fe, fh, Ee, Eh, gap, area, game, gamh, dcv):
         return _CalcXqw_fallback(iq, w, kr, fe, fh, Ee, Eh, gap, area, game, gamh, dcv)
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _CalcXqw_jit(
     iq,
     w,
@@ -3713,7 +3713,7 @@ def GetArrays(x, qx, kx):
     _NQ0 = GetArray0Index(qx)
 
 
-@jit(nopython=True, parallel=True)
+@jit(nopython=True, parallel=True, fastmath=True, cache=True)
 def _MakeKKP_jit(Nk, kr, dkr, NQ0, kkp):
     """JIT-compiled parallel version of MakeKKP computation."""
     from numba import prange

src/pulsesuite/PSTD3D/coulomb.py

@@ -35,7 +35,7 @@
 
 # JIT-compatible K0 approximation for modified Bessel function
 # nopython friendly version of K03 which is in usefulsubs.py
-@jit(nopython=True, cache=True)
+@jit(nopython=True, cache=True, fastmath=True)
 def _K03_jit(x):
     """
     JIT-compatible approximation of modified Bessel function K0(x).
@@ -66,7 +66,7 @@ def _K03_jit(x):
         return np.sqrt(np.pi / (2.0 * x)) * np.exp(-x)
 
 
-@jit(nopython=True, cache=True)
+@jit(nopython=True, cache=True, fastmath=True)
 def _Vint_jit(Qyk, y, alphae, alphah, Delta0, N1, N2, Ny):
     """JIT-compiled version of Vint calculation."""
     Vint_val = 0.0
@@ -782,7 +782,7 @@ def GetEps1Dqw(alphae, alphah, Delta0, epsr, me, mh, n1D, q, w):
 ###############################################################################
 
 
-@jit(nopython=True, parallel=True)
+@jit(nopython=True, parallel=True, fastmath=True, cache=True)
 def _CalcMVeh_core(p, Veh, MVeh, k3, UnDel):
     """Core JIT-compiled computation for CalcMVeh."""
     N = p.shape[0]
@@ -809,7 +809,7 @@ def _CalcMVeh_core(p, Veh, MVeh, k3, UnDel):
                             )
 
 
-@jit(nopython=True, parallel=True)
+@jit(nopython=True, parallel=True, fastmath=True, cache=True)
 def _BGRenorm_core(ne, nh, Vee, Vhh, BGR, UnDel):
     """Core JIT-compiled computation for BGRenorm."""
     N = len(ne)
@@ -830,7 +830,7 @@ def _BGRenorm_core(ne, nh, Vee, Vhh, BGR, UnDel):
             BGR[k, kp] = -sum_hh - sum_ee
 
 
-@jit(nopython=True, parallel=True)
+@jit(nopython=True, parallel=True, fastmath=True, cache=True)
 def _EeRenorm_core(ne, Vee, BGR, UnDel):
     """Core JIT-compiled computation for EeRenorm."""
     N = len(ne)
@@ -862,7 +862,7 @@ def _EeRenorm_core(ne, Vee, BGR, UnDel):
             )
 
 
-@jit(nopython=True, parallel=True)
+@jit(nopython=True, parallel=True, fastmath=True, cache=True)
 def _EhRenorm_core(nh, Vhh, BGR, UnDel):
     """Core JIT-compiled computation for EhRenorm."""
     N = len(nh)
@@ -894,7 +894,7 @@ def _EhRenorm_core(nh, Vhh, BGR, UnDel):
             )
 
 
-@jit(nopython=True, parallel=True)
+@jit(nopython=True, parallel=True, fastmath=True, cache=True)
 def _MBCE2_core(ne, nh, Veh2, Vee2, Win, Wout, k3, Ceh, Cee):
     """Core JIT-compiled computation for MBCE2."""
     Nk = len(ne) - 1
@@ -953,7 +953,7 @@ def _MBCE2_core(ne, nh, Veh2, Vee2, Win, Wout, k3, Ceh, Cee):
                     )
 
 
-@jit(nopython=True, parallel=True)
+@jit(nopython=True, parallel=True, fastmath=True, cache=True)
 def _MBCE_core(ne, nh, Veh2, Vee2, Win, Wout, k3, Ceh, Cee):
     """Core JIT-compiled computation for MBCE (identical to MBCE2)."""
     Nk = len(ne) - 1
@@ -1012,7 +1012,7 @@ def _MBCE_core(ne, nh, Veh2, Vee2, Win, Wout, k3, Ceh, Cee):
                     )
 
 
-@jit(nopython=True, parallel=True)
+@jit(nopython=True, parallel=True, fastmath=True, cache=True)
 def _MBCH_core(ne, nh, Veh2, Vhh2, Win, Wout, k3, Ceh, Chh):
     """Core JIT-compiled computation for MBCH."""
     Nk = len(ne) - 1

src/pulsesuite/PSTD3D/dcfield.py

@@ -335,7 +335,7 @@ def DC_Step_FD(ne, nh, nemid, nhmid, ky, Edc, dt, me, mh):
 # ============================================================================
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _EkReNorm_jit(n, En, V):
     Nk = len(n)
     Ec = np.zeros(Nk)
@@ -344,7 +344,7 @@ def _EkReNorm_jit(n, En, V):
     return Ec
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _DriftVt_jit(n, Ec, dkk, hbar_val):
     Nk = len(n)
     dEdk = np.zeros(Nk, dtype=np.complex128)
@@ -359,19 +359,19 @@ def _DriftVt_jit(n, Ec, dkk, hbar_val):
     return np.real(v)
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _Lrtz_jit(a, b):
     pi_val = 3.141592653589793
     return (b / pi_val) / (a**2 + b**2)
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _theta_jit(x):
     small = 1e-200
     return (np.abs(x) + x) / 2.0 / (np.abs(x) + small)
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _ThetaEM_jit(Ephn, m, g, ky, n, Cq2, v, N0, q, k, hbar_val, pi_val):
     q_idx = q - 1 if q > 0 else 0
     k_idx = k - 1 if k > 0 else 0
@@ -408,7 +408,7 @@ def _ThetaEM_jit(Ephn, m, g, ky, n, Cq2, v, N0, q, k, hbar_val, pi_val):
     return ThetaEM_val
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _ThetaABS_jit(Ephn, m, g, ky, n, Cq2, v, N0, q, k, hbar_val, pi_val):
     q_idx = q - 1 if q > 0 else 0
     k_idx = k - 1 if k > 0 else 0
@@ -445,7 +445,7 @@ def _ThetaABS_jit(Ephn, m, g, ky, n, Cq2, v, N0, q, k, hbar_val, pi_val):
     return ThetaABS_val
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _dndEk_jit(Ephn, m, q, dndq, hbar_val):
     x0 = Ephn * 2.0 * m / hbar_val**2
     N = len(q)
@@ -455,7 +455,7 @@ def _dndEk_jit(Ephn, m, q, dndq, hbar_val):
     return dndEk
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _ThetaEMABS_jit(Ephn, m, q, dndk, Cq2, v, hbar_val, pi_val):
     x0 = Ephn * 2.0 * m / hbar_val**2
     N = len(q)

src/pulsesuite/PSTD3D/dephasing.py

@@ -25,7 +25,7 @@
 # ── JIT kernels (module-level — Numba requirement) ──────────────────
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _CalcGammaE_jit(
     Vee2, Veh2, ne, nh, se, sh, k_p_q, k_m_q, k1p_m_q, k1p, xe, xh, pi_val, hbar_val
 ):
@@ -64,7 +64,7 @@ def _CalcGammaE_jit(
     return GammaE
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _CalcGammaH_jit(
     Vhh2, Veh2, ne, nh, se, sh, k_p_q, k_m_q, k1_m_q, k1, xe, xh, pi_val, hbar_val
 ):
@@ -103,7 +103,7 @@ def _CalcGammaH_jit(
     return GammaH
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _CalcOffDiagDeph_E_jit(
     ne, nh, Veh2, Vee2, Ee, Eh, k_p_q, k_m_q, gee, geh, pi_val, hbar_val
 ):
@@ -146,7 +146,7 @@ def _CalcOffDiagDeph_E_jit(
     return D * pi_val / hbar_val
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _CalcOffDiagDeph_H_jit(
     ne, nh, Veh2, Vhh2, Ee, Eh, k_p_q, k_m_q, ghh, geh, pi_val, hbar_val
 ):
@@ -189,7 +189,7 @@ def _CalcOffDiagDeph_H_jit(
     return D * pi_val / hbar_val
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _OffDiagDephasing_jit(Dh, De, pt, pp, k_p_q, undel, ii_val, hbar_val):
     """JIT-compiled version of OffDiagDephasing inner loops."""
     Nk = Dh.shape[0]
@@ -217,13 +217,13 @@ def _OffDiagDephasing_jit(Dh, De, pt, pp, k_p_q, undel, ii_val, hbar_val):
     return x * ii_val * hbar_val
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _Lrtz_jit(a, b, pi_val):
     """JIT-compiled version of Lrtz."""
     return (b / pi_val) / (a**2 + b**2)
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _CalcOffDiagDeph_E2_jit(ne, nh, Vsq, Ee, Eh, gee, geh, pi_val, hbar_val):
     """JIT-compiled version of CalcOffDiagDeph_E2."""
     Nk = len(ne)
@@ -270,7 +270,7 @@ def _CalcOffDiagDeph_E2_jit(ne, nh, Vsq, Ee, Eh, gee, geh, pi_val, hbar_val):
     return D * pi_val / hbar_val
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _CalcOffDiagDeph_H2_jit(ne, nh, Vsq, Ee, Eh, ghh, geh, pi_val, hbar_val):
     """JIT-compiled version of CalcOffDiagDeph_H2."""
     Nk = len(ne)
@@ -317,7 +317,7 @@ def _CalcOffDiagDeph_H2_jit(ne, nh, Vsq, Ee, Eh, ghh, geh, pi_val, hbar_val):
     return D * pi_val / hbar_val
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _OffDiagDephasing2_jit(Dh, De, p, pt, undel, ii_val, hbar_val):
     """JIT-compiled version of OffDiagDephasing2 inner loops."""
     Nk = p.shape[0]

src/pulsesuite/PSTD3D/epsrtl.py

@@ -142,7 +142,7 @@ def atanhc(x):
     return 0.5 * np.log((1.0 + x) / (1.0 - x))
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _PiT_jit(
     w, me, mh, Te, Th, dk, Ek, Ekq, n00_val, hbar_val, kB_val, twopi_val, pi_val
 ):
@@ -300,7 +300,7 @@ def GetEpsrLEpsrT(n1D, dcv0, Te, Th, me, mh, Eg, ky):
     # Note: stop statement removed in Python version
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _RecordEpsrT_loop_jit(
     ky_size,
     Nw,
@@ -694,7 +694,7 @@ def QqGq(ky, Nk, dk, dw, EpsR, EpsI, eh):
             f.write(f"{ky[q]} {Omega[q]} {Gam[q]}\n")
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _ZeroT_L_loop_jit(
     qy_size, Nw, dw, m, kf, qy, dq, Vc, hbar_val, e0_val, pi_val, ii_real, ii_imag
 ):
@@ -848,7 +848,7 @@ def ZeroT_L(B, m, qy, kf):
                 f.write(f"{eps[qq, ww_idx].real} {eps[qq, ww_idx].imag}\n")
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _ZeroT_T_loop_jit(
     qy_size,
     Nw,
@@ -1066,7 +1066,7 @@ def ZeroT_T(me, mh, Egap, dcv, qy, kf):
                 f.write(f"{Pi2[qq, ww_idx]} {-Pi3[qq, ww_idx].real}\n")
 
 
-@jit(nopython=True)
+@jit(nopython=True, fastmath=True, cache=True)
 def _RecordEpsrL_loop_jit(
     ky_size,
     Nw,