Redownload volumes (#28)

* bugfix for #27

* enhanced specificity, only removing corrupted file and not everything. locally tested.

* Added test

---------

Co-authored-by: Miles Wells <k1o0@3tk.co>

Author: grg2rsr

iblatlas/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.9.0'
+__version__ = '0.10.0'
 
 """A package for working with brain atlases.
 

iblatlas/atlas.py

@@ -1455,9 +1455,16 @@ def __init__(self, res_um=25, scaling=(1, 1, 1), mock=False, hist_path=None):
         super().__init__(image, label, dxyz, regions, ibregma, dims2xyz=dims2xyz, xyz2dims=xyz2dims)
 
     @staticmethod
-    def _read_volume(file_volume):
+    def _read_volume(file_volume: Path):
         if file_volume.suffix == '.nrrd':
-            volume, _ = nrrd.read(file_volume, index_order='C')  # ml, dv, ap
+            try:
+                volume, _ = nrrd.read(file_volume, index_order='C')  # ml, dv, ap
+            except nrrd.NRRDError as e:
+                _logger.error(f"An error occured when loading atlas volumes: {e}. \
+                                Deleting and redownloading the files.")
+                file_volume.unlink(missing_ok=True)
+                file_volume = _download_atlas_allen(file_volume)
+                volume, _ = nrrd.read(file_volume, index_order='C')  # ml, dv, ap
             # we want the coronal slice to be the most contiguous
             volume = np.transpose(volume, (2, 0, 1))  # image[iap, iml, idv]
         elif file_volume.suffix == '.npz':

iblatlas/tests/__init__.py

@@ -1,6 +1,12 @@
+import tempfile
+from pathlib import Path
+import unittest
+from unittest import mock
+from textwrap import dedent
+
 import numpy as np
 import matplotlib.pyplot as plt
-import unittest
+
 from iblatlas.atlas import AllenAtlas, BrainRegions
 from iblatlas.flatmaps import FlatMap
 from iblatlas.plots import plot_swanson, annotate_swanson
@@ -118,3 +124,44 @@ def test_load(self):
         df_genes, gene_expression, atlas_agea = agea.load()
         self.assertEqual(df_genes.shape[0], gene_expression.shape[0])
         self.assertEqual(gene_expression.shape[1:], (58, 41, 67))
+
+
+class TestReadVolume(unittest.TestCase):
+
+    @mock.patch('iblatlas.atlas._download_atlas_allen')
+    def test_read_volume(self, mock_download):
+        """Test that NRRD files are read corrrectly, and redownloaded when corrupted."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmpdir_path = Path(tmpdir)
+            file_image = tmpdir_path / 'annotation_25.nrrd'
+            # create dummy files
+            file_image.write_bytes(b'\xFF')  # invalid nrrd
+            # mock the download to just write a dummy file
+            mock_download.side_effect = self._side_effect
+            with self.assertLogs('iblatlas.atlas', level='ERROR') as cm:
+                volume = AllenAtlas._read_volume(file_image)
+            self.assertIn('An error occured when loading atlas volumes', cm.output[0])
+            # the mock should have been called twice, once for each file
+            self.assertEqual(mock_download.call_count, 1)
+        self.assertEqual(volume.shape, (3, 3, 3))
+
+    @staticmethod
+    def _side_effect(file_path):
+        """Write a dummy file with correct header for nrrd."""
+        s = dedent(
+            """\
+            NRRD0003
+            type: unsigned char
+            dimension: 3
+            sizes: 3 3 3
+            spacings: 1.0458000000000001 1.0458000000000001 1.0458000000000001
+            kinds: domain domain domain
+            encoding: ASCII\n\n
+            """)
+        s += '\n'.join(map(str, list(range(3**3))))
+        file_path.write_bytes(s.encode())
+        return file_path
+
+
+if __name__ == "__main__":
+    unittest.main(exit=False, verbosity=2)