Online linear regression using Kalman filtering#

We perform sequential (recursive) Bayesian inference for the parameters of a linear regression model using the Kalman filter. (This algorithm is also known as recursive least squares.) To do this, we treat the parameters of the model as the unknown hidden states. We assume that these are constant over time. The graphical model is shown below.

RLS

The model has the following form

\[\begin{align*} \theta_t &= \theta_{t-1} \\ y_t &= x_t^T \theta_t + r_t, \; r_t \sim N(0, \sigma^2) \end{align*}\]

This is a LG-SSM, where \(F=I\), \(Q=0\), \(H_t = x_t^T\) and \(R = \sigma^2\).

Setup#

from jax import numpy as jnp
from matplotlib import pyplot as plt
from dynamax.linear_gaussian_ssm import LinearGaussianSSM

Data#

Data is from https://github.com/probml/pmtk3/blob/master/demos/linregOnlineDemoKalman.m

n_obs = 21
x = jnp.linspace(0, 20, n_obs)
X = jnp.column_stack((jnp.ones_like(x), x))  # Design matrix.
y = jnp.array(
    [2.486, -0.303, -4.053, -4.336, -6.174, -5.604, -3.507, -2.326, -4.638, -0.233, -1.986, 1.028, -2.264,
     -0.451, 1.167, 6.652, 4.145, 5.268, 6.34, 9.626, 14.784])

Model#

F = jnp.eye(2)
Q = jnp.zeros((2, 2))  # No parameter drift.
obs_var = 1.0
R = jnp.ones((1, 1)) * obs_var
mu0 = jnp.zeros(2)
Sigma0 = jnp.eye(2) * 10.0


# the input_dim = 0 since we encode the covariates into the non-stationary emission matrix
lgssm = LinearGaussianSSM(state_dim = 2, emission_dim = 1, input_dim = 0)
params, _ = lgssm.initialize(
    initial_mean=mu0,
    initial_covariance=Sigma0,
    dynamics_weights=F,
    dynamics_covariance=Q,
    emission_weights=X[:, None, :], # (t, 1, D) where D = num input features
    emission_covariance=R,
    )

Online inference#

lgssm_posterior = lgssm.filter(params, y[:, None]) # reshape y to be (T,1)
kf_results = (lgssm_posterior.filtered_means, lgssm_posterior.filtered_covariances)

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[5], line 1
----> 1 lgssm_posterior = lgssm.filter(params, y[:, None]) # reshape y to be (T,1)
kf_results = (lgssm_posterior.filtered_means, lgssm_posterior.filtered_covariances)

File ~/work/dynamax/dynamax/dynamax/linear_gaussian_ssm/models.py:226, in LinearGaussianSSM.filter(self, params, emissions, inputs)
def filter(
   self,
   params: ParamsLGSSM,
   emissions: Float[Array, "ntime emission_dim"],
   inputs: Optional[Float[Array, "ntime input_dim"]] = None
) -> PosteriorGSSMFiltered:
--> 226     return lgssm_filter(params, emissions, inputs)

File ~/work/dynamax/dynamax/dynamax/linear_gaussian_ssm/inference.py:360, in preprocess_args.<locals>.wrapper(*args, **kwargs)
num_timesteps = len(emissions)
full_params, inputs = preprocess_params_and_inputs(params, num_timesteps, inputs)
--> 360 return f(full_params, emissions, inputs=inputs)

File ~/work/dynamax/dynamax/dynamax/linear_gaussian_ssm/inference.py:488, in lgssm_filter(params, emissions, inputs)
# Run the Kalman filter
carry = (0.0, params.initial.mean, params.initial.cov)
--> 488 (ll, _, _), (filtered_means, filtered_covs) = lax.scan(_step, carry, jnp.arange(num_timesteps))
return PosteriorGSSMFiltered(marginal_loglik=ll, filtered_means=filtered_means, filtered_covariances=filtered_covs)

    [... skipping hidden 9 frame]

File ~/work/dynamax/dynamax/dynamax/linear_gaussian_ssm/inference.py:476, in lgssm_filter.<locals>._step(carry, t)
y = emissions[t]
# Update the log likelihood
--> 476 ll += _log_likelihood(pred_mean, pred_cov, H, D, d, R, u, y)
# Condition on this emission
filtered_mean, filtered_cov = _condition_on(pred_mean, pred_cov, H, D, d, R, u, y)

