Remove hardcoded v_0 from earth_velocity

setup.py

@@ -19,7 +19,7 @@
     package_dir={'wimprates': 'wimprates'},
     package_data={'wimprates': [
         'data/bs/*', 'data/migdal/*', 'data/sd/*', 'data/dme/*']},
-    tests_require=requirements + ['pytest'],
+    tests_require=requirements + ['pytest', 'unittest'],
     keywords='wimp,spin-dependent,spin-independent,bremsstrahlung,migdal',
     classifiers=['Intended Audience :: Science/Research',
                  'Development Status :: 3 - Alpha',

tests/test_wimprates.py

@@ -8,60 +8,65 @@
 import numpy as np
 
 import wimprates as wr
+import unittest
 
 
-opts = dict(mw=50, sigma_nucleon=1e-45)
+class TestBenchmarks(unittest.TestCase):
+    opts = dict(mw=50, sigma_nucleon=1e-45)
+    def test_elastic(self):
+        ref = 33.052499179451
 
+        self.assertAlmostEqual(wr.rate_wimp_std(1, **self.opts), ref)
 
-def isclose(x, y):
-    assert np.abs(x - y)/x < 1e-5
+        # Test numericalunits.reset_units() does not affect results
+        nu.reset_units(123)
+        self.assertAlmostEqual(wr.rate_wimp_std(1, **self.opts), ref)
 
+        # Test vectorized call
+        energies = np.linspace(0.01, 40, 100)
+        dr = wr.rate_wimp_std(energies, **self.opts)
+        self.assertEqual(dr[0], wr.rate_wimp_std(0.01, **self.opts))
 
-def test_elastic():
-    ref = 33.19098343826968
 
-    isclose(wr.rate_wimp_std(1, **opts), ref)
+    def test_lightmediator(self):
+        self.assertAlmostEqual(wr.rate_wimp_std(1, m_med=1e-3, **self.opts),
+                0.0005479703883222002)
 
-    # Test numericalunits.reset_units() does not affect results
-    nu.reset_units(123)
-    isclose(wr.rate_wimp_std(1, **opts), ref)
 
-    # Test vectorized call
-    energies = np.linspace(0.01, 40, 100)
-    dr = wr.rate_wimp_std(energies, **opts)
-    assert dr[0] == wr.rate_wimp_std(0.01, **opts)
+    def test_spindependent(self):
+        self.assertAlmostEqual(wr.rate_wimp_std(1, interaction='SD_n_central', **self.opts),
+                0.00021688396095135553)
 
 
-def test_lightmediator():
-    isclose(wr.rate_wimp_std(1, m_med=1e-3, **opts),
-            0.0005502663384403058)
+    def test_migdal(self):
+        self.assertAlmostEqual(wr.rate_wimp_std(1, detection_mechanism='migdal', **self.opts),
+                0.2621719215956991)
 
 
-def test_spindependent():
-    isclose(wr.rate_wimp_std(1, interaction='SD_n_central', **opts),
-            0.00021779266679860948)
+    def test_brems(self):
+        self.assertAlmostEqual(wr.rate_wimp_std(1, detection_mechanism='bremsstrahlung', **self.opts),
+                0.00017137557193256555)
 
 
-def test_migdal():
-    isclose(wr.rate_wimp_std(1, detection_mechanism='migdal', **opts),
-            0.2610240963512907)
+    def test_dme(self):
+        self.assertAlmostEqual(
+            wr.rate_dme(100* nu.eV, 4, 'd',
+                        mw=nu.GeV/nu.c0**2, sigma_dme=4e-44 * nu.cm**2)
+                * nu.kg * nu.keV * nu.day,
+        2.232912243660405e-06)
 
+    def test_halo_scaling(self):
+        #check that passing rho multiplies the rate correctly:
+        ref = 33.052499179450834
+        halo_model = wr.StandardHaloModel(rho_dm=0.3 * nu.GeV / nu.c0 ** 2 / nu.cm ** 3)
+        self.assertAlmostEqual(wr.rate_wimp_std(1, halo_model=halo_model, **self.opts), ref)
 
-def test_brems():
-    isclose(wr.rate_wimp_std(1, detection_mechanism='bremsstrahlung', **opts),
-            0.00017062652972332665)
 
-
-def test_dme():
-    isclose(
-        wr.rate_dme(100* nu.eV, 4, 'd',
-                    mw=nu.GeV/nu.c0**2, sigma_dme=4e-44 * nu.cm**2)
-            * nu.kg * nu.keV * nu.day,
-    2.232912243660405e-06)
-
-def test_halo_scaling():
-    #check that passing rho multiplies the rate correctly:
-    ref = 33.19098343826968
-
-    isclose(wr.rate_wimp_std(1,halo_model = wr.StandardHaloModel(rho_dm = 0.3 * nu.GeV/nu.c0**2 / nu.cm**3) , **opts), ref)
+    def test_v_earth(self):
+        """Compare v_earth with and without specifying v_0"""
+        kms = nu.km/nu.s
+        v_0_default = 220 * kms
+        self.assertAlmostEqual(wr.v_earth(t=None, v_0=None)/kms,
+                               wr.v_earth(t=None, v_0=v_0_default)/kms,
+                               )
 

wimprates/halo.py

@@ -5,6 +5,7 @@
 import numpy as np
 import pandas as pd
 from scipy.special import erf
+from functools import lru_cache
 
 
 import wimprates as wr
@@ -59,14 +60,18 @@ def j2000_from_ymd(year, month, day_of_month):
 
 
 @export
-def earth_velocity(t):
+def earth_velocity(t, v_0 = None):
     """Returns 3d velocity of earth, in the galactic rest frame,
     in galactic coordinates.
     :param t: J2000.0 timestamp
+    :param v_0: Local standard of rest velocity
 
     Values and formula from https://arxiv.org/abs/1209.3339
     Assumes earth circular orbit.
     """
+    if v_0 is None :
+        v_0 = 220 * nu.km/nu.s
+
     # e_1 and e_2 are the directions of earth's velocity at t1
     # and t1 + 0.25 year.
     e_1 = np.array([0.9931, 0.1170, -0.01032])
@@ -84,36 +89,48 @@ def earth_velocity(t):
     v_earth_sun = v_orbit * (e_1 * np.cos(phi) + e_2 * np.sin(phi))
 
     # Velocity of Local Standard of Rest
-    v_lsr = np.array([0, 220, 0]) * nu.km/nu.s
+    v_lsr = np.array([0, v_0/(nu.km/nu.s), 0]) * nu.km/nu.s
     # Solar peculiar velocity
     v_pec = np.array([11, 12, 7]) * nu.km/nu.s
 
     return v_lsr + v_pec + v_earth_sun
 
 
 @export
-def v_earth(t=None):
+@lru_cache(None)
+def v_earth(t=None, v_0=None, _n_days_average=365):
     """Return speed of earth relative to galactic rest frame
     :param t: J2000 timestamp or None
+    :param v_0: Local standard of rest velocity
+    :param _n_days_average: if a v_0 is specified and t is not, average over
+    these many datapoints to get the year-averaged v_earth
     """
-    if t is None:
+    if t is None and v_0 is None:
         # Velocity of earth/sun relative to gal. center
-        # (eccentric orbit, so not equal to v_0)
+        # (eccentric orbit, so not equal to v_0). Only valid
+        # for the default v_0
         return 232 * nu.km / nu.s
-    else:
-        return np.sum(earth_velocity(t) ** 2) ** 0.5
+    if t is None:
+        # Lazy way of averaging over one-year period, the @lru_cache should
+        # cache the result
+        return np.mean([v_earth(x, v_0 = v_0)
+                        for x in np.linspace(0, 365.25, _n_days_average)]
+                       )
+    return np.sum(earth_velocity(t, v_0=v_0) ** 2) ** 0.5
 
 
 @export
-def v_max(t=None, v_esc=None):
+def v_max(t=None, v_esc=None, v_0=None):
     """Return maximum observable dark matter velocity on Earth."""
-    v_esc = 544 * nu.km/nu.s if v_esc is None else v_esc  # default
+    # defaults
+    v_esc = 544 * nu.km/nu.s if v_esc is None else v_esc
+    v_0 = 220 * nu.km / nu.s if v_0 is None else v_0
     # args do not change value when you do a
     # reset_unit so this is necessary to avoid errors
     if t is None:
-        return v_esc + v_earth(t)
+        return v_esc + v_earth(t, v_0=v_0)
     else:
-        return v_esc + np.sum(earth_velocity(t) ** 2) ** 0.5
+        return v_esc + np.sum(earth_velocity(t, v_0=v_0) ** 2) ** 0.5
 
 
 @export
@@ -132,7 +149,7 @@ def observed_speed_dist(v, t=None, v_0=None, v_esc=None):
     """
     v_0 = 220 * nu.km/nu.s if v_0 is None else v_0
     v_esc = 544 * nu.km/nu.s if v_esc is None else v_esc
-    v_earth_t = v_earth(t)
+    v_earth_t = v_earth(t, v_0=v_0)
 
     # Normalization constant, see Lewin&Smith appendix 1a
     _w = v_esc/v_0
@@ -155,13 +172,13 @@ def observed_speed_dist(v, t=None, v_0=None, v_esc=None):
         len(v)
     except TypeError:
         # Scalar argument
-        if v > v_max(t, v_esc):
+        if v > v_max(t, v_esc, v_0=v_0):
             return 0
         else:
             return y
     else:
         # Array argument
-        y[v > v_max(t, v_esc)] = 0
+        y[v > v_max(t, v_esc, v_0=v_0)] = 0
         return y
 
 
@@ -175,7 +192,7 @@ class used to pass a halo model to the rate computation
         giving normalised velocity distribution in earth rest-frame.
         :param rho_dm -- density in mass/volume of dark matter at the Earth
         The standard halo model also allows variation of v_0
-        :param v_0
+        :param v_0: Local standard of rest velocity
     """
 
     def __init__(self, v_0=None, v_esc=None, rho_dm=None):
@@ -186,5 +203,5 @@ def __init__(self, v_0=None, v_esc=None, rho_dm=None):
     def velocity_dist(self, v, t):
         # in units of per velocity,
         # v is in units of velocity
-        return observed_speed_dist(v, t, self.v_0, self.v_esc)
+        return observed_speed_dist(v, t, v_0=self.v_0, v_esc=self.v_esc)