def linearize(self, at_this_rv: randvars.Normal):
"""Linearize the dynamics function with a first order Taylor expansion."""
g = self.non_linear_model.state_trans_fun
dg = self.non_linear_model.jacob_state_trans_fun
x0 = at_this_rv.mean
def forcevecfun(t):
return g(t, x0) - dg(t, x0) @ x0
def dynamicsmatfun(t):
return dg(t, x0)
return statespace.DiscreteLinearGaussian(
input_dim=self.non_linear_model.input_dim,
output_dim=self.non_linear_model.output_dim,
state_trans_mat_fun=dynamicsmatfun,
shift_vec_fun=forcevecfun,
proc_noise_cov_mat_fun=self.non_linear_model.
proc_noise_cov_mat_fun,
proc_noise_cov_cholesky_fun=self.non_linear_model.
proc_noise_cov_cholesky_fun,
forward_implementation=self.forward_implementation,
backward_implementation=self.backward_implementation,
)
def _setup(
self,
test_ndim,
spdmat1,
spdmat2,
forw_impl_string_linear_gauss,
backw_impl_string_linear_gauss,
):
self.G = lambda t: spdmat1
self.S = lambda t: spdmat2
self.v = lambda t: np.arange(test_ndim)
self.transition = pnss.DiscreteLinearGaussian(
test_ndim,
test_ndim,
self.G,
self.v,
self.S,
forward_implementation=forw_impl_string_linear_gauss,
backward_implementation=backw_impl_string_linear_gauss,
)
self.g = lambda t, x: self.G(t) @ x + self.v(t)
self.dg = lambda t, x: self.G(t)