File ~/work/dynamax/dynamax/dynamax/linear_gaussian_ssm/inference.py:461, in lgssm_filter.<locals>._log_likelihood(pred_mean, pred_cov, H, D, d, R, u, y)
if R.ndim==2:
   S = R + H @ pred_cov @ H.T
--> 461     return MVN(m, S).log_prob(y)
else:
   L = H @ jnp.linalg.cholesky(pred_cov)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/decorator.py:232, in decorate.<locals>.fun(*args, **kw)
if not kwsyntax:
   args, kw = fix(args, kw, sig)
--> 232 return caller(func, *(extras + args), **kw)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py:342, in _DistributionMeta.__new__.<locals>.wrapped_init(***failed resolving arguments***)
# Note: if we ever want to have things set in `self` before `__init__` is
# called, here is the place to do it.
self_._parameters = None
--> 342 default_init(self_, *args, **kwargs)
# Note: if we ever want to override things set in `self` by subclass
# `__init__`, here is the place to do it.
if self_._parameters is None:
 # We prefer subclasses will set `parameters = dict(locals())` because
 # this has nearly zero overhead. However, failing to do this, we will
 # resolve the input arguments dynamically and only when needed.

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/mvn_full_covariance.py:191, in MultivariateNormalFullCovariance.__init__(self, loc, covariance_matrix, validate_args, allow_nan_stats, name)
     # No need to validate that covariance_matrix is non-singular.
     # LinearOperatorLowerTriangular has an assert_non_singular method that
     # is called by the Bijector.
     # However, cholesky() ignores the upper triangular part, so we do need
     # to separately assert symmetric.
     scale_tril = tf.linalg.cholesky(covariance_matrix)
--> 191     super(MultivariateNormalFullCovariance, self).__init__(
       loc=loc,
       scale_tril=scale_tril,
       validate_args=validate_args,
       allow_nan_stats=allow_nan_stats,
       name=name)
self._parameters = parameters

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/decorator.py:232, in decorate.<locals>.fun(*args, **kw)
if not kwsyntax:
   args, kw = fix(args, kw, sig)
--> 232 return caller(func, *(extras + args), **kw)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py:342, in _DistributionMeta.__new__.<locals>.wrapped_init(***failed resolving arguments***)
# Note: if we ever want to have things set in `self` before `__init__` is
# called, here is the place to do it.
self_._parameters = None
--> 342 default_init(self_, *args, **kwargs)
# Note: if we ever want to override things set in `self` by subclass
# `__init__`, here is the place to do it.
if self_._parameters is None:
 # We prefer subclasses will set `parameters = dict(locals())` because
 # this has nearly zero overhead. However, failing to do this, we will
 # resolve the input arguments dynamically and only when needed.

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/mvn_tril.py:228, in MultivariateNormalTriL.__init__(self, loc, scale_tril, validate_args, allow_nan_stats, experimental_use_kahan_sum, name)
 linop_cls = (KahanLogDetLinOpTriL if experimental_use_kahan_sum else
              tf.linalg.LinearOperatorLowerTriangular)
 scale = linop_cls(
     scale_tril,
     is_non_singular=True,
     is_self_adjoint=False,
     is_positive_definite=False)
--> 228 super(MultivariateNormalTriL, self).__init__(
   loc=loc,
   scale=scale,
   validate_args=validate_args,
   allow_nan_stats=allow_nan_stats,
   experimental_use_kahan_sum=experimental_use_kahan_sum,
   name=name)
self._parameters = parameters

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/decorator.py:232, in decorate.<locals>.fun(*args, **kw)
if not kwsyntax:
   args, kw = fix(args, kw, sig)
--> 232 return caller(func, *(extras + args), **kw)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py:342, in _DistributionMeta.__new__.<locals>.wrapped_init(***failed resolving arguments***)
# Note: if we ever want to have things set in `self` before `__init__` is
# called, here is the place to do it.
self_._parameters = None
--> 342 default_init(self_, *args, **kwargs)
# Note: if we ever want to override things set in `self` by subclass
# `__init__`, here is the place to do it.
if self_._parameters is None:
 # We prefer subclasses will set `parameters = dict(locals())` because
 # this has nearly zero overhead. However, failing to do this, we will
 # resolve the input arguments dynamically and only when needed.

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/mvn_linear_operator.py:205, in MultivariateNormalLinearOperator.__init__(self, loc, scale, validate_args, allow_nan_stats, experimental_use_kahan_sum, name)
if loc is not None:
 bijector = shift_bijector.Shift(
     shift=loc, validate_args=validate_args)(bijector)
