Feature/refactor tests

.githooks/pre-commit.sh → .githooks/pre-commit

@@ -60,7 +60,7 @@ fi
 
 # Format all staged files, then exit with an error code if any have uncommitted
 # changes.
-echo 'Formatting staged Python files . . .'
+echo 'Formatting staged Python files...'
 
 ########## PIP VERSION #############
 yapf -i -r "${PYTHON_FILES[@]}"

.github/workflows/test-package.yml

@@ -1,9 +1,9 @@
 name: Test package
 on:
   push:
-    branches: [main, dev]
+    branches: [main]
   pull_request:
-    branches: [main, dev]
+    branches: [main]
 jobs:
   build:
     runs-on: ubuntu-latest
@@ -20,8 +20,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install pytest
-        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+        python -m pip install -r requirements.txt
     - name: Test with pytest
       run: |
-        pytest -k 'not slow'
+        pytest --exitfirst --remote

README.rst

@@ -3,6 +3,9 @@
 pykoop
 ======
 
+.. image:: https://github.com/decarsg/pykoop/actions/workflows/test-package.yml/badge.svg
+    :target: https://github.com/decarsg/pykoop/actions/workflows/test-package.yml
+    :alt: Test package
 .. image:: https://zenodo.org/badge/DOI/10.5281/zenodo.5576490.svg
     :target: https://doi.org/10.5281/zenodo.5576490
     :alt: DOI
@@ -37,7 +40,6 @@ mass-spring-damper data. Using ``pykoop``, this can be implemented as:
 .. code-block:: python
 
     import pykoop
-    import numpy as np
     from sklearn.preprocessing import MaxAbsScaler, StandardScaler
 
     # Get sample mass-spring-damper data
@@ -59,7 +61,7 @@ mass-spring-damper data. Using ``pykoop``, this can be implemented as:
     # Predict using the pipeline
     X_initial = X_msd[[0], 1:3]
     U = X_msd[:, [3]]
-    X_pred = kp.predict_state(X_initial, U, episode_feature=False)
+    X_pred = kp.predict_trajectory(X_initial, U, episode_feature=False)
 
     # Score using the pipeline
     score = kp.score(X_msd)
@@ -116,6 +118,7 @@ Additional LMI solvers can be installed using
 .. code-block:: sh
 
     $ pip install mosek
+    $ pip install cvxopt
     $ pip install smcp
 
 Mosek is recommended, but is nonfree and requires a license.
@@ -129,11 +132,12 @@ The library can be tested using
 
 Note that ``pytest`` must be run from the repository's root directory.
 
-To skip slow unit tests, including all doctests and examples, run
+To skip unit tests that require a MOSEK license, including all doctests and
+examples, run
 
 .. code-block:: sh
 
-    $ pytest ./tests -k "not slow"
+    $ pytest ./tests -k "not mosek"
 
 The documentation can be compiled using
 

conftest.py

@@ -1,10 +1,15 @@
 """Pytest fixtures for doctests."""
+from typing import Any, Dict
+
 import numpy as np
 import pytest
 import sklearn
+from scipy import linalg, integrate
 
 import pykoop
 
+# Add remote option for MOSEK
+
 
 def pytest_addoption(parser) -> None:
     """Add ``--remote`` command line flag to run regressions remotely."""
@@ -30,6 +35,98 @@ def remote_url() -> str:
     return 'http://solve.mosek.com:30080'
 
 
