dapper.mods

  1"""Contains models included with DAPPER.
  2
  3.. include:: ./README.md
  4"""
  5
  6import copy as cp
  7import inspect
  8from pathlib import Path
  9
 10import numpy as np
 11import struct_tools
 12
 13# Imports used to set up HMMs
 14import dapper.tools.progressbar as pb
 15from dapper.dpr_config import rc
 16from dapper.mods.utils import Id_op
 17from dapper.tools.chronos import Chronology
 18from dapper.tools.localization import no_localization
 19from dapper.tools.matrices import CovMat
 20from dapper.tools.randvars import RV, GaussRV
 21from dapper.tools.seeding import set_seed
 22
 23from .integration import with_recursion, with_rk4
 24from .utils import Id_Obs, ens_compatible, linspace_int, partial_Id_Obs
 25
 26
 27class HiddenMarkovModel(struct_tools.NicePrint):
 28    """Container class (with some embellishments) for a Hidden Markov Model (HMM).
 29
 30    Should contain the details necessary to run synthetic DA experiments,
 31    also known as "twin experiment", or OSSE (observing system simulation experiment).
 32    The synthetic truth and observations may then be obtained by running
 33    `HiddenMarkovModel.simulate`.
 34
 35    .. note::
 36      Each model included with DAPPER comes with several examples
 37      of model settings from the literature.
 38      See, for example, `dapper.mods.Lorenz63.sakov2012`.
 39
 40    .. warning::
 41      These example configurations do not necessarily hold a high programming standard,
 42      as they may have been whipped up at short notice to replicate some experiments,
 43      and are not intended for re-use.
 44      Nevertheless, sometimes they are re-used by another configuration script,
 45      leading to a major gotcha/pitfall: changes made to the imported `HMM` (or
 46      the model's module itself) also impact the original object (since they
 47      are mutable and thereby referenced).  This *usually* isn't an issue, since
 48      one rarely imports two/more separate configurations. However, the test suite
 49      imports all configurations, which might then unintentionally interact.
 50      To avoid this, you should use the `copy` method of the `HMM`
 51      before making any changes to it.
 52
 53
 54    Parameters
 55    ----------
 56    Dyn: `Operator` or dict
 57        Operator for the dynamics.
 58    Obs: `Operator` or dict
 59        Operator for the observations
 60        Can also be time-dependent, ref `TimeDependentOperator`.
 61    tseq: `dapper.tools.chronos.Chronology`
 62        Time sequence of the HMM process.
 63    X0: `dapper.tools.randvars.RV`
 64        Random distribution of initial condition
 65    liveplotters: `list`, optional
 66        A list of tuples. See example use in function `LPs` of `dapper.mods.Lorenz63`.
 67        - The first element of the tuple determines if the liveplotter
 68        is shown by default. If `False`, the liveplotter is only shown when
 69        included among the `liveplots` argument of `assimilate`
 70        - The second element in the tuple gives the corresponding liveplotter
 71        function/class.
 72    sectors: `dict`, optional
 73        Labelled indices referring to parts of the state vector.
 74        When defined, field-mean statistics are computed for each sector.
 75        Example use can be found in  `examples/param_estim.py`
 76        and `dapper/mods/Lorenz96/miyoshi2011.py`
 77    name: str, optional
 78        Label for the `HMM`.
 79    """
 80
 81    def __init__(self, *args, **kwargs):
 82        # Expected args/kwargs, along with type and default.
 83        attrs = dict(Dyn=(Operator, None),
 84                     Obs=(TimeDependentOperator, None),
 85                     tseq=(Chronology, None),
 86                     X0=(RV, None),
 87                     liveplotters=(list, []),
 88                     sectors=(dict, {}),
 89                     name=(str, self._default_name()))
 90
 91        # Transfer args to kwargs
 92        for arg, kw in zip(args, attrs):
 93            assert (kw not in kwargs), "Could not sort out arguments."
 94            kwargs[kw] = arg
 95
 96        # Un-abbreviate
 97        abbrevs = {"LP": "liveplotters",
 98                   "loc": "localizer"}
 99        for key in list(kwargs):
