Merge pull request #641 from effigies/fix/deterministic_derivs

effigies · web-flow · commit 9e353501454f · 2021-06-08T16:15:18.000-04:00
RF: Write derivatives once, using deterministic gzip settings
diff --git a/niworkflows/interfaces/bids.py b/niworkflows/interfaces/bids.py
@@ -32,7 +32,7 @@
 from nipype.interfaces.io import add_traits
 from templateflow.api import templates as _get_template_list
 from ..utils.bids import _init_layout, relative_to_root
-from ..utils.images import overwrite_header
+from ..utils.images import set_consumables, unsafe_write_nifti_header_and_data
 from ..utils.misc import splitext as _splitext, _copy_any
 
 regz = re.compile(r"\.gz$")
@@ -576,16 +576,18 @@ def _run_interface(self, runtime):
             out_file = out_path / dest_file
             out_file.parent.mkdir(exist_ok=True, parents=True)
             self._results["out_file"].append(str(out_file))
-            self._results["compression"].append(_copy_any(orig_file, str(out_file)))
+            self._results["compression"].append(str(dest_file).endswith(".gz"))
+
+            # Set data and header iff changes need to be made. If these are
+            # still None when it's time to write, just copy.
+            new_data, new_header = None, None
 
             is_nifti = out_file.name.endswith(
                 (".nii", ".nii.gz")
             ) and not out_file.name.endswith((".dtseries.nii", ".dtseries.nii.gz"))
             data_dtype = self.inputs.data_dtype or DEFAULT_DTYPES[self.inputs.suffix]
             if is_nifti and any((self.inputs.check_hdr, data_dtype)):
-                # Do not use mmap; if we need to access the data at all, it will be to
-                # rewrite, risking a BusError
-                nii = nb.load(out_file, mmap=False)
+                nii = nb.load(orig_file)
 
                 if self.inputs.check_hdr:
                     hdr = nii.header
@@ -607,12 +609,10 @@ def _run_interface(self, runtime):
 
                     if curr_codes != xcodes or curr_units != units:
                         self._results["fixed_hdr"][i] = True
-                        hdr.set_qform(nii.affine, xcodes[0])
-                        hdr.set_sform(nii.affine, xcodes[1])
-                        hdr.set_xyzt_units(*units)
-
-                        # Rewrite file with new header
-                        overwrite_header(nii, out_file)
+                        new_header = hdr.copy()
+                        new_header.set_qform(nii.affine, xcodes[0])
+                        new_header.set_sform(nii.affine, xcodes[1])
+                        new_header.set_xyzt_units(*units)
 
                 if data_dtype == "source":  # match source dtype
                     try:
@@ -624,9 +624,6 @@ def _run_interface(self, runtime):
                         data_dtype = None
 
                 if data_dtype:
-                    if self.inputs.check_hdr:
-                        # load updated NIfTI
-                        nii = nb.load(out_file, mmap=False)
                     data_dtype = np.dtype(data_dtype)
                     orig_dtype = nii.get_data_dtype()
                     if orig_dtype != data_dtype:
@@ -639,9 +636,24 @@ def _run_interface(self, runtime):
                         else:
                             new_data = np.asanyarray(nii.dataobj, dtype=data_dtype)
                         # and set header to match
-                        nii.set_data_dtype(data_dtype)
-                        nii = nii.__class__(new_data, nii.affine, nii.header)
-                        nii.to_filename(out_file)
+                        if new_header is None:
+                            new_header = nii.header.copy()
+                        new_header.set_data_dtype(data_dtype)
+                del nii
+
+            if new_data is new_header is None:
+                _copy_any(orig_file, str(out_file))
+            else:
+                orig_img = nb.load(orig_file)
+                if new_data is None:
+                    set_consumables(new_header, orig_img.dataobj)
+                    new_data = orig_img.dataobj.get_unscaled()
+                unsafe_write_nifti_header_and_data(
+                    fname=out_file,
+                    header=new_header,
+                    data=new_data
+                )
+                del orig_img
 
         if len(self._results["out_file"]) == 1:
             meta_fields = self.inputs.copyable_trait_names()
diff --git a/niworkflows/interfaces/tests/test_bids.py b/niworkflows/interfaces/tests/test_bids.py
@@ -2,6 +2,7 @@
 import os
 from pathlib import Path
 import json
+from hashlib import sha1
 
 import numpy as np
 import nibabel as nb
