Ensemble kalman

docs/source/notebooks/Bayes_estimation_ssm.ipynb

@@ -307,7 +307,7 @@
     "  * the starting point of the chain is sampled from the prior; you may instead set it to a specific value using option `starting_point` (when instantiating `PMMH`); \n",
     "  * the proposal is an **adaptative** Gaussian random walk: this means that the covariance matrix of the random step is calibrated on the fly on past simulations (using vanishing adaptation). This may be disabled by setting option `adaptive=False`; \n",
     "  * a bootstrap filter is run to approximate the log-likelihood; you may use a different filter (e.g. a guided filter) by passing a `FeynmanKac` class to option `fk_cls`;\n",
-    "  * you may also want to pass various parameters to each call to ` SMC` through (dict) argument `smc_options`; e.g. `smc_options={'qmc': True}` will make each particle filter a SQMC algorithm. \n",
+    "  * you may also want to pass various parameters to each call to `SMC` through (dict) argument `smc_options`; e.g. `smc_options={'qmc': True}` will make each particle filter a SQMC algorithm. \n",
     "  \n",
     "Thus, by and large, quite a lot of flexibility is hidden behind this default behaviour."
    ]
@@ -444,7 +444,7 @@
     "Again, a few choices are made for you by default: \n",
     "\n",
     "* A bootstrap filter is run for each $\\theta-$particle; this may be changed by setting option `fk_class` while instantiating `SMC2`; e.g. ``fk_class=ssm.GuidedPF`` will run instead a guided filter. \n",
-    "* Option `init_Nx` determines the **initial** number of $x-$particles; the algorithm automatically increases $N_x$ each time the acceptance rate drops below $10%$ (as specified through option ``ar_to_increase=0.1``. Set this this option to `0.` if you do not want to increase $N_x$ in the course of t he algorithm.\n",
+    "* Option `init_Nx` determines the **initial** number of $x-$particles; the algorithm automatically increases $N_x$ each time the acceptance rate drops below $10\\%$ (as specified through option ``ar_to_increase=0.1``. Set this this option to `0.` if you do not want to increase $N_x$ in the course of t he algorithm.\n",
     "* The particle filters (in the $x-$dimension) are run with the default options of class `SMC`; e.g. resampling is set to `systematic` and so on; other options may be set by using option `smc_options`."
    ]
   },
@@ -473,7 +473,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -487,9 +487,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.12.7"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 1
+ "nbformat_minor": 4
 }

particles/distributions.py

@@ -840,7 +840,7 @@ def logpdf(self, x):
         return lp
 
     def rvs(self, size=None):
-        x = self.base_dist.rvs(size=size).astype(float)
+        x = self.base_dist.rvs(size=size)
         N = x.shape[0]
         is_missing = random.rand(N) < self.pmiss
         x[is_missing, ...] = np.nan

particles/ensemble_kalman.py