100            try:
101                full = abbrevs[key]
102            except KeyError:
103                pass
104            else:
105                assert (full not in kwargs), "Could not sort out arguments."
106                kwargs[full] = kwargs.pop(key)
107
108        # Convert dict to object
109        for key, (type_, default) in attrs.items():
110            val = kwargs.get(key, default)
111            if not isinstance(val, type_) and val is not None:
112                val = type_(**val)
113            kwargs[key] = val
114
115        # Transfer kwargs to self
116        for key in attrs:
117            setattr(self, key, kwargs.pop(key))
118        assert not kwargs, (
119            f"Arguments {list(kwargs)} not recognized. "
120            "If you want, you can still write them to the HMM, "
121            "but this must be done after initialisation.")
122
123        # Further defaults
124        # if not hasattr(self.Obs, "localizer"):
125        #     self.Obs.localizer = no_localization(self.Nx, self.Ny)
126
127        # Validation
128        # if self.Obs.noise.C == 0 or self.Obs.noise.C.rk != self.Obs.noise.C.M:
129        #     raise ValueError("Rank-deficient R not supported.")
130
131    # ndim shortcuts
132    @property
133    def Nx(self): return self.Dyn.M
134
135    printopts = {'ordering': ['Dyn', 'Obs', 'tseq', 'X0'], "indent": 4}
136
137    def simulate(self, desc='Truth & Obs'):
138        """Generate synthetic truth and observations."""
139        Dyn, Obs, tseq, X0 = self.Dyn, self.Obs, self.tseq, self.X0
140
141        # Init
142        xx = np.zeros((tseq.K + 1, Dyn.M))
143        yy = np.empty(tseq.Ko + 1, dtype=object)
144
145        x = X0.sample(1).squeeze()
146        xx[0] = x
147
148        # Loop
149        for k, ko, t, dt in pb.progbar(tseq.ticker, desc):
150            x = Dyn(x, t-dt, dt)
151            x = x + np.sqrt(dt)*Dyn.noise.sample(1).squeeze()
152            if ko is not None:
153                yy[ko] = Obs(ko)(x) + Obs(ko).noise.sample(1).squeeze()
154            xx[k] = x
155
156        return xx, yy
157
158    def copy(self):
159        return cp.deepcopy(self)
160
161    @staticmethod
162    def _default_name():
163        name = inspect.getfile(inspect.stack()[2][0])
164        try:
165            name = str(Path(name).relative_to(rc.dirs.dapper/'mods'))
166        except ValueError:
167            name = str(Path(name))
168        return name
169
170
171class TimeDependentOperator:
172    """Wrapper for `Operator` that enables time dependence.
173
174    The time instance should be specified by `ko`,
175    i.e. the index of an observation time.
176
177    Examples: `examples/time-dep-obs-operator.py` and `dapper/mods/QG/sakov2008.py`.
178    """
179    def __init__(self, **kwargs):
180        """Can be initialized like `Operator`, in which case the resulting
181        object will always return the same `Operator` nomatter the input time.
182
183        If initialized with 1 argument: `dict(time_dependent=func)`
184        then `func` must return an `Operator` object.
185        """
186        try:
187            fun = kwargs['time_dependent']
188            assert len(kwargs) == 1
189            assert callable(fun)
190            self.Ops = fun
191        except KeyError:
192            self.Op1 = Operator(**kwargs)
193
194    def __repr__(self):
195        return "<" + type(self).__name__ + '> ' + str(self)
196
197    def __str__(self):
198        if hasattr(self, 'Op1'):
199            return 'CONSTANT operator sepcified by .Op1:\n' + repr(self.Op1)
200        else:
201            return '.Ops: ' + repr(self.Ops)
202
203    def __call__(self, ko):
204        try:
205            return self.Ops(ko)
206        except AttributeError:
207            return self.Op1
208
209
210class Operator(struct_tools.NicePrint):
211    """Container for the dynamical and the observational maps.
212
213    Parameters
214    ----------
215    M: int
216        Length of output vectors.
217    model: function
218        The actual operator.
219    noise: RV, optional
220        The associated additive noise. The noise can also be a scalar or an
221        array, producing `GaussRV(C=noise)`.
222
223    Any remaining keyword arguments are written to the object as attributes.
224    """
225
226    def __init__(self, M, model=None, noise=None, **kwargs):
227        self.M = M
228
229        # Default to the Identity operator
230        if model is None:
231            model = Id_op()
232            kwargs['linear'] = lambda *args: np.eye(M)
233        # Assign
234        self.model = model
235
236        # None/0 => No noise
237        if isinstance(noise, RV):
238            self.noise = noise
239        else:
240            if noise is None:
241                noise = 0
242            if np.isscalar(noise):
243                self.noise = GaussRV(C=noise, M=M)
244            else:
245                self.noise = GaussRV(C=noise)
246
247        # Write attributes
248        for key, value in kwargs.items():
249            setattr(self, key, value)
250
251    def __call__(self, *args, **kwargs):
252        return self.model(*args, **kwargs)
253
254    printopts = {'ordering': ['M', 'model', 'noise'], "indent": 4}