Source code for probnum.randprocs.markov.integrator._preconditioner

"""Coordinate changes in state space models."""

import abc
from functools import cached_property

import numpy as np
import scipy.special  # for vectorised factorial

from probnum import config, linops, randvars

def apply_precon(precon, rv):
    # public (because it is needed in some integrator implementations),
    # but not exposed to the 'randprocs' namespace
    # (i.e. not imported in any

    # There is no way of checking whether `rv` has its Cholesky factor computed already
    # or not.
    # Therefore, since we need to update the Cholesky factor for square-root filtering,
    # we also update the Cholesky factor for non-square-root algorithms here,
    # which implies additional cost.
    # See Issues #319 and #329.
    # When they are resolved, this function here will hopefully be superfluous.

    new_mean = precon @ rv.mean
    new_cov_cholesky = precon @ rv.cov_cholesky  # precon is diagonal, so this is valid
    new_cov = new_cov_cholesky @ new_cov_cholesky.T

    return randvars.Normal(new_mean, new_cov, cov_cholesky=new_cov_cholesky)

class Preconditioner(abc.ABC):
    """Coordinate change transformations as preconditioners in state space models.

    For some models, this makes the filtering and smoothing steps more numerically

[docs] @abc.abstractmethod def __call__(self, step) -> np.ndarray: # if more than step is needed, add them into the signature in the future raise NotImplementedError
@cached_property def inverse(self) -> "Preconditioner": raise NotImplementedError class NordsieckLikeCoordinates(Preconditioner): """Nordsieck-like coordinates. Similar to Nordsieck coordinates (which store the Taylor coefficients instead of the derivatives), but better for ODE filtering and smoothing. Used in integrator-transitions, e.g. in :class:`IntegratedWienerTransition`. """ def __init__(self, powers, scales, dimension): # Clean way of assembling these coordinates cheaply, # because the powers and scales of the inverse # are better read off than inverted self.powers = powers self.scales = scales self.dimension = dimension
[docs] @classmethod def from_order(cls, order, dimension): # used to conveniently initialise in the beginning powers = np.arange(order, -1, -1) scales = scipy.special.factorial(powers) return cls( powers=powers + 0.5, scales=scales, dimension=dimension, )
[docs] def __call__(self, step): scaling_vector = np.abs(step) ** self.powers / self.scales if config.matrix_free: return linops.IdentityKronecker( num_blocks=self.dimension, B=linops.Scaling(factors=scaling_vector) ) return np.kron(np.eye(self.dimension), np.diag(scaling_vector))
@cached_property def inverse(self) -> "NordsieckLikeCoordinates": return NordsieckLikeCoordinates( powers=-self.powers, scales=1.0 / self.scales, dimension=self.dimension, )