New example: Lorenz model (#272)

* Add Lorenz weather forecast model example.
Add tests for Lorenz model.
Add description to the changelog

* Add Lorenz weather forecast model example.
Add tests.
Add description to the changelog

* Fix all comments from discussion

* Optimize forecast_lorenz method

* Update docstrings in Lorenz model

* Optimize cross-covariance function. Fix docstrings

* Update lorenz branch

* Fix some drawbacks in docstrings

* Fix confusing moment in cross-covariance function implementation

* Change forecast_lorenz implementation. Add all time series

* Change summary statistics implementation

* Fix docstrings

* Fix documentation of mean function of Lorenz model

* Fix axis in rolling function in covariance and cross-covariance functions

* Fix cross-covariance and auto-covariance functions

* Change initial state and put after running with 10000 total duration

* Fix docstring. Remove commented code

CHANGELOG.rst

@@ -5,6 +5,7 @@ dev
 ---
 - Add support for kwargs to elfi.set_client
 - Added new example: inference of transmission dynamics of bacteria in daycare centers
+- Added new example: Lorenz model
 
 0.7.1 (2018-04-11)
 ------------------

elfi/examples/lorenz.py

@@ -0,0 +1,320 @@
+"""Example implementation of the Lorenz forecast model.
+
+References
+----------
+- Dutta, R., Corander, J., Kaski, S. and Gutmann, M.U., 2016.
+  Likelihood-free inference by ratio estimation. arXiv preprint arXiv:1611.10242.
+  https://arxiv.org/abs/1611.10242
+
+"""
+
+from functools import partial
+
+import numpy as np
+
+import elfi
+
+
+def _lorenz_ode(y, params):
+    """Parametrized Lorenz 96 system defined by a coupled stochastic differential equations (SDE).
+
+    Parameters
+    ----------
+    y : numpy.ndarray of dimension (batch_size, n_obs)
+        Current state of the SDE.
+    params : list
+        The list of parameters needed to evaluate function. In this case it is
+        list of four elements - eta, theta1, theta2 and f.
+
+    Returns
+    -------
+    dy_dt : np.array
+        Rate of change of the SDE.
+
+    """
+    dy_dt = np.empty_like(y)
+
+    eta = params[0]
+    theta1 = params[1]
+    theta2 = params[2]
+
+    f = params[3]
+
+    g = theta1 + y * theta2
+
+    dy_dt[:, 0] = -y[:, -2] * y[:, -1] + y[:, -1] * y[:, 1] - y[:, 0] + f - g[:, 0] + eta[:, 0]
+
+    dy_dt[:, 1] = -y[:, -1] * y[:, 0] + y[:, 0] * y[:, 2] - y[:, 1] + f - g[:, 1] + eta[:, 1]
+
+    dy_dt[:, 2:-1] = (-y[:, :-3] * y[:, 1:-2] + y[:, 1:-2] * y[:, 3:] - y[:, 2:-1] + f - g[:, 2:-1]
+                      + eta[:, 2:-1])
+
+    dy_dt[:, -1] = (-y[:, -3] * y[:, -2] + y[:, -2] * y[:, 0] - y[:, -1] + f - g[:, -1]
+                    + eta[:, -1])
+
+    return dy_dt
+
+
+def runge_kutta_ode_solver(ode, time_step, y, params):
+    """4th order Runge-Kutta ODE solver.
+
+    Carnahan, B., Luther, H. A., and Wilkes, J. O. (1969).
+    Applied Numerical Methods. Wiley, New York.
+
+    Parameters
+    ----------
+    ode : function
+        Ordinary differential equation function. In the Lorenz model it is SDE.
+    time_step : float
+    y : np.ndarray of dimension (batch_size, n_obs)
+        Current state of the time-series.
+    params : list of parameters
+        The parameters needed to evaluate the ode. In this case it is
+        list of four elements - eta, theta1, theta2 and f.
+
+    Returns
+    -------
+    np.ndarray
+        Resulting state initiated at y and satisfying ode solved by this solver.
+
+    """
+    k1 = time_step * ode(y, params)
+
+    k2 = time_step * ode(y + k1 / 2, params)
+
+    k3 = time_step * ode(y + k2 / 2, params)
+
+    k4 = time_step * ode(y + k3, params)
+
+    y = y + (k1 + 2 * k2 + 2 * k3 + k4) / 6
+
+    return y
+
+
+def forecast_lorenz(theta1=None, theta2=None, f=10., phi=0.984, n_obs=40, n_timestep=160,
+                    batch_size=1, initial_state=None, random_state=None, total_duration=4):
+    """Forecast Lorenz model.
+
+    Wilks, D. S. (2005). Effects of stochastic parametrizations in the
+    Lorenz ’96 system. Quarterly Journal of the Royal Meteorological Society,
+    131(606), 389–407.
+
+    Parameters
+    ----------
+    theta1, theta2: list or numpy.ndarray
+        Closure parameters.
+    phi : float, optional
+        This value is used to express stochastic forcing term. It should be configured according
+        to force term and eventually impacts to the result of eta.
+        More details in Wilks (2005) et al.
+    initial_state: numpy.ndarray, optional
+        Initial state value of the time-series.
+    f : float, optional
+        Force term
+    n_obs : int, optional
+        Size of the observed 1D grid
+    n_timestep : int, optional
+        Number of the time step intervals
+    batch_size : int, optional
+    random_state : np.random.RandomState, optional
+    total_duration : float, optional
+
+    Returns
+    -------
+    np.ndarray of size (b, n, m) which is (batch_size, time, n_obs)
+        The computed SDE with time series.
+
+    """
+    if not initial_state:
+        initial_state = np.tile([2.40711741e-01, 4.75597337e+00, 1.19145654e+01, 1.31324866e+00,
+                                 2.82675744e+00, 3.96016971e+00, 2.10479504e+00, 5.47742826e+00,
+                                 5.42519447e+00, -1.45166074e+00, 2.01991521e+00, 3.93873313e+00,
+                                 8.22837848e+00, 4.89401702e+00, -5.66278973e+00, 1.58617220e+00,
+                                 -1.23849251e+00, -6.04649288e-01, 6.04132264e+00, 7.47588536e+00,
+                                 1.82761402e+00, 3.19209639e+00, -7.58539653e-02, -6.00928508e-03,
+                                 4.52902964e-01, 3.22063602e+00, 7.18613523e+00, 2.39210634e+00,
+                                 -2.65743666e+00, 2.32046235e-01, 1.28079141e+00, 4.23344286e+00,
+                                 6.94213238e+00, -1.15939497e+00, -5.23037351e-01, 1.54618811e+00,
+                                 1.77863869e+00, 3.30139201e+00, 7.47769309e+00, -3.91312909e-01],
+                                (batch_size, 1))
+
+    y = initial_state
+    eta = 0
+
+    theta1 = np.asarray(theta1).reshape(-1, 1)
+
+    theta2 = np.asarray(theta2).reshape(-1, 1)
+
+    time_step = total_duration / n_timestep
+
+    random_state = random_state or np.random
+
+    time_series = np.empty(shape=(batch_size, n_timestep, n_obs))
+    time_series[:, 0, :] = y
+
+    for i in range(1, n_timestep):
+        e = random_state.normal(0, 1, y.shape)
+        eta = phi * eta + e * np.sqrt(1 - pow(phi, 2))
+        params = (eta, theta1, theta2, f)
+
+        y = runge_kutta_ode_solver(ode=_lorenz_ode, time_step=time_step, y=y, params=params)
+        time_series[:, i, :] = y
+
+    return time_series
+
+
+def get_model(true_params=None, seed_obs=None, initial_state=None, n_obs=40, f=10., phi=0.984,
+              total_duration=4):
+    """Return a complete Lorenz model in inference task.
+
+    This is a simplified example that achieves reasonable predictions.
+
+    Hakkarainen, J., Ilin, A., Solonen, A., Laine, M., Haario, H., Tamminen,
+    J., Oja, E., and Järvinen, H. (2012). On closure parameter estimation in
+    chaotic systems. Nonlinear Processes in Geophysics, 19(1), 127–143.
+
+    Parameters
+    ----------
+    true_params : list, optional
+        Parameters with which the observed data is generated.
+    seed_obs : int, optional
+        Seed for the observed data generation.
+    initial_state : ndarray
+        Initial state value of the time-series.
+    n_obs : int, optional
+        Number of observed variables
+    f : float, optional
+        Force term
+    phi : float, optional
+        This value is used to express stochastic forcing term. It should be configured according
+        to force term and eventually impacts to the result of eta.
+        More details in Wilks (2005) et al.
+    total_duration : float, optional
+
+    Returns
+    -------
+    m : elfi.ElfiModel
+
+    """
+    simulator = partial(forecast_lorenz, initial_state=initial_state, f=f, n_obs=n_obs, phi=phi,
+                        total_duration=total_duration)
+
+    if not true_params:
+        true_params = [2.0, 0.1]
+
+    m = elfi.ElfiModel()
+
+    y_obs = simulator(*true_params, random_state=np.random.RandomState(seed_obs))
+    sumstats = []
+
+    elfi.Prior('uniform', 0.5, 3., model=m, name='theta1')
+    elfi.Prior('uniform', 0, 0.3, model=m, name='theta2')
+    elfi.Simulator(simulator, m['theta1'], m['theta2'], observed=y_obs, name='Lorenz')
+
+    sumstats.append(elfi.Summary(mean, m['Lorenz'], name='Mean'))
+
+    sumstats.append(elfi.Summary(var, m['Lorenz'], name='Var'))
+
+    sumstats.append(elfi.Summary(autocov, m['Lorenz'], name='Autocov'))
+
+    sumstats.append(elfi.Summary(cov, m['Lorenz'], name='Cov'))
+
+    sumstats.append(elfi.Summary(xcov, m['Lorenz'], True, name='CrosscovPrev'))
+
+    sumstats.append(elfi.Summary(xcov, m['Lorenz'], False, name='CrosscovNext'))
+
+    elfi.Distance('euclidean', *sumstats, name='d')
+
+    return m
+
+
+def mean(x):
+    """Return the mean of Y_{k}.
+
+    Parameters
+    ----------
+    x : np.array of size (b, n, m) which is (batch_size, time, n_obs)
+
+    Returns
+    -------
+    np.array of size (b,)
+        The computed mean of statistics over time and space.
+
+    """
+    return np.mean(x, axis=(1, 2))
+
+
+def var(x):
+    """Return the variance of Y_{k}.
+
+    Parameters
+    ----------
+    x : np.array of size (b, n, m) which is (batch_size, time, n_obs)
+
+    Returns
+    -------
+    np.array of size (b,)
+        The average over space of computed variance with respect to time.
+
+    """
+    return np.mean(np.var(x, axis=1), axis=1)
+
+
+def cov(x):
+    """Return the covariance of Y_{k} with its neighbour Y_{k+1}.
+
+    Parameters
+    ----------
+    x : np.array of size (b, n, m) which is (batch_size, time, n_obs)
+
+    Returns
+    -------
+    np.array of size (b,)
+        The average over space of computed covariance with respect to time.
+
+    """
+    x_next = np.roll(x, -1, axis=2)
+    return np.mean(np.mean((x - np.mean(x, keepdims=True, axis=1)) *
+                           (x_next - np.mean(x_next, keepdims=True, axis=1)),
+                           axis=1), axis=1)
+
+
+def xcov(x, prev=True):
+    """Return the cross-covariance of Y_{k} with its neighbours from previous time step.
+
+    Parameters
+    ----------
+    x : np.array of size (b, n, m) which is (batch_size, time, n_obs)
+    prev : bool
+        The side of previous neighbour. True for previous neighbour, False for next.
+
+    Returns
+    -------
+    np.array of size (b,)
+        The average over space of computed cross-covariance with respect to time.
+
+    """
+    x_lag = np.roll(x, 1, axis=2) if prev else np.roll(x, -1, axis=2)
+    return np.mean((x[:, :-1, :] - np.mean(x[:, :-1, :], keepdims=True, axis=1)) *
+                   (x_lag[:, 1:, :] - np.mean(x_lag[:, 1:, :], keepdims=True, axis=1)),
+                   axis=(1, 2))
+
+
+def autocov(x):
+    """Return the auto-covariance with time lag 1.
+
+    Parameters
+    ----------
+    x : np.array of size (b, n, m) which is (batch_size, time, n_obs)
+
+    Returns
+    -------
+    C : np.array of size (b,)
+        The average over space of computed auto-covariance with respect to time.
+
+    """
+    c = np.mean((x[:, :-1, :] - np.mean(x[:, :-1, :], keepdims=True, axis=1)) *
+                (x[:, 1:, :] - np.mean(x[:, 1:, :], keepdims=True, axis=1)),
+                axis=(1, 2))
+
+    return c

tests/unit/test_examples.py

@@ -5,7 +5,7 @@
 import pytest
 
 import elfi
-from elfi.examples import bdm, bignk, gauss, gnk, lotka_volterra, ricker, daycare
+from elfi.examples import bdm, bignk, gauss, gnk, lotka_volterra, ricker, daycare, lorenz
 
 
 def test_bdm():
@@ -72,6 +72,12 @@ def test_Ricker():
     rej.sample(20)
 
 
+def test_Lorenz():
+    m = lorenz.get_model()
+    rej = elfi.Rejection(m, m['d'], batch_size=10)
+    rej.sample(20)
+
+
 def test_gnk():
     m = gnk.get_model()
     rej = elfi.Rejection(m, m['d'], batch_size=10)