edmkit 0.0.12__tar.gz → 0.0.14__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (34) hide show
  1. {edmkit-0.0.12 → edmkit-0.0.14}/PKG-INFO +3 -3
  2. {edmkit-0.0.12 → edmkit-0.0.14}/pyproject.toml +4 -3
  3. {edmkit-0.0.12 → edmkit-0.0.14}/pyproject.toml.orig +4 -3
  4. edmkit-0.0.14/src/edmkit/ccm.py +274 -0
  5. edmkit-0.0.14/src/edmkit/embedding.py +85 -0
  6. edmkit-0.0.14/src/edmkit/generate/__init__.py +10 -0
  7. edmkit-0.0.14/src/edmkit/generate/ctrnn.py +25 -0
  8. {edmkit-0.0.12 → edmkit-0.0.14}/src/edmkit/generate/double_pendulum.py +6 -30
  9. edmkit-0.0.14/src/edmkit/generate/duffing.py +31 -0
  10. edmkit-0.0.14/src/edmkit/generate/lorenz.py +27 -0
  11. edmkit-0.0.14/src/edmkit/generate/lorenz96.py +25 -0
  12. edmkit-0.0.14/src/edmkit/generate/rk45.py +85 -0
  13. edmkit-0.0.14/src/edmkit/generate/rossler.py +27 -0
  14. edmkit-0.0.14/src/edmkit/metrics/__init__.py +35 -0
  15. edmkit-0.0.14/src/edmkit/metrics/mae.py +77 -0
  16. edmkit-0.0.14/src/edmkit/metrics/pearson_correlation.py +96 -0
  17. edmkit-0.0.14/src/edmkit/metrics/rmse.py +79 -0
  18. {edmkit-0.0.12 → edmkit-0.0.14}/src/edmkit/simplex_projection/__init__.py +1 -2
  19. {edmkit-0.0.12 → edmkit-0.0.14}/src/edmkit/simplex_projection/simplex_projection.py +26 -15
  20. {edmkit-0.0.12 → edmkit-0.0.14}/src/edmkit/simplex_projection/soft.py +13 -7
  21. {edmkit-0.0.12 → edmkit-0.0.14}/src/edmkit/smap.py +21 -15
  22. {edmkit-0.0.12 → edmkit-0.0.14}/src/edmkit/splits.py +3 -2
  23. edmkit-0.0.14/src/edmkit/theiler_window.py +34 -0
  24. {edmkit-0.0.12 → edmkit-0.0.14}/src/edmkit/util.py +1 -1
  25. edmkit-0.0.12/src/edmkit/ccm.py +0 -508
  26. edmkit-0.0.12/src/edmkit/embedding.py +0 -246
  27. edmkit-0.0.12/src/edmkit/generate/__init__.py +0 -6
  28. edmkit-0.0.12/src/edmkit/generate/lorenz.py +0 -49
  29. edmkit-0.0.12/src/edmkit/metrics.py +0 -132
  30. edmkit-0.0.12/src/edmkit/simplex_projection/loo.py +0 -91
  31. {edmkit-0.0.12 → edmkit-0.0.14}/README.md +0 -0
  32. {edmkit-0.0.12 → edmkit-0.0.14}/src/edmkit/generate/mackey_glass.py +0 -0
  33. {edmkit-0.0.12 → edmkit-0.0.14}/src/edmkit/simplex_projection/knn.py +0 -0
  34. {edmkit-0.0.12 → edmkit-0.0.14}/src/edmkit/types.py +0 -0
@@ -1,14 +1,14 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: edmkit
3
- Version: 0.0.12
3
+ Version: 0.0.14
4
4
  Summary: Simple EDM (Empirical Dynamic Modeling) library
5
5
  Author: FUJISHIGE TEMMA
6
6
  Author-email: FUJISHIGE TEMMA <tenma.x0@gmail.com>
7
7
  Requires-Dist: kdtree-rs>=0.3.1
8
8
  Requires-Dist: numpy>=2.5.1
9
9
  Requires-Dist: scipy>=1.17.1
10
- Requires-Dist: tinygrad>=0.13.0
11
- Requires-Dist: tinysolve>=0.2.0
10
+ Requires-Dist: tinygrad>=0.14.0
11
+ Requires-Dist: tinysolve>=0.2.1
12
12
  Requires-Python: >=3.13
