Add ultranest sampler

.coveragerc

@@ -4,5 +4,7 @@
 exclude_lines =
     if __name__ == .__main__.:
     raise
+    print
+    warn
 
 ignore_errors = True

.github/workflows/test_package.yml

@@ -70,8 +70,9 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         run: |
           coverage run --source=dddm -m pytest -vsx --durations 0
-          coverage run --append --source=dddm -m pytest --nbmake -n=auto ./notebooks
           coveralls --service=github
+          # Failfast since we don't using coveralls
+          pytest -vsx
       - name: Test package (linux)
         if: matrix.os == 'ubuntu-latest' && matrix.test == 'pytest'
         env:

.sourcery.yml

@@ -8,6 +8,7 @@ refactor:
   skip:
     inline-immediately-returned-variable
     use-named-expression
+    use-dictionary-union
 
 github:
   request_review: owner

dddm/context.py

@@ -60,8 +60,10 @@ class Context:
         'nestle': dddm.samplers.nestle.NestleSampler,
         'multinest': dddm.samplers.pymultinest.MultiNestSampler,
         'emcee': dddm.samplers.emcee.MCMCStatModel,
+        'ultranest': dddm.samplers.ultranest.UltraNestSampler,
         'multinest_combined': dddm.samplers.multi_detectors.CombinedMultinest,
         'nestle_combined': dddm.samplers.multi_detectors.CombinedNestle,
