add ultranest import

README.rst

@@ -2,7 +2,7 @@
 anesthetic: nested sampling post-processing
 ===========================================
 :Authors: Will Handley and Lukas Hergt
-:Version: 2.0.0
+:Version: 2.1.0
 :Homepage: https://github.com/handley-lab/anesthetic
 :Documentation: http://anesthetic.readthedocs.io/
 

anesthetic/_version.py

@@ -1 +1 @@
-__version__ = '2.0.0'
+__version__ = '2.1.0'

anesthetic/read/chain.py

@@ -3,6 +3,7 @@
 from anesthetic.read.getdist import read_getdist
 from anesthetic.read.cobaya import read_cobaya
 from anesthetic.read.multinest import read_multinest
+from anesthetic.read.ultranest import read_ultranest
 
 
 def read_chains(root, *args, **kwargs):
@@ -12,6 +13,7 @@ def read_chains(root, *args, **kwargs):
 
         * `PolyChord <https://github.com/PolyChord/PolyChordLite>`_,
         * `MultiNest <https://github.com/farhanferoz/MultiNest>`_,
+        * `UltraNest <https://github.com/JohannesBuchner/UltraNest>`_,
         * `CosmoMC <https://github.com/cmbant/CosmoMC>`_,
         * `Cobaya <https://github.com/CobayaSampler/cobaya>`_,
         * or anything `GetDist <https://github.com/cmbant/getdist>`_
@@ -45,7 +47,11 @@ def read_chains(root, *args, **kwargs):
             "anesthetic.html#anesthetic.samples.MCMCSamples.remove_burn_in"
         )
     errors = []
-    for read in [read_polychord, read_multinest, read_cobaya, read_getdist]:
+    readers = [
+        read_polychord, read_multinest, read_cobaya,
+        read_ultranest, read_getdist
+    ]
+    for read in readers:
         try:
             samples = read(root, *args, **kwargs)
             samples.root = root

anesthetic/read/cobaya.py

@@ -1,4 +1,4 @@
-"""Read MCMCSamples from cobaya chains."""
+"""Read MCMCSamples from Cobaya chains."""
 import os
 import re
 import numpy as np
@@ -13,8 +13,8 @@ def read_paramnames(root):
     columns as there are parameters (sampled and derived) plus an
     additional two corresponding to the weights (first column) and the
     log-posterior (second column). The first line should start with a # and
-    should list the parameter names corresponding to the columns. This
-    will be used as label in the pandas array.
+    should list the parameter names corresponding to the columns. These
+    will be used as handles in the pandas array.
     """
     with open(root + ".1.txt") as f:
         header = f.readline()[1:]
@@ -34,6 +34,12 @@ def read_cobaya(root, *args, **kwargs):
     Note that in order to optimally read chains from Cobaya you need to have
     `GetDist <https://getdist.readthedocs.io/en/latest/>`__ installed.
 
+    Parameters
+    ----------
+    root : str
+        root name for reading files in Cobaya format, i.e. the files
+        ``<root>.*.txt`` and ``<root>.updated.yaml``.
+
     Returns
     -------
     :class:`anesthetic.samples.MCMCSamples`

anesthetic/read/multinest.py

@@ -1,12 +1,22 @@
-"""Read NestedSamples from multinest chains."""
+"""Read NestedSamples from MultiNest chains."""
 import os
 import numpy as np
 from anesthetic.read.getdist import read_paramnames
 from anesthetic.samples import NestedSamples
 
 
 def read_multinest(root, *args, **kwargs):
-    """Read <root>ev.dat and <root>phys_live.points in MultiNest format."""
+    """Read MultiNest chain files.
+
+    Parameters
+    ----------
+    root : str
+        root name for reading files in MultiNest format, i.e. the files
+        ``<root>ev.dat`` and ``<root>phys_live.points`` in the old format, and
+        ``<root>dead-birth.txt`` and ``<root>phys_live-birth.txt`` in the new
+        format.
+
+    """
     try:
         data = np.loadtxt(root + 'dead-birth.txt')
         samples, logL, logL_birth, _ = np.split(data, [-4, -3, -2], axis=1)

anesthetic/read/polychord.py

@@ -1,12 +1,20 @@
-"""Read NestedSamples from polychord chains."""
+"""Read NestedSamples from PolyChord chains."""
 import os
 import numpy as np
 from anesthetic.read.getdist import read_paramnames
 from anesthetic.samples import NestedSamples
 
 
 def read_polychord(root, *args, **kwargs):
-    """Read ``<root>_dead-birth.txt`` in PolyChord format."""
+    """Read PolyChord chain files.
+
+    Parameters
+    ----------
+    root : str
+        root name for reading files in PolyChord format, i.e. the files
+        ``<root>_dead-birth.txt`` and ``<root>_phys_live-birth.txt``.
+
+    """
     birth_file = root + '_dead-birth.txt'
     birth_file
     data = np.loadtxt(birth_file)

anesthetic/read/ultranest.py

@@ -0,0 +1,40 @@
+"""Read NestedSamples from UltraNest results."""
+import os
+import json
+from anesthetic.samples import NestedSamples
+try:
+    import h5py
+except ImportError:
+    pass
+
+
+def read_ultranest(root, *args, **kwargs):
+    """Read UltraNest files.
+
+    Parameters
+    ----------
+    root : str
+        root name for reading files in UltraNest format, i.e. the files
+        ``<root>/info/results.json`` and ``<root>/results/points.hdf5``.
+
+    """
+    with open(os.path.join(root, 'info', 'results.json')) as infofile:
+        labels = json.load(infofile)['paramnames']
+    num_params = len(labels)
+
+    filepath = os.path.join(root, 'results', 'points.hdf5')
+    with h5py.File(filepath, 'r') as fileobj:
+        points = fileobj['points']
+        _, ncols = points.shape
+        x_dim = ncols - 3 - num_params
+        logL_birth = points[:, 0]
+        logL = points[:, 1]
+        samples = points[:, 3+x_dim:3+x_dim+num_params]
+
+    kwargs['label'] = kwargs.get('label', os.path.basename(root))
+    columns = kwargs.pop('columns', labels)
+    labels = kwargs.pop('labels', labels)
+    data = samples
+
+    return NestedSamples(data=data, logL=logL, logL_birth=logL_birth,
+                         columns=columns, labels=labels, *args, **kwargs)

anesthetic/samples.py

@@ -359,7 +359,7 @@ def plot_2d(self, axes=None, *args, **kwargs):
             axes = self.drop_labels().columns
 
         if not isinstance(axes, AxesDataFrame):
-            _, axes = make_2d_axes(axes, labels=self.get_labels(),
+            _, axes = make_2d_axes(axes, labels=self.get_labels_map(),
                                    upper=('upper' in kind),
                                    lower=('lower' in kind),
                                    diagonal=('diagonal' in kind))