@@ -0,0 +1,170 @@
+
+import collections
+
+import numpy as np
+from scipy.linalg import solve
+
+from particles import distributions as dists
+from particles import state_space_models as ssms
+from particles import utils
+import particles
+
+
+
+
+def EnK_step(ssm, t, xp, x_prop, y, return_weights = False):  
+  flag_reshape = False
+  if x_prop.ndim == 1:
+    flag_reshape = True
+    dx = 1
+    J = len(x_prop)
+    mapped_X_prop = ssm.PY(t, xp, x_prop).rvs()
+    x_prop = np.reshape(x_prop, (J,1))
+  else:
+    J, dx = x_prop.shape    
+    mapped_X_prop = ssm.PY(t, xp, x_prop).rvs()
+  # weights = ssm.PY(t, xp, x_prop).logpdf(mapped_X_prop)
+  mapped_X_prop = np.reshape(mapped_X_prop, (J,1))
+  full_cov = np.cov(x_prop,mapped_X_prop, rowvar=False)
+  Cuu = np.atleast_2d(full_cov[0:dx, 0:dx])
+  Cup = np.atleast_2d(full_cov[0:dx, dx:])
+  CppGamma = np.atleast_2d(full_cov[dx:, dx:])
+  K = Cup@np.linalg.inv(CppGamma)
+  cov_shrinkage = (np.eye(dx) - K@Cup.T)
+  Qhat = cov_shrinkage@Cuu@cov_shrinkage.T
+  new_filt = x_prop - (K@(mapped_X_prop.T - y.T)).T
+
+  if flag_reshape: # cast back to original form
+    new_filt = np.reshape(new_filt, (J,))
+
+  if return_weights: # possibly remove all this
+    if t == 0:        
+      Pxweights = ssm.PX0().logpdf(x_prop)
+      Qxweights = ssm.PX0().logpdf(new_filt)  
+    else:        
+      Pxweights = ssm.PX(t,xp).logpdf(x_prop)
+      Qxweights = ssm.PX(t,xp).logpdf(new_filt)
+    correction_logweights = Pxweights - Qxweights
+    return new_filt, K, correction_logweights
+  else:
+    return new_filt, K, None
+
+
+error_msg = "arguments of MVNonlinearGauss.__init__ have inconsistent shapes"
+
+class MVNonlinearGauss(ssms.StateSpaceModel):
+    r"""Multivariate nonlinear additive Gaussian model.
+
+    .. math::
+        X_0 & \sim N(\mu_0, cov_0) \\
+        X_t & = F(X_{t-1}) + U_t, \quad   U_t\sim N(0, cov_X) \\
+        Y_t & = G(X_t) + V_t,     \quad   V_t \sim N(0, cov_Y)
+
+    The only mandatory parameters are `covX` and `covY` (from which the
+    dimensions dx and dy of, respectively, X_t, and Y_t, are deduced). The
+    default values for the other parameters are:
+
+    * `mu0` : an array of zeros (of size dx)
+    * `cov0`: cov_X
+    * `F` : Identity mapping  of shape (dx, dx)
+    * `G` : (dy, dx) matrix such that G[i, j] = 1[i=j]
+
+    Note
+    ----
+    The Kalman filter takes as an input an instance of this class (or one of
+    its subclasses).
+    """
+
+    def __init__(self, F=None, G=None, covX=None, covY=None, mu0=None, cov0=None, **kwargs):
+        self.covX, self.covY = np.atleast_2d(covX), np.atleast_2d(covY)
+        self.dx, self.dy = self.covX.shape[0], self.covY.shape[0]
+        self.mu0 = np.zeros(self.dx) if mu0 is None else mu0
+        self.cov0 = self.covX if cov0 is None else np.atleast_2d(cov0)
+        # Assign F and G only if they are not already defined as methods
+        if not hasattr(self, "F"):
+            if F is None:
+                self.F = lambda t, x: x
+            else:
+                self.F = F
+
+        if not hasattr(self, "G"):
+            if G is None:
+                self.G = lambda t, x: x[0:self.dy]
+            else:
+                self.G = G
+        self.check_shapes()
+        if hasattr(self, "default_params"):
+            self.__dict__.update(self.default_params)
+        self.__dict__.update(kwargs)
+
+    def check_shapes(self):
+        """
+        Check all dimensions are correct.
+        """
+        assert self.covX.shape == (self.dx, self.dx), error_msg
+        assert self.covY.shape == (self.dy, self.dy), error_msg
+        # assert self.F.shape == (self.dx, self.dx), error_msg
+        # assert self.G.shape == (self.dy, self.dx), error_msg
+        assert self.mu0.shape == (self.dx,), error_msg
+        assert self.cov0.shape == (self.dx, self.dx), error_msg
+
+    def PX0(self):
+        return dists.MvNormal(loc=self.mu0, cov=self.cov0)
+
+    def PX(self, t, xp):
+        return dists.MvNormal(loc=self.F(t, xp), cov=self.covX)
+
+    def PY(self, t, xp, x):
+        return dists.MvNormal(loc=self.G(t, x), cov=self.covY)
+
+
+
+class EnsembleKalman(particles.FeynmanKac):
+  """ Ensemble Kalman formalism of a given state-space model.
+
+  Note that the EnKF is an approximate algorithm! This means that the Feynman-Kac
+  model is only an exact representation of the state space model if the setting is linear and Gaussian
+
+  Parameters
+  ----------
+
+  ssm: `StateSpaceModel` object
+      the considered state-space model
+  data: list-like
+      the data
+
+  Returns
+  -------
+  `FeynmanKac` object
+      the Feynman-Kac representation of the Ensemble Kalman filter for the
+      considered state-space model
+  """
+  def __init__(self, ssm=None, data=None):
+      self.ssm = ssm
+      self.data = data
+      self.du = self.ssm.PX0().dim
+
+  @property
+  def T(self):
+      return 0 if self.data is None else len(self.data)
+
+  def M0(self, N):
+      x_prop = self.ssm.PX0().rvs(size=N)
+      if np.isnan(self.data[0]):
+        new_filt = x_prop
+      else:
+        new_filt, _, _ = EnK_step(self.ssm, 0, None, x_prop, self.data[0])
+      return new_filt
+
+  def M(self, t, xp):
+      x_prop = self.ssm.PX(t, xp).rvs()
+      if np.isnan(self.data[t]):
+        new_filt = x_prop
+      else:
+        new_filt, _, _ = EnK_step(self.ssm, t, xp, x_prop, self.data[t])
+      return new_filt
+
+  def logG(self, t, xp, x):
+      return np.zeros(x.shape[0])
+
+

