Refactor pre/post processing implementation (PyDMD#554)

Refactor pre/post processing implementation

Author: fandreuz

pydmd/bopdmd.py

@@ -1369,7 +1369,7 @@ def _initialize_alpha(self):
         # rank-truncated SVD of Y.
         Y = (ux1 + ux2) / 2
         # Element-wise division by time differences. w/o large T
-        Z = (ux2 - ux1) / (self._time[1:] - self._time[:-1])  
+        Z = (ux2 - ux1) / (self._time[1:] - self._time[:-1])
         U, s, V = compute_svd(Y, self._svd_rank)
         S = np.diag(s)
 
@@ -1712,16 +1712,16 @@ def plot_eig_uq(
 
         if flip_axes:
             eigs = self.eigs.imag + 1j * self.eigs.real
-            plt.xlabel("$Im(\omega)$")
-            plt.ylabel("$Re(\omega)$")
+            plt.xlabel(r"$Im(\omega)$")
+            plt.ylabel(r"$Re(\omega)$")
 
             if eigs_true is not None:
                 eigs_true = eigs_true.imag + 1j * eigs_true.real
 
         else:
             eigs = self.eigs
-            plt.xlabel("$Re(\omega)$")
-            plt.ylabel("$Im(\omega)$")
+            plt.xlabel(r"$Re(\omega)$")
+            plt.ylabel(r"$Im(\omega)$")
 
         for e, std in zip(eigs, self.eigenvalues_std):
             # Plot 2 standard deviations.

pydmd/preprocessing/hankel.py

@@ -2,22 +2,32 @@
 Hankel pre-processing.
 """
 
-from functools import partial
-from typing import Dict, List, Tuple, Union
+from __future__ import annotations
+
+import sys
+from typing import Tuple, Union
 
 import numpy as np
 
 from pydmd.dmdbase import DMDBase
-from pydmd.preprocessing.pre_post_processing import PrePostProcessingDMD
+from pydmd.preprocessing.pre_post_processing import (
+    PrePostProcessing,
+    PrePostProcessingDMD,
+)
 from pydmd.utils import pseudo_hankel_matrix
 
-_reconstruction_method_type = Union[str, np.ndarray, List, Tuple]
+if sys.version_info >= (3, 12):
+    from typing import override
+else:
+    from typing_extensions import override
+
+_ReconstructionMethodType = Union[str, np.ndarray, list, tuple]
 
 
 def hankel_preprocessing(
     dmd: DMDBase,
     d: int,
-    reconstruction_method: _reconstruction_method_type = "first",
+    reconstruction_method: _ReconstructionMethodType = "first",
 ):
     """
     Hankel pre-processing.
@@ -26,19 +36,13 @@ def hankel_preprocessing(
     :param d: Hankel matrix rank
     :param reconstruction_method: Reconstruction method.
     """
-    return PrePostProcessingDMD(
-        dmd,
-        partial(_preprocessing, d=d),
-        partial(
-            _hankel_post_processing,
-            d=d,
-            reconstruction_method=reconstruction_method,
-        ),
-    )
+    if isinstance(reconstruction_method, (tuple, list)):
+        reconstruction_method = np.asarray(reconstruction_method)
 
-
-def _preprocessing(_: Dict, X: np.ndarray, d: int, **kwargs):
-    return (pseudo_hankel_matrix(X, d),) + tuple(kwargs.values())
+    pre_post_processing = _HankelPrePostProcessing(
+        d=d, reconstruction_method=reconstruction_method
+    )
+    return PrePostProcessingDMD(dmd, pre_post_processing)
 
 
 def _reconstructions(rec: np.ndarray, d: int):
@@ -92,24 +96,34 @@ def _first_reconstructions(reconstructions: np.ndarray, d: int) -> np.ndarray:
     return reconstructions[..., time_idxes, d_idxes, :].swapaxes(-1, -2)
 
 
-def _hankel_post_processing(
-    _: Dict,  # No state
-    X: np.ndarray,
-    d: int,
-    reconstruction_method: _reconstruction_method_type,
-) -> np.ndarray:
-    rec = _reconstructions(X, d=d)
-    rec = np.ma.array(rec, mask=np.isnan(rec))
-
-    if reconstruction_method == "first":
-        result = _first_reconstructions(rec, d=d)
-    elif reconstruction_method == "mean":
-        result = np.nanmean(rec, axis=1).T
-    elif isinstance(reconstruction_method, (np.ndarray, list, tuple)):
-        result = np.ma.average(rec, axis=1, weights=reconstruction_method).T
-    else:
-        raise ValueError(
-            f"The reconstruction method wasn't recognized: {reconstruction_method}"
-        )
-
-    return result.filled(fill_value=0)
+class _HankelPrePostProcessing(PrePostProcessing):
+    def __init__(
+        self, *, d: int, reconstruction_method: _ReconstructionMethodType
+    ):
+        self._d = d
+        self._reconstruction_method = reconstruction_method
+
+    @override
+    def pre_processing(self, X: np.ndarray) -> Tuple[None, np.ndarray]:
+        return None, pseudo_hankel_matrix(X, self._d)
+
+    @override
+    def post_processing(
+        self, pre_processing_output: None, Y: np.ndarray
+    ) -> np.ndarray:
+        rec = _reconstructions(Y, d=self._d)
+        rec = np.ma.array(rec, mask=np.isnan(rec))
+
+        if isinstance(self._reconstruction_method, str):
+            if self._reconstruction_method == "first":
+                result = _first_reconstructions(rec, d=self._d)
+            elif self._reconstruction_method == "mean":
+                result = np.nanmean(rec, axis=1).T
+        elif isinstance(self._reconstruction_method, np.ndarray):
+            result = np.ma.average(
+                rec, axis=1, weights=self._reconstruction_method
+            ).T
+        else:
+            raise ValueError(f"{self._reconstruction_method=} not recognized")
+
+        return result.filled(fill_value=0)

pydmd/preprocessing/pre_post_processing.py

@@ -2,27 +2,30 @@
 Pre/post-processing capability for DMD instances.
 """
 
-from typing import Callable, Dict
+from __future__ import annotations
 
-from pydmd.dmdbase import DMDBase
+from inspect import isroutine
+from typing import Any, Dict, Generic, Tuple, TypeVar
 
+import numpy as np
 
-def _shallow_preprocessing(_: Dict, *args, **kwargs):
-    return args + tuple(kwargs.values())
+from pydmd.dmdbase import DMDBase
 
+# Pre-processing output type
+S = TypeVar("S")
 
-def _shallow_postprocessing(_: Dict, *args):
-    # The first item of args is always the output of dmd.reconstructed_data
-    return args[0]
 
+class PrePostProcessing(Generic[S]):
+    def pre_processing(self, X: np.ndarray) -> Tuple[S, np.ndarray]:
+        return None, X
 
-def _tuplify(value):
-    if isinstance(value, tuple):
-        return value
-    return (value,)
+    def post_processing(
+        self, pre_processing_output: S, Y: np.ndarray
+    ) -> np.ndarray:
+        return Y
 
 
-class PrePostProcessingDMD:
+class PrePostProcessingDMD(Generic[S]):
     """
     Pre/post-processing decorator. This class is not thread-safe in case of
     stateful transformations.
@@ -40,20 +43,14 @@ class PrePostProcessingDMD:
     def __init__(
         self,
         dmd: DMDBase,
-        pre_processing: Callable = _shallow_preprocessing,
-        post_processing: Callable = _shallow_postprocessing,
+        pre_post_processing: PrePostProcessing[S] = PrePostProcessing(),
     ):
         if dmd is None:
             raise ValueError("DMD instance cannot be None")
-        if pre_processing is None:
-            pre_processing = _shallow_preprocessing
-        if post_processing is None:
-            post_processing = _shallow_postprocessing
+        self._pre_post_processing = pre_post_processing
 
-        self._pre_post_processed_dmd = dmd
-        self._pre_processing = pre_processing
-        self._post_processing = post_processing
-        self._state_holder = None
+        self._dmd = dmd
+        self._pre_processing_output: S | None = None
 
     def __getattribute__(self, name):
         try:
@@ -65,18 +62,14 @@ def __getattribute__(self, name):
             return self._pre_processing_fit
 
         if "reconstructed_data" == name:
-            output = self._post_processing(
-                self._state_holder,
-                self._pre_post_processed_dmd.reconstructed_data,
-            )
-            return output
+            return self._reconstructed_data_with_post_processing()
 
         # This check is needed to allow copy/deepcopy
-        if name != "_pre_post_processed_dmd":
-            sub_dmd = self._pre_post_processed_dmd
+        if name != "_dmd":
+            sub_dmd = self._dmd
             if isinstance(sub_dmd, PrePostProcessingDMD):
                 return PrePostProcessingDMD.__getattribute__(sub_dmd, name)
-            return object.__getattribute__(self._pre_post_processed_dmd, name)
+            return object.__getattribute__(self._dmd, name)
         return None
 
     @property
@@ -87,19 +80,61 @@ def pre_post_processed_dmd(self):
         :return: decorated DMD instance.
         :rtype: pydmd.DMDBase
         """
-        return self._pre_post_processed_dmd
+        return self._dmd
 
     @property
     def modes_activation_bitmask(self):
-        return self._pre_post_processed_dmd.modes_activation_bitmask
+        return self._dmd.modes_activation_bitmask
 
     @modes_activation_bitmask.setter
     def modes_activation_bitmask(self, value):
-        self._pre_post_processed_dmd.modes_activation_bitmask = value
+        self._dmd.modes_activation_bitmask = value
 
     def _pre_processing_fit(self, *args, **kwargs):
-        self._state_holder = dict()
-        pre_processing_output = _tuplify(
-            self._pre_processing(self._state_holder, *args, **kwargs)
+        X = PrePostProcessingDMD._extract_training_data(*args, **kwargs)
+        self._pre_processing_output, pre_processed_training_data = (
+            self._pre_post_processing.pre_processing(X)
+        )
+        new_args, new_kwargs = PrePostProcessingDMD._replace_training_data(
+            pre_processed_training_data, *args, **kwargs
         )
-        return self._pre_post_processed_dmd.fit(*pre_processing_output)
+        return self._dmd.fit(*new_args, **new_kwargs)
+
+    def _reconstructed_data_with_post_processing(self) -> np.ndarray:
+        data = self._dmd.reconstructed_data
+
+        if not isroutine(data):
+            return self._pre_post_processing.post_processing(
+                self._pre_processing_output,
+                data,
+            )
+
+        # e.g. DMDc
+        def output(*args, **kwargs) -> np.ndarray:
+            return self._pre_post_processing.post_processing(
+                self._pre_processing_output,
+                data(*args, **kwargs),
+            )
+
+        return output
+
+    @staticmethod
+    def _extract_training_data(*args, **kwargs):
+        if len(args) >= 1:
+            return args[0]
+        elif "X" in kwargs:
+            return kwargs["X"]
+        raise ValueError(
+            f"Could not extract training data from {args=}, {kwargs=}"
+        )
+
+    @staticmethod
+    def _replace_training_data(
+        new_training_data: Any, *args, **kwargs
+    ) -> [Tuple[Any, ...], Dict[str, Any]]:
+        if len(args) >= 1:
+            return (new_training_data,) + args[1:], kwargs
+        elif "X" in kwargs:
+            new_kwargs = dict(kwargs)
+            new_kwargs["X"] = new_training_data
+            return args, new_kwargs