Seed dithering

CHANGES.md

@@ -1,6 +1,11 @@
 version 1.2.5 (not yet released)
 -------------
 
+New Features
+
+    - writing images supports the dither_seed keyword, to seed
+      the random number generator for subtractive dithering
+
 Bug Fixes
 
     - Fig bug slicing tables that have TBIT columns

fitsio/fitsio_pywrap.c

@@ -1395,6 +1395,7 @@ PyFITSObject_create_image_hdu(struct PyFITSObject* self, PyObject* args, PyObjec
     int extver=0;
     float qlevel=0;
     int qmethod=0;
+    int dither_seed=0;
     float hcomp_scale=0;
     int hcomp_smooth=0;
 
@@ -1411,6 +1412,7 @@ PyFITSObject_create_image_hdu(struct PyFITSObject* self, PyObject* args, PyObjec
 
          "qlevel",
          "qmethod",
+         "dither_seed",
 
          "hcomp_scale",
          "hcomp_smooth",
@@ -1419,14 +1421,15 @@ PyFITSObject_create_image_hdu(struct PyFITSObject* self, PyObject* args, PyObjec
          "extver",
          NULL,
     };
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "Oi|OiOfifisi", kwlist,
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "Oi|OiOfiifisi", kwlist,
                           &array_obj, &nkeys,
                           &dims_obj,
                           &comptype,
                           &tile_dims_obj,
 
                           &qlevel,
                           &qmethod,
+                          &dither_seed,
 
                           &hcomp_scale,
                           &hcomp_smooth,
@@ -1495,6 +1498,12 @@ PyFITSObject_create_image_hdu(struct PyFITSObject* self, PyObject* args, PyObjec
                 goto create_image_hdu_cleanup;
             }
 
+            if (dither_seed >= 0) {
+                if (fits_set_dither_seed(self->fits, dither_seed, &status)) {
+                    goto create_image_hdu_cleanup;
+                }
+            }
+
             if (comptype == HCOMPRESS_1) {
 
                 if (fits_set_hcomp_scale(self->fits, hcomp_scale, &status)) {

fitsio/fitslib.py

@@ -288,6 +288,7 @@ def write(filename, data, extname=None, extver=None, header=None,
           names=None, write_bitcols=False, compress=None, tile_dims=None,
           qlevel=DEFAULT_QLEVEL,
           qmethod=DEFAULT_QMETHOD,
+          dither_seed=None,
           hcomp_scale=DEFAULT_HCOMP_SCALE,
           hcomp_smooth=False,
           **keys):
@@ -376,6 +377,9 @@ def write(filename, data, extname=None, extver=None, header=None,
                 Preserves zeros
 
         Defaults to 'SUBTRACTIVE_DITHER_1' which follows the fpack defaults
+    dither_seed: int
+        Seed for the subtractive dither.  Seeding makes the lossy
+        compression reproducible.
 
     hcomp_scale: float
         Scale value for HCOMPRESS, 0.0 means lossless compression. Default is
@@ -409,6 +413,7 @@ def write(filename, data, extname=None, extver=None, header=None,
             tile_dims=tile_dims,
             qlevel=qlevel,
             qmethod=qmethod,
+            dither_seed=dither_seed,
             hcomp_scale=hcomp_scale,
             hcomp_smooth=hcomp_smooth,
         )
@@ -624,6 +629,7 @@ def write(self, data, units=None, extname=None, extver=None,
               tile_dims=None,
               qlevel=DEFAULT_QLEVEL,
               qmethod=DEFAULT_QMETHOD,
+              dither_seed=None,
               hcomp_scale=DEFAULT_HCOMP_SCALE,
               hcomp_smooth=False,
               header=None, names=None,
@@ -686,6 +692,9 @@ def write(self, data, units=None, extname=None, extver=None,
                     Preserves zeros
 
             Defaults to 'SUBTRACTIVE_DITHER_1' which follows the fpack defaults
+        dither_seed: int
+            Seed for the subtractive dither.  Seeding makes the lossy
+            compression reproducible.
 
         hcomp_scale: float
             Scale value for HCOMPRESS, 0.0 means lossless compression. Default
@@ -731,6 +740,7 @@ def write(self, data, units=None, extname=None, extver=None,
                              tile_dims=tile_dims,
                              qlevel=qlevel,
                              qmethod=qmethod,
+                             dither_seed=dither_seed,
                              hcomp_scale=hcomp_scale,
                              hcomp_smooth=hcomp_smooth,
                              header=header)
@@ -745,6 +755,7 @@ def write_image(self, img, extname=None, extver=None,
                     compress=None, tile_dims=None,
                     qlevel=DEFAULT_QLEVEL,
                     qmethod=DEFAULT_QMETHOD,
+                    dither_seed=None,
                     hcomp_scale=DEFAULT_HCOMP_SCALE,
                     hcomp_smooth=False,
                     header=None):
@@ -792,6 +803,9 @@ def write_image(self, img, extname=None, extver=None,
                     Preserves zeros
 
             Defaults to 'SUBTRACTIVE_DITHER_1' which follows the fpack defaults
+        dither_seed: int
+            Seed for the subtractive dither.  Seeding makes the lossy
+            compression reproducible.
 
         hcomp_scale: float
             Scale value for HCOMPRESS, 0.0 means lossless compression. Default
@@ -823,6 +837,7 @@ def write_image(self, img, extname=None, extver=None,
             tile_dims=tile_dims,
             qlevel=qlevel,
             qmethod=qmethod,
+            dither_seed=dither_seed,
             hcomp_scale=hcomp_scale,
             hcomp_smooth=hcomp_smooth,
         )
@@ -844,6 +859,7 @@ def create_image_hdu(self,
                          tile_dims=None,
                          qlevel=DEFAULT_QLEVEL,
                          qmethod=DEFAULT_QMETHOD,
+                         dither_seed=None,
                          hcomp_scale=DEFAULT_HCOMP_SCALE,
                          hcomp_smooth=False,
                          header=None):
@@ -917,7 +933,9 @@ def create_image_hdu(self,
                     Preserves zeros
 
             Defaults to 'SUBTRACTIVE_DITHER_1' which follows the fpack defaults
-
+        dither_seed: int
+            Seed for the subtractive dither.  Seeding makes the lossy
+            compression reproducible.
         hcomp_scale: float
             Scale value for HCOMPRESS, 0.0 means lossless compression. Default
             is 0.0 following the fpack defaults.
@@ -1002,6 +1020,7 @@ def create_image_hdu(self,
 
         comptype = get_compress_type(compress)
         qmethod = get_qmethod(qmethod)
+        dither_seed = get_dither_seed(dither_seed)
 
         tile_dims = get_tile_dims(tile_dims, dims)
         if qlevel is None:
@@ -1032,6 +1051,7 @@ def create_image_hdu(self,
 
             qlevel=qlevel,
             qmethod=qmethod,
+            dither_seed=dither_seed,
 
             hcomp_scale=hcomp_scale,
             hcomp_smooth=hcomp_smooth,
@@ -1800,6 +1820,15 @@ def get_qmethod(qmethod):
     return _qmethod_map[qmethod]
 
 
+def get_dither_seed(dither_seed):
+    if dither_seed is None:
+        # -1 means do not set the seed
+        return -1
+    else:
+        # must fit in an int
+        return numpy.int32(dither_seed)
+
+
 def check_comptype_img(comptype, dtype_str):
 
     if comptype == NOCOMPRESS:

fitsio/tests/test_image_compression.py

@@ -219,3 +219,101 @@ def test_compress_preserve_zeros():
 
                     for zind in zinds:
                         assert rdata[zind[0], zind[1]] == 0.0
+
+
+@pytest.mark.parametrize(
+    'compress',
+    [
+        'rice',
+        'hcompress',
+        'plio',
+    ]
+)
+@pytest.mark.parametrize(
+    'match_seed',
+    [False, True],
+)
+@pytest.mark.parametrize(
+    'use_fits_object',
+    [False, True],
+)
+@pytest.mark.parametrize(
+    'dtype',
+    ['f4', 'f8'],
+)
+def test_compressed_seed(compress, match_seed, use_fits_object, dtype):
+    """
+    Test writing and reading a rice compressed image
+    """
+    nrows = 5
+    ncols = 20
+
+    qlevel = 16
+
+    seed = 1919
+    rng = np.random.RandomState(seed)
+
+    if match_seed:
+        # dither_seed = 9881
+        dither_seed1 = 9881
+        dither_seed2 = 9881
+    else:
+        # dither_seed = None
+        dither_seed1 = 3
+        dither_seed2 = 4
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        fname1 = os.path.join(tmpdir, 'test1.fits')
+        fname2 = os.path.join(tmpdir, 'test2.fits')
+
+        data = rng.normal(size=(nrows, ncols))
+        if compress == 'plio':
+            data = data.clip(min=0)
+        data = data.astype(dtype)
+
+        if use_fits_object:
+            with FITS(fname1, 'rw') as fits1:
+                fits1.write(
+                    data, compress=compress, qlevel=qlevel,
+                    # dither_seed=dither_seed,
+                    dither_seed=dither_seed1,
+                )
+                rdata1 = fits1[-1].read()
+
+            with FITS(fname2, 'rw') as fits2:
+                fits2.write(
+                    data, compress=compress, qlevel=qlevel,
+                    # dither_seed=dither_seed,
+                    dither_seed=dither_seed2,
+                )
+                rdata2 = fits2[-1].read()
+        else:
+            write(
+                fname1, data, compress=compress, qlevel=qlevel,
+                # dither_seed=dither_seed,
+                dither_seed=dither_seed1,
+            )
+            rdata1 = read(fname1)
+
+            write(
+                fname2, data, compress=compress, qlevel=qlevel,
+                # dither_seed=dither_seed,
+                dither_seed=dither_seed2,
+            )
+            rdata2 = read(fname2)
+
+        mess = "%s compressed images ('%s')" % (compress, dtype)
+
+        if match_seed:
+            assert np.all(rdata1 == rdata2), mess
+        else:
+            assert np.all(rdata1 != rdata2), mess
+
+
+if __name__ == '__main__':
+    test_compressed_seed(
+        compress='rice',
+        match_seed=False,
+        use_fits_object=True,
+        dtype='f4',
+    )