Extension functionality: Merge aggregations

cfapyx/creator.py

@@ -41,6 +41,7 @@ def _first_pass(self, agg_dims: list = None) -> tuple:
 
         ## First Pass - Determine dimensions
         for x, file in enumerate(self.files):
+            is_aggregated = False
             logger.info(f'First pass: File {x+1}/{len(self.files)}')
 
             ds = self._call_file(file)
@@ -51,6 +52,18 @@ def _first_pass(self, agg_dims: list = None) -> tuple:
             all_dims = ds.dimensions.keys()
             all_vars = ds.variables.keys()
 
+            # Determine if standard aggregation or fragment extension
+            for v in all_vars:
+                if hasattr(ds[v],'aggregated_dimensions'):
+                    is_aggregated = True
+                    if self.agg_extend is None:
+                        self.agg_extend = True
+                    elif not self.agg_extend:
+                        raise ValueError(
+                            'Mixed file types not allowed. Can only extend'
+                            ' fragment sets or aggregation sets.'
+                        )
+
             coord_variables = []
             pure_dimensions = []
             variables       = []
@@ -71,6 +84,7 @@ def _first_pass(self, agg_dims: list = None) -> tuple:
             if not var_info:
                 var_info = {v: {} for v in variables}
 
+            # Should be consistent across fragments or aggregation files
             logger.info(f'Coordinate variables: {coord_variables}')
             logger.info(f'Pure dimensions: {pure_dimensions}')
             logger.info(f'Variables: {variables}')
@@ -105,8 +119,21 @@ def _first_pass(self, agg_dims: list = None) -> tuple:
                 if new_info['type'] == 'coord':
                     # Only coordinate dimensions can have attributes
                     dim_info = self._update_info(ds[d], dim_info, new_info)
+                    if is_aggregated:
+                        if 'sizes' not in dim_info[d]:
+                            dim_info[d]['sizes'] = []
+                        if 'starts' not in dim_info[d]:
+                            dim_info[d]['starts'] = []
+                        #dim_info[d]['sizes'].append(arr_components['sizes'])
+                        # This is to prevent the line 141 below triggering again.
+                        #dim_info[d]['starts'].append(arr_components['starts'])
                 else:
-                    dim_info[d] = new_info
+                    if is_aggregated and (dim_info[d] != {} or isinstance(dim_info[d], list)):
+                        if not isinstance(dim_info[d], list):
+                            dim_info[d] = [dim_info[d]]
+                        dim_info[d].append(new_info)
+                    else:
+                        dim_info[d] = new_info
 
                 if arr_components is not None:
                     if first_time:
@@ -124,7 +151,7 @@ def _first_pass(self, agg_dims: list = None) -> tuple:
             for v in variables:
 
                 try:
-                    fill = ds[v].getncattr('_FillValue')
+                    fill = ds[v].get_fill_value()
                 except:
                     fill = None
 
@@ -141,8 +168,28 @@ def _first_pass(self, agg_dims: list = None) -> tuple:
                     '_FillValue': fill,
                 }
 
+                if is_aggregated:
+                    # Special handling for extraction of string variables
+                    # Easier to keep numpy arrays separate if they are wrapped in lists
+                    if ds[v].size == 1:
+                        new_info['arr'] = [np.array(ds[v][0], dtype=ds[v].dtype)]
+                    else:
+                        new_info['arr'] = [np.array(list(ds[v]), dtype=ds[v].dtype)]
+
                 var_info = self._update_info(ds[v], var_info, new_info)
 
+            # No variables in current file are aggregations
+            if self.agg_extend is None and not is_aggregated:
+                self.agg_extend = is_aggregated
+
+            # Any variables in current file are aggregated, while previous 
+            # files were not.
+            if self.agg_extend != is_aggregated:
+                raise ValueError(
+                    'Mixed file types not allowed. Can only extend'
+                    ' fragment sets or aggregation sets.'
+                )
+
             arranged_files[tuple(fcoord)] = file
 
         return arranged_files, global_attrs, var_info, dim_info