--> 205 super(MultivariateNormalLinearOperator, self).__init__(
   # TODO(b/137665504): Use batch-adding meta-distribution to set the batch
   # shape instead of tf.zeros.
   # We use `Sample` instead of `Independent` because `Independent`
   # requires concatenating `batch_shape` and `event_shape`, which loses
   # static `batch_shape` information when `event_shape` is not statically
   # known.
   distribution=sample.Sample(
       normal.Normal(
           loc=tf.zeros(batch_shape, dtype=dtype),
           scale=tf.ones([], dtype=dtype)),
       event_shape,
       experimental_use_kahan_sum=experimental_use_kahan_sum),
   bijector=bijector,
   validate_args=validate_args,
   name=name)
self._parameters = parameters

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/decorator.py:232, in decorate.<locals>.fun(*args, **kw)
if not kwsyntax:
   args, kw = fix(args, kw, sig)
--> 232 return caller(func, *(extras + args), **kw)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py:342, in _DistributionMeta.__new__.<locals>.wrapped_init(***failed resolving arguments***)
# Note: if we ever want to have things set in `self` before `__init__` is
# called, here is the place to do it.
self_._parameters = None
--> 342 default_init(self_, *args, **kwargs)
# Note: if we ever want to override things set in `self` by subclass
# `__init__`, here is the place to do it.
if self_._parameters is None:
 # We prefer subclasses will set `parameters = dict(locals())` because
 # this has nearly zero overhead. However, failing to do this, we will
 # resolve the input arguments dynamically and only when needed.

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/transformed_distribution.py:244, in _TransformedDistribution.__init__(self, distribution, bijector, kwargs_split_fn, validate_args, parameters, name)
self._zero = tf.constant(0, dtype=tf.int32, name='zero')
# We don't just want to check isinstance(JointDistribution) because
# TransformedDistributions with multipart bijectors are effectively
# joint but don't inherit from JD. The 'duck-type' test is that
# JDs have a structured dtype.
--> 244 dtype = self.bijector.forward_dtype(self.distribution.dtype)
self._is_joint = tf.nest.is_nested(dtype)
super(_TransformedDistribution, self).__init__(
   dtype=dtype,
   reparameterization_type=self._distribution.reparameterization_type,
   (...)
   parameters=parameters,
   name=name)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/bijectors/bijector.py:1705, in Bijector.forward_dtype(self, dtype, name, **kwargs)
 input_dtype = nest_util.broadcast_structure(
     self.forward_min_event_ndims, self.dtype)
else:
 # Make sure inputs are compatible with statically-known dtype.
-> 1705   input_dtype = nest.map_structure_up_to(
     self.forward_min_event_ndims,
     lambda x: dtype_util.convert_to_dtype(x, dtype=self.dtype),
     nest_util.coerce_structure(self.forward_min_event_ndims, dtype),
     check_types=False)
output_dtype = self._forward_dtype(input_dtype, **kwargs)
try:
 # kwargs may alter dtypes themselves, but we currently require
 # structure to be statically known.

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/python/internal/backend/jax/nest.py:324, in map_structure_up_to(shallow_structure, func, *structures, **kwargs)
def map_structure_up_to(shallow_structure, func, *structures, **kwargs):
--> 324   return map_structure_with_tuple_paths_up_to(
     shallow_structure,
     lambda _, *args: func(*args),  # Discards path.
     *structures,
     **kwargs)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/python/internal/backend/jax/nest.py:353, in map_structure_with_tuple_paths_up_to(shallow_structure, func, expand_composites, *structures, **kwargs)
for input_tree in structures:
 assert_shallow_structure(
     shallow_structure, input_tree, check_types=check_types)
--> 353 return dm_tree.map_structure_with_path_up_to(
   shallow_structure, func, *structures, **kwargs)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tree/__init__.py:778, in map_structure_with_path_up_to(***failed resolving arguments***)
results = []
for path_and_values in _multiyield_flat_up_to(shallow_structure, *structures):
--> 778   results.append(func(*path_and_values))
return unflatten_as(shallow_structure, results)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/python/internal/backend/jax/nest.py:326, in map_structure_up_to.<locals>.<lambda>(_, *args)
def map_structure_up_to(shallow_structure, func, *structures, **kwargs):
 return map_structure_with_tuple_paths_up_to(
     shallow_structure,
--> 326       lambda _, *args: func(*args),  # Discards path.
     *structures,
     **kwargs)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/bijectors/bijector.py:1707, in Bijector.forward_dtype.<locals>.<lambda>(x)
 input_dtype = nest_util.broadcast_structure(
     self.forward_min_event_ndims, self.dtype)
