1
1
"""Pytest fixtures for doctests."""
2
+ from typing import Any , Dict
3
+
2
4
import numpy as np
3
5
import pytest
4
6
import sklearn
7
+ from scipy import linalg , integrate
5
8
6
9
import pykoop
7
10
11
+ # Add remote option for MOSEK
12
+
8
13
9
14
def pytest_addoption (parser ) -> None :
10
15
"""Add ``--remote`` command line flag to run regressions remotely."""
@@ -30,6 +35,98 @@ def remote_url() -> str:
30
35
return 'http://solve.mosek.com:30080'
31
36
32
37
38
+ # Add mass-spring-damper reference trajectory for all tests
39
+
40
+
41
+ @pytest .fixture
42
+ def mass_spring_damper_sine_input () -> Dict [str , Any ]:
43
+ """Compute mass-spring-damper response with sine input."""
44
+ # Set up problem
45
+ t_range = (0 , 10 )
46
+ t_step = 0.1
47
+ msd = pykoop .dynamic_models .MassSpringDamper (0.5 , 0.7 , 0.6 )
48
+
49
+ def u (t ):
50
+ return 0.1 * np .sin (t )
51
+
52
+ # Initial condition
53
+ x0 = np .array ([0 , 0 ])
54
+ # Solve ODE for training data
55
+ t , x = msd .simulate (
56
+ t_range ,
57
+ t_step ,
58
+ x0 ,
59
+ u ,
60
+ rtol = 1e-8 ,
61
+ atol = 1e-8 ,
62
+ )
63
+ # Compute discrete-time A and B matrices
64
+ Ad = linalg .expm (msd .A * t_step )
65
+
66
+ def integrand (s ):
67
+ return linalg .expm (msd .A * (t_step - s )).ravel ()
68
+
69
+ Bd = integrate .quad_vec (integrand , 0 , t_step )[0 ].reshape ((2 , 2 )) @ msd .B
70
+ U_valid = np .hstack ((Ad , Bd ))
71
+ # Split the data
72
+ y_train , y_valid = np .split (x , 2 , axis = 0 )
73
+ u_train , u_valid = np .split (np .reshape (u (t ), (- 1 , 1 )), 2 , axis = 0 )
74
+ X_train = np .hstack ((y_train [:- 1 , :], u_train [:- 1 , :]))
75
+ Xp_train = y_train [1 :, :]
76
+ X_valid = np .hstack ((y_valid [:- 1 , :], u_valid [:- 1 , :]))
77
+ Xp_valid = y_valid [1 :, :]
78
+ return {
79
+ 'X_train' : X_train ,
80
+ 'Xp_train' : Xp_train ,
81
+ 'X_valid' : X_valid ,
82
+ 'Xp_valid' : Xp_valid ,
83
+ 'n_inputs' : 1 ,
84
+ 'episode_feature' : False ,
85
+ 'U_valid' : U_valid ,
86
+ 't_step' : t_step ,
87
+ }
88
+
89
+
90
+ @pytest .fixture
91
+ def mass_spring_damper_no_input () -> Dict [str , Any ]:
92
+ """Compute mass-spring-damper response with no input."""
93
+ # Set up problem
94
+ t_range = (0 , 10 )
95
+ t_step = 0.1
96
+ msd = pykoop .dynamic_models .MassSpringDamper (0.5 , 0.7 , 0.6 )
97
+ # Initial condition
98
+ x0 = np .array ([1 , 0 ])
99
+ # Solve ODE for training data
100
+ t , x = msd .simulate (
101
+ t_range ,
102
+ t_step ,
103
+ x0 ,
104
+ lambda t : np .zeros ((1 , )),
105
+ rtol = 1e-8 ,
106
+ atol = 1e-8 ,
107
+ )
108
+ U_valid = linalg .expm (msd .A * t_step )
109
+ # Split the data
110
+ y_train , y_valid = np .split (x , 2 , axis = 0 )
111
+ X_train = y_train [:- 1 , :]
112
+ Xp_train = y_train [1 :, :]
113
+ X_valid = y_valid [:- 1 , :]
114
+ Xp_valid = y_valid [1 :, :]
115
+ return {
116
+ 'X_train' : X_train ,
117
+ 'Xp_train' : Xp_train ,
118
+ 'X_valid' : X_valid ,
119
+ 'Xp_valid' : Xp_valid ,
120
+ 'n_inputs' : 0 ,
121
+ 'episode_feature' : False ,
122
+ 'U_valid' : U_valid ,
123
+ 't_step' : t_step ,
124
+ }
125
+
126
+
127
+ # Add common packages and data to doctest namespace
128
+
129
+
33
130
@pytest .fixture (autouse = True )
34
131
def add_np (doctest_namespace ):
35
132
"""Add numpy to namespace."""
0 commit comments