optimized-svm 0.1.0__py3-none-any.whl

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.
@@ -0,0 +1,153 @@
1
+ """Cross-validation error and its gradient w.r.t. the hyperparameters (Strasdat, Sec. 9.2).
2
+
3
+ Hyperparameters are ``theta = (log lam, log s_1, ..., log s_K)`` (regularization and feature
4
+ block scalings). For each fold the training problem is solved and the derivative of the
5
+ solution is obtained from the implicit function theorem applied to the (semismooth) optimality
6
+ conditions ``a_l + phi_l'(ka_l + b + s_l c) = 0``, ``sum(a) = 0`` (and ``s^T a = A``). Only the
7
+ terms with positive curvature ``h_l = phi_l''`` react to parameter changes. Instead of the
8
+ ``K`` forward sensitivities, a single *adjoint* system is solved per fold:
9
+
10
+ ``[[K_PP / lam + diag(1 / h_P), 1, s_P], [1^T, 0, 0], [s_P^T, 0, 0]] [u; .]``
11
+ ``= [K_VP^T dE / lam; sum dE; 0]``
12
+
13
+ after which ``dE/dtheta = -u^T (d ka_P / dtheta) + dE^T (d f_V / dtheta)`` with ``a`` fixed.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ from collections.abc import Sequence
19
+ from dataclasses import dataclass, field
20
+ from typing import Any
21
+
22
+ import numpy as np
23
+ from numpy.typing import NDArray
24
+
25
+ from .kernels import Blocks, n_blocks, rbf, rbf_with_gradient, scale_features
26
+ from .linalg import solve_bordered
27
+ from .losses import Loss
28
+ from .problem import TrainingProblem
29
+ from .training import Solution, train
30
+ from .validation import ValidationLoss
31
+
32
+ Array = NDArray[np.float64]
33
+ Fold = tuple[NDArray[np.intp], NDArray[np.intp]]
34
+
35
+
36
+ @dataclass
37
+ class FoldResult:
38
+ error: float
39
+ gradient: Array
40
+ solution: Solution
41
+ n_support: int
42
+ time: float
43
+
44
+
45
+ def fold_error_and_gradient(
46
+ x_train: Array,
47
+ y_train: Array,
48
+ w_train: Array,
49
+ x_val: Array,
50
+ y_val: Array,
51
+ loss: Loss,
52
+ validation: ValidationLoss,
53
+ lam: float,
54
+ blocks: Blocks,
55
+ max_asum: float | None = None,
56
+ smo_options: dict[str, Any] | None = None,
57
+ newton_options: dict[str, Any] | None = None,
58
+ alpha0: Array | None = None,
59
+ ) -> FoldResult:
60
+ """Validation error of one fold and its gradient w.r.t. ``(log lam, log s_k)``.
61
+
62
+ ``x_train``/``x_val`` must already be scaled.
63
+ """
64
+ prob = TrainingProblem(x_train, y_train, loss, lam, w_train, max_asum)
65
+ sol = train(prob, smo_options, newton_options, alpha0)
66
+ a = sol.a
67
+ sv = np.flatnonzero(a)
68
+ a_sv = a[sv] / lam
69
+ x_terms = prob.x[prob.terms.kernel_index]
70
+
71
+ # validation predictions and explicit derivatives (a fixed)
72
+ k_val_sv, dk_val_sv = rbf_with_gradient(x_val, x_terms[sv], blocks)
73
+ ka_val = k_val_sv @ a_sv
74
+ f_val = ka_val + sol.b
75
+ err, d_err = validation(f_val, y_val)
76
+ grad = np.zeros(1 + n_blocks(blocks))
77
+ grad[0] = -d_err @ ka_val
78
+ grad[1:] = a_sv @ np.tensordot(d_err, dk_val_sv, axes=(0, 0))
79
+
80
+ # implicit part: adjoint system on terms with positive curvature
81
+ dec = prob.decision(sol.b, sol.c, sol.ka)
82
+ h = prob.loss.d2(dec, prob.terms)
83
+ p = np.flatnonzero(h > 0.0)
84
+ if p.size > 0:
85
+ _, dk_p_sv = rbf_with_gradient(x_terms[p], x_terms[sv], blocks)
86
+ q = prob.kernel(p, p) / lam + np.diag(1.0 / h[p])
87
+ cols = [np.ones(p.size)]
88
+ rhs0 = [d_err.sum()]
89
+ if max_asum is not None and sol.c > 0.0:
90
+ cols.append(prob.sign[p])
91
+ rhs0.append(0.0)
92
+ k_val_p = rbf(x_val, x_terms[p])
93
+ u, _ = solve_bordered(q, np.stack(cols, axis=1), k_val_p.T @ d_err / lam, np.array(rhs0))
94
+ # d ka_P / d log lam = -ka_P, d ka_P / d log s_k = dK_{P,SV} a_SV / lam
95
+ grad[0] += u @ sol.ka[p]
96
+ grad[1:] -= a_sv @ np.tensordot(u, dk_p_sv, axes=(0, 0))
97
+ return FoldResult(err, grad, sol, int(sv.size), sol.time_smo + sol.time_newton)
98
+
99
+
100
+ @dataclass
101
+ class CVResult:
102
+ error: float
103
+ gradient: Array
104
+ folds: list[FoldResult] = field(default_factory=list)
105
+
106
+
107
+ def cv_error_and_gradient(
108
+ theta: Array,
109
+ x: Array,
110
+ y: Array,
111
+ w: Array,
112
+ folds: Sequence[Fold],
113
+ loss: Loss,
114
+ validation: ValidationLoss,
115
+ blocks: Blocks = None,
116
+ max_asum: float | None = None,
117
+ smo_options: dict[str, Any] | None = None,
118
+ newton_options: dict[str, Any] | None = None,
119
+ warm_start: list[Array | None] | None = None,
120
+ ) -> CVResult:
121
+ """Mean validation error over all folds and its gradient w.r.t. ``theta``.
122
+
123
+ ``warm_start`` is a list with one (possibly ``None``) start vector per fold; it is updated
124
+ in place with the new solutions.
125
+ """
126
+ lam = float(np.exp(theta[0]))
127
+ xs = scale_features(x, np.exp(theta[1:]), blocks)
128
+ total, grad, results = 0.0, np.zeros_like(theta, dtype=np.float64), []
129
+ n_val = 0
130
+ for i, (train_idx, val_idx) in enumerate(folds):
131
+ start = warm_start[i] if warm_start is not None else None
132
+ res = fold_error_and_gradient(
133
+ xs[train_idx],
134
+ y[train_idx],
135
+ w[train_idx],
136
+ xs[val_idx],
137
+ y[val_idx],
138
+ loss,
139
+ validation,
140
+ lam,
141
+ blocks,
142
+ max_asum,
143
+ smo_options,
144
+ newton_options,
145
+ start,
146
+ )
147
+ if warm_start is not None:
148
+ warm_start[i] = res.solution.a
149
+ total += res.error
150
+ grad += res.gradient
151
+ n_val += len(val_idx)
152
+ results.append(res)
153
+ return CVResult(total / n_val, grad / n_val, results)
@@ -0,0 +1,63 @@
1
+ """Gaussian kernel with feature scaling and its derivatives w.r.t. log-scalings.
2
+
3
+ The hyperparameters of the kernel are scalings ``s_k > 0`` of feature blocks ``B_k``:
4
+ ``k(x, z) = exp(-sum_k s_k^2 |x_{B_k} - z_{B_k}|^2)``. With a single block this is the RBF
5
+ kernel with ``gamma = s^2``; with one block per feature it is an ARD kernel. The training code
6
+ works on *scaled* features ``x_{B_k} * s_k`` and an RBF kernel with ``gamma = 1``.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ from collections.abc import Sequence
12
+
13
+ import numpy as np
14
+ from numpy.typing import ArrayLike, NDArray
15
+
16
+ Array = NDArray[np.float64]
17
+ Blocks = Sequence[Sequence[int]] | None
18
+
19
+
20
+ def squared_distances(x: Array, z: Array) -> Array:
21
+ """Matrix of squared Euclidean distances ``|x_i - z_j|^2`` (clipped at zero)."""
22
+ d = (x * x).sum(axis=1)[:, None] + (z * z).sum(axis=1)[None, :] - 2.0 * x @ z.T
23
+ return np.maximum(d, 0.0)
24
+
25
+
26
+ def rbf(x: Array, z: Array) -> Array:
27
+ """``exp(-|x - z|^2)`` for already scaled features."""
28
+ return np.exp(-squared_distances(x, z))
29
+
30
+
31
+ def rbf_with_gradient(x: Array, z: Array, blocks: Blocks = None) -> tuple[Array, Array]:
32
+ """Kernel matrix and its derivatives w.r.t. the log-scalings of all blocks.
33
+
34
+ Returns ``K`` of shape ``(n, m)`` and ``dK`` of shape ``(n, m, n_blocks)`` with
35
+ ``dK[..., k] = d K / d log s_k = -2 K * |x_{B_k} - z_{B_k}|^2`` (features already scaled).
36
+ """
37
+ if blocks is None:
38
+ d = squared_distances(x, z)
39
+ k = np.exp(-d)
40
+ return k, (-2.0 * k * d)[:, :, None]
41
+ parts = [squared_distances(x[:, list(b)], z[:, list(b)]) for b in blocks]
42
+ k = np.exp(-sum(parts))
43
+ return k, -2.0 * k[:, :, None] * np.stack(parts, axis=-1)
44
+
45
+
46
+ def n_blocks(blocks: Blocks) -> int:
47
+ return 1 if blocks is None else len(blocks)
48
+
49
+
50
+ def scale_features(x: ArrayLike, scaling: ArrayLike, blocks: Blocks = None) -> Array:
51
+ """Multiplies the features of each block by its scaling."""
52
+ x = np.asarray(x, dtype=np.float64)
53
+ scaling = np.atleast_1d(np.asarray(scaling, dtype=np.float64))
54
+ if blocks is None:
55
+ if scaling.size != 1:
56
+ raise ValueError("a single scaling is required without blocks")
57
+ return x * scaling[0]
58
+ if scaling.size != len(blocks):
59
+ raise ValueError(f"expected {len(blocks)} scalings, got {scaling.size}")
60
+ xs = x.copy()
61
+ for s, block in zip(scaling, blocks, strict=True):
62
+ xs[:, list(block)] *= s
63
+ return xs
@@ -0,0 +1,24 @@
1
+ """Bordered symmetric positive definite systems."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import numpy as np
6
+ from numpy.typing import NDArray
7
+ from scipy.linalg import cho_factor, cho_solve
8
+
9
+ Array = NDArray[np.float64]
10
+
11
+
12
+ def solve_bordered(q: Array, a: Array, r: Array, r0: Array) -> tuple[Array, Array]:
13
+ """Solves ``[[Q, A], [A^T, 0]] [u; v] = [r; r0]`` for SPD ``Q`` (Schur complement).
14
+
15
+ ``A`` has shape ``(n, k)`` with ``k`` small (one column per equality constraint).
16
+ """
17
+ k = a.shape[1]
18
+ if q.shape[0] == 0:
19
+ return np.zeros(0), np.zeros(k)
20
+ q_cho = cho_factor(q)
21
+ qia = cho_solve(q_cho, a)
22
+ qir = cho_solve(q_cho, r)
23
+ v = np.linalg.solve(a.T @ qia, a.T @ qir - r0)
24
+ return qir - qia @ v, v
@@ -0,0 +1,176 @@
1
+ """Loss functions of the training problem, expressed per *loss term*.
2
+
3
+ A training problem with ``n`` samples consists of ``m`` loss terms: one per sample, or two for
4
+ the epsilon-insensitive loss (one for each side of the tube). Term ``l`` belongs to sample
5
+ ``kernel_index[l]``, has weight ``w_l``, label ``y_l`` and sign ``s_l`` (used by the optional
6
+ 1-norm constraint, ``0`` if not applicable). The decision value of a term is
7
+ ``f(x_l) + s_l c``.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from dataclasses import dataclass
13
+
14
+ import numpy as np
15
+ from numpy.typing import NDArray
16
+ from scipy.special import gammaln
17
+
18
+ from .smooth_max import Degree, smooth_max
19
+
20
+ Array = NDArray[np.float64]
21
+
22
+
23
+ @dataclass(frozen=True)
24
+ class Terms:
25
+ """Term structure of a training problem."""
26
+
27
+ kernel_index: NDArray[np.intp]
28
+ y: Array
29
+ w: Array
30
+ sign: Array
31
+
32
+ @property
33
+ def size(self) -> int:
34
+ return self.y.size
35
+
36
+
37
+ class Loss:
38
+ """Base class: per-term loss ``phi_l(t)`` and its derivatives."""
39
+
40
+ #: whether the dual of this loss can be solved by ``smo_solver`` (as warm start)
41
+ smo_kind: str | None = None
42
+ #: whether the 1-norm constraint is meaningful (non-zero signs)
43
+ signed: bool = False
44
+
45
+ def terms(self, y: Array, w: Array) -> Terms:
46
+ n = y.size
47
+ return Terms(np.arange(n), y, w, np.zeros(n))
48
+
49
+ def value(self, t: Array, terms: Terms) -> Array:
50
+ raise NotImplementedError
51
+
52
+ def d1(self, t: Array, terms: Terms) -> Array:
53
+ raise NotImplementedError
54
+
55
+ def d2(self, t: Array, terms: Terms) -> Array:
56
+ raise NotImplementedError
57
+
58
+
59
+ @dataclass(frozen=True)
60
+ class HingeLoss(Loss):
61
+ """Smoothed hinge loss ``w * m_delta(shift - y t)`` for labels ``y`` in ``{-1, +1}``."""
62
+
63
+ shift: float = 1.0
64
+ smoothing: float = 0.5
65
+ degree: Degree = "2"
66
+ smo_kind = "classification"
67
+ signed = True
68
+
69
+ def terms(self, y: Array, w: Array) -> Terms:
70
+ return Terms(np.arange(y.size), y, w, y.astype(float))
71
+
72
+ def value(self, t: Array, terms: Terms) -> Array:
73
+ m = smooth_max(self.smoothing, self.degree)[0]
74
+ return terms.w * m(self.shift - terms.y * t)
75
+
76
+ def d1(self, t: Array, terms: Terms) -> Array:
77
+ dm = smooth_max(self.smoothing, self.degree)[1]
78
+ return -terms.w * terms.y * dm(self.shift - terms.y * t)
79
+
80
+ def d2(self, t: Array, terms: Terms) -> Array:
81
+ ddm = smooth_max(self.smoothing, self.degree)[2]
82
+ return terms.w * ddm(self.shift - terms.y * t)
83
+
84
+
85
+ @dataclass(frozen=True)
86
+ class EpsilonInsensitiveLoss(Loss):
87
+ """Smoothed epsilon-insensitive loss ``w * (m(y - t - eps) + m(t - y - eps))``.
88
+
89
+ Uses two terms per sample: sign ``+1`` for ``m(y - t - eps)`` and ``-1`` for the other side.
90
+ """
91
+
92
+ epsilon: float = 0.1
93
+ smoothing: float = 0.5
94
+ degree: Degree = "2"
95
+ smo_kind = "regression"
96
+ signed = True
97
+
98
+ def terms(self, y: Array, w: Array) -> Terms:
99
+ n = y.size
100
+ idx = np.arange(n)
101
+ return Terms(
102
+ np.concatenate((idx, idx)),
103
+ np.concatenate((y, y)),
104
+ np.concatenate((w, w)),
105
+ np.concatenate((np.ones(n), -np.ones(n))),
106
+ )
107
+
108
+ def _arg(self, t: Array, terms: Terms) -> Array:
109
+ return -terms.sign * (t - terms.y) - self.epsilon
110
+
111
+ def value(self, t: Array, terms: Terms) -> Array:
112
+ m = smooth_max(self.smoothing, self.degree)[0]
113
+ return terms.w * m(self._arg(t, terms))
114
+
115
+ def d1(self, t: Array, terms: Terms) -> Array:
116
+ dm = smooth_max(self.smoothing, self.degree)[1]
117
+ return -terms.sign * terms.w * dm(self._arg(t, terms))
118
+
119
+ def d2(self, t: Array, terms: Terms) -> Array:
120
+ ddm = smooth_max(self.smoothing, self.degree)[2]
121
+ return terms.w * ddm(self._arg(t, terms))
122
+
123
+
124
+ @dataclass(frozen=True)
125
+ class SquaredLoss(Loss):
126
+ """Least-squares loss ``w / 2 * (t - y)^2`` (LS-SVM)."""
127
+
128
+ def value(self, t: Array, terms: Terms) -> Array:
129
+ return 0.5 * terms.w * (t - terms.y) ** 2
130
+
131
+ def d1(self, t: Array, terms: Terms) -> Array:
132
+ return terms.w * (t - terms.y)
133
+
134
+ def d2(self, t: Array, terms: Terms) -> Array:
135
+ return terms.w * np.ones_like(t)
136
+
137
+
138
+ @dataclass(frozen=True)
139
+ class PoissonLoss(Loss):
140
+ """Poisson negative log-likelihood ``w * (exp(t) - y t + log(y!))`` for counts ``y``."""
141
+
142
+ def value(self, t: Array, terms: Terms) -> Array:
143
+ with np.errstate(over="ignore"): # inf is handled by the line search
144
+ return terms.w * (np.exp(t) - terms.y * t + gammaln(terms.y + 1))
145
+
146
+ def d1(self, t: Array, terms: Terms) -> Array:
147
+ return terms.w * (np.exp(t) - terms.y)
148
+
149
+ def d2(self, t: Array, terms: Terms) -> Array:
150
+ return terms.w * np.exp(t)
151
+
152
+
153
+ @dataclass(frozen=True)
154
+ class ApproxPoissonLoss(Loss):
155
+ """Poisson-type loss with a linear left tail (from the original experiments):
156
+
157
+ ``w * (exp(t) - (y + 1) t - 1 + log(y!))`` for ``t > 0`` and ``w * (-y t + log(y!))`` else.
158
+
159
+ It is convex and continuously differentiable; for ``t <= 0`` its curvature vanishes.
160
+ """
161
+
162
+ def value(self, t: Array, terms: Terms) -> Array:
163
+ pos = t > 0.0
164
+ with np.errstate(over="ignore"):
165
+ r = np.where(
166
+ pos, np.exp(np.where(pos, t, 0.0)) - (terms.y + 1.0) * t - 1.0, -terms.y * t
167
+ )
168
+ return terms.w * (r + gammaln(terms.y + 1))
169
+
170
+ def d1(self, t: Array, terms: Terms) -> Array:
171
+ pos = t > 0.0
172
+ return terms.w * np.where(pos, np.exp(np.where(pos, t, 0.0)) - (terms.y + 1.0), -terms.y)
173
+
174
+ def d2(self, t: Array, terms: Terms) -> Array:
175
+ pos = t > 0.0
176
+ return terms.w * np.where(pos, np.exp(np.where(pos, t, 0.0)), 0.0)
@@ -0,0 +1,99 @@
1
+ """Training problem in the dual parametrisation used by the Newton method.
2
+
3
+ For loss terms ``phi_l`` (see :mod:`optimized_svm.losses`), the primal training problem
4
+ ``min lam/2 |w|^2 + sum_l phi_l(<w, x_l> + b + s_l c) + A c`` is written in terms of
5
+ ``a`` with ``w = 1/lam sum_l a_l phi(x_l)`` (representer theorem). With
6
+ ``ka = K a / lam`` the objective reads ``1/2 a^T ka + sum_l phi_l(ka_l + b + s_l c) + A c``
7
+ (Strasdat, Ch. 6 and 8).
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import numpy as np
13
+ from numpy.typing import NDArray
14
+
15
+ from .kernels import rbf
16
+ from .losses import Loss
17
+
18
+ Array = NDArray[np.float64]
19
+ Index = NDArray[np.intp] | NDArray[np.bool_] | slice
20
+
21
+ #: kernel matrices up to this number of samples are computed once and kept in memory
22
+ PRECOMPUTE_LIMIT = 8000
23
+
24
+
25
+ class TrainingProblem:
26
+ """Training problem on (already scaled) features ``x`` with an RBF kernel (``gamma = 1``)."""
27
+
28
+ def __init__(
29
+ self,
30
+ x: Array,
31
+ y: Array,
32
+ loss: Loss,
33
+ lam: float,
34
+ weights: Array | None = None,
35
+ max_asum: float | None = None,
36
+ ) -> None:
37
+ if lam <= 0:
38
+ raise ValueError(f"lam must be positive, got {lam}")
39
+ if max_asum is not None and not loss.signed:
40
+ raise ValueError(f"max_asum is not supported for {type(loss).__name__}")
41
+ self.x = np.ascontiguousarray(x, dtype=np.float64)
42
+ self.y = np.asarray(y, dtype=np.float64)
43
+ self.w = np.ones_like(self.y) if weights is None else np.asarray(weights, dtype=np.float64)
44
+ self.loss = loss
45
+ self.lam = float(lam)
46
+ self.max_asum = max_asum
47
+ self.terms = loss.terms(self.y, self.w)
48
+ n = self.x.shape[0]
49
+ # the matrix is stored per *sample*, not per loss term
50
+ self._k = rbf(self.x, self.x) if n <= PRECOMPUTE_LIMIT else None
51
+
52
+ @property
53
+ def n_samples(self) -> int:
54
+ return self.x.shape[0]
55
+
56
+ @property
57
+ def size(self) -> int:
58
+ """Number of loss terms (= dual variables)."""
59
+ return self.terms.size
60
+
61
+ @property
62
+ def sign(self) -> Array:
63
+ return self.terms.sign
64
+
65
+ def kernel(self, rows: Index = slice(None), cols: Index = slice(None)) -> Array:
66
+ """Kernel matrix between loss terms ``rows`` and ``cols``."""
67
+ ki = self.terms.kernel_index
68
+ r, c = ki[rows], ki[cols]
69
+ if self._k is not None:
70
+ return self._k[np.ix_(r, c)]
71
+ return rbf(self.x[r], self.x[c])
72
+
73
+ def kernel_to(self, z: Array, rows: Index = slice(None)) -> Array:
74
+ """Kernel matrix between loss terms ``rows`` and new (scaled) points ``z``."""
75
+ return rbf(self.x[self.terms.kernel_index[rows]], z)
76
+
77
+ def sample_coefficients(self, a: Array) -> Array:
78
+ """Coefficients per sample ``beta_k = sum_{l: k(l) = k} a_l``."""
79
+ return np.bincount(self.terms.kernel_index, weights=a, minlength=self.n_samples)
80
+
81
+ def ka(self, a: Array) -> Array:
82
+ """``K a / lam`` per loss term (computed via the sample coefficients)."""
83
+ beta = self.sample_coefficients(a)
84
+ sv = np.flatnonzero(beta)
85
+ if self._k is not None:
86
+ per_sample = self._k[:, sv] @ beta[sv]
87
+ else:
88
+ per_sample = rbf(self.x, self.x[sv]) @ beta[sv]
89
+ return per_sample[self.terms.kernel_index] / self.lam
90
+
91
+ def decision(self, b: float, c: float, ka: Array) -> Array:
92
+ return ka + b + self.sign * c
93
+
94
+ def objective(self, a: Array, b: float, c: float, ka: Array) -> float:
95
+ dec = self.decision(b, c, ka)
96
+ obj = 0.5 * float(a @ ka) + float(self.loss.value(dec, self.terms).sum())
97
+ if self.max_asum is not None:
98
+ obj += self.max_asum * c
99
+ return obj
@@ -0,0 +1,90 @@
1
+ """Smooth approximations of the maximum function ``max(0, t)`` (Strasdat, Sec. 4.4).
2
+
3
+ Each approximation ``m_delta`` coincides with ``max(0, t)`` for ``|t| >= delta`` and is a
4
+ polynomial on ``(-delta, delta)``. The degree determines the smoothness:
5
+
6
+ * ``"2"``: ``C^1``, piecewise quadratic (its conjugate is quadratic, so SMO applies),
7
+ * ``"3"``: ``C^2``, piecewise cubic,
8
+ * ``"6"``: ``C^2``, a single sextic polynomial on ``(-delta, delta)``.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ from collections.abc import Callable
14
+ from typing import Literal
15
+
16
+ import numpy as np
17
+ from numpy.typing import NDArray
18
+
19
+ Array = NDArray[np.float64]
20
+ Degree = Literal["2", "3", "6"]
21
+ SmoothMax = tuple[Callable[[Array], Array], Callable[[Array], Array], Callable[[Array], Array]]
22
+
23
+
24
+ def _deg2(d: float) -> SmoothMax:
25
+ def m(t: Array) -> Array:
26
+ return np.where(t >= d, t, np.where(t <= -d, 0.0, 0.25 / d * (t + d) ** 2))
27
+
28
+ def dm(t: Array) -> Array:
29
+ return np.where(t >= d, 1.0, np.where(t <= -d, 0.0, 0.5 / d * (t + d)))
30
+
31
+ def ddm(t: Array) -> Array:
32
+ return np.where(np.abs(t) < d, 0.5 / d, 0.0)
33
+
34
+ return m, dm, ddm
35
+
36
+
37
+ def _deg3(d: float) -> SmoothMax:
38
+ def m(t: Array) -> Array:
39
+ inner = np.where(t <= 0.0, (d + t) ** 3 / (6 * d**2), t + (d - t) ** 3 / (6 * d**2))
40
+ return np.where(t >= d, t, np.where(t <= -d, 0.0, inner))
41
+
42
+ def dm(t: Array) -> Array:
43
+ inner = np.where(t <= 0.0, (d + t) ** 2 / (2 * d**2), 1.0 - (d - t) ** 2 / (2 * d**2))
44
+ return np.where(t >= d, 1.0, np.where(t <= -d, 0.0, inner))
45
+
46
+ def ddm(t: Array) -> Array:
47
+ return np.where(np.abs(t) < d, (d - np.abs(t)) / d**2, 0.0)
48
+
49
+ return m, dm, ddm
50
+
51
+
52
+ def _deg6(d: float) -> SmoothMax:
53
+ def m(t: Array) -> Array:
54
+ u = t / d
55
+ inner = d * (5 + u * (16 + u * (15 + u * u * (-5 + u * u)))) / 32
56
+ return np.where(t >= d, t, np.where(t <= -d, 0.0, inner))
57
+
58
+ def dm(t: Array) -> Array:
59
+ u = t / d
60
+ inner = (8 + u * (15 + u * u * (-10 + 3 * u * u))) / 16
61
+ return np.where(t >= d, 1.0, np.where(t <= -d, 0.0, inner))
62
+
63
+ def ddm(t: Array) -> Array:
64
+ u = t / d
65
+ return np.where(np.abs(t) < d, (15 + u * u * (-30 + 15 * u * u)) / (16 * d), 0.0)
66
+
67
+ return m, dm, ddm
68
+
69
+
70
+ def _exact() -> SmoothMax:
71
+ return (
72
+ lambda t: np.maximum(t, 0.0),
73
+ lambda t: (t > 0.0).astype(float),
74
+ lambda t: np.zeros_like(t, dtype=float),
75
+ )
76
+
77
+
78
+ def smooth_max(delta: float, degree: Degree = "2") -> SmoothMax:
79
+ """Returns ``(m, dm, ddm)``: the smoothed maximum and its first two derivatives.
80
+
81
+ ``delta = 0`` yields the (non-smooth) maximum itself.
82
+ """
83
+ if delta < 0:
84
+ raise ValueError(f"delta must be non-negative, got {delta}")
85
+ if delta == 0:
86
+ return _exact()
87
+ try:
88
+ return {"2": _deg2, "3": _deg3, "6": _deg6}[str(degree)](float(delta))
89
+ except KeyError:
90
+ raise ValueError(f"unknown degree {degree!r}, use '2', '3' or '6'") from None