else:
 # Make sure inputs are compatible with statically-known dtype.
 input_dtype = nest.map_structure_up_to(
     self.forward_min_event_ndims,
-> 1707       lambda x: dtype_util.convert_to_dtype(x, dtype=self.dtype),
     nest_util.coerce_structure(self.forward_min_event_ndims, dtype),
     check_types=False)
output_dtype = self._forward_dtype(input_dtype, **kwargs)
try:
 # kwargs may alter dtypes themselves, but we currently require
 # structure to be statically known.

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/internal/dtype_util.py:247, in convert_to_dtype(tensor_or_dtype, dtype, dtype_hint)
elif isinstance(tensor_or_dtype, np.ndarray):
 dt = base_dtype(dtype or dtype_hint or tensor_or_dtype.dtype)
--> 247 elif np.issctype(tensor_or_dtype):
 dt = base_dtype(dtype or dtype_hint or tensor_or_dtype)
else:
 # If this is a Python object, call `convert_to_tensor` and grab the dtype.
 # Note that this will add ops in graph-mode; we may want to consider
 # other ways to handle this case.

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/numpy/__init__.py:397, in __getattr__(attr)
   raise AttributeError(__former_attrs__[attr])
if attr in __expired_attributes__:
--> 397     raise AttributeError(
       f"`np.{attr}` was removed in the NumPy 2.0 release. "
       f"{__expired_attributes__[attr]}"
   )
if attr == "chararray":
   warnings.warn(
       "`np.chararray` is deprecated and will be removed from "
       "the main namespace in the future. Use an array with a string "
       "or bytes dtype instead.", DeprecationWarning, stacklevel=2)

AttributeError: `np.issctype` was removed in the NumPy 2.0 release. Use `issubclass(rep, np.generic)` instead.

Offline inference#

We compute the offline posterior given all the data using Bayes rule for linear regression. This should give the same results as the final step of online inference.

posterior_prec = jnp.linalg.inv(Sigma0) + X.T @ X / obs_var
b = jnp.linalg.inv(Sigma0) @ mu0 + X.T @ y / obs_var
posterior_mean = jnp.linalg.solve(posterior_prec, b)
batch_results = (posterior_mean, posterior_prec)

Plot results#

# Unpack kalman filter results
post_weights_kf, post_sigma_kf = kf_results
w0_kf_hist, w1_kf_hist = post_weights_kf.T
w0_kf_err, w1_kf_err = jnp.sqrt(post_sigma_kf[:, [0, 1], [0, 1]].T)

# Unpack batch results
post_weights_batch, post_prec_batch = batch_results
w0_post_batch, w1_post_batch = post_weights_batch
Sigma_post_batch = jnp.linalg.inv(post_prec_batch)
w0_std_batch, w1_std_batch = jnp.sqrt(Sigma_post_batch[[0, 1], [0, 1]])

fig, ax = plt.subplots()
timesteps = jnp.arange(len(w0_kf_hist))

# Plot online kalman filter posterior.
ax.errorbar(timesteps, w0_kf_hist, w0_kf_err, fmt="-o", label="$w_0$", color="black", fillstyle="none")
ax.errorbar(timesteps, w1_kf_hist, w1_kf_err, fmt="-o", label="$w_1$", color="tab:red")

# Plot batch posterior.
ax.hlines(y=w0_post_batch, xmin=timesteps[0], xmax=timesteps[-1], color="black", label="$w_0$ batch")
ax.hlines(
    y=w1_post_batch, xmin=timesteps[0], xmax=timesteps[-1], color="tab:red", linestyle="--", label="$w_1$ batch"
)
ax.fill_between(timesteps, w0_post_batch - w0_std_batch, w0_post_batch + w0_std_batch, color="black", alpha=0.4)
ax.fill_between(timesteps, w1_post_batch - w1_std_batch, w1_post_batch + w1_std_batch, color="tab:red", alpha=0.4)

ax.set_xlabel("time")
ax.set_ylabel("weights")
ax.legend()

<matplotlib.legend.Legend at 0x7ff4ba3527f0>

../../_images/db54f96b1c55f1a98b577c924ea584a23a41100dc800d8618923b7842fa2c075.png

Online linear regression using Kalman filtering

Contents