@@ -164,7 +211,7 @@ def _second_pass(
         for x, file in enumerate(second_set):
             logger.info(f'Second pass: File {x+1}/{len(self.files)}')
 
-            ds = self._call_file(file) # Ideally don't want to do this twice.
+            ds = self._call_file(file) # Typically have to do this twice.
 
             for v in non_aggregated:
                 new_values = np.array(ds.variables[v])
@@ -261,7 +308,11 @@ def _update_info(
             info[id]['attrs'] = self._accumulate_attrs(info[id]['attrs'], attrs)
 
             for attr, value in new_info.items():
-                if value != info[id][attr]:
+                if attr == 'arr':
+                    if 'arr' not in info[id]:
+                        info[id]['arr'] = []
+                    info[id]['arr'] += value
+                elif value != info[id][attr]:
                     if np.isnan(value) and np.isnan(info[id][attr]):
                         pass
                     else:
@@ -461,6 +512,7 @@ def _accumulate_attrs(self, attrs: dict, ncattrs: dict) -> dict:
         for correct values.
         """
 
+        first_time = False
         if attrs is None:
             first_time = True
             attrs = {}
@@ -510,18 +562,19 @@ def _write_dimensions(self):
 
         for dim, di in self.dim_info.items():
 
-            f_size   = di['f_size']
+            f_size   = di.get('f_size', None)
             dim_size = di['size']
 
             real_part = self.ds.createDimension(
                 dim,
                 dim_size
             )
 
-            frag_part = self.ds.createDimension(
-                f'f_{dim}',
-                f_size,
-            )
+            if f_size is not None:
+                frag_part = self.ds.createDimension(
+                    f'f_{dim}',
+                    f_size,
+                )
 
             f_dims[f'f_{dim}'] = f_size
 
@@ -748,7 +801,10 @@ def _write_nonagg_variable(
                 logger.warning(err)
 
         if 'data' in meta:
-            var_arr[:] = meta['data']
+            if meta['dtype'] == str:
+                var_arr[:] = np.array(meta['data'], dtype=meta['dtype'])
+            else:
+                var_arr[:] = meta['data']
 
 class CFANetCDF(CFACreateMixin, CFAWriteMixin):
 
@@ -768,6 +824,8 @@ def __init__(self, files: list, concat_msg : str = CONCAT_MSG):
         the provided set of files. A custom concat message can also be set
         here if needed."""
 
+        self.agg_extend = None
+
         if isinstance(files, str):
             fileset = glob.glob(files)
             self.files = self._filter_files(fileset)
@@ -816,6 +874,15 @@ def create(
         # First pass collect info
         arranged_files, global_attrs, var_info, dim_info = self._first_pass(agg_dims=agg_dims)
 
+        if self.agg_extend:
+            self._extend(
+                arranged_files,
+                global_attrs,
+                var_info,
+                dim_info
+            )
+            return
+
         global_attrs, var_info, dim_info = self._apply_filters(updates, removals, global_attrs, var_info, dim_info)
 
         # Arrange aggregation dimensions
@@ -873,6 +940,14 @@ def write(
 
             f_dims['versions'] = self.max_files
 
+        if self.agg_extend:
+
+            # Skip writing shapes
+
+            self._write_variables()
+            self.ds.close()
+            return
+
         self._write_shape_dims(f_dims)
         self._write_fragment_shapes()
         self._write_fragment_addresses()
@@ -912,6 +987,171 @@ def handle_conventions(self, value) -> str:
             )
         return delim.join(updated_conventions)
 
+    def _extend(self, arranged_files: tuple, global_attrs: dict, var_info: dict, dim_info: dict):
+        """
+        Extend arranged files according to their coordinates.
+
+        NOTE: Aggregation will be limited to dimensions that are already aggregated. i.e This 
+        feature will only work on files where aggregation dimensions are consistent (i.e extending dimensions).
+
+        These functions must:
+        - Identify the dimension(s) that are being extended and identify the affected variables.
+        - Extend scalar dimensions `f_` if those dimensions are already greater than 1.
+        - Identify the affected fragment constructors (map, uris) and extend along the extending dimensions.
+        """
+        st_dim_info = {}
+
+        # Identify Extending dimensions
+        ext_dims = []
+        for d, info in dim_info.items():
+
+            if not isinstance(info, list):
+                # Coordinate variables are not collected into a list
+                # Default value may be overridden by future construction
+
+                #st_dim_info[d] = info
+                #st_dim_info[d]['arrays'] = info['arrays'][0] 
+                continue
+
+            if 'f_' not in d:
+                if d not in st_dim_info:
+
+                    # Default value - reset using f_ numbers if needed
+                    if isinstance(info, list):
+                        if info[0]['size'] != info[-1]['size']:
+                            raise ValueError(
+                                'Aggregation not possible for differing non-aggregated values.'
+                            )
+                        st_dim_info[d] = info[0]
+                    else:
+                        st_dim_info[d] = info
+                    st_dim_info[d]['array'] = st_dim_info[d].get('arrays',None)
+
+                if 'map_' in d:
+                    # Constructor dimension - don't add more f_dims
+                    st_dim_info[d].pop('f_size')
+
+                continue
+
+            for fileinst in info:
+                if fileinst['size'] != 1:
+                    ext_dims.append(d)
+                    break
+
+            if d not in ext_dims:
+                # Non extending aggregation dimension
+                rd = d.split('_')[-1]
+
+                # Skip coordinate dimensions here
+                if not isinstance(dim_info[rd], dict):
+                    if len(set([i['size'] for i in dim_info[rd]])) != 1:
+                        raise ValueError(
+                            f'Non-extending dimension {rd} differs in size between files'
+                        )
+                    st_dim_info[rd] = dim_info[rd][0]
+                else:
+                    st_dim_info[rd] = dim_info[rd]
+                    st_dim_info[rd]['array'] = st_dim_info[rd]['arrays'][0]
+                    if len(set(dim_info[rd]['sizes'])) != 1:
+                        raise ValueError(
+                            f'Non-extending dimension {rd} differs in size between files'
+                        )
+                    st_dim_info[rd]['size'] = dim_info[rd]['sizes'][0]
+                st_dim_info[rd]['f_size'] = info[0]['size']
+
+        new_dim_sizes = {}
+        for d in ext_dims:
+            real_dim = d.split('_')[-1]
+            new_dim_sizes[d] = sum([i['size'] for i in dim_info[d]])
+
+            # Dim info transformations for writing
+            st_dim_info[real_dim] = dim_info[real_dim]
+            st_dim_info[real_dim]['f_size'] = new_dim_sizes[d]
+            st_dim_info[real_dim]['size']   = sum(dim_info[real_dim]['sizes'])
+
+            # Sort array components using index list
+            sorted_a = [[] for a in dim_info[real_dim]['arrays']]
+            sorted_starts = np.argsort(dim_info[real_dim]['starts'])
+            for x, a in enumerate(dim_info[real_dim]['arrays']):
+
+                sorted_a[sorted_starts[x]] = a
+
+            array = np.concatenate(
+                sorted_a,
+                axis=0
+            )
+
+            st_dim_info[real_dim]['array'] = array
+
+
+        sorted_dim_sizes = sorted([
+            (v, d)  for d, v in new_dim_sizes.items()],
+            key = lambda x: x[0]
+        )
+
+        orders = [np.argsort(dim_info[d.split('_')[-1]]['starts']) for d in ext_dims]
+
+        # Identify the affected fragment components
+        extension_vars = []
+        for v, info in var_info.items():
+            extended = False
+            for d in ext_dims:
+                if d in info['dims']:
+                    extension_vars.append(v)
+
+            if not extended:
+                var_info[v]['data'] = var_info[v]['arr'][0]
+
+        # Perform ordering/concatenation for fragment components
+        for ev in extension_vars:
+
+            # Sort array components using index list
+            sorted_a = [[] for a in var_info[ev]['arr']]
+            for x, a in enumerate(var_info[ev]['arr']):
+
+                for nd in sorted_dim_sizes:
+                    if nd[1] in var_info[ev]['dims']:
+                        dominant_dim = nd[1]
+                        break
+
+                sorted_a[
+                    orders[ext_dims.index(dominant_dim)][x]
+                ] = a
+
+            # Assumes they have been sorted - not necessarily the case
+            array = np.concatenate(
+                sorted_a,
+                axis=var_info[ev]['dims'].index(dominant_dim)
+            )
+
+            if 'fragment_map_' in ev:
+
+                array = np.ma.masked_values(array, var_info[ev]['_FillValue'])
+                # Smooth paddings
+                for i, row in enumerate(array):
+                    premask = None
+                    gap = False
+                    for j, value in enumerate(row):
+                        if premask is None and not np.ma.is_masked(value):
+                            premask = value
+                        elif np.ma.is_masked(value):
+                            gap = True
+                        elif not np.ma.is_masked(value) and gap:
+                            if value == premask:
+                                array[i][j] = np.ma.masked
+
+            var_info[ev]['data'] = array
+
+        # Writing part done as a second step as before.
+        # Need to perform any transformations to get to that stage.
+
+        # Define the fragment space
+        self.fragment_space = [v['f_size'] for v in dim_info.values() if 'f_size' in v]
+
+        self.global_attrs = global_attrs
+        self.dim_info   = st_dim_info
+        self.var_info   = var_info
+
     def display_attrs(self):
         """
         Display the global attributes consolidated in the