Changes to HDF5 processing

mosaic/file_manipulation/h5.py

@@ -123,7 +123,7 @@ def _write_dataset(name, obj, group):
         dataset.attrs['is_str'] = isinstance(flat_obj[0], str)
 
 
-def read(obj, lazy=True, filter=None):
+def read(obj, lazy=True, filter=None, only=None):
     if isinstance(obj, h5py.Group):
         if filter is None:
             filter = {}
@@ -137,6 +137,8 @@ def read(obj, lazy=True, filter=None):
         if obj.attrs.get('is_array'):
             data = []
             for key in sorted(obj.keys()):
+                if only is not None and key not in only:
+                    continue
                 try:
                     value = read(obj[key], lazy=lazy, filter=filter)
                 except FilterException:
@@ -145,6 +147,8 @@ def read(obj, lazy=True, filter=None):
         else:
             data = {}
             for key in obj.keys():
+                if only is not None and key not in only:
+                    continue
                 try:
                     value = read(obj[key], lazy=lazy, filter=filter)
                 except FilterException:
@@ -301,9 +305,9 @@ def file(self):
     def close(self):
         self._file.close()
 
-    def load(self, lazy=True, filter=None):
+    def load(self, lazy=True, filter=None, only=None):
         group = self._file['/']
-        description = read(group, lazy=lazy, filter=filter)
+        description = read(group, lazy=lazy, filter=filter, only=only)
         return Struct(description)
 
     def dump(self, description):

stride/problem/acquisitions.py

@@ -1187,7 +1187,11 @@ def __get_desc__(self, **kwargs):
         return description
 
     def __set_desc__(self, description):
-        for shot_desc in description.shots:
+        if 'shots' in description:
+            shots = description.shots
+        else:
+            shots = description.values()
+        for shot_desc in shots:
             if shot_desc.id not in self._shots:
                 shot = Shot(shot_desc.id,
                             geometry=self._geometry,

stride/problem/base.py

@@ -146,7 +146,7 @@ def load(self, *args, **kwargs):
         """
         kwargs['parameter'] = self.name
         with h5.HDF5(*args, **kwargs, mode='r') as file:
-            description = file.load(filter=kwargs.pop('filter', None))
+            description = file.load(filter=kwargs.pop('filter', None), only=kwargs.pop('only', None))
 
             self.__set_desc__(description)
 
@@ -241,7 +241,7 @@ def load(self, *args, **kwargs):
         """
         kwargs['parameter'] = self.name
         with h5.HDF5(*args, **kwargs, mode='r') as file:
-            description = file.load(filter=kwargs.pop('filter', None))
+            description = file.load(filter=kwargs.pop('filter', None), only=kwargs.pop('only', None))
 
             # TODO If there's already a grid and they don't match, resample instead
             if 'space' in description and self._grid.space is None: