Add method for DMD modes selection

pydmd/dmdbase.py

@@ -376,6 +376,41 @@ def _compute_amplitudes(self):
 
         return a
 
+
+    def select_modes(self, func):
+        """
+        Select the DMD modes by using the given `func`.
+        `func` has to be a callable function which takes as input the DMD
+        object itself and return a numpy.ndarray of boolean where `False`
+        indicates that the corresponding mode will be discarded.
+
+        :param callable func: the function to select the modes
+
+        Example:
+
+        >>> def stable_modes(dmd_object):
+        >>>    toll = 1e-3
+        >>>    return np.abs(np.abs(dmd.eigs) - 1) < toll
+        >>> dmd = DMD(svd_rank=10)
+        >>> dmd.fit(sample_data)
+        >>> dmd.select_modes(stable_modes)
+        """
+        selected_indeces = func(self)
+
+        self.operator._eigenvalues = self.operator._eigenvalues[selected_indeces]
+        self.operator._Lambda = self.operator._Lambda[selected_indeces]
+
+        self.operator._eigenvectors = self.operator._eigenvectors[:, selected_indeces]
+        self.operator._modes = self.operator._modes[:, selected_indeces]
+
+        self.operator._Atilde = np.linalg.multi_dot([
+            self.operator._eigenvectors,
+            np.diag(self.operator._eigenvalues),
+            np.linalg.pinv(self.operator._eigenvectors)])
+
+        self._b = self._compute_amplitudes()
+
+
     def plot_eigs(self,
                   show_axes=True,
                   show_unit_circle=True,

tests/test_dmdbase.py

@@ -1,5 +1,6 @@
 from unittest import TestCase
 from pydmd.dmdbase import DMDBase
+from pydmd import DMD
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -103,3 +104,13 @@ def test_sorted_eigs_default(self):
     def test_sorted_eigs_param(self):
         dmd = DMDBase(sorted_eigs='real')
         assert dmd.operator._sorted_eigs == 'real'
+
+    def test_select_modes(self):
+        def stable_modes(dmd_object):
+            toll = 1e-3
+            return np.abs(np.abs(dmd_object.eigs) - 1) < toll
+        dmd = DMD(svd_rank=10)
+        dmd.fit(sample_data)
+        exp = dmd.reconstructed_data
+        dmd.select_modes(stable_modes)
+        np.testing.assert_array_almost_equal(exp, dmd.reconstructed_data)