Resistance to reset_units

README.md

@@ -14,9 +14,7 @@ Installation and usage
  - [See this basic example for usage.](https://github.com/JelleAalbers/wimprates/blob/master/notebooks/Example.ipynb)
 
 The package uses numericalunits (https://pypi.python.org/pypi/numericalunits); all function inputs
-are expected to have proper units. 
-
-**Do NOT call reset_units in your own code without reloading this module!**
+are expected to have proper units.
 
 
 Features

tests/test_wimprates.py

@@ -4,6 +4,7 @@
 values here.
 
 """
+import numericalunits as nu
 import numpy as np
 
 import wimprates as wr
@@ -17,8 +18,13 @@ def isclose(x, y):
 
 
 def test_elastic():
-    isclose(wr.rate_wimp_std(1, **opts),
-            33.19098343826968)
+    ref = 33.19098343826968
+
+    isclose(wr.rate_wimp_std(1, **opts), ref)
+
+    # Test numericalunits.reset_units() does not affect results
+    nu.reset_units()
+    isclose(wr.rate_wimp_std(1, **opts), ref)
 
     # Test vectorized call
     energies = np.linspace(0.01, 40, 100)
@@ -44,5 +50,3 @@ def test_migdal():
 def test_brems():
     isclose(wr.rate_wimp_std(1, detection_mechanism='bremsstrahlung', **opts),
             0.00017062652972332665)
-
-

wimprates/bremsstrahlung.py

@@ -29,7 +29,7 @@ def vmin_w(w, mw):
 
     From Kouvaris/Pradler [arxiv:1607.01789v2], equation in text below eq. 10
     """
-    return (2 * w / wr.reduced_mass(wr.mn, mw))**0.5
+    return (2 * w / wr.mu_nucleus(mw))**0.5
 
 
 def erec_bound(sign, w, v, mw):
@@ -42,7 +42,7 @@ def erec_bound(sign, w, v, mw):
     :param mw: WIMP mass
     :param v: WIMP speed (earth/detector frame)
     """
-    return (wr.reduced_mass(wr.mn, mw)**2 * v**2 / wr.mn
+    return (wr.mu_nucleus(mw)**2 * v**2 / wr.mn()
             * (1
                - vmin_w(w, mw)**2 / (2 * v**2)
                + sign * (1 - vmin_w(w, mw)**2 / v**2)**0.5))
@@ -69,7 +69,7 @@ def sigma_w_erec(w, erec, v, mw, sigma_nucleon,
 
     # Note mn -> mn c^2, Kouvaris/Pradtler and McCabe use natural units
     return (4 * nu.alphaFS / (3 * np.pi * w) *
-            erec / (wr.mn * nu.c0**2) *
+            erec / (wr.mn() * nu.c0**2) *
             form_atomic**2 *
             wr.sigma_erec(erec, v, mw, sigma_nucleon, interaction, m_med))
 
@@ -129,7 +129,7 @@ def integrand(v):
         return (sigma_w(w, v, mw, sigma_nucleon, interaction, m_med) *
                 v * wr.observed_speed_dist(v, t))
 
-    return wr.rho_dm / mw * (1 / wr.mn) * quad(
+    return wr.rho_dm() / mw * (1 / wr.mn()) * quad(
         integrand,
         vmin,
         wr.v_max(t),

wimprates/elastic_nr.py

@@ -8,22 +8,30 @@
 
 import wimprates as wr
 export, __all__ = wr.exporter()
-__all__ += 'An mn'.split()
 
-# Standard atomic weight of target (averaged across all isotopes)
-An = 131.293
-mn = An * nu.amu    # Mass of nucleus (not nucleon!)
+
+@export
+def an():
+    """Standard atomic weight of target (averaged across all isotopes)"""
+    return 131.293
+
+
+@export
+def mn():
+    """Mass of nucleus (not nucleon!)"""
+    return an() * nu.amu
+
 
 spin_isotopes = [
     # A, mass, J (nuclear spin), abundance
     # Data from Wikipedia (Jelle, 12 January 2018)
-    (129, 128.9047794 * nu.amu, 1/2, 26.401e-2),
-    (131, 130.9050824 * nu.amu, 3/2, 21.232e-2),
+    (129, 128.9047794, 1/2, 26.401e-2),
+    (131, 130.9050824, 3/2, 21.232e-2),
 ]
 
 # Load spin-dependent structure functions
 s_data = wr.load_pickle('sd/structure_f_erec_xe.pkl')
-s_energies = s_data['_energies'] * nu.keV
+s_energies = s_data['_energies']
 structure_functions = {}
 # Don't use k, v; dangerous globals...
 for _k, _v in s_data.items():
@@ -39,13 +47,21 @@ def reduced_mass(m1, m2):
 
 
 @export
-def e_max(mw, v, mn=mn):
+def mu_nucleus(mw):
+    """DM-nucleus reduced mass"""
+    return reduced_mass(mw, mn())
+
+
+@export
+def e_max(mw, v, m_nucleus=None):
     """Kinematic nuclear recoil energy maximum
     :param mw: Wimp mass
-    :param mn: Nucleus mass. Defaults to standard atomic mass.
+    :param m_nucleus: Nucleus mass. Defaults to standard atomic mass.
     :param v: Wimp speed
     """
-    return 2 * reduced_mass(mw, mn)**2 * v**2 / mn
+    if m_nucleus is None:
+        m_nucleus = mn()
+    return 2 * reduced_mass(mw, m_nucleus)**2 * v**2 / m_nucleus
 
 
 @export
@@ -56,14 +72,16 @@ def spherical_bessel_j1(x):
 
 @export
 @wr.vectorize_first
-def helm_form_factor_squared(erec, anucl=An):
+def helm_form_factor_squared(erec, anucl=None):
     """Return Helm form factor squared from Lewin & Smith
 
     Lifted from Andrew Brown's code with minor edits
 
     :param erec: nuclear recoil energy
     :param anucl: Nuclear mass number
     """
+    if anucl is None:
+        anucl = an()
     en = erec / nu.keV
     if anucl <= 0:
         raise ValueError("Invalid value of A!")
@@ -109,17 +127,18 @@ def sigma_erec(erec, v, mw, sigma_nucleon,
     """
     if interaction == 'SI':
         sigma_nucleus = (sigma_nucleon
-                         * (reduced_mass(mn, mw) / reduced_mass(nu.amu, mw))**2
-                         * An**2)
+                         * (mu_nucleus(mw) / reduced_mass(nu.amu, mw))**2
+                         * an()**2)
         result = (sigma_nucleus
                   / e_max(mw, v)
-                  * helm_form_factor_squared(erec, anucl=An))
+                  * helm_form_factor_squared(erec, anucl=an()))
 
     elif interaction.startswith('SD'):
         _, coupling, s_assumption = interaction.split('_')
 
         result = np.zeros_like(erec)
         for A, mn_isotope, J, abundance in spin_isotopes:
