edmkit 0.0.6__tar.gz → 0.0.8__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.
@@ -1,13 +1,13 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: edmkit
3
- Version: 0.0.6
3
+ Version: 0.0.8
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
+ Requires-Dist: kdtree-rs>=0.1.0
7
8
  Requires-Dist: numpy>=2.4.3
8
9
  Requires-Dist: scipy>=1.17.1
9
10
  Requires-Dist: tinygrad>=0.11.0
10
- Requires-Dist: usearch>=2.23.0
11
11
  Requires-Python: >=3.13
12
12
  Description-Content-Type: text/markdown
13
13
 
@@ -1,19 +1,20 @@
1
1
  [project]
2
2
  name = "edmkit"
3
- version = "0.0.6"
3
+ version = "0.0.8"
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"
7
7
  requires-python = ">= 3.13"
8
8
  dependencies = [
9
+ "kdtree-rs>=0.1.0",
9
10
  "numpy>=2.4.3",
10
11
  "scipy>=1.17.1",
11
12
  "tinygrad>=0.11.0",
12
- "usearch>=2.23.0",
13
13
  ]
14
14
 
15
15
  [dependency-groups]
16
16
  dev = [
17
+ "griffe2md>=1.5.0",
17
18
  "hypothesis>=6.151.9",
18
19
  "pytest>=9.0.2",
19
20
  "ruff>=0.15.6",
@@ -1,6 +1,5 @@
1
1
  from collections.abc import Callable
2
2
  from functools import partial
3
- from typing import TypeAlias
4
3
 
5
4
  import numpy as np
6
5
 
@@ -8,9 +7,9 @@ from edmkit.simplex_projection import simplex_projection
8
7
  from edmkit.smap import smap
9
8
  from edmkit.types import PredictFunc
10
9
 
11
- SampleFunc: TypeAlias = Callable[[np.ndarray, int], np.ndarray]
10
+ type SampleFunc = Callable[[np.ndarray, int], np.ndarray]
12
11
  """SampleFunc is a function that takes (pool, size) and returns a sampled array."""
13
- AggregateFunc: TypeAlias = Callable[[np.ndarray], float]
12
+ type AggregateFunc = Callable[[np.ndarray], float]
14
13
  """AggregateFunc is a function that takes an array of values and returns a single value."""
15
14
 
16
15
 
@@ -50,7 +49,7 @@ def bootstrap(
50
49
  Target time series (potential driver)
51
50
  lib_sizes : np.ndarray
52
51
  Array of library sizes to test convergence.
53
- predict_func : :type: `PredictFunc`
52
+ predict_func : PredictFunc
54
53
  Prediction function with signature (X, Y, Q) -> predictions.
55
54
  n_samples : int, default 20
56
55
  Number of random samples per library size for bootstrapping.
@@ -58,16 +57,16 @@ def bootstrap(
58
57
  1-D array of integer indices from which library members are sampled.
59
58
  prediction_pool : np.ndarray
60
59
  1-D array of integer indices that are predicted.
61
- sample_func : :type: `SampleFunc` | None, default None
60
+ sample_func : SampleFunc or None, default None
62
61
  Function responsible for drawing a library sample of a given size.
63
62
  When None, a fresh RNG-backed sampler is created per call.
64
- batch_size : int | None, default 10
63
+ batch_size : int or None, default 10
65
64
  If specified, predictions are made in batches to limit memory usage.
66
65
 
67
66
  Returns
68
67
  -------
69
- samples : np.ndarray of shape (n_samples, len(lib_sizes))
70
- Per-sample correlation coefficients.
68
+ samples : np.ndarray
69
+ Per-sample correlation coefficients of shape ``(n_samples, len(lib_sizes))``.
71
70
  """
72
71
  if sample_func is None:
73
72
  sample_func = make_sample_func()
@@ -142,7 +141,7 @@ def ccm(
142
141
  Target time series (potential driver)
143
142
  lib_sizes : np.ndarray
144
143
  Array of library sizes to test convergence.
145
- predict_func : :type: `PredictFunc`
144
+ predict_func : PredictFunc
146
145
  Prediction function with signature (X, Y, Q) -> predictions.
147
146
  Can be `simplex_projection`, `smap` with partial application, or a custom function.
148
147
  n_samples : int, default 100
@@ -151,13 +150,13 @@ def ccm(
151
150
  1-D array of integer indices from which library members are sampled.
152
151
  prediction_pool : np.ndarray
153
152
  1-D array of integer indices that are predicted.
154
- sample_func : :type: `SampleFunc` | None, default None
153
+ sample_func : SampleFunc or None, default None
155
154
  Function responsible for drawing a library sample of a given size.
156
- It receives `(pool, size)` and returns an array of indices.
155
+ It receives ``(pool, size)`` and returns an array of indices.
157
156
  When None, a fresh RNG-backed sampler is created per call.
158
- aggregate_func : :type: `AggregateFunc`, default `np.mean`
157
+ aggregate_func : AggregateFunc, default np.mean
159
158
  Reducer applied to the correlation samples for each library size.
160
- batch_size : int | None, default None
159
+ batch_size : int or None, default None
161
160
  If not specified, batch_size == n_samples.
162
161
  If specified, predictions are made in batches to limit memory usage.
163
162
  Returns
@@ -14,13 +14,17 @@ def lagged_embed(x: np.ndarray, tau: int, e: int):
14
14
 
15
15
  Parameters
16
16
  ----------
17
- `x` : `np.ndarray` of shape `(N,)`
18
- `tau` : `int`
19
- `e` : `int`
17
+ x : np.ndarray
18
+ 1D time series of shape ``(N,)``.
19
+ tau : int
20
+ Time delay.
21
+ e : int
22
+ Embedding dimension.
20
23
 
21
24
  Returns
22
25
  -------
23
- `np.ndarray` of shape `(N - (e - 1) * tau, e)`
26
+ np.ndarray
27
+ Embedded array of shape ``(N - (e - 1) * tau, e)``.
24
28
 
25
29
  Raises
26
30
  ------
@@ -2,3 +2,5 @@
2
2
  from .double_pendulum import double_pendulum, to_xy
3
3
  from .lorenz import lorenz
4
4
  from .mackey_glass import mackey_glass
5
+
6
+ __all__ = ["double_pendulum", "to_xy", "lorenz", "mackey_glass"]
@@ -0,0 +1,105 @@
1
+ import numpy as np
2
+
3
+
4
+ def double_pendulum(
5
+ m1: float,
6
+ m2: float,
7
+ L1: float,
8
+ L2: float,
9
+ g: float,
10
+ X0: np.ndarray,
11
+ dt: float,
12
+ t_max: int,
13
+ ):
14
+ """Generate double pendulum dynamics via forward Euler integration.
15
+
16
+ Parameters
17
+ ----------
18
+ m1 : float
19
+ Mass of first pendulum.
20
+ m2 : float
21
+ Mass of second pendulum.
22
+ L1 : float
23
+ Length of first pendulum.
24
+ L2 : float
25
+ Length of second pendulum.
26
+ g : float
27
+ Gravitational acceleration.
28
+ X0 : np.ndarray
29
+ Initial state ``(theta1, theta2, omega1, omega2)`` of shape ``(4,)``.
30
+ dt : float
31
+ Integration time step.
32
+ t_max : int
33
+ Maximum time.
34
+
35
+ Returns
36
+ -------
37
+ t : np.ndarray
38
+ Time array.
39
+ X : np.ndarray
40
+ State trajectory of shape ``(N, 4)``.
41
+ """
42
+
43
+ def f(x: np.ndarray):
44
+ theta1, theta2, omega1, omega2 = x
45
+ delta = theta1 - theta2
46
+
47
+ denom = 2 * m1 + m2 - m2 * np.cos(2 * delta)
48
+
49
+ dtheta1_dt = omega1
50
+ dtheta2_dt = omega2
51
+
52
+ domega1_dt = (
53
+ -g * (2 * m1 + m2) * np.sin(theta1)
54
+ - m2 * g * np.sin(theta1 - 2 * theta2)
55
+ - 2 * np.sin(delta) * m2 * (omega2**2 * L2 + omega1**2 * L1 * np.cos(delta))
56
+ ) / (L1 * denom)
57
+
58
+ domega2_dt = (2 * np.sin(delta) * (omega1**2 * L1 * (m1 + m2) + g * (m1 + m2) * np.cos(theta1) + omega2**2 * L2 * m2 * np.cos(delta))) / (
59
+ L2 * denom
60
+ )
61
+
62
+ return np.array([dtheta1_dt, dtheta2_dt, domega1_dt, domega2_dt])
63
+
64
+ t = np.arange(0, t_max, dt)
65
+ X = np.zeros((len(t), 4))
66
+ X[0] = X0
67
+
68
+ for i in range(1, len(t)):
69
+ X[i] = X[i - 1] + dt * f(X[i - 1])
70
+
71
+ return t, X
72
+
73
+
74
+ def to_xy(L1: float, L2: float, theta1: np.ndarray, theta2: np.ndarray):
75
+ """Convert double pendulum angles to Cartesian coordinates.
76
+
77
+ Parameters
78
+ ----------
79
+ L1 : float
80
+ Length of first pendulum.
81
+ L2 : float
82
+ Length of second pendulum.
83
+ theta1 : np.ndarray
84
+ Angle of first pendulum.
85
+ theta2 : np.ndarray
86
+ Angle of second pendulum.
87
+
88
+ Returns
89
+ -------
90
+ x1 : np.ndarray
91
+ x-coordinate of first pendulum.
92
+ y1 : np.ndarray
93
+ y-coordinate of first pendulum.
94
+ x2 : np.ndarray
95
+ x-coordinate of second pendulum.
96
+ y2 : np.ndarray
97
+ y-coordinate of second pendulum.
98
+ """
99
+ x1 = L1 * np.sin(theta1)
100
+ y1 = -L1 * np.cos(theta1)
101
+
102
+ x2 = x1 + L2 * np.sin(theta2)
103
+ y2 = y1 - L2 * np.cos(theta2)
104
+
105
+ return x1, y1, x2, y2
@@ -0,0 +1,47 @@
1
+ import numpy as np
2
+
3
+
4
+ def lorenz(
5
+ sigma: float,
6
+ rho: float,
7
+ beta: float,
8
+ X0: np.ndarray,
9
+ dt: float,
10
+ t_max: int,
11
+ ):
12
+ """Generate a Lorenz system trajectory via forward Euler integration.
13
+
14
+ Parameters
15
+ ----------
16
+ sigma : float
17
+ Prandtl number (typical: 10).
18
+ rho : float
19
+ Rayleigh number (typical: 28).
20
+ beta : float
21
+ Geometric factor (typical: 8/3).
22
+ X0 : np.ndarray
23
+ Initial condition of shape ``(3,)``.
24
+ dt : float
25
+ Integration time step.
26
+ t_max : int
27
+ Maximum time.
28
+
29
+ Returns
30
+ -------
31
+ t : np.ndarray
32
+ Time array.
33
+ X : np.ndarray
34
+ Trajectory of shape ``(N, 3)`` for ``(x, y, z)``.
35
+ """
36
+
37
+ def f(x: np.ndarray):
38
+ return np.array([[-sigma, sigma, 0], [rho, -1, -x[0]], [0, x[0], -beta]]) @ x
39
+
40
+ t = np.arange(0, t_max, dt)
41
+ X = np.zeros((len(t), 3))
42
+ X[0] = X0
43
+
44
+ for i in range(1, len(t)):
45
+ X[i] = X[i - 1] + dt * f(X[i - 1])
46
+
47
+ return t, X
@@ -0,0 +1,44 @@
1
+ import numpy as np
2
+
3
+
4
+ def mackey_glass(tau: float, n: int, beta: float, gamma: float, x0: float, dt: float, t_max: int):
5
+ """Generate a Mackey-Glass chaotic time series via forward Euler integration.
6
+
7
+ Parameters
8
+ ----------
9
+ tau : float
10
+ Delay parameter (typical: 17 for chaos).
11
+ n : int
12
+ Nonlinearity exponent (typical: 10).
13
+ beta : float
14
+ Feedback strength (typical: 0.2).
15
+ gamma : float
16
+ Decay rate (typical: 0.1).
17
+ x0 : float
18
+ Initial condition.
19
+ dt : float
20
+ Integration time step.
21
+ t_max : int
22
+ Maximum time.
23
+
24
+ Returns
25
+ -------
26
+ t : np.ndarray
27
+ Time array.
28
+ x : np.ndarray
29
+ 1D time series.
30
+ """
31
+
32
+ def f(x, x_tau):
33
+ return beta * x_tau / (1 + x_tau**n) - gamma * x
34
+
35
+ t = np.arange(0, t_max, dt)
36
+ x = np.zeros_like(t)
37
+
38
+ tau_idx = int(tau / dt)
39
+ x[:tau_idx] = x0
40
+
41
+ for i in range(tau_idx, len(t)):
42
+ x[i] = x[i - 1] + dt * f(x[i - 1], x[i - tau_idx])
43
+
44
+ return t, x
@@ -1,8 +1,8 @@
1
- from typing import TYPE_CHECKING, Callable, TypeAlias
1
+ from typing import TYPE_CHECKING, Callable
2
2
 
3
3
  import numpy as np
4
4
 
5
- MetricFunc: TypeAlias = Callable[[np.ndarray, np.ndarray], np.ndarray]
5
+ type MetricFunc = Callable[[np.ndarray, np.ndarray], np.ndarray]
6
6
  """MetricFunc is a function that takes (predictions, observations) and returns a metric value."""
7
7
 
8
8
 
@@ -0,0 +1,31 @@
1
+ import numpy as np
2
+ from kdtree import KDTree
3
+
4
+
5
+ def knn(X: np.ndarray, Q: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
6
+ """Find the k-nearest neighbors of `Q` in `X` using `kdtree.KDTree`.
7
+
8
+ Parameters
9
+ ----------
10
+ X : np.ndarray
11
+ The input data (N, E)
12
+ Q : np.ndarray
13
+ The query points (M, E)
14
+ k : int
15
+ The number of nearest neighbors to find (typically E+1 for simplex projection).
16
+
17
+ Returns
18
+ -------
19
+ distances : np.ndarray
20
+ The distances from each query point in `Q` to its k nearest neighbors in `X` (M, k)
21
+ indices : np.ndarray
22
+ The indices of the k nearest neighbors in `X` for each query point in `Q` (M, k)
23
+ """
24
+
25
+ N = X.shape[0]
26
+
27
+ if N < k:
28
+ raise ValueError(f"Not enough points in X to find {k} neighbors, got N={N}")
29
+
30
+ tree = KDTree(X)
31
+ return tree.query(Q, k=k)
@@ -17,18 +17,18 @@ def loo(
17
17
 
18
18
  Parameters
19
19
  ----------
20
- `X` : `np.ndarray`
20
+ X : np.ndarray
21
21
  The input data of shape (N,) or (N, E) or (B, N, E).
22
- `Y` : `np.ndarray`
22
+ Y : np.ndarray
23
23
  The target data of shape (N,) or (N, E') or (B, N, E').
24
- `theiler_window` : `int`
24
+ theiler_window : int
25
25
  Theiler window half-width. Library points ``j`` where
26
26
  ``|i - j| <= theiler_window`` are excluded when predicting point ``i``.
27
- For lagged embedding, use ``(E - 1) * tau + n_ahead``.
27
+ For lagged embedding, use ``(E - 1) * tau``.
28
28
 
29
29
  Returns
30
30
  -------
31
- predictions : `np.ndarray`
31
+ predictions : np.ndarray
32
32
  The predicted values of shape (N,) or (N, E') or (B, N, E').
33
33
 
34
34
  Raises
@@ -20,21 +20,21 @@ def simplex_projection(
20
20
 
21
21
  Parameters
22
22
  ----------
23
- `X` : `np.ndarray`
23
+ X : np.ndarray
24
24
  The input data of shape (N,) or (N, E) or (B, N, E)
25
- `Y` : `np.ndarray`
25
+ Y : np.ndarray
26
26
  The target data of shape (N,) or (N, E') or (B, N, E')
27
- `Q` : `np.ndarray`
27
+ Q : np.ndarray
28
28
  The query points of shape (M,) or (M, E) or (B, M, E) for which to find the nearest neighbors in `X`.
29
- `mask` : `np.ndarray | None`
29
+ mask : np.ndarray or None
30
30
  Boolean mask of shape (N,) or (B, N) indicating which library points to include when finding nearest neighbors for the queries in `Q`.
31
- `use_tensor` : `bool`, default `False`
31
+ use_tensor : bool, default False
32
32
  Whether to use `tinygrad.Tensor` for computation.
33
33
  **This may be slower than the NumPy implementation in most cases for now.**
34
34
 
35
35
  Returns
36
36
  -------
37
- predictions : `np.ndarray`
37
+ predictions : np.ndarray
38
38
  The predicted values based on the weighted mean of the nearest neighbors in `Y`.
39
39
 
40
40
  Raises
@@ -21,23 +21,23 @@ def smap(
21
21
 
22
22
  Parameters
23
23
  ----------
24
- `X` : `np.ndarray`
24
+ X : np.ndarray
25
25
  The input data
26
- `Y` : `np.ndarray`
26
+ Y : np.ndarray
27
27
  The target data
28
- `Q` : `np.ndarray`
28
+ Q : np.ndarray
29
29
  The query points for which to make predictions.
30
- `theta` : `float`
30
+ theta : float
31
31
  Locality parameter. (0: global linear, >0: local linear)
32
- `alpha` : `float`, default `1e-10`
32
+ alpha : float, default 1e-10
33
33
  Regularization parameter to stabilize the inversion.
34
- `use_tensor` : `bool`, default `False`
34
+ use_tensor : bool, default False
35
35
  Whether to use `tinygrad.Tensor` for computation.
36
36
  **This may be slower than the NumPy implementation in most cases for now.**
37
37
 
38
38
  Returns
39
39
  -------
40
- predictions : `np.ndarray`
40
+ predictions : np.ndarray
41
41
  The predicted values based on the weighted linear regression.
42
42
 
43
43
  Raises
@@ -137,21 +137,21 @@ def _numpy(
137
137
 
138
138
  Parameters
139
139
  ----------
140
- `X` : `np.ndarray`
140
+ X : np.ndarray
141
141
  (N,) or (N, E) or (B, N, E)
142
- `Y` : `np.ndarray`
142
+ Y : np.ndarray
143
143
  (N,) or (N, E') or (B, N, E')
144
- `Q` : `np.ndarray`
144
+ Q : np.ndarray
145
145
  The query points for which to make predictions.
146
146
  (M,) or (M, E) or (B, M, E)
147
- `theta` : `float`
147
+ theta : float
148
148
  Locality parameter. (0: global linear, >0: local linear)
149
- `alpha` : `float`, default `1e-10`
149
+ alpha : float, default 1e-10
150
150
  Regularization parameter to stabilize the inversion.
151
151
 
152
152
  Returns
153
153
  -------
154
- predictions : `np.ndarray`
154
+ predictions : np.ndarray
155
155
  The predicted values based on the weighted linear regression.
156
156
  (M, E') or (B, M, E')
157
157
 
@@ -245,20 +245,20 @@ def _tensor(
245
245
 
246
246
  Parameters
247
247
  ----------
248
- `X` : `np.ndarray`
248
+ X : np.ndarray
249
249
  The input data
250
- `Y` : `np.ndarray`
250
+ Y : np.ndarray
251
251
  The target data
252
- `Q` : `np.ndarray`
252
+ Q : np.ndarray
253
253
  The query points for which to make predictions.
254
- `theta` : `float`
254
+ theta : float
255
255
  Locality parameter. (0: global linear, >0: local linear)
256
- `alpha` : `float`, default `1e-10`
256
+ alpha : float, default 1e-10
257
257
  Regularization parameter to stabilize the inversion.
258
258
 
259
259
  Returns
260
260
  -------
261
- predictions : `np.ndarray`
261
+ predictions : np.ndarray
262
262
  The predicted values based on the weighted linear regression.
263
263
 
264
264
  Raises
@@ -1,4 +1,4 @@
1
- from typing import TYPE_CHECKING, Callable, NamedTuple, TypeAlias
1
+ from typing import TYPE_CHECKING, Callable, NamedTuple
2
2
 
3
3
  import numpy as np
4
4
 
@@ -10,7 +10,7 @@ class Fold(NamedTuple):
10
10
  validation: np.ndarray # index array
11
11
 
12
12
 
13
- SplitFunc: TypeAlias = Callable[[int], list[Fold]]
13
+ type SplitFunc = Callable[[int], list[Fold]]
14
14
 
15
15
 
16
16
  def temporal_fold(
@@ -8,11 +8,13 @@ def pad(As: list[np.ndarray]):
8
8
 
9
9
  Parameters
10
10
  ----------
11
- `As` : `list` of `np.ndarray` of shape `(N, D_i)`
11
+ As : list[np.ndarray]
12
+ List of arrays of shape ``(N, D_i)``.
12
13
 
13
14
  Returns
14
15
  -------
15
- Single `np.ndarray` of shape `(B, N, max(D))` where `B` is `len(As)`
16
+ np.ndarray
17
+ Single array of shape ``(B, N, max(D))`` where B is ``len(As)``.
16
18
 
17
19
  Raises
18
20
  ------
@@ -41,21 +43,18 @@ def pairwise_distance(A: Tensor, B: Tensor | None = None) -> Tensor:
41
43
 
42
44
  Parameters
43
45
  ----------
44
- `A` : `Tensor` of shape `(N, D)` or `(B, N, D)`
45
- - `B`: batch size
46
- - `N`: number of points
47
- - `D`: dimension of each point
48
- `B` : `Tensor` of shape `(M, D)` or `(B, M, D)`
49
- - `B`: batch size
50
- - `M`: number of points
51
- - `D`: dimension of each point
46
+ A : Tensor
47
+ Shape ``(N, D)`` or ``(B, N, D)``.
48
+ ``B`` is batch size, ``N`` is number of points, ``D`` is dimension of each point.
49
+ B : Tensor
50
+ Shape ``(M, D)`` or ``(B, M, D)``.
51
+ ``B`` is batch size, ``M`` is number of points, ``D`` is dimension of each point.
52
52
 
53
53
  Returns
54
54
  -------
55
- When `A` is of shape `(N, D)`:
56
- `Tensor` of shape `(N, N)` [or `(N, M)`] where the element at position `(i, j)` is the squared Euclidean distance between `A[i]` and `A[j]` [or between `A[i]` and `B[j]`].
57
- When `A` is of shape `(B, N, D)`:
58
- `Tensor` of shape `(B, N, N)` [or `(B, N, M)`] where the element at position `(b, i, j)` is the squared Euclidean distance between `A[b, i]` and `A[b, j]`.
55
+ Tensor
56
+ When `A` is of shape ``(N, D)``: shape ``(N, N)`` [or ``(N, M)``] where the element at position ``(i, j)`` is the squared Euclidean distance between ``A[i]`` and ``A[j]`` [or between ``A[i]`` and ``B[j]``].
57
+ When `A` is of shape ``(B, N, D)``: shape ``(B, N, N)`` [or ``(B, N, M)``] where the element at position ``(b, i, j)`` is the squared Euclidean distance between ``A[b, i]`` and ``A[b, j]``.
59
58
 
60
59
  Raises
61
60
  ------
@@ -84,21 +83,18 @@ def pairwise_distance_np(A: np.ndarray, B: np.ndarray | None = None) -> np.ndarr
84
83
 
85
84
  Parameters
86
85
  ----------
87
- `A` : `np.ndarray` of shape `(N, D)` or `(B, N, D)`
88
- - `B`: batch size
89
- - `N`: number of points
90
- - `D`: dimension of each point
91
- `B` : `np.ndarray` of shape `(M, D)` or `(B, M, D)`
92
- - `B`: batch size
93
- - `M`: number of points
94
- - `D`: dimension of each point
86
+ A : np.ndarray
87
+ Shape ``(N, D)`` or ``(B, N, D)``.
88
+ ``B`` is batch size, ``N`` is number of points, ``D`` is dimension of each point.
89
+ B : np.ndarray
90
+ Shape ``(M, D)`` or ``(B, M, D)``.
91
+ ``B`` is batch size, ``M`` is number of points, ``D`` is dimension of each point.
95
92
 
96
93
  Returns
97
94
  -------
98
- When `A` is of shape `(N, D)`:
99
- `np.ndarray` of shape `(N, N)` [or `(N, M)`] where the element at position `(i, j)` is the squared Euclidean distance between `A[i]` and `A[j]` [or between `A[i]` and `B[j]`].
100
- When `A` is of shape `(B, N, D)`:
101
- `np.ndarray` of shape `(B, N, N)` [or `(B, N, M)`] where the element at position `(b, i, j)` is the squared Euclidean distance between `A[b, i]` and `A[b, j]`.
95
+ np.ndarray
96
+ When `A` is of shape ``(N, D)``: shape ``(N, N)`` [or ``(N, M)``] where the element at position ``(i, j)`` is the squared Euclidean distance between ``A[i]`` and ``A[j]`` [or between ``A[i]`` and ``B[j]``].
97
+ When `A` is of shape ``(B, N, D)``: shape ``(B, N, N)`` [or ``(B, N, M)``] where the element at position ``(b, i, j)`` is the squared Euclidean distance between ``A[b, i]`` and ``A[b, j]``.
102
98
 
103
99
  Raises
104
100
  ------
@@ -128,12 +124,15 @@ def dtw(A: np.ndarray, B: np.ndarray):
128
124
 
129
125
  Parameters
130
126
  ----------
131
- `A` : Tensor of shape `(N,D)`
132
- `B` : Tensor of shape `(M,D)`
127
+ A : np.ndarray
128
+ Sequence of shape ``(N, D)``.
129
+ B : np.ndarray
130
+ Sequence of shape ``(M, D)``.
133
131
 
134
132
  Returns
135
133
  -------
136
- distance : float
134
+ distance : float
135
+ The DTW distance between the two sequences.
137
136
  """
138
137
  N: int = A.shape[0]
139
138
  M: int = B.shape[0]
@@ -169,13 +168,17 @@ def autocorrelation(x: np.ndarray, max_lag: int, step: int = 1):
169
168
 
170
169
  Parameters
171
170
  ----------
172
- `x` : `np.ndarray` The input array for which to compute the autocorrelation.
173
- `max_lag` : `int` The maximum lag up to which the autocorrelation is computed.
174
- `step` : `int`, `optional` The step size for the lag. Default is 1.
171
+ x : np.ndarray
172
+ The input array for which to compute the autocorrelation.
173
+ max_lag : int
174
+ The maximum lag up to which the autocorrelation is computed.
175
+ step : int, optional
176
+ The step size for the lag. Default is 1.
175
177
 
176
178
  Returns
177
179
  -------
178
- `np.ndarray` of shape `(max_lag // step + 1,)` containing the autocorrelation values.
180
+ np.ndarray
181
+ Array of shape ``(max_lag // step + 1,)`` containing the autocorrelation values.
179
182
  """
180
183
  x = x - np.mean(x)
181
184
 
@@ -1,52 +0,0 @@
1
- import numpy as np
2
-
3
-
4
- def double_pendulum(
5
- m1: float,
6
- m2: float,
7
- L1: float,
8
- L2: float,
9
- g: float,
10
- X0: np.ndarray, # (theta1, theta2, omega1, omega2)
11
- dt: float,
12
- t_max: int,
13
- ):
14
- def f(x: np.ndarray):
15
- theta1, theta2, omega1, omega2 = x
16
- delta = theta1 - theta2
17
-
18
- denom = 2 * m1 + m2 - m2 * np.cos(2 * delta)
19
-
20
- dtheta1_dt = omega1
21
- dtheta2_dt = omega2
22
-
23
- domega1_dt = (
24
- -g * (2 * m1 + m2) * np.sin(theta1)
25
- - m2 * g * np.sin(theta1 - 2 * theta2)
26
- - 2 * np.sin(delta) * m2 * (omega2**2 * L2 + omega1**2 * L1 * np.cos(delta))
27
- ) / (L1 * denom)
28
-
29
- domega2_dt = (2 * np.sin(delta) * (omega1**2 * L1 * (m1 + m2) + g * (m1 + m2) * np.cos(theta1) + omega2**2 * L2 * m2 * np.cos(delta))) / (
30
- L2 * denom
31
- )
32
-
33
- return np.array([dtheta1_dt, dtheta2_dt, domega1_dt, domega2_dt])
34
-
35
- t = np.arange(0, t_max, dt)
36
- X = np.zeros((len(t), 4))
37
- X[0] = X0
38
-
39
- for i in range(1, len(t)):
40
- X[i] = X[i - 1] + dt * f(X[i - 1])
41
-
42
- return t, X
43
-
44
-
45
- def to_xy(L1: float, L2: float, theta1: np.ndarray, theta2: np.ndarray):
46
- x1 = L1 * np.sin(theta1)
47
- y1 = -L1 * np.cos(theta1)
48
-
49
- x2 = x1 + L2 * np.sin(theta2)
50
- y2 = y1 - L2 * np.cos(theta2)
51
-
52
- return x1, y1, x2, y2
@@ -1,22 +0,0 @@
1
- import numpy as np
2
-
3
-
4
- def lorenz(
5
- sigma: float,
6
- rho: float,
7
- beta: float,
8
- X0: np.ndarray,
9
- dt: float,
10
- t_max: int,
11
- ):
12
- def f(x: np.ndarray):
13
- return np.array([[-sigma, sigma, 0], [rho, -1, -x[0]], [0, x[0], -beta]]) @ x
14
-
15
- t = np.arange(0, t_max, dt)
16
- X = np.zeros((len(t), 3))
17
- X[0] = X0
18
-
19
- for i in range(1, len(t)):
20
- X[i] = X[i - 1] + dt * f(X[i - 1])
21
-
22
- return t, X
@@ -1,17 +0,0 @@
1
- import numpy as np
2
-
3
-
4
- def mackey_glass(tau: float, n: int, beta: float, gamma: float, x0: float, dt: float, t_max: int):
5
- def f(x, x_tau):
6
- return beta * x_tau / (1 + x_tau**n) - gamma * x
7
-
8
- t = np.arange(0, t_max, dt)
9
- x = np.zeros_like(t)
10
-
11
- tau_idx = int(tau / dt)
12
- x[:tau_idx] = x0
13
-
14
- for i in range(tau_idx, len(t)):
15
- x[i] = x[i - 1] + dt * f(x[i - 1], x[i - tau_idx])
16
-
17
- return t, x
@@ -1,40 +0,0 @@
1
- import numpy as np
2
- from scipy.spatial import KDTree
3
- from usearch.index import Index
4
-
5
-
6
- def knn(X: np.ndarray, Q: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
7
- """Find the k-nearest neighbors of `Q` in `X` using either `usearch` or `scipy.spatial.KDTree` depending on the size and dimensionality of the data.
8
-
9
- Parameters
10
- ----------
11
- `X` : `np.ndarray`
12
- The input data (N, E)
13
- `Q` : `np.ndarray`
14
- The query points (M, E)
15
- `k` : `int`
16
- The number of nearest neighbors to find (typically E+1 for simplex projection).
17
-
18
- Returns
19
- -------
20
- distances : `np.ndarray`
21
- The distances from each query point in `Q` to its k nearest neighbors in `X` (M, k)
22
- indices : `np.ndarray`
23
- The indices of the k nearest neighbors in `X` for each query point in `Q` (M, k)
24
- """
25
-
26
- N, E = X.shape
27
-
28
- if N < k:
29
- raise ValueError(f"Not enough points in X to find {k} neighbors, got N={N}")
30
-
31
- if E >= 15 and N >= 10_000:
32
- index = Index(ndim=E, metric="l2sq")
33
- index.add(np.arange(len(X)), np.ascontiguousarray(X, dtype=np.float32))
34
- matches = index.search(np.ascontiguousarray(Q, dtype=np.float32), k)
35
- distances = np.atleast_2d(np.sqrt(np.asarray(matches.distances)))
36
- indices = np.atleast_2d(np.asarray(matches.keys).astype(np.intp))
37
- return distances, indices
38
- else:
39
- tree = KDTree(X)
40
- return tree.query(Q, k=k)
File without changes
File without changes