tests/ensemble_kalman/test_EnKF.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python
+
+""" Checks that the Ensemble Kalman filter works on a simple Lorenz 63 model """
+
+
+
+from matplotlib import pyplot as plt
+import numpy as np
+import seaborn as sb
+from scipy import stats
+
+import sys
+sys.path.insert(0, "..\\..")
+import particles
+from particles import state_space_models as ssm
+from particles import distributions as dists
+from particles import ensemble_kalman as enk
+from particles.collectors import Moments
+np.random.seed(1)
+# setup of filtering problem
+def atleast_2d_last(x):
+    """Ensure x has at least 2 dims, adding a new axis at the end if needed."""
+    x = np.array(x, copy=False)
+    if x.ndim < 2:
+        return x[..., np.newaxis]
+    return x
+colors = ["tab:blue", "tab:green", "tab:orange"]
+
+class Lorenz_63(enk.MVNonlinearGauss):
+  def F(self, t, xp):  # Distribution of X_t given X_{t-1}=xp (p=past)
+      x = xp[:,[0]]
+      y = xp[:,[1]]
+      z = xp[:,[2]]
+      return xp + self.dt*np.hstack([self.sigma*(y-x), x*(self.rho - z)-y, x*y-self.beta*z])
+  def G(self, t, x):  # Distribution of Y_t given X_t=x (and possibly X_{t-1}=xp)
+      return x[:,[0]]#return dists.Normal(loc=x[:,0], scale=self.obsnoise)
+
+
+
+my_model = Lorenz_63(rho=28., sigma=10., beta=8./3, dt=0.01, covX = np.eye(3), covY = 1.*np.eye(1))  # actual model
+
+#my_model = enk.MVNonlinearGauss(F=lambda t, x: 0.8*x, G=lambda t, x: np.exp(x), covX=0.1, covY=0.05, mu0=None, cov0=None)
+# my_model = ssm.StochVol()
+# toymodel = enk.MVNonlinearGauss(F=lambda x: x, G=lambda x: np.exp(x), covX=1., covY=.1, mu0=None, cov0=None)
+true_states, data = my_model.simulate(500) #200  # we simulate from the model 100 data points
+
+J = 30 # size of ensembles
+
+plt.figure()
+plt.subplot(211)
+plt.plot(np.vstack(true_states))
+plt.subplot(212)
+plt.plot(np.vstack(data))
+
+
+#%% 
+
+algorithms = [enk.EnsembleKalman, ssm.Bootstrap]
+alg_titles = ["EnKF", "Bootstrap"]
+plt.figure(figsize=(6,6))
+N_MC = 25
+MSEs_mean = [[None  for m in range(N_MC)] for n in range(len(algorithms))]
+coverage = [[None  for m in range(N_MC)] for n in range(len(algorithms))]
+for n_alg, alg in enumerate(algorithms):
+  for nMC in range(N_MC):
+    fk = alg(my_model, data)
+    filt = particles.SMC(fk=fk, N=J, collect=[Moments()], store_history=True) 
+    filt.run()
+
+    filt_path = np.stack(filt.hist.X)
+
+
+
+    means =  atleast_2d_last(np.stack([m['mean'] for m in filt.summaries.moments]))
+    var = atleast_2d_last(np.stack([m['var'] for m in filt.summaries.moments]))
+    if nMC == 0:
+      plt.subplot(len(algorithms),1,n_alg+1)
+      for m in range(means.shape[1]):
+        if means.shape[0] > 1:
+          plt.plot(means[:,m], color=colors[m])
+          plt.fill_between(range(len(data)), y1=means[:,m]-2*np.sqrt(var[:,m]), y2=means[:,m]+2*np.sqrt(var[:,m]), color=colors[m], alpha=0.3)
+          plt.plot(np.vstack(true_states)[:,m], "--", color=colors[m])
+        else:
+          plt.errorbar(y=0, x=true_states[0], xerr=[means[:,m]-2*np.sqrt(var[:,m]),means[:,m]+2*np.sqrt(var[:,m])],
+            capsize=5,
+            ecolor="lightgrey",
+            markerfacecolor="black",
+            markeredgecolor="black",
+            marker="o",
+            linestyle="none",)
+          plt.ylim([-0.5,0.5])
+      plt.title(alg_titles[n_alg])
+
+    MSEs_mean[n_alg][nMC] = np.sqrt(np.sum((np.vstack(true_states) - means)**2))
+
+
+    coverage[n_alg][nMC] = [np.mean((np.vstack(true_states) >= means-nn*np.sqrt(var)) & (np.vstack(true_states) <= means+nn*np.sqrt(var))) for nn in [1,2,3]]
+
+MSEs_mean = np.array(MSEs_mean)
+coverage = np.array(coverage)
+plt.tight_layout()
+
+
+#%%
+plt.figure()
+color_quantile = ["tab:green", "tab:orange", "tab:red"]
+
+for n_quantile in range(3):
+  parts = plt.violinplot(coverage[:,:,n_quantile].T, positions=range(len(algorithms)))
+  for pc in parts['bodies']:
+    pc.set_facecolor(color_quantile[n_quantile])
+    pc.set_edgecolor(color_quantile[n_quantile])
+
+  parts['cmaxes'].set_colors("black")
+  parts['cmaxes'].set_alpha(0.2)
+  parts['cmins'].set_colors("black")
+  parts['cmins'].set_alpha(0.2)
+  parts['cbars'].set_colors("black")
+  parts['cbars'].set_alpha(0.2)
+  plt.plot(coverage[:,:,n_quantile], '.k')
+plt.axhline(y = 0.68, label="$\\pm 1 \\sigma$", color = "tab:green")
+plt.axhline(y = 0.95, label="$\\pm 2 \\sigma$", color = "tab:orange")
+plt.axhline(y = 0.997, label="$\\pm 3 \\sigma$", color = "tab:red")
+plt.xticks(range(len(algorithms)), alg_titles)
+plt.legend()
+
+
+
+plt.figure()
+plt.title("MSE")
+plt.boxplot(MSEs_mean.T, positions=range(len(algorithms)))
+plt.xticks(range(len(algorithms)), alg_titles)