Feat/kalman with fitting (#692)

* First proof of concept

* Clean up and add checks

* Remove support for multiple timeseries

* Add tests and fixes

* Add test given KF and fixes

* Add docstrings

* Add and adapt example

* Improve covariates (error) handling

* Remove unused imports

Co-authored-by: Dennis Bader <[email protected]>

* Improve docstring line length

Co-authored-by: Dennis Bader <[email protected]>

* Fix typos

* Remove filterpy from requirements

* Remove unnecessary parentheses

* Add num_samples to docstring

* Add num_block_rows to docstring

* Clarify covariance formula

* Update to nfoursid with relaxed requirements

* Bump nfoursid requirement

Co-authored-by: Dennis Bader <[email protected]>

Author: tomasvanpottelbergh

darts/models/filtering/filtering_model.py

@@ -8,7 +8,7 @@
 from abc import ABC, abstractmethod
 
 from darts.timeseries import TimeSeries
-from darts.logging import get_logger
+from darts.logging import get_logger, raise_if_not
 
 logger = get_logger(__name__)
 
@@ -35,4 +35,5 @@ def filter(self, series: TimeSeries) -> TimeSeries:
         TimeSeries
             A time series containing the filtered values.
         """
-        pass
+        raise_if_not(series.is_deterministic, 'The input series must be '
+                                              'deterministic (observations).')

darts/models/filtering/gaussian_process_filter.py

@@ -59,8 +59,6 @@ def filter(self,
             A stochastic ``TimeSeries`` sampled from the Gaussian Process, or its mean
             if `num_samples` is set to 1.
         """
-        raise_if_not(series.is_deterministic, 'The input series for the Gaussian Process filter must be '
-                                              'deterministic (observations).')
         super().filter(series)
 
         values = series.values(copy=False)

darts/models/filtering/kalman_filter.py

@@ -4,150 +4,204 @@
 """
 
 from abc import ABC
-
-from typing import Optional
-from filterpy.kalman import KalmanFilter as FpKalmanFilter
 from copy import deepcopy
+from typing import Optional
+
 import numpy as np
+import pandas as pd
+from nfoursid.kalman import Kalman
+from nfoursid.nfoursid import NFourSID
 
 from darts.models.filtering.filtering_model import FilteringModel
 from darts.timeseries import TimeSeries
-from darts.utils.utils import raise_if_not
+from darts.logging import raise_if, raise_if_not
 
 
 class KalmanFilter(FilteringModel, ABC):
     def __init__(
             self, 
             dim_x: int = 1,
-            x_init: Optional[np.array] = None,
-            P: Optional[np.array] = None,
-            Q: Optional[np.array] = None,
-            R: Optional[np.array] = None,
-            H: Optional[np.array] = None,
-            F: Optional[np.array] = None,
-            kf: Optional[FpKalmanFilter] = None
+            kf: Optional[Kalman] = None
             ):
         """
-        This model implements a Kalman filter over a time series (without control signal).
+        This model implements a Kalman filter over a time series.
 
         The key method is `KalmanFilter.filter()`.
         It considers the provided time series as containing (possibly noisy) observations z obtained from a
         (possibly noisy) linear dynamical system with hidden state x. The function `filter(series)` returns a new
-        `TimeSeries` describing the distribution of the state x, as inferred by the Kalman filter from
-        sequentially observing z from `series`.
-        Depending on the use case, this can be used to de-noise a series or infer the underlying hidden state of the
-        data generating process (assuming notably that the dynamical system generating the data is known, as captured
-        by the `F` matrix.).
+        `TimeSeries` describing the distribution of the output z (without noise), as inferred by the Kalman filter from
+        sequentially observing z from `series`, and the dynamics of the linear system of order dim_x.
 
-        This implementation wraps around filterpy.kalman.KalmanFilter, so more information the parameters can be found
-        here: https://filterpy.readthedocs.io/en/latest/kalman/KalmanFilter.html
+        The method `KalmanFilter.fit()` is used to initialize the Kalman filter by estimating the state space model of 
+        a linear dynamical system and the covariance matrices of the process and measurement noise using the N4SID 
+        algorithm.
 
-        The dimensionality of the measurements z is automatically inferred upon calling `filter()`.
-        This implementation doesn't include control signal.
+        This implementation uses Kalman from the NFourSID package. More information can be found here:
+        https://nfoursid.readthedocs.io/en/latest/source/kalman.html.
+
+        The dimensionality of the measurements z and optional control signal (covariates) u is automatically inferred upon
+        calling `filter()`.
 
         Parameters
         ----------
         dim_x : int
-            Size of the Kalman filter state vector. It determines the dimensionality of the `TimeSeries`
-            returned by the `filter()` function.
-        x_init : ndarray (dim_x, 1), default: [0, 0, ..., 0]
-            Initial state; will be updated at each time step.
-        P : ndarray (dim_x, dim_x), default: identity matrix
-            initial covariance matrix; will be update at each time step
-        Q : ndarray (dim_x, dim_x), default: identity matrix
-            Process noise covariance matrix
-        R : ndarray (dim_z, dim_z), default: identity matrix
-            Measurement noise covariance matrix. `dim_z` must match the dimensionality (width) of the `TimeSeries`
-            used with `filter()`.
-        H : ndarray (dim_z, dim_x), default: all-ones matrix
-            measurement function; describes how the measurement z is obtained from the state vector x
-        F : ndarray (dim_x, dim_x), default: identity matrix
-            State transition matrix; describes how the state evolves from one time step to the next
-            in the underlying dynamical system.
-        kf : filterpy.kalman.KalmanFilter
-            Optionally, an instance of `filterpy.kalman.KalmanFilter`.
-            If this is provided, the other parameters are ignored. This instance will be copied for every
+            Size of the Kalman filter state vector.
+        kf : nfoursid.kalman.Kalman
+            Optionally, an instance of `nfoursid.kalman.Kalman`.
+            If this is provided, the parameter dim_x is ignored. This instance will be copied for every
             call to `filter()`, so the state is not carried over from one time series to another across several
             calls to `filter()`.
-            The various dimensionality in the filter must match those in the `TimeSeries` used when calling `filter()`.
+            The various dimensionalities of the filter must match those of the `TimeSeries` used when calling `filter()`.
         """
+        # TODO: Add support for x_init. Needs reimplementation of NFourSID.
+
         super().__init__()
+        self._expect_covariates = False
+
         if kf is None:
-            self.dim_x = dim_x
-            self.x_init = x_init if x_init is not None else np.zeros(self.dim_x,)
-            self.P = P if P is not None else np.eye(self.dim_x)
-            self.Q = Q if Q is not None else np.eye(self.dim_x)
-            self.R = R
-            self.H = H
-            self.F = F if F is not None else np.eye(self.dim_x)
             self.kf = None
-            self.kf_provided = False
+            self.dim_x = dim_x
+            self._kf_provided = False
         else:
             self.kf = kf
-            self.kf_provided = True
+            self.dim_u = kf.state_space.u_dim
+            self.dim_x = kf.state_space.x_dim
+            self.dim_y = kf.state_space.y_dim
+            self._kf_provided = True
+            if self.dim_u > 0:
+                self._expect_covariates = True
 
     def __str__(self):
         return 'KalmanFilter(dim_x={})'.format(self.dim_x)
 
+    def fit(self,
+            series: TimeSeries,
+            covariates: Optional[TimeSeries] = None,
+            num_block_rows: Optional[int] = None) -> None:
+        """
+        Initializes the Kalman filter using the N4SID algorithm.
+
+        Parameters
+        ----------
+        series : TimeSeries
+            The series of outputs (observations) used to infer the underlying state space model.
+            This must be a deterministic series (containing one sample).
+        covariates : Optional[TimeSeries]
+            An optional series of inputs (control signal) that will also be used to infer the underlying state space model.
+            This must be a deterministic series (containing one sample).
+        num_block_rows : Optional[int]
+            The number of block rows to use in the block Hankel matrices used in the N4SID algorithm. 
+            See the documentation of nfoursid.nfoursid.NFourSID for more information.
+            If not provided, the dimensionality of the state space model will be used, with a maximum of 10.
+        """
+        if covariates is not None:
+            self._expect_covariates = True
+            covariates = covariates.slice_intersect(series)
+            raise_if_not(series.has_same_time_as(covariates),
+                         'The number of timesteps in the series and the covariates must match.')
+        
+        # TODO: Handle multiple timeseries. Needs reimplementation of NFourSID?
+        self.dim_y = series.width
+        outputs = series.pd_dataframe()
+        outputs.columns = [f'y_{i}' for i in outputs.columns]
+
+        if covariates is not None:
+            self.dim_u = covariates.width
+            inputs = covariates.pd_dataframe()
+            inputs.columns = [f'u_{i}' for i in inputs.columns]
+            input_columns = list(inputs.columns)
+            measurements = pd.concat([outputs, inputs], axis=1)
+        else:
+            measurements = outputs
+            input_columns = None
+
+        if num_block_rows is None:
+            num_block_rows = max(10, self.dim_x)
+        nfoursid = NFourSID(measurements,
+                            output_columns=list(outputs.columns),
+                            input_columns=input_columns,
+                            num_block_rows=num_block_rows)
+        nfoursid.subspace_identification()
+        state_space_identified, covariance_matrix = nfoursid.system_identification(
+            rank=self.dim_x
+        )
+
+        self.kf = Kalman(state_space_identified, covariance_matrix)
+
+
     def filter(self,
                series: TimeSeries,
-               num_samples: int = 1):
+               covariates: Optional[TimeSeries] = None,
+               num_samples: int = 1) -> TimeSeries:
         """
         Sequentially applies the Kalman filter on the provided series of observations.
 
         Parameters
         ----------
         series : TimeSeries
-            The series of observations used to infer the state values according to the specified Kalman process.
+            The series of outputs (observations) used to infer the underlying outputs according to the specified Kalman process.
+            This must be a deterministic series (containing one sample).
+        covariates : Optional[TimeSeries]
+            An optional series of inputs (control signal), necessary if the Kalman filter was initialized with covariates.
             This must be a deterministic series (containing one sample).
+        num_samples : int, default: 1
+            The number of samples to generate from the inferred distribution of the output z. If this is set to 1, the
+            output is a `TimeSeries` containing a single sample using the mean of the distribution.
 
         Returns
         -------
         TimeSeries
-            A stochastic `TimeSeries` of state values, of dimension `dim_x`.
+            A (stochastic) `TimeSeries` of the inferred output z, of the same width as the input series.
         """
+        super().filter(series)
 
-        raise_if_not(series.is_deterministic, 'The input series for the Kalman filter must be '
-                                              'deterministic (observations).')
+        raise_if(self.kf is None, 'The Kalman filter has not been fitted yet. Call `fit()` first '
+                                  'or provide Kalman filter in constructor.')
+                                  
+        raise_if_not(series.width == self.dim_y, 'The provided TimeSeries dimensionality does not match '
+                                                 'the output dimensionality of the Kalman filter.')
 
-        dim_z = series.width
+        raise_if(covariates is not None and not self._expect_covariates,
+                 'Covariates were provided, but the Kalman filter was not fitted with covariates.')
 
-        if not self.kf_provided:
-            kf = FpKalmanFilter(dim_x=self.dim_x, dim_z=dim_z)
-            kf.x = self.x_init
-            kf.P = self.P
-            kf.Q = self.Q
-            kf.R = self.R if self.R is not None else np.eye(dim_z)
-            kf.H = self.H if self.H is not None else np.ones((dim_z, self.dim_x))
-            kf.F = self.F
-        else:
-            raise_if_not(dim_z == self.kf.dim_z, 'The provided TimeSeries dimensionality does not match '
-                                                 'the filter observation dimensionality dim_z.')
-            kf = deepcopy(self.kf)
+        if self._expect_covariates:
+            raise_if(covariates is None,
+                     'The Kalman filter was fitted with covariates, but these were not provided.')
 
-        super().filter(series)
-        values = series.values(copy=False)
+            raise_if_not(covariates.is_deterministic,
+                         'The covariates must be deterministic (observations).')
+
+            covariates = covariates.slice_intersect(series)
+            raise_if_not(series.has_same_time_as(covariates),
+                         'The number of timesteps in the series and the covariates must match.')
 
+        kf = deepcopy(self.kf)
+
+        y_values = series.values(copy=False)
+        if self._expect_covariates:
+            u_values = covariates.values(copy=False)
+        else:
+            u_values = np.zeros((len(y_values), 0))
+        
         # For each time step, we'll sample "n_samples" from a multivariate Gaussian
         # whose mean vector and covariance matrix come from the Kalman filter.
         if num_samples == 1:
-            sampled_states = np.zeros(((len(values)), self.dim_x, ))
+            sampled_states = np.zeros((len(y_values), self.dim_y, ))
         else:
-            sampled_states = np.zeros(((len(values)), self.dim_x, num_samples))
+            sampled_states = np.zeros((len(y_values), self.dim_y, num_samples))
 
-        # process_means = np.zeros((len(values), self.dim_x))  # mean values
-        # process_covariances = ...                            # covariance matrices; TODO
-        for i in range(len(values)):
-            obs = values[i, :]
-            kf.predict()
-            kf.update(obs)
-            mean_vec = kf.x.reshape(self.dim_x,)
+        for i in range(len(y_values)):
+            y = y_values[i, :].reshape(-1, 1)
+            u = u_values[i, :].reshape(-1, 1)
+            kf.step(y, u)
+            mean_vec = kf.y_filtereds[-1].reshape(self.dim_y,)
 
             if num_samples == 1:
-                # It's actually not sampled in this case
                 sampled_states[i, :] = mean_vec
             else:
-                cov_matrix = kf.P
+                # The measurement covariance matrix is given by the sum of the covariance matrix of the
+                # state estimate (transformed by C) and the covariance matrix of the measurement noise.
+                cov_matrix = kf.state_space.c @ kf.p_filtereds[-1] @ kf.state_space.c.T + kf.r
                 sampled_states[i, :, :] = np.random.multivariate_normal(mean_vec, cov_matrix, size=num_samples).T
 
         # TODO: later on for a forecasting model we'll have to do something like