+        'ultranest_combined': dddm.samplers.multi_detectors.CombinedUltraNest,
     })
     _halo_classes = immutabledict({
         'shm': dddm.SHM,

dddm/samplers/__init__.py

@@ -9,3 +9,6 @@
 
 from .multi_detectors import *
 from . import multi_detectors
+
+from .ultranest import *
+from . import ultranest

dddm/samplers/multi_detectors.py

@@ -6,6 +6,7 @@
 import dddm
 from .pymultinest import MultiNestSampler, convert_dic_to_savable
 from .nestle import NestleSampler
+from .ultranest import UltraNestSampler
 import typing as ty
 
 export, __all__ = dddm.exporter()
@@ -47,9 +48,9 @@ def save_sub_configs(self, force_index=False):
             if 'logging' not in c.config:
                 raise ValueError(f'{c} does not have logging in config ({list(c.config.keys())})')
             save_as = os.path.join(f'{save_dir}', f'{c.config["detector"]}_')
-            with open(save_as + 'config.json', 'w') as file:
+            with open(f'{save_as}config.json', 'w') as file:
                 json.dump(convert_dic_to_savable(c.config), file, indent=4)
-            np.save(save_as + 'config.npy', convert_dic_to_savable(c.config))
+            np.save(f'{save_as}config.npy', convert_dic_to_savable(c.config))
             shutil.copy(c.config['logging'], save_as +
                         c.config['logging'].split('/')[-1])
             self.log.info('save_sub_configs::\tdone_saving')
@@ -169,3 +170,61 @@ def __init__(
             for one_class in self.sub_detectors
         ]
         self.log.debug(f'Sub detectors are set: {self.sub_classes}')
+
+
+@export
+class CombinedUltraNest(_CombinedInference, UltraNestSampler):
+    def __init__(
+            self,
+            wimp_mass: ty.Union[float, int],
+            cross_section: ty.Union[float, int],
+            spectrum_class: ty.List[ty.Union[dddm.DetectorSpectrum, dddm.GenSpectrum]],
+            prior: dict,
+            tmp_folder: str,
+            results_dir: str = None,
+            fit_parameters=('log_mass', 'log_cross_section', 'v_0', 'v_esc', 'density', 'k'),
+
+            detector_name=None,
+            verbose=False,
+            notes='default',
+            nlive=1024,
+            tol=0.1,
+    ):
+        assert detector_name is not None
+        # Make list explicit
+        spectrum_classes = spectrum_class
+        del spectrum_class
+
+        UltraNestSampler.__init__(self,
+                                  wimp_mass=wimp_mass,
+                                  cross_section=cross_section,
+                                  spectrum_class=spectrum_classes,
+                                  prior=prior,
+                                  tmp_folder=tmp_folder,
+                                  fit_parameters=fit_parameters,
+                                  detector_name=detector_name,
+                                  verbose=verbose,
+                                  results_dir=results_dir,
+                                  notes=notes,
+                                  nlive=nlive,
+                                  tol=tol,
+                                  )
+        if len(spectrum_classes) < 2:
+            self.log.warning(
+                "Don't use this class for single experiments! Use NestedSamplerStatModel instead")
+        self.sub_detectors = spectrum_classes
+        self.config['sub_sets'] = [str(sp) for sp in spectrum_classes]
+        self.sub_classes = [
+            UltraNestSampler(wimp_mass=wimp_mass,
+                             cross_section=cross_section,
+                             spectrum_class=one_class,
+                             prior=prior,
+                             tmp_folder=tmp_folder,
+                             fit_parameters=fit_parameters,
+                             detector_name=one_class.detector_name,
+                             verbose=verbose,
+                             notes=notes,
+                             )
+            for one_class in self.sub_detectors
+        ]
+        self.log.debug(f'Sub detectors are set: {self.sub_classes}')

dddm/samplers/ultranest.py

@@ -0,0 +1,65 @@
+from .pymultinest import MultiNestSampler
+import ultranest
+import datetime
+import os
+import dddm
+import warnings
+
+export, __all__ = dddm.exporter()
+
+
+@export
+class UltraNestSampler(MultiNestSampler):
+    """
+    Skeleton of ultranest sampler
+    """
+
+    def run(self):
+        warnings.warn('Ultranest sampler is not completely implemented yet')
+        self._fix_parameters()
+        self._print_before_run()
+
+        n_dims = len(self.config["fit_parameters"])
+        tol = self.config['tol']  # the stopping criterion
+        save_at = self.get_save_dir()
+
+        self.log.warning(f'start_fit for {n_dims} parameters')
+
+        start = datetime.datetime.now()
+
+        # Multinest saves output to a folder. First write to the tmp folder,
+        # move it to the results folder later
+        _tmp_folder = self.get_save_dir()
+        save_at_temp = os.path.join(_tmp_folder, 'ultra_nest')
+
+        sampler = ultranest.ReactiveNestedSampler(
+            param_names=list(self.config["fit_parameters"]),
+            loglike=self._log_probability_nested,
+            transform=self._log_prior_transform_nested,  # SafePrior,
+            log_dir=_tmp_folder,
+            resume='resume',
+        )
+        result = sampler.run(
+            min_num_live_points=self.config['nlive'],
+            min_ess=self.config['nlive'],
+            dlogz=tol,
+        )
+        sampler.print_results()
+        self.result_file = save_at_temp
+
+        # Open a save-folder after successful running multinest. Move the
+        # multinest results there.
+        dddm.utils.check_folder_for_file(save_at)
+        end = datetime.datetime.now()
+        dt = (end - start).total_seconds()
+        self.log.info(f'fit_done in {dt} s ({dt / 3600} h)')
+        self.log_dict['did_run'] = True
+        # release the config
+        self.config = dddm.utils._immutable_to_dict(self.config)
+        self.config['fit_time'] = dt
+
+        self.log.info('Finished with running Multinest!')
+        return result, sampler
+
+    def get_summary(self):
+        raise NotImplementedError

extra_requirements/requirements-tests.txt

@@ -3,6 +3,7 @@ coveralls==3.3.1
 emcee==3.1.2
 gitpython==3.1.27
 gvar==11.10
+h5py==3.6.0
 hypothesis==6.21.1
 immutabledict==2.2.1
 matplotlib==3.5.2
@@ -22,11 +23,12 @@ PyYAML==6.0
 scipy==1.8.1
 seaborn==0.11.2
 tqdm==4.64.0
+ultranest==3.4.6
 vegas==5.1.1
 
 
-# These two packages require metadata not shipped via pip+git
 git+https://github.com/jorana/wimprates
+# These two packages require metadata not shipped via pip+git
 # git+https://github.com/JoranAngevaare/DarkELF
 # git+https://github.com/jorana/verne
 

requirements.txt

@@ -1,6 +1,7 @@
 corner
 coveralls
 emcee
+h5py
 hypothesis
 immutabledict
 matplotlib
@@ -15,7 +16,7 @@ pytest-runner
 scipy
 seaborn~=0.11.2
 tqdm
-
+ultranest
 
 # These two packages require metadata not shipped via pip+git
 # git+https://github.com/jorana/wimprates

tests/test_samplers/test_ultranest.py

@@ -0,0 +1,62 @@
+from unittest import TestCase, skipIf
+import dddm
+import numpy as np
+
+
+# @skipIf(dddm.utils.is_windows(), "Ultranest only works on linux")
+class PymultinestTest(TestCase):
+    def setUp(self) -> None:
+        self.ct = dddm.test_context()
+
+    def test(self, halo_name='shm', max_sigma_off=5, **kwargs, ):
+        base_config = dict(wimp_mass=50,
+                           cross_section=1e-45,
+                           sampler_name='ultranest',
+                           detector_name='Xe_simple',
+                           prior="Pato_2010",
+                           halo_name=halo_name,
+                           detector_kwargs=None,
+                           halo_kwargs=None if halo_name == 'shm' else dict(location='XENON'),
+                           sampler_kwargs=dict(nlive=100, tol=0.9, verbose=1),
+                           fit_parameters=('log_mass', 'log_cross_section',),
+                           )
+        config = {**base_config, **kwargs}  # noqa
+        sampler = self.ct.get_sampler_for_detector(**config)
+
+        results, _ = sampler.run()
+
+        fails = []
+        for thing, expected, avg, std in zip(
+                    base_config.get('fit_parameters'),
+                    [getattr(sampler, f) for f in base_config.get('fit_parameters')],
+                    results['posterior']['mean'],
+                    results['posterior']['stdev']
+                ):
+            nsigma_off = np.abs(expected - avg) / std
+            # assert False, dict(thing=thing, expected=expected, avg=avg, nsigma_off=nsigma_off)
+            message = (f'For {thing}: expected {expected:.2f} yielded '
+                       f'{avg:.2f} +/- {std:.2f}. Off '
+                       f'by {nsigma_off:.1f} sigma')
+            if nsigma_off > max_sigma_off:
+                fails += [message]
+            print(message)
+
+        self.assertFalse(fails, fails)
+        return sampler
+
+    @skipIf(*dddm.test_utils.skip_long_test())
+    def test_combined(self,
+                      halo_name='shm',
+                      fit_parameters=('log_mass', 'log_cross_section',)):
+        self.test(
+            wimp_mass=50,
+            cross_section=1e-45,
+            sampler_name='ultranest_combined',
+            detector_name=['Xe_simple', 'Ar_simple', 'Ge_simple'],
+            prior="Pato_2010",
+            halo_name=halo_name,
+            detector_kwargs=None,
+            halo_kwargs=None if halo_name == 'shm' else dict(location='XENON'),
+            sampler_kwargs=dict(nlive=100, tol=0.9, verbose=1, detector_name='test_combined'),
+            fit_parameters=fit_parameters,
+        )