13
13
  Description-Content-Type: text/markdown
14
14
 
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "edmkit"
3
- version = "0.0.12"
3
+ version = "0.0.14"
4
4
  description = "Simple EDM (Empirical Dynamic Modeling) library"
5
5
  readme = "README.md"
6
6
  requires-python = ">= 3.13"
@@ -8,8 +8,8 @@ dependencies = [
8
8
  "kdtree-rs>=0.3.1",
9
9
  "numpy>=2.5.1",
10
10
  "scipy>=1.17.1",
11
- "tinygrad>=0.13.0",
12
- "tinysolve>=0.2.0",
11
+ "tinygrad>=0.14.0",
12
+ "tinysolve>=0.2.1",
13
13
  ]
14
14
 
15
15
  [[project.authors]]
@@ -20,6 +20,7 @@ email = "tenma.x0@gmail.com"
20
20
  dev = [
21
21
  "griffe2md>=1.5.0",
22
22
  "hypothesis>=6.151.9",
23
+ "matplotlib>=3.11.2",
23
24
  "pytest>=9.1.1",
24
25
  "ruff>=0.15.20",
25
26
  "ty>=0.0.56",
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "edmkit"
3
- version = "0.0.12"
3
+ version = "0.0.14"
4
4
  description = "Simple EDM (Empirical Dynamic Modeling) library"
5
5
  authors = [{ name = "FUJISHIGE TEMMA", email = "tenma.x0@gmail.com" }]
6
6
  readme = "README.md"
@@ -9,14 +9,15 @@ dependencies = [
9
9
  "kdtree-rs>=0.3.1",
10
10
  "numpy>=2.5.1",
11
11
  "scipy>=1.17.1",
12
- "tinygrad>=0.13.0",
13
- "tinysolve>=0.2.0",
12
+ "tinygrad>=0.14.0",
13
+ "tinysolve>=0.2.1",
14
14
  ]
15
15
 
16
16
  [dependency-groups]
17
17
  dev = [
18
18
  "griffe2md>=1.5.0",
19
19
  "hypothesis>=6.151.9",
20
+ "matplotlib>=3.11.2",
20
21
  "pytest>=9.1.1",
21
22
  "ruff>=0.15.20",
22
23
  "ty>=0.0.56",
@@ -0,0 +1,274 @@
1
+ from collections.abc import Callable
2
+
3
+ import numpy as np
4
+
5
+ from edmkit.metrics import MetricFunc
6
+ from edmkit.types import PredictFunc
7
+
8
+ __all__ = [
9
+ "AggregateFunc",
10
+ "SampleFunc",
11
+ "bootstrap",
12
+ "ccm",
13
+ "make_sample_func",
14
+ ]
15
+
16
+ type SampleFunc = Callable[[np.ndarray, int], np.ndarray]
17
+ """SampleFunc is a function that takes (pool, size) and returns a sampled array."""
18
+ type AggregateFunc = Callable[[np.ndarray], float]
19
+ """AggregateFunc is a function that takes an array of values and returns a single value."""
20
+
21
+
22
+ def make_sample_func(seed: int | None = 42) -> SampleFunc:
23
+ """Create a sample function with its own independent RNG."""
24
+ rng = np.random.default_rng(seed)
25
+
26
+ def sample_func(pool: np.ndarray, size: int) -> np.ndarray:
27
+ return rng.choice(pool, size=size, replace=True)
28
+
29
+ return sample_func
30
+
31
+
32
+ def bootstrap(
33
+ X: np.ndarray,
34
+ Y: np.ndarray,
35
+ lib_sizes: np.ndarray,
36
+ predict_func: PredictFunc,
37
+ metric_func: MetricFunc,
38
+ n_samples: int = 20,
39
+ *,
40
+ library_pool: np.ndarray,
41
+ prediction_pool: np.ndarray,
42
+ sample_func: SampleFunc | None = None,
43
+ batch_size: int | None = 20,
44
+ ) -> np.ndarray:
45
+ """
46
+ Perform Convergent Cross Mapping and return per-sample scores.
47
+
48
+ Same as :func:`ccm` but returns the raw per-sample scores instead of
49
+ aggregating them.
50
+
51
+ Parameters
52
+ ----------
53
+ X : np.ndarray
54
+ Library delay vectors of shape ``(T,)`` or ``(T, E)`` (potential response).
55
+ Y : np.ndarray
56
+ Target values of shape ``(T,)`` or ``(T, E')`` (potential driver).
57
+ A 2D `Y` cross-maps to a delay vector; the metric scores all columns jointly.
58
+ lib_sizes : np.ndarray
59
+ Array of library sizes to test convergence.
60
+ predict_func : PredictFunc
61
+ Prediction function with signature (X, Y, Q) -> predictions.
62
+ metric_func : MetricFunc
63
+ Metric function with signature (predictions, observations) -> metric value.
64
+ n_samples : int, default 20
65
+ Number of random samples per library size for bootstrapping.
66
+ library_pool : np.ndarray
67
+ 1-D array of integer indices from which library members are sampled.
68
+ prediction_pool : np.ndarray
69
+ 1-D array of integer indices that are predicted.
70
+ sample_func : SampleFunc or None, default None
71
+ Function responsible for drawing a library sample of a given size.
72
+ When None, a fresh RNG-backed sampler is created per call.
73
+ batch_size : int or None, default 20
74
+ If specified, predictions are made in batches to limit memory usage.
75
+
76
+ Returns
77
+ -------
78
+ samples : np.ndarray
79
+ Per-sample skill scores of shape ``(n_samples, len(lib_sizes))``.
80
+ """
81
+ if sample_func is None:
82
+ sample_func = make_sample_func()
83
+
84
+ if X.shape[0] != Y.shape[0]:
85
+ raise ValueError(f"X and Y must have same length, got {X.shape[0]} and {Y.shape[0]}")
86
+ if predict_func is None or not callable(predict_func):
87
+ raise ValueError(f"predict_func must be callable, got {type(predict_func)}")
88
+ if n_samples <= 0:
89
+ raise ValueError(f"n_samples must be positive, got {n_samples}")
90
+
91
+ if batch_size is None:
92
+ batch_size = n_samples
93
+ else:
94
+ batch_size = min(batch_size, n_samples)
95
+
96
+ if X.ndim == 1:
97
+ X = X[:, None]
98
+ if Y.ndim == 1:
99
+ Y = Y[:, None]
100
+
101
+ prediction_indices = np.tile(prediction_pool, (batch_size, 1))
102
+ Q = X[prediction_indices]
103
+ actual = Y[prediction_indices]
104
+
105
+ samples = np.zeros((n_samples, len(lib_sizes)))
106
+
107
+ for i, lib_size in enumerate(lib_sizes):
108
+ remaining = n_samples
109
+ while remaining > 0:
110
+ batch = min(batch_size, remaining)
111
+
112
+ library_indices = np.vstack([sample_func(library_pool, lib_size) for _ in range(batch)])
113
+
114
+ lib_X = X[library_indices]
115
+ lib_Y = Y[library_indices]
116
+
117
+ predictions = predict_func(lib_X, lib_Y, Q[:batch])
118
+
119
+ offset = n_samples - remaining
120
+ samples[offset : offset + batch, i] = metric_func(predictions, actual[:batch])
121
+ remaining -= batch
122
+
123
+ return samples
124
+
125
+
126
+ def ccm(
127
+ X: np.ndarray,
128
+ Y: np.ndarray,
129
+ lib_sizes: np.ndarray,
130
+ predict_func: PredictFunc,
131
+ metric_func: MetricFunc,
132
+ n_samples: int = 20,
133
+ *,
134
+ library_pool: np.ndarray,
135
+ prediction_pool: np.ndarray,
136
+ sample_func: SampleFunc | None = None,
137
+ aggregate_func: AggregateFunc = np.mean,
138
+ batch_size: int | None = 20,
139
+ ) -> np.ndarray:
140
+ """
141
+ Perform Convergent Cross Mapping using a custom prediction function.
142
+
143
+ CCM tests for causality from X to Y by using the attractor reconstructed from Y
144
+ to predict values of X. If X causes Y, then Y's attractor contains information
145
+ about X, allowing cross-mapping from Y to X.
146
+
147
+ Parameters
148
+ ----------
149
+ X : np.ndarray
150
+ Library delay vectors of shape ``(T,)`` or ``(T, E)`` (potential response).
151
+ Y : np.ndarray
152
+ Target values of shape ``(T,)`` or ``(T, E')`` (potential driver).
153
+ A 2D `Y` cross-maps to a delay vector; the metric scores all columns jointly.
154
+ lib_sizes : np.ndarray
155
+ Array of library sizes to test convergence.
156
+ predict_func : PredictFunc
157
+ Prediction function with signature (X, Y, Q) -> predictions.
158
+ Can be `simplex_projection`, `smap` with partial application, or a custom function.
159
+ metric_func : MetricFunc
160
+ Metric function with signature (predictions, observations) -> metric value.
161
+ n_samples : int, default 100
162
+ Number of random samples per library size for bootstrapping.
163
+ library_pool : np.ndarray
164
+ 1-D array of integer indices from which library members are sampled.
165
+ prediction_pool : np.ndarray
166
+ 1-D array of integer indices that are predicted.
167
+ sample_func : SampleFunc or None, default None
168
+ Function responsible for drawing a library sample of a given size.
169
+ It receives ``(pool, size)`` and returns an array of indices.
170
+ When None, a fresh RNG-backed sampler is created per call.
171
+ aggregate_func : AggregateFunc, default np.mean
172
+ Reducer applied to the skill samples for each library size.
173
+ batch_size : int or None, default 20
174
+ If not specified, batch_size == n_samples.
175
+ If specified, predictions are made in batches to limit memory usage.
176
+ Returns
177
+ -------
178
+ correlations : np.ndarray
179
+ Mean correlation coefficient for each library size.
180
+
181
+ Raises
182
+ ------
183
+ ValueError
184
+ - If `X` and `Y` have different lengths.
185
+ - If `lib_sizes` contains non-positive values.
186
+ - If `predict_func` is not callable.
187
+ - If `n_samples` is not positive.
188
+ - If `aggregate_func` is not callable.
189
+ - If `library_pool` or `prediction_pool` is invalid.
190
+
191
+ Notes
192
+ -----
193
+ - Higher correlation at larger library sizes indicates convergence and suggests X influences Y (X -> Y causality)
194
+ - Convergence is the key signature of causality in CCM
195
+ - The method uses Y's attractor to predict X (cross-mapping)
196
+
197
+ Examples
198
+ --------
199
+ ```python
200
+ from functools import partial
201
+
202
+ import numpy as np
203
+
204
+ from edmkit.ccm import ccm
205
+ from edmkit.embedding import embed
206
+ from edmkit.metrics import pearson_correlation
207
+ from edmkit.simplex_projection import simplex_projection
208
+ from edmkit.smap import smap
209
+
210
+ # Generate coupled logistic maps (X drives Y)
211
+ N = 1000
212
+ rx, ry, Bxy = 3.8, 3.5, 0.02
213
+ X = np.zeros(N)
214
+ Y = np.zeros(N)
215
+ X[0], Y[0] = 0.4, 0.2
216
+ for i in range(1, N):
217
+ X[i] = X[i - 1] * (rx - rx * X[i - 1])
218
+ Y[i] = Y[i - 1] * (ry - ry * Y[i - 1]) + Bxy * X[i - 1]
219
+
220
+ tau = 1
221
+ E = 2
222
+
223
+ # To test X -> Y causality, cross-map from Y's attractor to X
224
+ coordinates = np.array([[0, tau * j] for j in range(E)])
225
+ Y_embedding, _ = embed(coordinates, Y)
226
+ shift = tau * (E - 1)
227
+ X_aligned = X[shift:]
228
+
229
+ library_pool = np.arange(Y_embedding.shape[0] // 2)
230
+ prediction_pool = np.arange(Y_embedding.shape[0] // 2, Y_embedding.shape[0])
231
+
232
+ # logarithmic within range 10 to max library size
233
+ lib_sizes = np.logspace(np.log10(10), np.log10(library_pool[-1]), num=5, dtype=int)
234
+
235
+ # Using simplex projection
236
+ correlations = ccm(
237
+ Y_embedding,
238
+ X_aligned,
239
+ lib_sizes=lib_sizes,
240
+ predict_func=simplex_projection,
241
+ metric_func=pearson_correlation,
242
+ library_pool=library_pool,
243
+ prediction_pool=prediction_pool,
244
+ )
245
+
246
+ # Using S-Map with partial application
247
+ correlations = ccm(
248
+ Y_embedding,
249
+ X_aligned,
250
+ lib_sizes=lib_sizes,
251
+ predict_func=partial(smap, theta=2.0, alpha=1e-10),
252
+ metric_func=pearson_correlation,
253
+ library_pool=library_pool,
254
+ prediction_pool=prediction_pool,
255
+ )
256
+ ```
257
+ """
258
+ if aggregate_func is None or not callable(aggregate_func):
259
+ raise ValueError("aggregate_func must be a callable")
260
+
261
+ samples = bootstrap(
262
+ X,
263
+ Y,
264
+ lib_sizes,
265
+ predict_func,
266
+ metric_func,
267
+ n_samples,
268
+ library_pool=library_pool,
269
+ prediction_pool=prediction_pool,
270
+ sample_func=sample_func,
271
+ batch_size=batch_size,
272
+ )
273
+
274
+ return np.array([aggregate_func(samples[:, i]) for i in range(samples.shape[1])])
@@ -0,0 +1,85 @@
1
+ import numpy as np
2
+
3
+ __all__ = ["embed"]
4
+
5
+
6
+ def embed(coordinates: np.ndarray, x: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
7
+ """Delay vectors for an arbitrary choice of variables and lags.
8
+
9
+ Parameters
10
+ ----------
11
+ coordinates : np.ndarray
12
+ Integer array of shape ``(E, 2)``. Row ``j`` is ``(variable, lag)``, selecting
13
+ observation ``x[t - lag, variable]`` for the delay vector at time ``t``.
14
+ A negative lag selects an observation taken after time ``t``.
15
+ x : np.ndarray
16
+ Observations of shape ``(T, d)``, or ``(T,)``.
17
+
18
+ Returns
19
+ -------
20
+ X : np.ndarray
21
+ Delay vectors of shape ``(N, E)``, one row per time in `t`.
22
+ t : np.ndarray
23
+ Times of shape ``(N,)`` at which every coordinate is available, that is
24
+ ``lags.max() <= t <= T - 1 + lags.min()``. A time outside ``[0, T - 1]`` is
25
+ included when every coordinate still reads an observation inside `x`.
26
+
27
+ Raises
28
+ ------
29
+ ValueError
30
+ - If `coordinates` is not an integer array of shape ``(E, 2)`` with ``E >= 1``.
31
+ - If `x` is not a 1D or 2D array.
32
+ - If a variable index lies outside `x`.
33
+ - If no time has every coordinate available.
34
+
35
+ Examples
36
+ --------
37
+ ```python
38
+ import numpy as np
39
+
40
+ from edmkit.embedding import embed
41
+
42
+ x = np.arange(20).reshape(10, 2) # ten times, two variables
43
+
44
+ # Classical univariate delay coordinates.
45
+ tau, E = 2, 3
46
+ coordinates = np.array([[0, tau * j] for j in range(E)])
47
+ X, t = embed(coordinates, x)
48
+ print(X)
49
+ print(t)
50
+ # [[ 8 4 0]
51
+ # [10 6 2]
52
+ # [12 8 4]
53
+ # [14 10 6]
54
+ # [16 12 8]
55
+ # [18 14 10]]
56
+ # [4 5 6 7 8 9]
57
+
58
+ # One step ahead of the same coordinates, as iterated forecasting needs.
59
+ Y, t_ahead = embed(coordinates - [0, 1], x)
60
+ ```
61
+ """
62
+ coordinates = np.asarray(coordinates)
63
+ if coordinates.ndim != 2 or coordinates.shape[-1] != 2:
64
+ raise ValueError(f"coordinates must have shape (E, 2), got coordinates.shape={coordinates.shape}")
65
+ if coordinates.shape[0] < 1:
66
+ raise ValueError("coordinates must select at least one delay coordinate")
67
+ if not np.issubdtype(coordinates.dtype, np.integer):
68
+ raise ValueError(f"coordinates must be an integer array, got coordinates.dtype={coordinates.dtype}")
69
+ if x.ndim == 1:
70
+ x = x[:, None]
71
+ if x.ndim != 2:
72
+ raise ValueError(f"x must have shape (T,) or (T, d), got x.shape={x.shape}")
73
+
74
+ variables, lags = coordinates[:, 0], coordinates[:, 1]
75
+ if variables.min() < 0 or variables.max() >= x.shape[1]:
76
+ raise ValueError(f"variable indices must lie in [0, {x.shape[1] - 1}], got {variables.min()} to {variables.max()}")
77
+
78
+ # A delay vector at time t needs 0 <= t - lag <= T - 1 for every coordinate.
79
+ first = int(lags.max())
80
+ last = x.shape[0] - 1 + int(lags.min())
81
+ if last < first:
82
+ raise ValueError(f"no time has every coordinate available: lags span {int(lags.min())} to {int(lags.max())} for T={x.shape[0]}")
83
+
84
+ t = np.arange(first, last + 1)
85
+ return x[t[:, None] - lags[None, :], variables[None, :]], t
@@ -0,0 +1,10 @@
1
+ from .ctrnn import ctrnn
2
+ from .double_pendulum import double_pendulum, to_xy
3
+ from .duffing import duffing
4
+ from .lorenz import lorenz
5
+ from .lorenz96 import lorenz96
6
+ from .mackey_glass import mackey_glass
7
+ from .rk45 import rk45
8
+ from .rossler import rossler
9
+
10
+ __all__ = ["ctrnn", "double_pendulum", "duffing", "lorenz", "lorenz96", "mackey_glass", "rk45", "rossler", "to_xy"]
@@ -0,0 +1,25 @@
1
+ import numpy as np
2
+
3
+ __all__ = ["ctrnn"]
4
+
5
+
6
+ def ctrnn(g: float, J: np.ndarray):
7
+ """Vector field of a continuous-time recurrent neural network ``dx/dt = -x + g J tanh(x)``.
8
+
9
+ Parameters
10
+ ----------
11
+ g : float
12
+ Gain (chaos for ``g > 1`` when ``J`` is large and random with variance ``1 / N``).
13
+ J : np.ndarray
14
+ Connectivity of shape ``(N, N)``; ``J[i, j]`` is the weight with which unit ``j`` drives unit ``i``.
15
+
16
+ Returns
17
+ -------
18
+ f : Callable[[float, np.ndarray], np.ndarray]
19
+ Right-hand side ``f(t, x)`` for a state of shape ``(N,)``.
20
+ """
21
+
22
+ def f(t: float, x: np.ndarray):
23
+ return -x + g * J @ np.tanh(x)
24
+
25
+ return f
@@ -3,17 +3,8 @@ import numpy as np
3
3
  __all__ = ["double_pendulum", "to_xy"]
4
4
 
5
5
 
6
- def double_pendulum(
7
- m1: float,
8
- m2: float,
9
- L1: float,
10
- L2: float,
11
- g: float,
12
- X0: np.ndarray,
13
- dt: float,
14
- t_max: int,
15
- ):
16
- """Generate double pendulum dynamics via forward Euler integration.
6
+ def double_pendulum(m1: float, m2: float, L1: float, L2: float, g: float):
7
+ """Vector field of the double pendulum.
17
8
 
18
9
  Parameters
19
10
  ----------
@@ -27,22 +18,14 @@ def double_pendulum(
27
18
  Length of second pendulum.
28
19
  g : float
29
20
  Gravitational acceleration.
30
- X0 : np.ndarray
31
- Initial state ``(theta1, theta2, omega1, omega2)`` of shape ``(4,)``.
32
- dt : float
33
- Integration time step.
34
- t_max : int
35
- Maximum time.
36
21
 
37
22
  Returns
38
23
  -------
39
- t : np.ndarray
40
- Time array.
41
- X : np.ndarray
42
- State trajectory of shape ``(N, 4)``.
24
+ f : Callable[[float, np.ndarray], np.ndarray]
25
+ Right-hand side ``f(t, x)`` for a state ``(theta1, theta2, omega1, omega2)`` of shape ``(4,)``.
43
26
  """
44
27
 
45
- def f(x: np.ndarray):
28
+ def f(t: float, x: np.ndarray):
46
29
  theta1, theta2, omega1, omega2 = x
47
30
  delta = theta1 - theta2
48
31
 
@@ -63,14 +46,7 @@ def double_pendulum(
63
46
 
64
47
  return np.array([dtheta1_dt, dtheta2_dt, domega1_dt, domega2_dt])
65
48
 
66
- t = np.arange(0, t_max, dt)
67
- X = np.zeros((len(t), 4))
68
- X[0] = X0
69
-
70
- for i in range(1, len(t)):
71
- X[i] = X[i - 1] + dt * f(X[i - 1])
72
-
73
- return t, X
49
+ return f
74
50
 
75
51
 
76
52
  def to_xy(L1: float, L2: float, theta1: np.ndarray, theta2: np.ndarray):
@@ -0,0 +1,31 @@
1
+ import numpy as np
2
+
3
+ __all__ = ["duffing"]
4
+
5
+
6
+ def duffing(alpha: float, beta: float, delta: float, gamma: float, omega: float):
7
+ """Vector field of the forced Duffing oscillator.
8
+
9
+ Parameters
10
+ ----------
11
+ alpha : float
12
+ Linear stiffness (typical: -1).
13
+ beta : float
14
+ Cubic stiffness (typical: 1).
15
+ delta : float
16
+ Damping coefficient (typical: 0.3).
17
+ gamma : float
18
+ Forcing amplitude (typical: 0.5 for chaos).
19
+ omega : float
20
+ Forcing angular frequency (typical: 1.2).
21
+
22
+ Returns
23
+ -------
24
+ f : Callable[[float, np.ndarray], np.ndarray]
25
+ Right-hand side ``f(t, x)`` for a state ``(x, v)`` of shape ``(2,)``.
26
+ """
27
+
28
+ def f(t: float, x: np.ndarray):
29
+ return np.array([[0, 1], [-alpha - beta * x[0] ** 2, -delta]]) @ x + np.array([0, gamma * np.cos(omega * t)])
30
+
31
+ return f
@@ -0,0 +1,27 @@
1
+ import numpy as np
2
+
3
+ __all__ = ["lorenz"]
4
+
5
+
6
+ def lorenz(sigma: float, rho: float, beta: float):
7
+ """Vector field of the Lorenz system.
8
+
9
+ Parameters
10
+ ----------
11
+ sigma : float
12
+ Prandtl number (typical: 10).
13
+ rho : float
14
+ Rayleigh number (typical: 28).
15
+ beta : float
16
+ Geometric factor (typical: 8/3).
17
+
18
+ Returns
19
+ -------
20
+ f : Callable[[float, np.ndarray], np.ndarray]
21
+ Right-hand side ``f(t, x)`` for a state ``(x, y, z)`` of shape ``(3,)``.
22
+ """
23
+
24
+ def f(t: float, x: np.ndarray):
25
+ return np.array([[-sigma, sigma, 0], [rho, -1, -x[0]], [0, x[0], -beta]]) @ x
26
+
27
+ return f
@@ -0,0 +1,25 @@
1
+ import numpy as np
2
+
3
+ __all__ = ["lorenz96"]
4
+
5
+
6
+ def lorenz96(F: float):
7
+ """Vector field of the Lorenz 96 system.
8
+
9
+ Parameters
10
+ ----------
11
+ F : float
12
+ Forcing constant (typical: 8 for chaos).
13
+
14
+ Returns
15
+ -------
16
+ f : Callable[[float, np.ndarray], np.ndarray]
17
+ Right-hand side ``f(t, x)`` for a state of shape ``(D,)`` with ``D >= 4``.
18
+ """
19
+
20
+ def f(t: float, x: np.ndarray):
21
+ # p[i + 2] is x[i] on the ring, so this is (x[i + 1] - x[i - 2]) * x[i - 1] - x[i] + F
22
+ p = np.concatenate([x[-2:], x, x[:1]])
23
+ return (p[3:] - p[:-3]) * p[1:-2] - x + F
24
+
25
+ return f