@@ -24,61 +25,70 @@
 
 @pytest.mark.parametrize("out_path_base", [None, "fmriprep"])
 @pytest.mark.parametrize(
-    "source,input_files,entities,expectation",
+    "source,input_files,entities,expectation,checksum",
     [
         (
             T1W_PATH,
             ["anat.nii.gz"],
             {"desc": "preproc"},
             "sub-100185/anat/sub-100185_desc-preproc_T1w.nii.gz",
+            "7c047921def32da260df4a985019b9f5231659fa",
         ),
         (
             T1W_PATH,
             ["anat.nii.gz"],
             {"desc": "preproc", "space": "MNI"},
             "sub-100185/anat/sub-100185_space-MNI_desc-preproc_T1w.nii.gz",
+            "b22399f50ce454049d5d074457a92ab13e7fdf8c",
         ),
         (
             T1W_PATH,
             ["anat.nii.gz"],
             {"desc": "preproc", "space": "MNI", "resolution": "native"},
             "sub-100185/anat/sub-100185_space-MNI_desc-preproc_T1w.nii.gz",
+            "b22399f50ce454049d5d074457a92ab13e7fdf8c",
         ),
         (
             T1W_PATH,
             ["anat.nii.gz"],
             {"desc": "preproc", "space": "MNI", "resolution": "high"},
             "sub-100185/anat/sub-100185_space-MNI_res-high_desc-preproc_T1w.nii.gz",
+            "b22399f50ce454049d5d074457a92ab13e7fdf8c",
         ),
         (
             T1W_PATH,
             ["tfm.txt"],
             {"from": "fsnative", "to": "T1w", "suffix": "xfm"},
             "sub-100185/anat/sub-100185_from-fsnative_to-T1w_mode-image_xfm.txt",
+            "da39a3ee5e6b4b0d3255bfef95601890afd80709",
         ),
         (
             T1W_PATH,
             ["tfm.h5"],
             {"from": "MNI152NLin2009cAsym", "to": "T1w", "suffix": "xfm"},
             "sub-100185/anat/sub-100185_from-MNI152NLin2009cAsym_to-T1w_mode-image_xfm.h5",
+            "da39a3ee5e6b4b0d3255bfef95601890afd80709",
         ),
         (
             T1W_PATH,
             ["anat.nii.gz"],
             {"desc": "brain", "suffix": "mask"},
             "sub-100185/anat/sub-100185_desc-brain_mask.nii.gz",
+            "d425f0096b6b6d1252973e48b31d760c0b1bdc11",
         ),
         (
             T1W_PATH,
             ["anat.nii.gz"],
             {"desc": "brain", "suffix": "mask", "space": "MNI"},
             "sub-100185/anat/sub-100185_space-MNI_desc-brain_mask.nii.gz",
+            "a2a6efa16eb23173d0ee64779de879711bc74643",
         ),
         (
             T1W_PATH,
             ["anat.surf.gii"],
             {"suffix": "pial", "hemi": "L"},
             "sub-100185/anat/sub-100185_hemi-L_pial.surf.gii",
+            "da39a3ee5e6b4b0d3255bfef95601890afd80709",
         ),
         (
             T1W_PATH,
@@ -88,6 +98,8 @@
                 f"sub-100185/anat/sub-100185_desc-{s}_dseg.nii"
                 for s in ("aseg", "aparcaseg")
             ],
+            ["a235cdf59f9bf077ba30bf2523a56508e3a5aabb",
+             "a235cdf59f9bf077ba30bf2523a56508e3a5aabb"],
         ),
         (
             T1W_PATH,
@@ -97,6 +109,8 @@
                 f"sub-100185/anat/sub-100185_desc-preproc_T1w.{ext}"
                 for ext in ("nii", "json")
             ],
+            ["25c107d4a3e6f98e48aa752c5bbd88ab8e8d069f",
+             "da39a3ee5e6b4b0d3255bfef95601890afd80709"],
         ),
         (
             T1W_PATH,
@@ -106,63 +120,73 @@
                 f"sub-100185/anat/sub-100185_label-{lab}_probseg.nii.gz"
                 for lab in ("GM", "WM", "CSF")
             ],