+# Add mass-spring-damper reference trajectory for all tests
+
+
+@pytest.fixture
+def mass_spring_damper_sine_input() -> Dict[str, Any]:
+    """Compute mass-spring-damper response with sine input."""
+    # Set up problem
+    t_range = (0, 10)
+    t_step = 0.1
+    msd = pykoop.dynamic_models.MassSpringDamper(0.5, 0.7, 0.6)
+
+    def u(t):
+        return 0.1 * np.sin(t)
+
+    # Initial condition
+    x0 = np.array([0, 0])
+    # Solve ODE for training data
+    t, x = msd.simulate(
+        t_range,
+        t_step,
+        x0,
+        u,
+        rtol=1e-8,
+        atol=1e-8,
+    )
+    # Compute discrete-time A and B matrices
+    Ad = linalg.expm(msd.A * t_step)
+
+    def integrand(s):
+        return linalg.expm(msd.A * (t_step - s)).ravel()
+
+    Bd = integrate.quad_vec(integrand, 0, t_step)[0].reshape((2, 2)) @ msd.B
+    U_valid = np.hstack((Ad, Bd))
+    # Split the data
+    y_train, y_valid = np.split(x, 2, axis=0)
+    u_train, u_valid = np.split(np.reshape(u(t), (-1, 1)), 2, axis=0)
+    X_train = np.hstack((y_train[:-1, :], u_train[:-1, :]))
+    Xp_train = y_train[1:, :]
+    X_valid = np.hstack((y_valid[:-1, :], u_valid[:-1, :]))
+    Xp_valid = y_valid[1:, :]
+    return {
+        'X_train': X_train,
+        'Xp_train': Xp_train,
+        'X_valid': X_valid,
+        'Xp_valid': Xp_valid,
+        'n_inputs': 1,
+        'episode_feature': False,
+        'U_valid': U_valid,
+        't_step': t_step,
+    }
+
+
+@pytest.fixture
+def mass_spring_damper_no_input() -> Dict[str, Any]:
+    """Compute mass-spring-damper response with no input."""
+    # Set up problem
+    t_range = (0, 10)
+    t_step = 0.1
+    msd = pykoop.dynamic_models.MassSpringDamper(0.5, 0.7, 0.6)
+    # Initial condition
+    x0 = np.array([1, 0])
+    # Solve ODE for training data
+    t, x = msd.simulate(
+        t_range,
+        t_step,
+        x0,
+        lambda t: np.zeros((1, )),
+        rtol=1e-8,
+        atol=1e-8,
+    )
+    U_valid = linalg.expm(msd.A * t_step)
+    # Split the data
+    y_train, y_valid = np.split(x, 2, axis=0)
+    X_train = y_train[:-1, :]
+    Xp_train = y_train[1:, :]
+    X_valid = y_valid[:-1, :]
+    Xp_valid = y_valid[1:, :]
+    return {
+        'X_train': X_train,
+        'Xp_train': Xp_train,
+        'X_valid': X_valid,
+        'Xp_valid': Xp_valid,
+        'n_inputs': 0,
+        'episode_feature': False,
+        'U_valid': U_valid,
+        't_step': t_step,
+    }
+
+
+# Add common packages and data to doctest namespace
+
+
 @pytest.fixture(autouse=True)
 def add_np(doctest_namespace):
     """Add numpy to namespace."""

constraints.txt

@@ -7,7 +7,10 @@ optht>=0.2.0
 Deprecated>=1.2.13
 
 pytest
+pytest-regressions
 matplotlib
 
 sphinx
 sphinx-rtd-theme
+
+mosek>=9.2.49

doc/conf.py

@@ -49,10 +49,10 @@
 autosummary_generate = True
 
 intersphinx_mapping = {
-    'python': ('https://docs.python.org/3', None),
-    'numpy': ('https://numpy.org/doc/stable', None),
-    'scipy': ('https://docs.scipy.org/doc/scipy/reference', None),
-    'matplotlib': ('https://matplotlib.org/', None),
+    'python': ('https://docs.python.org/3/', None),
+    'numpy': ('https://numpy.org/doc/stable/', None),
+    'scipy': ('https://docs.scipy.org/doc/scipy/', None),
+    'matplotlib': ('https://matplotlib.org/stable/', None),
     'scikit-learn': ('https://scikit-learn.org/stable/', None),
 }
 

examples/example_eigenvalue_comparison.py

@@ -1,24 +1,26 @@
 """Example comparing eigenvalues of different versions of EDMD."""
 
 import numpy as np
-import pykoop
-import pykoop.lmi_regressors
 from matplotlib import pyplot as plt
 from scipy import integrate, linalg
 
+import pykoop
+import pykoop.lmi_regressors
+
 plt.rc('lines', linewidth=2)
 plt.rc('axes', grid=True)
 plt.rc('grid', linestyle='--')
 
 
-def main() -> None:
+def example_eigenvalue_comparison() -> None:
     """Compare eigenvalues of different versions of EDMD."""
     # Get example data
     X = pykoop.example_data_msd()
+    # Set solver (you can switch to ``'mosek'`` if you have a license).
+    solver = 'cvxopt'
 
     # Regressor with no constraint
-    reg_no_const = pykoop.KoopmanPipeline(
-        regressor=pykoop.lmi_regressors.LmiEdmd())
+    reg_no_const = pykoop.KoopmanPipeline(regressor=pykoop.Edmd())
     reg_no_const.fit(X, n_inputs=1, episode_feature=True)
     U_no_const = reg_no_const.regressor_.coef_.T
 
