Make it easy to estimate linear parameters with multitones (#32)

Author: fhchl

dynax/estimation.py

@@ -457,57 +457,68 @@ def H(s: complex):
 
 
 def estimate_spectra(
-    u: np.ndarray, y: np.ndarray, sr: float, nperseg: int
+    u: ArrayLike,
+    y: ArrayLike,
+    sr: float,
+    with_dc: bool = False,
+    **kwargs,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Estimate cross- and autospectral densities using Welch's method.
 
     Args:
         u: Input signal.
         y: Output signal.
         sr: Sampling rate.
-        nperseg: Number of samples per segment.
+        with_dc: Wether or not to include zero frequency term.
+        kwargs: Passed to `scipy.signal.csd`.
 
     Returns:
         Tuple `(f, S_yu, S_uu)` of frequencies, cross- and autospectral densities.
 
     """
     u_ = np.asarray(u)
     y_ = np.asarray(y)
+
     # Prep for correct broadcasting in sig.csd
     if u_.ndim == 1:
         u_ = u_[:, None]
     if y_.ndim == 1:
         y_ = y_[:, None]
-    f, S_yu = sig.csd(u_[:, None, :], y_[:, :, None], fs=sr, nperseg=nperseg, axis=0)
-    _, S_uu = sig.welch(u, fs=sr, nperseg=nperseg, axis=0)
+    f, Syu = sig.csd(u_[:, None, :], y_[:, :, None], fs=sr, **kwargs, axis=0)
+    _, Suu = sig.welch(u, fs=sr, **kwargs, axis=0)
+
     # Reshape back with dimensions of arguments
-    S_yu = S_yu.reshape((nperseg // 2 + 1,) + y.shape[1:] + u.shape[1:])
-    S_uu = S_uu.reshape((nperseg // 2 + 1,) + u.shape[1:])
-    return f, S_yu, S_uu
+    Syu = Syu.reshape((Syu.shape[0],) + y.shape[1:] + u.shape[1:])
+    Suu = Suu.reshape((Suu.shape[0],) + u.shape[1:])
+
+    if not with_dc:
+        # remove dc term
+        f = f[1:]
+        Syu = Syu[1:]
+        Suu = Suu[1:]
+
+    return f, Syu, Suu
 
 
 def fit_csd_matching(
     sys: AbstractSystem,
-    u: ArrayLike,
-    y: ArrayLike,
-    sr: float = 1.0,
-    nperseg: int = 1024,
+    f: ArrayLike,
+    Syu: ArrayLike,
+    Suu: ArrayLike,
     reg: float = 0,
     x_scale: bool = True,
     verbose_mse: bool = True,
     absolute_sigma: bool = False,
-    fit_dc: bool = False,
     linearize_kwargs: dict | None = None,
     **kwargs,
 ) -> OptimizeResult:
     """Estimate parameters of linearized system by matching cross-spectral densities.
 
     Args:
         sys: Concrete dynamical system.
-        u: Input signal.
-        y: Output signal.
-        sr: Sampling rate.
-        nperseg: Number of samples per segment.
+        f: Frequencies.
+        S_yu: Cross-spectral density.
+        S_uu: Auto-spectral density.
         reg: Weight of the :math:`L_2` regularization.
         x_scale: Whether parameters are scaled by the initial values.
         verbose_mse: Whether the cost is scaled to the mean-squared error during logging
@@ -516,7 +527,6 @@ def fit_csd_matching(
             parameter covariance reflects these absolute values. Otherwise, only
             the relative magnitudes of the sigma values matter. See also
             :func:`scipy.optimize.curve_fit`.
-        fit_dc: Whether to fit the DC term.
         linearize_kwargs: Arguments passed to
             :py:meth:`~dynax.system.AbstractSystem.linearize`.
         kwargs: Optional parameters for `scipy.optimize.least_squares`.
@@ -536,14 +546,6 @@ def fit_csd_matching(
     if linearize_kwargs is None:
         linearize_kwargs = {}
 
-    f, Syu, Suu = estimate_spectra(u, y, sr=sr, nperseg=nperseg)
-
-    if not fit_dc:
-        # remove dc term
-        f = f[1:]
-        Syu = Syu[1:]
-        Suu = Suu[1:]
-
     s = 2 * np.pi * f * 1j
     weight = 1 / np.std(Syu, axis=0)
     init_params, bounds, unravel = ravel_and_bounds(sys)

dynax/util.py

@@ -4,8 +4,9 @@
 import equinox
 import jax
 import jax.numpy as jnp
+import numpy as np
 
-from .custom_types import Array, Scalar
+from .custom_types import Array, ArrayLike, Scalar
 
 
 def value_and_jacfwd(fun: Callable, x: Array) -> tuple[Array, Callable]:
@@ -38,17 +39,17 @@ def value_and_jacrev(fun: Callable, x: Array) -> tuple[Array, Callable]:
     return y, jac
 
 
-def mse(target: Array, prediction: Array, axis: int = 0) -> Scalar:
+def mse(target: ArrayLike, prediction: ArrayLike, axis: int = 0) -> Scalar:
     """Compute mean-squared error."""
     return jnp.mean(jnp.abs(target - prediction) ** 2, axis=axis)
 
 
-def nmse(target: Array, prediction: Array, axis: int = 0) -> Scalar:
+def nmse(target: ArrayLike, prediction: ArrayLike, axis: int = 0) -> Scalar:
     """Compute normalized mean-squared error."""
     return mse(target, prediction, axis) / jnp.mean(jnp.abs(target) ** 2, axis=axis)
 
 
-def nrmse(target: Array, prediction: Array, axis: int = 0) -> Scalar:
+def nrmse(target: ArrayLike, prediction: ArrayLike, axis: int = 0) -> Scalar:
     """Compute normalized root mean-squared error."""
     return jnp.sqrt(nmse(target, prediction, axis))
 
@@ -96,3 +97,33 @@ def broadcast_right(arr, target):
 
 def dim2shape(x: int | Literal["scalar"]) -> tuple:
     return () if x == "scalar" else (x,)
+
+
+def multitone(length: int, num_tones: int, first_tone: int = 0) -> Array:
+    """Create a periodic bandpassed multitone signal.
+
+    Has flat flat spectrum and crest factor <= 2.
+
+    From: S. Boyd, "Multitone Signals with Low Crest Factor"
+          IEEE Transactions on Circuits and Systems, CAS-33(10):1018-1022, October 1986
+    """
+    if np.log2(num_tones) % 1 != 0:
+        raise ValueError("Number of tones must be power of 2")
+
+    def rudin_sign(k: int) -> int:
+        previous_bit, sign = 1, 1
+        k = k - 1
+        while k > 0:
+            previous_bit = k % 2
+            if (k := k // 2) == 0:
+                break
+            if previous_bit == 1 and k % 2 == 1:
+                sign = -sign
+        return sign
+
+    t = np.linspace(0, 2 * np.pi, length, endpoint=False)
+    u = 0
+    for k in range(1, num_tones + 1):
+        u += rudin_sign(k) * np.sin((k + first_tone) * t)
+    u *= np.sqrt(2 / num_tones)
+    return u

examples/fit_long_input.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 from dynax import fit_csd_matching, fit_least_squares, Flow, pretty
+from dynax.estimation import estimate_spectra
 from dynax.example_models import NonlinearDrag
 
 
@@ -27,9 +28,8 @@
 # If we have long-duration, wide-band input data we can fit the linear
 # parameters first by matching the transfer-functions. In this example the result is
 # not very good.
-initial_sys = fit_csd_matching(
-    initial_sys, u_train, y_train, samplerate, nperseg=500
-).result
+f, Syu, Suu = estimate_spectra(u_train, y_train, samplerate, nperseg=512)
+initial_sys = fit_csd_matching(initial_sys, f, Syu, Suu).result
 print("linear params fitted:", pretty(initial_sys))
 
 # Combine the fitted ODE with an ODE solver

tests/test_estimation.py

@@ -17,7 +17,9 @@
     non_negative_field,
     transfer_function,
 )
+from dynax.estimation import estimate_spectra
 from dynax.example_models import LotkaVolterra, NonlinearDrag, SpringMassDamper
+from dynax.util import multitone
 
 
 tols = {"rtol": 1e-02, "atol": 1e-04}
@@ -222,9 +224,10 @@ def test_csd_matching():
     # output
     _, y = model(t, u)
     # fit
-    init_sys = SpringMassDamper(1.0, 1.0, 1.0)
-    fitted_sys = fit_csd_matching(init_sys, u, y, sr, nperseg=1024, verbose=1).result
+    f, Syu, Suu = estimate_spectra(u, y, sr, nperseg=1024)
+    fitted_sys = fit_csd_matching(sys, f, Syu, Suu, verbose=1).result
 
+    # poor fit due to random input signal that requires windowing
     assert eqx.tree_equal(
         fitted_sys,
         sys,
@@ -233,6 +236,37 @@ def test_csd_matching():
     )
 
 
+def test_csd_matching_multitone():
+    np.random.seed(123)
+    # model
+    sys = SpringMassDamper(1.0, 1.0, 1.0)
+    model = Flow(sys, stepsize_controller=PIDController(rtol=1e-6, atol=1e-6))
+    # input
+    nperseg = 1024
+    segs = 10
+    num_total_samples = nperseg * segs
+    sr = 50
+    t = np.arange(num_total_samples) / sr
+    u = np.tile(multitone(nperseg, num_tones=nperseg // 2), segs)
+    # output
+    _, y = model(t, u)
+    # fit
+    warmup_seps = 5
+    warmup = warmup_seps * nperseg
+    f, Syu, Suu = estimate_spectra(
+        u[warmup:], y[warmup:], sr, nperseg=nperseg, window="boxcar", noverlap=0
+    )
+    fitted_sys = fit_csd_matching(sys, f, Syu, Suu, verbose=1).result
+
+    # good fit with much less data!
+    assert eqx.tree_equal(
+        fitted_sys,
+        sys,
+        rtol=1e-2,
+        atol=1e-2,
+    )
+
+
 def test_estimate_initial_state():
     class NonlinearDragFreeInitialState(NonlinearDrag):
         initial_state: Array = field(init=False)

tests/test_util.py

@@ -0,0 +1,17 @@
+import numpy as np
+
+from dynax.util import multitone
+
+
+def test_multitone():
+    def crest_factor(sig):
+        return (
+            np.linalg.norm(sig, np.inf) / (np.linalg.norm(sig, 2) / np.sqrt(len(sig)))
+        ).item()
+
+    length = 1000
+    for num_tones in np.pow(2, np.arange(8)):
+        for first_tone in [0, 1, 10]:
+            u = multitone(length, num_tones, first_tone)
+            cr = crest_factor(u)
+            assert cr < 2 or np.allclose(cr, 2)