Repository: dfm/emcee

Issue:
Enable HDF5 compression in the backend
**General information:**

- emcee version: 3.0.2
- platform: (Arch) Linux
- installation method (pip/conda/source/other?): AUR repository

**Problem description:**
It would be nice to be able to enable compression when using HDF5 files via the emcee backends. A quick skim of the documentation and `backends/hdf.py` file suggests this isn't currently possible. Perhaps the h5py `compression=` argument could be exposed in the `emcee.backends.HDFBackend` class?

Discussion:

Patch:
diff --git a/src/emcee/backends/hdf.py b/src/emcee/backends/hdf.py
index 9503512..0f0c0de 100644
--- a/src/emcee/backends/hdf.py
+++ b/src/emcee/backends/hdf.py
@@ -22,9 +22,9 @@ except ImportError:
 def does_hdf5_support_longdouble():
     if h5py is None:
         return False
-    with NamedTemporaryFile(prefix="emcee-temporary-hdf5",
-                            suffix=".hdf5",
-                            delete=False) as f:
+    with NamedTemporaryFile(
+        prefix="emcee-temporary-hdf5", suffix=".hdf5", delete=False
+    ) as f:
         f.close()
 
         with h5py.File(f.name, "w") as hf:
@@ -53,12 +53,23 @@ class HDFBackend(Backend):
             ``RuntimeError`` if the file is opened with write access.
 
     """
-    def __init__(self, filename, name="mcmc", read_only=False, dtype=None):
+
+    def __init__(
+        self,
+        filename,
+        name="mcmc",
+        read_only=False,
+        dtype=None,
+        compression=None,
+        compression_opts=None,
+    ):
         if h5py is None:
             raise ImportError("you must install 'h5py' to use the HDFBackend")
         self.filename = filename
         self.name = name
         self.read_only = read_only
+        self.compression = compression
+        self.compression_opts = compression_opts
         if dtype is None:
             self.dtype_set = False
             self.dtype = np.float64
@@ -109,18 +120,27 @@ class HDFBackend(Backend):
             g.attrs["ndim"] = ndim
             g.attrs["has_blobs"] = False
             g.attrs["iteration"] = 0
-            g.create_dataset("accepted", data=np.zeros(nwalkers))
+            g.create_dataset(
+                "accepted",
+                data=np.zeros(nwalkers),
+                compression=self.compression,
+                compression_opts=self.compression_opts,
+            )
             g.create_dataset(
                 "chain",
                 (0, nwalkers, ndim),
                 maxshape=(None, nwalkers, ndim),
                 dtype=self.dtype,
+                compression=self.compression,
+                compression_opts=self.compression_opts,
             )
             g.create_dataset(
                 "log_prob",
                 (0, nwalkers),
                 maxshape=(None, nwalkers),
                 dtype=self.dtype,
+                compression=self.compression,
+                compression_opts=self.compression_opts,
             )
 
     def has_blobs(self):
@@ -206,6 +226,8 @@ class HDFBackend(Backend):
                         (ntot, nwalkers),
                         maxshape=(None, nwalkers),
                         dtype=dt,
+                        compression=self.compression,
+                        compression_opts=self.compression_opts,
                     )
                 else:
                     g["blobs"].resize(ntot, axis=0)
@@ -239,18 +261,25 @@ class HDFBackend(Backend):
 
 
 class TempHDFBackend(object):
-
-    def __init__(self, dtype=None):
+    def __init__(self, dtype=None, compression=None, compression_opts=None):
         self.dtype = dtype
         self.filename = None
+        self.compression = compression
+        self.compression_opts = compression_opts
 
     def __enter__(self):
-        f = NamedTemporaryFile(prefix="emcee-temporary-hdf5",
-                               suffix=".hdf5",
-                               delete=False)
+        f = NamedTemporaryFile(
+            prefix="emcee-temporary-hdf5", suffix=".hdf5", delete=False
+        )
         f.close()
         self.filename = f.name
-        return HDFBackend(f.name, "test", dtype=self.dtype)
+        return HDFBackend(
+            f.name,
+            "test",
+            dtype=self.dtype,
+            compression=self.compression,
+            compression_opts=self.compression_opts,
+        )
 
     def __exit__(self, exception_type, exception_value, traceback):
         os.remove(self.filename)


Test patch:
diff --git a/src/emcee/tests/unit/test_backends.py b/src/emcee/tests/unit/test_backends.py
index 370008e..b7a01d3 100644
--- a/src/emcee/tests/unit/test_backends.py
+++ b/src/emcee/tests/unit/test_backends.py
@@ -1,8 +1,8 @@
 # -*- coding: utf-8 -*-
 
 import os
+from os.path import join
 from itertools import product
-from tempfile import NamedTemporaryFile
 
 import numpy as np
 import pytest
@@ -263,3 +263,14 @@ def test_longdouble_preserved(backend):
 
             assert s.log_prob.dtype == np.longdouble
             assert np.all(s.log_prob == lp)
+
+
+@pytest.mark.skipif(h5py is None, reason="HDF5 not available")
+def test_hdf5_compression():
+    with backends.TempHDFBackend(compression="gzip") as b:
+        run_sampler(b, blobs=True)
+        # re-open and read
+        b.get_chain()
+        b.get_blobs()
+        b.get_log_prob()
+        b.accepted