@@ -30,20 +32,27 @@ def main() -> None:
     ])  # yapf: disable
     reg_diss_const = pykoop.KoopmanPipeline(
         regressor=pykoop.lmi_regressors.LmiEdmdDissipativityConstr(
-            supply_rate=Xi))
+            supply_rate=Xi,
+            solver_params={'solver': solver},
+        ))
     reg_diss_const.fit(X, n_inputs=1, episode_feature=True)
     U_diss_const = reg_diss_const.regressor_.coef_.T
 
     # Regressor with H-infinity regularization
     reg_hinf_reg = pykoop.KoopmanPipeline(
-        regressor=pykoop.lmi_regressors.LmiEdmdHinfReg(alpha=5))
+        regressor=pykoop.lmi_regressors.LmiEdmdHinfReg(
+            alpha=5,
+            solver_params={'solver': solver},
+        ))
     reg_hinf_reg.fit(X, n_inputs=1, episode_feature=True)
     U_hinf_reg = reg_hinf_reg.regressor_.coef_.T
 
     # Regressor with spectral radius constraint
     reg_sr_const = pykoop.KoopmanPipeline(
         regressor=pykoop.lmi_regressors.LmiEdmdSpectralRadiusConstr(
-            spectral_radius=0.8))
+            spectral_radius=0.8,
+            solver_params={'solver': solver},
+        ))
     reg_sr_const.fit(X, n_inputs=1, episode_feature=True)
     U_sr_const = reg_sr_const.regressor_.coef_.T
 
@@ -72,5 +81,5 @@ def plt_eig(A: np.ndarray,
 
 
 if __name__ == '__main__':
-    main()
+    example_eigenvalue_comparison()
     plt.show()

examples/example_pipeline_cv.py

@@ -8,7 +8,7 @@
 import pykoop.dynamic_models
 
 
-def main() -> None:
+def example_pipeline_cv() -> None:
     """Cross-validate regressor parameters."""
     # Get sample mass-spring-damper data
     X_msd = pykoop.example_data_msd()
@@ -96,4 +96,4 @@ def main() -> None:
 
 
 if __name__ == '__main__':
-    main()
+    example_pipeline_cv()

examples/example_pipeline_simple.py

@@ -1,12 +1,11 @@
 """Example of how to use the Koopman pipeline."""
 
-import numpy as np
 from sklearn.preprocessing import MaxAbsScaler, StandardScaler
 
 import pykoop
 
 
-def main() -> None:
+def example_pipeline_simple() -> None:
     """Demonstrate how to use the Koopman pipeline."""
     # Get sample mass-spring-damper data
     X_msd = pykoop.example_data_msd()
@@ -27,11 +26,11 @@ def main() -> None:
     # Predict using the pipeline
     X_initial = X_msd[[0], 1:3]
     U = X_msd[:, [3]]
-    X_pred = kp.predict_state(X_initial, U, episode_feature=False)
+    X_pred = kp.predict_trajectory(X_initial, U, episode_feature=False)
 
     # Score using the pipeline
     score = kp.score(X_msd)
 
 
 if __name__ == '__main__':
-    main()
+    example_pipeline_simple()

examples/example_pipeline_vdp.py

@@ -7,7 +7,7 @@
 import pykoop.dynamic_models
 
 
-def main() -> None:
+def example_pipeline_vdp() -> None:
     """Demonstrate how to use the Koopman pipeline."""
     # Get sample data
     X_vdp = pykoop.example_data_vdp()
@@ -38,15 +38,15 @@ def main() -> None:
     u = X_valid[:, 3:]
 
     # Predict with re-lifting between timesteps (default)
-    X_pred_local = kp.predict_state(
+    X_pred_local = kp.predict_trajectory(
         x0,
         u,
         relift_state=True,
         episode_feature=False,
     )
 
     # Predict without re-lifting between timesteps
-    X_pred_global = kp.predict_state(
+    X_pred_global = kp.predict_trajectory(
         x0,
         u,
         relift_state=False,
@@ -85,7 +85,7 @@ def main() -> None:
     Psi_valid = kp.lift(X_valid[:, 1:], episode_feature=False)
 
     # Predict lifted state with re-lifting between timesteps (default)
-    Psi_pred_local = kp.predict_state(
+    Psi_pred_local = kp.predict_trajectory(
         x0,
         u,
         relift_state=True,
@@ -95,7 +95,7 @@ def main() -> None:
     )
 
     # Predict lifted state without re-lifting between timesteps
-    Psi_pred_global = kp.predict_state(
+    Psi_pred_global = kp.predict_trajectory(
         x0,
         u,
         relift_state=False,
@@ -141,8 +141,7 @@ def main() -> None:
     ax[0, 0].set_ylabel('$u$')
     ax[0, 0].set_title('Exogenous input')
 
-    plt.show()
-
 
 if __name__ == '__main__':
-    main()
+    example_pipeline_vdp()
+    plt.show()