Repository: terrapower/armi
Pull request #1778 (2024-07-12 17:45:38)

Make new jaggedArray backwards compatible
## What is the change?

The recent new jagged array implementation broke backwards compatibility with old databases. Upon closer inspection, the two formats for storing the jagged array are pretty closely related, and it is possible to preserve backwards compatibility with just a little bit of extra work in the `unpack` routine.

## Why is the change being made?

<!-- MANDATORY: Explain why the change is necessary -->
<!-- Optional: Link to any related GitHub Issues -->

This change is being made to maintain backwards compatibility with ARMI databases created before PR #1726 was merged. This will be extremely valuable to users, who will be able to load from old databases without having to revert to an older version of their app.


---

## Checklist

<!--
    You (the pull requester) should put an `x` in the boxes below you have completed.

    (If a checkbox doesn't apply to your PR, check it anyway.)
-->

- [x] This PR has only [one purpose or idea](https://terrapower.github.io/armi/developer/tooling.html#one-idea-one-pr).
- [x] [Tests](https://terrapower.github.io/armi/developer/tooling.html#test-it) have been added/updated to verify any new/changed code.

<!-- Check the code quality -->

- [x] The code style follows [good practices](https://terrapower.github.io/armi/developer/standards_and_practices.html).
- [x] The commit message(s) follow [good practices](https://terrapower.github.io/armi/developer/tooling.html).

<!-- Check the project-level cruft -->

- [x] The [release notes](https://terrapower.github.io/armi/developer/tooling.html#add-release-notes) have been updated if necessary.
- [x] The [documentation](https://terrapower.github.io/armi/developer/tooling.html#document-it) is still up-to-date in the `doc` folder.
- [x] The dependencies are still up-to-date in `pyproject.toml`.


Patch:
diff --git a/armi/bookkeeping/db/jaggedArray.py b/armi/bookkeeping/db/jaggedArray.py
index 3135f7da..c28b9874 100644
--- a/armi/bookkeeping/db/jaggedArray.py
+++ b/armi/bookkeeping/db/jaggedArray.py
@@ -173,17 +173,19 @@ class JaggedArray:
             List of numpy arrays with varying dimensions (i.e., jagged arrays)
         """
         unpackedJaggedData: List[Optional[np.ndarray]] = []
-        numElements = len(self.offsets) + len(self.nones)
+        shapeIndices = [i for i, x in enumerate(self.shapes) if sum(x) != 0]
+        numElements = len(shapeIndices) + len(self.nones)
         j = 0  # non-None element counter
         for i in range(numElements):
             if i in self.nones:
                 unpackedJaggedData.append(None)
             else:
+                k = shapeIndices[j]
                 unpackedJaggedData.append(
                     np.ndarray(
-                        self.shapes[j],
+                        self.shapes[k],
                         dtype=self.dtype,
-                        buffer=self.flattenedArray[self.offsets[j] :],
+                        buffer=self.flattenedArray[self.offsets[k] :],
                     )
                 )
                 j += 1


Test patch:
diff --git a/armi/bookkeeping/db/tests/test_jaggedArray.py b/armi/bookkeeping/db/tests/test_jaggedArray.py
index 89d6c53f..d519bc7f 100644
--- a/armi/bookkeeping/db/tests/test_jaggedArray.py
+++ b/armi/bookkeeping/db/tests/test_jaggedArray.py
@@ -47,6 +47,48 @@ class TestJaggedArray(unittest.TestCase):
         flatArray = JaggedArray.flatten(testdata)
         self.assertEqual(flatArray, [1, 2, 3, 4, 5, None, 6, 7, 8, 9])
 
+    def test_backwardsCompatible(self):
+        """
+        Test that the new JaggedArray can unpack the old database jagged data format.
+
+        The "old" database format contains shapes and offsets for locations that have None.
+        The "new" database format only contains shapes and offsets for non-None values.
+        The "new" unpacking routine is able to read either format.
+        """
+        paramName = "test_old"
+        data = [[1, 2], None, [3, 4, 5], None, None, [6, 7, 8, 9]]
+        flattenedArray = numpy.array([1, 2, 3, 4, 5, 6, 7, 8, 9])
+        shapes = [(2,), (0,), (3,), (0,), (0,), (4,)]
+        offsets = [0, 2, 2, 5, 5, 5, 5]
+        nones = [1, 3, 4]
+        h5file = "test_oldFormat.h5"
+        with h5py.File(h5file, "w") as hf:
+            dset = hf.create_dataset(
+                data=flattenedArray,
+                name=paramName,
+            )
+            dset.attrs["jagged"] = True
+            dset.attrs["offsets"] = offsets
+            dset.attrs["shapes"] = shapes
+            dset.attrs["noneLocations"] = nones
+
+        with h5py.File(h5file, "r") as hf:
+            dataset = hf[paramName]
+            values = dataset[()]
+            offsets = dataset.attrs["offsets"]
+            shapes = dataset.attrs["shapes"]
+            nones = dataset.attrs["noneLocations"]
+
+        roundTrip = JaggedArray.fromH5(
+            values,
+            offsets,
+            shapes,
+            nones,
+            dtype=flattenedArray.dtype,
+            paramName=paramName,
+        )
+        self._compareArrays(data, roundTrip)
+
     def _compareRoundTrip(self, data, paramName):
         """Make sure that data is unchanged by packing/unpacking."""
         jaggedArray = JaggedArray(data, paramName)