+            mn_isotope *= nu.amu
             s = structure_functions[(A, coupling, s_assumption)]
             # x isn't quite sigma_nucleus:
             # you'd have to multilpy by structure(0),
@@ -128,7 +147,9 @@ def sigma_erec(erec, v, mw, sigma_nucleon,
             x = (sigma_nucleon * 4 * np.pi
                  * reduced_mass(mw, mn_isotope)**2
                  / (3 * reduced_mass(mw, nu.mp)**2 * (2 * J + 1)))
-            result += abundance * x / e_max(mw, v, mn=mn_isotope) * s(erec)
+            result += (abundance
+                       * x / e_max(mw, v, mn_isotope)
+                       * s(erec / nu.keV))
 
     else:
         raise ValueError("Unsupported DM-nucleus interaction '%s'"
@@ -141,7 +162,7 @@ def sigma_erec(erec, v, mw, sigma_nucleon,
 def mediator_factor(erec, m_med):
     if m_med == float('inf'):
         return 1
-    q = (2 * mn * erec)**0.5
+    q = (2 * mn() * erec)**0.5
     return m_med**4 / (m_med**2 + (q/nu.c0)**2)**2
 
 
@@ -151,7 +172,7 @@ def vmin_elastic(erec, mw):
     :param erec: recoil energy
     :param mw: Wimp mass
     """
-    return np.sqrt(mn * erec / (2 * reduced_mass(mw, mn)**2))
+    return np.sqrt(mn() * erec / (2 * mu_nucleus(mw)**2))
 
 
 @export
@@ -186,7 +207,7 @@ def integrand(v):
         return (sigma_erec(erec, v, mw, sigma_nucleon, interaction, m_med)
                 * v * wr.observed_speed_dist(v, t))
 
-    return wr.rho_dm / mw * (1 / mn) * quad(
+    return wr.rho_dm() / mw * (1 / mn()) * quad(
         integrand,
         v_min, wr.v_max(t),
         **kwargs)[0]

wimprates/halo.py

@@ -8,18 +8,26 @@
 
 import wimprates as wr
 export, __all__ = wr.exporter()
-__all__ += 'rho_dm v_0 v_esc'.split()
 
 
-# Local dark matter density
-rho_dm = 0.3 * nu.GeV/nu.c0**2 / nu.cm**3
+@export
+def rho_dm():
+    """Local dark matter density"""
+    return 0.3 * nu.GeV/nu.c0**2 / nu.cm**3
 
-# Most common velocity of WIMPs in the halo,
-# relative to galactic center (asymptotic)
-v_0 = 220 * nu.km/nu.s
 
-# Galactic escape velocity
-v_esc = 544 * nu.km/nu.s
+@export
+def v_0():
+    """Most common velocity of WIMPs in the halo,
+    relative to galactic center (asymptotic)
+    """
+    return 220 * nu.km/nu.s
+
+
+@export
+def v_esc():
+    """Galactic escape velocity"""
+    return 544 * nu.km/nu.s
 
 
 # J2000.0 epoch conversion (converts datetime to days since epoch)
@@ -105,9 +113,9 @@ def v_earth(t=None):
 def v_max(t=None):
     """Return maximum observable dark matter velocity on Earth."""
     if t is None:
-        return v_esc + v_earth(t)
+        return v_esc() + v_earth(t)
     else:
-        return v_esc + np.sum(earth_velocity(t) ** 2) ** 0.5
+        return v_esc() + np.sum(earth_velocity(t) ** 2) ** 0.5
 
 
 @export
@@ -125,16 +133,18 @@ def observed_speed_dist(v, t=None):
     v_earth_t = v_earth(t)
 
     # Normalization constant, see Lewin&Smith appendix 1a
-    _w = v_esc/v_0
+    _w = v_esc()/v_0()
     k = erf(_w) - 2/np.pi**0.5 * _w * np.exp(-_w**2)
 
     # Maximum cos(angle) for this velocity, otherwise v0
-    xmax = np.minimum(1, (v_esc**2 - v_earth_t**2 - v**2)/(2 * v_earth_t * v))
-
-    y = (k * v / (np.pi**0.5 * v_0 * v_earth_t)
-         * (np.exp(-((v-v_earth_t)/v_0)**2)
-            - np.exp(-1/v_0**2 * (v**2 + v_earth_t**2
-                                  + 2 * v * v_earth_t * xmax))))
+    xmax = np.minimum(1,
+                      (v_esc()**2 - v_earth_t**2 - v**2)
+                      / (2 * v_earth_t * v))
+
+    y = (k * v / (np.pi**0.5 * v_0() * v_earth_t)
+         * (np.exp(-((v-v_earth_t)/v_0())**2)
+            - np.exp(-1/v_0()**2 * (v**2 + v_earth_t**2
+                                    + 2 * v * v_earth_t * xmax))))
 
     # Zero if v > v_max
     try:

wimprates/migdal.py

@@ -1,7 +1,6 @@
 """Migdal effect
 
 """
-
 import numericalunits as nu
 import numpy as np
 import pandas as pd
@@ -26,15 +25,15 @@
               5.4e3, 4.9e3,
               1.1e3, 9.3e2, 6.6e2,
               2e2, 1.4e2, 6.1e1,
-              2.1e1, 9.8]) * nu.eV))
+              2.1e1, 9.8])))
 
 
 def vmin_migdal(w, erec, mw):
     """Return minimum WIMP velocity to make a Migdal signal with energy w,
     given elastic recoil energy erec and WIMP mass mw.
     """
-    y = (wr.mn * erec / (2 * wr.reduced_mass(wr.mn, mw) ** 2))**0.5
-    y += w / (2 * wr.mn * erec)**0.5
+    y = (wr.mn() * erec / (2 * wr.mu_nucleus(mw) ** 2))**0.5
+    y += w / (2 * wr.mn() * erec)**0.5
     return np.maximum(0, y)
 
 
@@ -69,6 +68,7 @@ def rate_migdal(w, mw, sigma_nucleon, interaction='SI', m_med=float('inf'),
 
     result = 0
     for state, binding_e in binding_es_for_migdal.items():
+        binding_e *= nu.eV
         # Only consider n=3 and n=4
         # n=5 is the valence band so unreliable in in liquid
         # n=1,2 contribute very little
@@ -97,7 +97,7 @@ def diff_rate(v, erec):
 
                 # Migdal effect |Z|^2
                 # TODO: ?? what is explicit (eV/c)**2 doing here?
-                * (nu.me * (2 * erec / wr.mn)**0.5 / (nu.eV / nu.c0))**2
+                * (nu.me * (2 * erec / wr.mn())**0.5 / (nu.eV / nu.c0))**2
                 / (2 * np.pi)
                 * p(eelec))
 
@@ -113,4 +113,4 @@ def diff_rate(v, erec):
 
         result += r
 
-    return wr.rho_dm / mw * (1 / wr.mn) * result
+    return wr.rho_dm() / mw * (1 / wr.mn()) * result