Add support to TiledDataset for missing, irregular or overlapping tiles

dkist/conftest.py

@@ -299,16 +299,23 @@ def eit_dataset():
         return f.tree["dataset"]
 
 
-@pytest.fixture
-def simple_tiled_dataset(dataset):
+@pytest.fixture(params=[False,
+                        [[False, False],
+                         [True, False]]],
+                ids=["simple-nomask", "simple-masked"])
+def simple_tiled_dataset(dataset, request):
     datasets = [copy.deepcopy(dataset) for i in range(4)]
     for ds in datasets:
         ds.meta["inventory"] = dataset.meta["inventory"]
     dataset_array = np.array(datasets).reshape((2,2))
-    return TiledDataset(dataset_array, dataset.meta["inventory"])
+    return TiledDataset(dataset_array, dataset.meta["inventory"], mask=request.param)
 
 
-@pytest.fixture
+@pytest.fixture(params=[False,
+                        [[False, True, False],
+                         [True, False, True],
+                         [False, True, False]]],
+                ids=["large-nomask", "large-masked"])
 def large_tiled_dataset(tmp_path_factory):
     vbidir = tmp_path_factory.mktemp("data")
     with gzip.open(Path(rootdir) / "large_vbi.asdf.gz", mode="rb") as gfo:

dkist/dataset/tiled_dataset.py

@@ -82,8 +82,8 @@ def _from_components(cls, shape, file_managers, wcses, header_tables, inventory)
 
         return cls(datasets, inventory)
 
-    def __init__(self, dataset_array, inventory=None):
-        self._data = np.array(dataset_array, dtype=object)
+    def __init__(self, dataset_array, inventory=None, mask=False):
+        self._data = np.ma.masked_array(dataset_array, dtype=object, mask=mask)
         self._inventory = inventory or {}
         self._validate_component_datasets(self._data, inventory)
 
@@ -105,7 +105,7 @@ def __getitem__(self, aslice):
 
     @staticmethod
     def _validate_component_datasets(datasets, inventory):
-        datasets = datasets.flat
+        datasets = datasets.compressed()
         inv_1 = datasets[0].meta["inventory"]
         if inv_1 and inv_1 is not inventory:
             raise ValueError("The inventory record of the first dataset does not match the one passed to TiledDataset")
@@ -122,7 +122,7 @@ def flat(self):
         """
         Represent this `.TiledDataset` as a 1D array.
         """
-        return type(self)(self._data.flat, self.inventory)
+        return type(self)(self._data.compressed(), self.inventory)
 
     @property
     def inventory(self):
@@ -137,7 +137,7 @@ def combined_headers(self):
         A single `astropy.table.Table` containing all the FITS headers for all
         files in this dataset.
         """
-        return vstack([ds.meta["headers"] for ds in self._data.flat])
+        return vstack([ds.meta["headers"] for ds in self._data.compressed()])
 
     @property
     def shape(self):

dkist/io/asdf/converters/tiled_dataset.py

@@ -5,16 +5,23 @@ class TiledDatasetConverter(Converter):
     tags = [
         "tag:dkist.nso.edu:dkist/tiled_dataset-0.1.0",
         "asdf://dkist.nso.edu/tags/tiled_dataset-1.0.0",
+        "asdf://dkist.nso.edu/tags/tiled_dataset-1.1.0",
     ]
     types = ["dkist.dataset.tiled_dataset.TiledDataset"]
 
     def from_yaml_tree(cls, node, tag, ctx):
         from dkist.dataset.tiled_dataset import TiledDataset
 
-        return TiledDataset(node["datasets"], node["inventory"])
+        try:
+            mask = node["mask"]
+        except KeyError:
+            mask = None
+
+        return TiledDataset(node["datasets"], node["inventory"], mask)
 
     def to_yaml_tree(cls, tiled_dataset, tag, ctx):
         tree = {}
         tree["inventory"] = tiled_dataset._inventory
         tree["datasets"] = tiled_dataset._data.tolist()
+        tree["mask"] = tiled_dataset._data.mask.tolist()
         return tree

dkist/io/asdf/resources/schemas/tiled_dataset-1.1.0.yaml

@@ -0,0 +1,35 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "asdf://dkist.nso.edu/schemas/tiled_dataset-1.0.0"
+
+title: |
+  A DKIST Tiled Dataset object.
+description:
+  The container for a set of Dataset objects.
+
+type: object
+properties:
+  datasets:
+    description: A nested structure of Dataset objects
+    type: array
+    items:
+      type: array
+      items:
+        anyOf:
+          - tag: "asdf://dkist.nso.edu/tags/dataset-1.*"
+          - null: true
+  inventory:
+    description: A copy of the inventory record for this dataset.
+    type: object
+  mask:
+    description: A mask to indicate if invalid or missing Datasets should be ignored.
+    type: array
+    items:
+      type: array
+      items:
+        - type: boolean
+
+required: [datasets, inventory, mask]
+additionalProperties: false
+...

dkist/io/asdf/tests/test_dataset.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import pytest
+from pytest_lazy_fixtures import lf
 
 import asdf
 import astropy.table
@@ -68,7 +69,7 @@ def test_roundtrip_tiled_dataset(simple_tiled_dataset):
 @pytest.mark.parametrize("tagobj",
                          [
                              "dataset",
-                             "simple_tiled_dataset",
+                             lf("simple_tiled_dataset"),
                          ],
                          indirect=True)
 def test_save_dataset_without_file_schema(tagobj, tmp_path):

pyproject.toml

@@ -55,6 +55,7 @@ tests = [
   "pytest-filter-subpackage",
   "pytest-benchmark",
   "pytest-xdist",
+  "pytest-lazy-fixtures",
   "hypothesis",
   "tox",
   "pydot",