Refactoring and Multiprocessing

.gitignore

@@ -1,3 +1,6 @@
+# Notebooks in root dir
+/*.ipynb
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -72,4 +75,6 @@ target/
 *.pdf
 
 # Temp
-notebooks/dm_e
+notebooks/dm_e
+
+!notebooks/*.png

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "wimprates/data/migdal/Cox"]
+	path = wimprates/data/migdal/Cox/cox_submodule
+	url = https://github.com/petercox/Migdal.git

pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"

requirements.txt

@@ -4,3 +4,4 @@ pandas
 tqdm
 boltons
 numericalunits
+importlib_metadata

setup.py

@@ -1,28 +1,37 @@
 import setuptools
 
-readme = open('README.md').read()
-history = open('HISTORY.md').read().replace('.. :changelog:', '')
-requirements = open('requirements.txt').read().splitlines()
+readme = open("README.md").read()
+history = open("HISTORY.md").read().replace(".. :changelog:", "")
+requirements = open("requirements.txt").read().splitlines()
 
 setuptools.setup(
-    name='wimprates',
-    version='0.5.0',
-    description='Differential rates of WIMP-nucleus scattering',
-    long_description=readme + '\n\n' + history,
-    long_description_content_type='text/markdown',
-    author='Jelle Aalbers',
-    url='https://github.com/jelleaalbers/wimprates',
-    license='MIT',
+    name="wimprates",
+    version="0.5.0",
+    description="Differential rates of WIMP-nucleus scattering",
+    long_description=readme + "\n\n" + history,
+    long_description_content_type="text/markdown",
+    author="Jelle Aalbers",
+    url="https://github.com/jelleaalbers/wimprates",
+    license="MIT",
     packages=setuptools.find_packages(),
-    setup_requires=['pytest-runner'],
+    setup_requires=["pytest-runner"],
     install_requires=requirements,
-    package_dir={'wimprates': 'wimprates'},
-    package_data={'wimprates': [
-        'data/bs/*', 'data/migdal/*', 'data/sd/*', 'data/dme/*']},
-    tests_require=requirements + ['pytest', 'unittest'],
-    keywords='wimp,spin-dependent,spin-independent,bremsstrahlung,migdal',
-    classifiers=['Intended Audience :: Science/Research',
-                 'Development Status :: 3 - Alpha',
-                 'Programming Language :: Python',
-                 'Programming Language :: Python :: 3'],
-    zip_safe=False)
+    package_dir={"wimprates": "wimprates"},
+    package_data={
+        "wimprates": ["data/bs/*", "data/migdal/**", "data/sd/*", "data/dme/*"]
+    },
+    tests_require=requirements + ["pytest", "unittest"],
+    keywords="wimp,spin-dependent,spin-independent,bremsstrahlung,migdal",
+    classifiers=[
+        "Intended Audience :: Science/Research",
+        "Development Status :: 3 - Alpha",
+        "Programming Language :: Python",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+    ],
+    python_requires=">=3.9",
+    zip_safe=False,
+)

tests/test_halo.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 
 import pandas as pd
-from wimprates import j2000, StandardHaloModel, j2000_from_ymd
+from wimprates import j2000, StandardHaloModel, j2000_from_ymd, j2000_to_datetime
 import numericalunits as nu
 import numpy as np
 
@@ -21,6 +21,11 @@ def test_j2000_datetime():
     assert j2000(date) == 3318.25
 
 
+def test_datetime_j2000():
+    date = 3318.25
+    assert j2000_to_datetime(date) == datetime(year=2009, month=1, day=31, hour=18)
+
+
 def test_j2000_ns_int():
     date = datetime(year=2009, month=1, day=31, hour=18)
     value = pd.to_datetime(date).value

tests/test_wimprates.py

@@ -14,6 +14,8 @@
 class TestBenchmarks(unittest.TestCase):
     opts = dict(mw=50,
                 sigma_nucleon=1e-45,
+                save_cache=False,
+                load_cache=False,
                 )
     def test_elastic(self):
         ref = 30.39515403337126
@@ -40,9 +42,19 @@ def test_spindependent(self):
                 0.00019944698779638946)
 
 
-    def test_migdal(self):
-        self.assertAlmostEqual(wr.rate_wimp_std(1, detection_mechanism='migdal', **self.opts),
+    def test_migdal_Ibe(self):
+        self.assertAlmostEqual(wr.rate_wimp_std(1, detection_mechanism='migdal', migdal_model="Ibe", **self.opts),
                 0.27459766238555017)
+
+
+    def test_migdal_Cox(self):
+        self.assertAlmostEqual(wr.rate_wimp_std(1, detection_mechanism='migdal', migdal_model="Cox", **self.opts),
+                0.2843514286729741)
+
+
+    def test_migdal_Cox_dipole(self):
+        self.assertAlmostEqual(wr.rate_wimp_std(1, detection_mechanism='migdal', migdal_model="Cox", dipole=True, **self.opts),
+                0.30438231874513705)
 
 
     def test_brems(self):
@@ -104,3 +116,6 @@ def test_average_v_earth(self):
             # places=1 means that we get the same results at the first decimal (fine for 500.0<?>)
             places=1
         )
+
+if __name__ == "__main__":
+    unittest.main()

wimprates/__init__.py

@@ -14,4 +14,4 @@
 from .migdal import *
 from .electron import *
 from .summary import *
-
+from .data.migdal.Cox.cox_wrapper import *

wimprates/data/migdal/Cox/cox_wrapper.py

@@ -0,0 +1,52 @@
+"""
+The paths in the code in Peter Cox's package assume the working directory to be
+the root of its package. With this wrapper we change the working dir when computing
+the interpolators, then returning the Migdal class instance once they've been instantiated.
+The working directory is then reset.
+"""
+
+from functools import lru_cache
+import os
+import sys
+
+from .cox_submodule.Migdal import Migdal
+import wimprates as wr
+
+export, __all__ = wr.exporter()
+
+
+@export
+@lru_cache
+def cox_migdal_model(element: str, **kwargs) -> Migdal:
+    """
+    This function creates a Cox Migdal model for a given element.
+
+    Parameters:
+    - element (str): The element for which the Cox Migdal model is created.
+    - **kwargs: Additional keyword arguments for loading probabilities and total probabilities.
+
+    Returns:
+    - material: The Cox Migdal material object.
+
+    Example usage:
+    cox_migdal_model("carbon", arg1=value1, arg2=value2)
+
+    Note: The Cox's model assumes that the main process is running in its root directory and uses
+        relative paths. Therefore, we need to switch the working directory to the root of the package
+        when computing the interpolators. 
+        This wrapper function changes the working directory temporarily, instantiates the Migdal class, 
+        and then resets the working directory back to its original state.
+    """
+    original_cwd = os.getcwd()
+
+    try:
+        migdal_directory = os.path.join(os.path.dirname(__file__), "cox_submodule")
+        os.chdir(migdal_directory)
+
+        material = Migdal(element)
+        material.load_probabilities(**kwargs)
+        material.load_total_probabilities(**kwargs)
+    finally:
+        os.chdir(original_cwd)
+
+    return material

wimprates/electron.py

@@ -1,5 +1,6 @@
 """Dark matter - electron scattering
 """
+from functools import lru_cache
 import numericalunits as nu
 import numpy as np
 from scipy.interpolate import RegularGridInterpolator, interp1d
@@ -9,13 +10,17 @@
 export, __all__ = wr.exporter()
 __all__ += ['dme_shells', 'l_to_letter', 'l_to_number']
 
-# Load form factor and construct interpolators
-shell_data = wr.load_pickle('dme/dme_ionization_ff.pkl')
-for _shell, _sd in shell_data.items():
-    _sd['log10ffsquared_itp'] = RegularGridInterpolator(
-        (_sd['lnks'], _sd['lnqs']),
-        np.log10(_sd['ffsquared']),
-        bounds_error=False, fill_value=-float('inf'),)
+
+@lru_cache()
+def get_shell_data():
+    """Load form factor and construct interpolators"""
+    shell_data = wr.load_pickle('dme/dme_ionization_ff.pkl')
+    for _shell, _sd in shell_data.items():
+        _sd['log10ffsquared_itp'] = RegularGridInterpolator(
+            (_sd['lnks'], _sd['lnqs']),
+            np.log10(_sd['ffsquared']),
+            bounds_error=False, fill_value=-float('inf'),)
+    return shell_data
 
 
 dme_shells = [(5, 1), (5, 0), (4, 2), (4, 1), (4, 0)]
@@ -54,6 +59,9 @@ def dme_ionization_ff(shell, e_er, q):
     # Ry = rydberg = 13.6 eV
     ry = nu.me * nu.e ** 4 / (8 * nu.eps0 ** 2 * nu.hPlanck ** 2)
     lnk = np.log(e_er / ry) / 2
+
+    shell_data = get_shell_data()
+
     return 10**(shell_data[shell]['log10ffsquared_itp'](
         np.vstack([lnk, lnq]).T))
 
@@ -84,8 +92,8 @@ def v_min_dme(eb, erec, q, mw):
     return (erec + eb) / q + q / (2 * mw)
 
 
-# Precompute velocity integrals for t=None
 @export 
+@lru_cache()
 def velocity_integral_without_time(halo_model=None):
     halo_model = wr.StandardHaloModel() if halo_model is None else halo_model
     _v_mins = np.linspace(0, 1, 1000) * wr.v_max(None, halo_model.v_esc)
@@ -105,7 +113,6 @@ def velocity_integral_without_time(halo_model=None):
         fill_value=0, bounds_error=False)
     return inverse_mean_speed_kms
 
-inverse_mean_speed_kms = velocity_integral_without_time()
 
 
 @export
@@ -141,9 +148,12 @@ def rate_dme(erec, n, l, mw, sigma_dme,
 
     # No bounds are given for the q integral
     # but the form factors are only specified in a limited range of q
+    shell_data = get_shell_data()
     qmax = (np.exp(shell_data[shell]['lnqs'].max())
             * (nu.me * nu.c0 * nu.alphaFS))
 
+    # Precompute velocity integrals for t=None
+    inverse_mean_speed_kms = velocity_integral_without_time()
     if t is None:
         # Use precomputed inverse mean speed,
         # so we only have to do a single integral

wimprates/halo.py

@@ -67,6 +67,17 @@ def j2000_from_ymd(year, month, day_of_month):
             + np.floor(30.61 * (m + 1))
             + day_of_month - 730563.5)
 
+@export
+def j2000_to_datetime(j2000_date):
+    """
+    Returns date in np.datetime64 instance from the fractional number of days since J2000.0 epoch.
+    It is effectively the inverse of the j2000 function.
+    """
+    zero_value = pd.to_datetime("2000-01-01T12:00").value
+
+    nanoseconds_per_day = nu.day / nu.ns
+    _date = pd.to_datetime(j2000_date * nanoseconds_per_day).value
+    return pd.to_datetime(_date + zero_value).round("s")
 
 @export
 def earth_velocity(t, v_0 = None):