+            ["7c047921def32da260df4a985019b9f5231659fa"] * 3,
         ),
         # BOLD data
         (
             BOLD_PATH,
             ["aroma.csv"],
             {"suffix": "AROMAnoiseICs"},
             "sub-100185/func/sub-100185_task-machinegame_run-1_AROMAnoiseICs.csv",
+            "da39a3ee5e6b4b0d3255bfef95601890afd80709",
         ),
         (
             BOLD_PATH,
             ["confounds.tsv"],
             {"suffix": "regressors", "desc": "confounds"},
             "sub-100185/func/sub-100185_task-machinegame_run-1_desc-confounds_regressors.tsv",
+            "da39a3ee5e6b4b0d3255bfef95601890afd80709",
         ),
         (
             BOLD_PATH,
             ["mixing.tsv"],
             {"suffix": "mixing", "desc": "MELODIC"},
             "sub-100185/func/sub-100185_task-machinegame_run-1_desc-MELODIC_mixing.tsv",
+            "da39a3ee5e6b4b0d3255bfef95601890afd80709",
         ),
         (
             BOLD_PATH,
             ["lh.func.gii"],
             {"space": "fsaverage", "density": "10k", "hemi": "L"},
             "sub-100185/func/sub-100185_task-machinegame_run-1_"
             "space-fsaverage_den-10k_hemi-L_bold.func.gii",
+            "da39a3ee5e6b4b0d3255bfef95601890afd80709",
         ),
         (
             BOLD_PATH,
             ["hcp.dtseries.nii"],
             {"space": "fsLR", "density": "91k"},
             "sub-100185/func/sub-100185_task-machinegame_run-1_"
             "space-fsLR_den-91k_bold.dtseries.nii",
+            "53d9b486d08fec5a952f68fcbcddb38a72818d4c",
         ),
         (
             BOLD_PATH,
             ["ref.nii"],
             {"space": "MNI", "suffix": "boldref"},
             "sub-100185/func/sub-100185_task-machinegame_run-1_space-MNI_boldref.nii",
+            "53d9b486d08fec5a952f68fcbcddb38a72818d4c",
         ),
         (
             BOLD_PATH,
             ["dseg.nii"],
             {"space": "MNI", "suffix": "dseg", "desc": "aseg"},
             "sub-100185/func/sub-100185_task-machinegame_run-1_space-MNI_desc-aseg_dseg.nii",
+            "6d2cae7f56c246d7934e2e21e7b472ecc63a4257",
         ),
         (
             BOLD_PATH,
             ["mask.nii"],
             {"space": "MNI", "suffix": "mask", "desc": "brain"},
             "sub-100185/func/sub-100185_task-machinegame_run-1_space-MNI_desc-brain_mask.nii",
+            "c365991854931181a1444d6803f5289448e7e266",
         ),
         (
             BOLD_PATH,
             ["bold.nii"],
             {"space": "MNI", "desc": "preproc"},
             "sub-100185/func/sub-100185_task-machinegame_run-1_space-MNI_desc-preproc_bold.nii",
+            "aa1eed935e6a8dcca646b0c78ee57218e30e2974",
         ),
         # Nondeterministic order - do we really need this to work, or we can stay safe with
         # MapNodes?
@@ -175,30 +199,35 @@
             ["anat.html"],
             {"desc": "conform", "datatype": "figures"},
             "sub-100185/figures/sub-100185_desc-conform_T1w.html",
+            "da39a3ee5e6b4b0d3255bfef95601890afd80709",
         ),
         (
             BOLD_PATH,
             ["aroma.csv"],
             {"suffix": "AROMAnoiseICs", "extension": "h5"},
             ValueError,
+            None,
         ),
         (
             T1W_PATH,
             ["anat.nii.gz"] * 3,
             {"desc": "preproc", "space": "MNI"},
             ValueError,
+            None,
         ),
         (
             "sub-07/ses-preop/anat/sub-07_ses-preop_T1w.nii.gz",
             ["tfm.h5"],
             {"from": "orig", "to": "target", "suffix": "xfm"},
             "sub-07/ses-preop/anat/sub-07_ses-preop_from-orig_to-target_mode-image_xfm.h5",
+            "da39a3ee5e6b4b0d3255bfef95601890afd80709",
         ),
         (
             "sub-07/ses-preop/anat/sub-07_ses-preop_run-1_T1w.nii.gz",
             ["tfm.txt"],
             {"from": "orig", "to": "T1w", "suffix": "xfm"},
             "sub-07/ses-preop/anat/sub-07_ses-preop_run-1_from-orig_to-T1w_mode-image_xfm.txt",
+            "da39a3ee5e6b4b0d3255bfef95601890afd80709",
         ),
     ],
 )
