Implement Debye Huckel Screening (pynucastro#870)

Implements weak screening as per equation A1 of Chugunov 2009 (see below). Also adds a decorator for screening functions that performs a pre-screening check (e.g. comparing debye_huckel against a threshold) and only computes the full screening factor if that check passes, which can significantly improve runtimes.

pynucastro/screening/__init__.py

@@ -3,6 +3,7 @@
 __all__ = ["screen", "screening_util"]
 
 from .screen import (NseState, PlasmaState, ScreenFactors, chugunov_2007,
-                     chugunov_2009, make_plasma_state, make_screen_factors,
-                     potekhin_1998, screen5)
+                     chugunov_2009, debye_huckel, make_plasma_state,
+                     make_screen_factors, potekhin_1998, screen5,
+                     screening_check)
 from .screening_util import ScreeningPair, get_screening_map

pynucastro/screening/screen.py

@@ -8,8 +8,8 @@
 from pynucastro.numba_util import jitclass, njit
 
 __all__ = ["PlasmaState", "ScreenFactors", "chugunov_2007", "chugunov_2009",
-           "make_plasma_state", "make_screen_factors", "potekhin_1998",
-           "screen5"]
+           "debye_huckel", "make_plasma_state", "make_screen_factors",
+           "potekhin_1998", "screen5", "screening_check"]
 
 
 @jitclass()
@@ -194,6 +194,29 @@ def make_screen_factors(n1, n2):
     return ScreenFactors(n1.Z, n1.A, n2.Z, n2.A)
 
 
+@njit
+def debye_huckel(state, scn_fac) -> float:
+    """Calculates the Debye-Huckel enhancement factor for weak Coloumb coupling,
+    following the appendix of :cite:t:`chugunov:2009`.
+
+    :param PlasmaState state:     the precomputed plasma state factors
+    :param ScreenFactors scn_fac: the precomputed ion pair factors
+    :returns: screening correction factor
+    """
+    z1z2 = scn_fac.z1 * scn_fac.z2
+
+    # Gamma_e from eq. 6
+    Gamma_e = state.gamma_e_fac / state.temp
+
+    # eq. A1
+    h_DH = z1z2 * np.sqrt(3 * Gamma_e**3 * state.z2bar / state.zbar)
+
+    # machine limit the output
+    h_max = 300
+    h = min(h_DH, h_max)
+    return np.exp(h)
+
+
 @njit
 def screen5(state: PlasmaState, scn_fac):
     """Calculates screening factors following the appendix of :cite:t:`Wallace:1982`.
@@ -586,3 +609,32 @@ def potekhin_1998(state, scn_fac):
     scor = np.exp(h12)
 
     return scor
+
+
+def screening_check(check_func=debye_huckel, threshold: float = 1.01):
+    """A decorator factory that wraps a screening function with a
+    check that determines whether that function can be skipped for a
+    given plasma state and screening pair.
+
+    :param func: the function to check against the threshold
+    :param threshold: the threshold to check against. If screen_check
+                      is less than the threshold, skip screen_func
+    :returns: a decorator for wrapping screening functions
+    """
+
+    def screening_decorator(screen_func):
+        """Decorates a screening function with a computation
+        that determines whether the check is skippable.
+
+        :param screen_func: the screening function being wrapped
+        :returns: a wrapped screening function that performs this check
+        """
+
+        @njit
+        def screening_wrapper(state, scn_fac):
+            F0 = check_func(state, scn_fac)
+            if F0 <= threshold:
+                return F0
+            return screen_func(state, scn_fac)
+        return screening_wrapper
+    return screening_decorator

pynucastro/screening/tests/test_screen.py

@@ -2,9 +2,9 @@
 from pytest import approx
 
 import pynucastro as pyna
-from pynucastro.screening import (chugunov_2007, chugunov_2009,
+from pynucastro.screening import (chugunov_2007, chugunov_2009, debye_huckel,
                                   make_plasma_state, make_screen_factors,
-                                  potekhin_1998, screen5)
+                                  potekhin_1998, screen5, screening_check)
 
 
 class TestScreen:
@@ -63,10 +63,21 @@ def test_chugunov_2009(self, plasma_state, scn_fac):
         scor = chugunov_2009(plasma_state, scn_fac)
         assert scor == approx(2.87983449091315e+33)
 
+    def test_debye_huckel(self, plasma_state, scn_fac):
+        scor = debye_huckel(plasma_state, scn_fac)
+        assert scor == approx(1.9424263952412558e+130)
+
     def test_potekhin_1998(self, plasma_state, scn_fac):
         scor = potekhin_1998(plasma_state, scn_fac)
         assert scor == approx(1.0508243810383098e+36)
 
     def test_screen5(self, plasma_state, scn_fac):
         scor = screen5(plasma_state, scn_fac)
         assert scor == approx(4.049488384394272e+33)
+
+    @pytest.mark.parametrize("screen_func", [chugunov_2007, chugunov_2009, potekhin_1998, screen5])
+    def test_screening_check(self, plasma_state, scn_fac, screen_func):
+        wrapped_screen_func = screening_check()(screen_func)
+        scor_wrapped = wrapped_screen_func(plasma_state, scn_fac)
+        scor = screen_func(plasma_state, scn_fac)
+        assert scor_wrapped == approx(scor)