Merge pull request #24 from wpreimes/postprocess

Add postprocess option to ts2img

src/repurpose/img2ts.py

@@ -439,6 +439,8 @@ def calc(self):
 
  for img_stack_dict, timestamps in self.img_bulk():
  # =================================================================
+ logging.info(f"Finished reading bulk with {len(timestamps)} images")
+
  start_time = datetime.now()
 
  # temporally drop grids, due to issue when pickling them...
@@ -451,6 +453,7 @@ def calc(self):
 
  keys = list(img_stack_dict.keys())
  for key in keys:
+ #print(key)
  # rename variable in output dataset
  if self.variable_rename is None:
  var_new_name = str(key)
@@ -473,6 +476,7 @@ def calc(self):
  ITER_KWARGS = {'cell': [], 'celldata': []}
 
  for cell in np.unique(target_grid.activearrcell):
+ #print(cell)
  cell_idx = np.where(cells == cell)[0]
 
  if len(cell_idx) == 0:
@@ -486,7 +490,6 @@ def calc(self):
  np.atleast_2d(img_stack_dict[k])[:, cell_idx], 0, 1)
  img_stack_dict[k] = np.delete(img_stack_dict[k], cell_idx,
  axis=1)
-
  cells = np.delete(cells, cell_idx)
 
  ITER_KWARGS['celldata'].append(celldata)

src/repurpose/ts2img.py

@@ -42,8 +42,15 @@ def _convert(converter: 'Ts2Img',
  if ts is None:
  continue
  if preprocess_func is not None:
- preprocess_kwargs = preprocess_kwargs or {}
- ts = preprocess_func(ts, **preprocess_kwargs)
+ preprocess_func = np.atleast_1d(preprocess_func)
+ if preprocess_kwargs is None:
+ preprocess_kwargs = [{}] * len(preprocess_func)
+ preprocess_kwargs = np.atleast_1d(preprocess_kwargs)
+ if len(preprocess_func) != len(preprocess_kwargs):
+ raise ValueError("Length of preprocess_func and "
+ "preprocess_kwargs is different")
+ for func, kwargs in zip(preprocess_func, preprocess_kwargs):
+ ts = func(ts, **kwargs)
  if np.any(np.isin(ts.columns, Ts2Img._protected_vars)):
  raise ValueError(
  f"Time series contains protected variables. "
@@ -263,7 +270,7 @@ def calc(self, path_out, format_out='slice', preprocess=None,
  - stack: write all time steps into one file. In this case if there
  is a {datetime} placeholder in the fn_template, then the time
  range is inserted.
- preprocess: callable, optional (default: None)
+ preprocess: callable or list[Callable], optional (default: None)
  Function that is applied to each time series before converting it.
  The first argument is the data frame that the reader returns.
  Additional keyword arguments can be passed via `preprocess_kwargs`.
@@ -279,25 +286,25 @@ def preprocess_add(df: pd.DataFrame, **preprocess_kwargs) \
  df['var3'] = df['var1'] + df['var2']
  return df
  ```
- preprocess_kwargs: dict, optional (default: None)
+ preprocess_kwargs: dict or list[dict], optional (default: None)
  Keyword arguments for the preprocess function. If None are given,
  then the preprocessing function is is called with only the input
  data frame and no additional arguments (see example above).
- postprocess: Callable, optional (default: None)
- Function that is applied to the image stack after loading the data
- and before writing it to disk. The function must take xarray
+ postprocess: Callable or list[Callable], optional (default: None)
+ Function(s) applied to the image stack after loading the data
+ and before writing it to disk. The function must take an xarray
  Dataset as the first argument and return an xarray Dataset of the
- same form as the input data.
+ same form.
  A simple example for a preprocessing function to add a new variable
  from the sum of two existing variables:
  ```
- def preprocess_add(stack: xr.Dataset, **postprocess_kwargs) \
+ def postprocess_add(stack: xr.Dataset, **postprocess_kwargs) \
  -> xr.Dataset
  stack = stack.assign(var3=lambda x: x['var0'] + x['var2'])
  return stack
  ```
- postprocess_kwargs: dict, optional (default: None)
- Keyword arguments for the postprocess function. If None are given,
+ postprocess_kwargs: dict or list[dict], optional (default: None)
+ Keyword arguments for the postprocess function(s). If None are given,
  then the postprocess function is called with only the input
  image stack and no additional arguments (see example above).
  fn_template: str, optional (default: "{datetime}.nc")
@@ -381,10 +388,6 @@ def preprocess_add(stack: xr.Dataset, **postprocess_kwargs) \
 
  self.stack = self.stack.drop_isel(time=idx_empty)
 
- if postprocess is not None:
- postprocess_kwargs = postprocess_kwargs or {}
- self.stack = postprocess(self.stack, **postprocess_kwargs)
-
  if var_fillvalues is not None:
  for var, fillvalue in var_fillvalues.items():
  self.stack[var].values = np.nan_to_num(
@@ -410,6 +413,19 @@ def preprocess_add(stack: xr.Dataset, **postprocess_kwargs) \
  for var in var_attrs:
  self.stack[var].attrs.update(var_attrs[var])
 
+ if postprocess is not None:
+ # this is done as late as possible
+ postprocess = np.atleast_1d(postprocess)
+ if postprocess_kwargs is None:
+ postprocess_kwargs = [{}] * len(postprocess)
+ postprocess_kwargs = np.atleast_1d(postprocess_kwargs)
+ if len(postprocess) != len(postprocess_kwargs):
+ raise ValueError("postprocess and postprocess_kwargs "
+ "have different lengths")
+
+ for func, kwargs in zip(postprocess, postprocess_kwargs):
+ self.stack = func(self.stack, **kwargs)
+
  if self.stack['time'].size == 0:
  warnings.warn("No images in stack to write to disk.")
  self.stack = None