@@ -210,6 +239,7 @@ def test_DerivativesDataSink_build_path(
     input_files,
     entities,
     expectation,
+    checksum,
     dismiss_entities,
 ):
     """Check a few common derivatives generated by NiPreps."""
@@ -274,9 +304,12 @@ def test_DerivativesDataSink_build_path(
     output = dds.run().outputs.out_file
     if isinstance(output, str):
         output = [output]
+        checksum = [checksum]
 
     for out, exp in zip(output, expectation):
         assert Path(out).relative_to(tmp_path) == Path(base) / exp
+    for out, chksum in zip(output, checksum):
+        assert sha1(Path(out).read_bytes()).hexdigest() == chksum
 
 
 def test_DerivativesDataSink_dtseries_json_hack(tmp_path):
@@ -423,6 +456,52 @@ def test_DerivativesDataSink_t1w(tmp_path, space, size, units, xcodes, fixed):
     assert nii.header.get_xyzt_units() == ("mm", "unknown")
 
 
+@pytest.mark.parametrize("dtype", ("i2", "u2", "f4"))
+def test_DerivativesDataSink_values(tmp_path, dtype):
+    # We use static checksums above, which ensures we don't break things, but
+    # pins the tests to specific values.
+    # Here we use random values, check that the values are preserved, and then
+    # the checksums are unchanged across two runs.
+    fname = str(tmp_path / "source.nii.gz")
+    rng = np.random.default_rng()
+    hdr = nb.Nifti1Header()
+    hdr.set_qform(np.eye(4), code=1)
+    hdr.set_sform(np.eye(4), code=1)
+    nb.Nifti1Image(rng.uniform(500, 2000, (5, 5, 5)), np.eye(4), hdr).to_filename(fname)
+
+    orig_data = np.asanyarray(nb.load(fname).dataobj)
+    expected = np.rint(orig_data) if dtype[0] in "iu" else orig_data
+
+    dds = bintfs.DerivativesDataSink(
+        base_directory=str(tmp_path),
+        keep_dtype=True,
+        data_dtype=dtype,
+        desc="preproc",
+        source_file=T1W_PATH,
+        in_file=fname,
+    ).run()
+
+    out_file = Path(dds.outputs.out_file)
+
+    nii = nb.load(out_file)
+    assert np.allclose(nii.dataobj, expected)
+
+    checksum = sha1(out_file.read_bytes()).hexdigest()
+    out_file.unlink()
+
+    # Rerun to ensure determinism with non-zero data
+    dds = bintfs.DerivativesDataSink(
+        base_directory=str(tmp_path),
+        keep_dtype=True,
+        data_dtype=dtype,
+        desc="preproc",
+        source_file=T1W_PATH,
+        in_file=fname,
+    ).run()
+
+    assert sha1(out_file.read_bytes()).hexdigest() == checksum
+
+
 @pytest.mark.parametrize("field", ["RepetitionTime", "UndefinedField"])
 def test_ReadSidecarJSON_connection(testdata_dir, field):
     """
diff --git a/niworkflows/utils/images.py b/niworkflows/utils/images.py
@@ -1,6 +1,7 @@
 """Utilities to manipulate images."""
 import nibabel as nb
 import numpy as np
+from gzip import GzipFile
 
 
 def rotation2canonical(img):
@@ -33,12 +34,16 @@ def unsafe_write_nifti_header_and_data(fname, header, data):
     If you're not using this for NIfTI files specifically, you're playing
     with Fortran-ordered fire.
     """
-    # ImageOpener handles zips transparently
-    with nb.openers.ImageOpener(fname, mode="wb") as fobj:
+    with open(fname, "wb") as fobj:
+        # Avoid setting fname or mtime, for deterministic outputs
+        if str(fname).endswith(".gz"):
+            fobj = GzipFile("", "wb", 9, fobj, 0.0)
         header.write_to(fobj)
         # This function serializes one block at a time to reduce memory usage a bit
         # It assumes Fortran-ordered data.
         nb.volumeutils.array_to_file(data, fobj, offset=header.get_data_offset())
+        if str(fname).endswith(".gz"):
+            fobj.close()
 
 
 def set_consumables(header, dataobj):
