dapper.mods.QG

Quasi-geostraphic 2D flow. Described in detail by bib.sakov2008deterministic.

Adapted from Pavel Sakov's enkf-matlab package.

More info:

  • governing_eqn.png
  • demo.py
  • ψ (psi) is the stream function (i.e. surface elevation)
  • Doubling time "between 25 and 50"
  • Note Sakov's trick of increasing RKH2 from 2.0e-12 to 2.0e-11 to stabilize the ensemble integration, which may be necessary for EnKF's with small N. See example in counillon2009.
  1"""Quasi-geostraphic 2D flow. Described in detail by `bib.sakov2008deterministic`.
  2
  3Adapted from Pavel Sakov's enkf-matlab package.
  4
  5More info:
  6
  7- `governing_eqn.png`
  8- `demo.py`
  9- ψ (psi) is the stream function (i.e. surface elevation)
 10- Doubling time "between 25 and 50"
 11- Note Sakov's trick of increasing RKH2 from 2.0e-12 to 2.0e-11 to stabilize
 12  the ensemble integration, which may be necessary for EnKF's with small N.
 13  See example in `counillon2009`.
 14"""
 15
 16import sys
 17from pathlib import Path
 18
 19import matplotlib as mpl
 20import numpy as np
 21
 22import dapper.mods as modelling
 23import dapper.tools.liveplotting as LP
 24
 25#########################
 26# Model
 27#########################
 28default_prms = dict(
 29    # These parameters may be interesting to change.
 30    dtout=5.0,  # dt for output to DAPPER.
 31    dt=1.25,  # dt used internally by Fortran. CFL = 2.0
 32    RKB=0,  # bottom     friction
 33    RKH=0,  # horizontal friction
 34    RKH2=2.0e-12,  # horizontal friction, biharmonic
 35    F=1600,  # Froud number
 36    R=1.0e-5,  # ≈ Rossby number
 37    scheme="'rk4'",  # One of (2ndorder, rk4, dp5)
 38    # Do not change the following:
 39    tend=0,  # Only used by standalone QG
 40    verbose=0,  # Turn off
 41    rstart=0,  # Restart: switch
 42    restartfname="''",  # Restart: read file
 43    outfname="''",  # Restart: write file
 44)
 45
 46
 47class model_config:
 48    """Define model.
 49
 50    Helps ensure consistency between prms file (that Fortran module reads)
 51    and Python calls to step(), for example for dt.
 52    """
 53
 54    def __init__(self, name, prms, mp=True):
 55        """Use `prms={}` to get the default configuration."""
 56        # Insert prms. Assert key is present in defaults.
 57        D = default_prms.copy()
 58        for key in prms:
 59            assert key in D
 60            D[key] = prms[key]
 61
 62        # Fortran code does not adjust its dt to divide dtout.
 63        # Nor is it worth implementing -- just assert:
 64        assert D["dtout"] % D["dt"] == 0, "Must be integer multiple"
 65
 66        self.prms = D
 67        self.mp = mp
 68        self.name = name
 69        self.fname = Path(__file__).parent / "f90" / f"prms_{name}.txt"
 70
 71        # Create string
 72        text = [f"  {key.ljust(20)!s} = {D[key]!s}" for key in D]
 73        text = (
 74            f"""! Parameter namelist ("{name}") generated via Python
 75        &parameters\n"""
 76            + "\n".join(text)
 77            + """\n/\n"""
 78        )
 79
 80        # Write string to file
 81        with open(self.fname, "w") as f:
 82            f.write(text)
 83
 84    @property
 85    def f90(self):
 86        try:
 87            from .f90.py_mod import interface_mod
 88
 89            return interface_mod
 90        except ImportError as error:
 91            error.msg = error.msg + (
 92                "\nHave you compiled the (Fortran) model?\n"
 93                f"See README in {__name__.replace('.', '/')}/f90"
 94            )
 95            raise
 96
 97    def step_1(self, x0, t, dt):
 98        """Step a single state vector."""
 99        # Coz fortran.step() reads dt (dtout) from prms file:
100        assert self.prms["dtout"] == dt
101        # Coz Fortran is typed.
102        assert isinstance(t, float)
103        # QG is autonomous (⇒ t should not matter), but Fortran doesn't like nan/inf.
104        assert np.isfinite(t)
105        # Copy coz Fortran will modify in-place.
106        psi = py2f(x0.copy())
107        # Call Fortran model.
108        self.f90.step(t, psi, self.fname)
109        # Flattening
110        x = f2py(psi)
111        return x
112
113    def step(self, E, t, dt):
114        """Vector and 2D-array (ens) input, with multiproc for ens case."""
115        if E.ndim == 1:
116            return self.step_1(E, t, dt)
117        if E.ndim == 2:
118            if self.mp:  # PARALLELIZED:
119                # Note: the relative overhead for parallelization decreases
120                # as the ratio dtout/dt increases.
121                # But the overhead is already negligible with a ratio of 4.
122                import dapper.tools.multiproc as multiproc
123
124                with multiproc.Pool(self.mp) as pool:
125                    E = pool.map(lambda x: self.step_1(x, t=t, dt=dt), E)
126                E = np.array(E)
127            else:  # NON-PARALLELIZED:
128                for n, x in enumerate(E):
129                    E[n] = self.step_1(x, t, dt)
130            return E
131
132
133#########################
134# Domain management
135#########################
136# Domain size "hardcoded" in f90/parameters.f90.
137
138# "Physical" domain length -- copied from f90/parameters.f90.
139# In my tests, only square domains generate any dynamics of interest.
140NX1 = 2
141NY1 = 2
142# Resolution level -- copied MREFIN from parameters.f90
143res = 7
144# Grid lengths.
145nx = NX1 * 2 ** (res - 1) + 1  # (axis=1)
146ny = NY1 * 2 ** (res - 1) + 1  # (axis=0)
147# Actually, the BCs are psi = nabla psi = nabla^2 psi = 0,
148# => psi should always be zero on the boundries.
149# => it'd be safer to rm boundries from the DA state vector,
150#    yielding ndim(state)=(nx-2)*(ny-2), but this is not done here.
151
152# Fortran model (e.g. f90/interface.f90) requires orientation: X[ix,iy].
153shape = (nx, ny)
154# Passing arrays to/from Fortran requries that flags['F_CONTIGUOUS']==True.
155order = "F"
156
157
158def py2f(x):
159    return x.reshape(shape, order=order)
160
161
162def f2py(X):
163    return X.flatten(order=order)
164
165
166# However, FOR PRINTING/PLOTTING PURPOSES, the y-axis should be vertical
167# [imshow(mat) uses the same orientation as print(mat)].
168def square(x):
169    return x.reshape(shape[::-1])
170
171
172def ind2sub(ind):
173    return np.unravel_index(ind, shape[::-1])
174
175
176#########################
177# Free run
178#########################
179def gen_sample(model, nSamples, SpinUp, Spacing):
180    simulator = modelling.with_recursion(model.step, prog="Simulating")
181    K = SpinUp + nSamples * Spacing
182    Nx = np.prod(shape)  # total state length
183    sample = simulator(np.zeros(Nx), K, 0.0, model.prms["dtout"])
184    return sample[SpinUp::Spacing]
185
186
187sample_filename = modelling.rc.dirs.samples / "QG_samples.npz"
188if (not sample_filename.is_file()) and ("pdoc" not in sys.modules):
189    print(
190        "Did not find sample file",
191        sample_filename,
192        "for experiment initialization. Generating...",
193    )
194    sample = gen_sample(model_config("sample_generation", {}), 400, 700, 10)
195    np.savez(sample_filename, sample=sample)
196
197
198#########################
199# Liveplotting
200#########################
201cm = mpl.colors.ListedColormap(0.85 * mpl.cm.jet(np.arange(256)))
202center = nx * int(ny / 2) + int(0.5 * nx)
203
204
205def LP_setup(jj):
206    return [
207        (1, LP.spatial2d(square, ind2sub, jj, cm)),
208        (0, LP.spectral_errors),
209        (0, LP.sliding_marginals(dims=center + np.arange(4))),
210    ]
default_prms = {'dtout': 5.0, 'dt': 1.25, 'RKB': 0, 'RKH': 0, 'RKH2': 2e-12, 'F': 1600, 'R': 1e-05, 'scheme': "'rk4'", 'tend': 0, 'verbose': 0, 'rstart': 0, 'restartfname': "''", 'outfname': "''"}
class model_config: 
 48class model_config:
 49    """Define model.
 50
 51    Helps ensure consistency between prms file (that Fortran module reads)
 52    and Python calls to step(), for example for dt.
 53    """
 54
 55    def __init__(self, name, prms, mp=True):
 56        """Use `prms={}` to get the default configuration."""
 57        # Insert prms. Assert key is present in defaults.
 58        D = default_prms.copy()
 59        for key in prms:
 60            assert key in D
 61            D[key] = prms[key]
 62
 63        # Fortran code does not adjust its dt to divide dtout.
 64        # Nor is it worth implementing -- just assert:
 65        assert D["dtout"] % D["dt"] == 0, "Must be integer multiple"
 66
 67        self.prms = D
 68        self.mp = mp
 69        self.name = name
 70        self.fname = Path(__file__).parent / "f90" / f"prms_{name}.txt"
 71
 72        # Create string
 73        text = [f"  {key.ljust(20)!s} = {D[key]!s}" for key in D]
 74        text = (
 75            f"""! Parameter namelist ("{name}") generated via Python
 76        &parameters\n"""
 77            + "\n".join(text)
 78            + """\n/\n"""
 79        )
 80
 81        # Write string to file
 82        with open(self.fname, "w") as f:
 83            f.write(text)
 84
 85    @property
 86    def f90(self):
 87        try:
 88            from .f90.py_mod import interface_mod
 89
 90            return interface_mod
 91        except ImportError as error:
 92            error.msg = error.msg + (
 93                "\nHave you compiled the (Fortran) model?\n"
 94                f"See README in {__name__.replace('.', '/')}/f90"
 95            )
 96            raise
 97
 98    def step_1(self, x0, t, dt):
 99        """Step a single state vector."""
100        # Coz fortran.step() reads dt (dtout) from prms file:
101        assert self.prms["dtout"] == dt
102        # Coz Fortran is typed.
103        assert isinstance(t, float)
104        # QG is autonomous (⇒ t should not matter), but Fortran doesn't like nan/inf.
105        assert np.isfinite(t)
106        # Copy coz Fortran will modify in-place.
107        psi = py2f(x0.copy())
108        # Call Fortran model.
109        self.f90.step(t, psi, self.fname)
110        # Flattening
111        x = f2py(psi)
112        return x
113
114    def step(self, E, t, dt):
115        """Vector and 2D-array (ens) input, with multiproc for ens case."""
116        if E.ndim == 1:
117            return self.step_1(E, t, dt)
118        if E.ndim == 2:
119            if self.mp:  # PARALLELIZED:
120                # Note: the relative overhead for parallelization decreases
121                # as the ratio dtout/dt increases.
122                # But the overhead is already negligible with a ratio of 4.
123                import dapper.tools.multiproc as multiproc
124
125                with multiproc.Pool(self.mp) as pool:
126                    E = pool.map(lambda x: self.step_1(x, t=t, dt=dt), E)
127                E = np.array(E)
128            else:  # NON-PARALLELIZED:
129                for n, x in enumerate(E):
130                    E[n] = self.step_1(x, t, dt)
131            return E

Define model.

Helps ensure consistency between prms file (that Fortran module reads) and Python calls to step(), for example for dt.

model_config(name, prms, mp=True) 
55    def __init__(self, name, prms, mp=True):
56        """Use `prms={}` to get the default configuration."""
57        # Insert prms. Assert key is present in defaults.
58        D = default_prms.copy()
59        for key in prms:
60            assert key in D
61            D[key] = prms[key]
62
63        # Fortran code does not adjust its dt to divide dtout.
64        # Nor is it worth implementing -- just assert:
65        assert D["dtout"] % D["dt"] == 0, "Must be integer multiple"
66
67        self.prms = D
68        self.mp = mp
69        self.name = name
70        self.fname = Path(__file__).parent / "f90" / f"prms_{name}.txt"
71
72        # Create string
73        text = [f"  {key.ljust(20)!s} = {D[key]!s}" for key in D]
74        text = (
75            f"""! Parameter namelist ("{name}") generated via Python
76        &parameters\n"""
77            + "\n".join(text)
78            + """\n/\n"""
79        )
80
81        # Write string to file
82        with open(self.fname, "w") as f:
83            f.write(text)

Use prms={} to get the default configuration.

prms
mp
name
fname
f90 
85    @property
86    def f90(self):
87        try:
88            from .f90.py_mod import interface_mod
89
90            return interface_mod
91        except ImportError as error:
92            error.msg = error.msg + (
93                "\nHave you compiled the (Fortran) model?\n"
94                f"See README in {__name__.replace('.', '/')}/f90"
95            )
96            raise
def step_1(self, x0, t, dt): 
 98    def step_1(self, x0, t, dt):
 99        """Step a single state vector."""
100        # Coz fortran.step() reads dt (dtout) from prms file:
101        assert self.prms["dtout"] == dt
102        # Coz Fortran is typed.
103        assert isinstance(t, float)
104        # QG is autonomous (⇒ t should not matter), but Fortran doesn't like nan/inf.
105        assert np.isfinite(t)
106        # Copy coz Fortran will modify in-place.
107        psi = py2f(x0.copy())
108        # Call Fortran model.
109        self.f90.step(t, psi, self.fname)
110        # Flattening
111        x = f2py(psi)
112        return x

Step a single state vector.

def step(self, E, t, dt): 
114    def step(self, E, t, dt):
115        """Vector and 2D-array (ens) input, with multiproc for ens case."""
116        if E.ndim == 1:
117            return self.step_1(E, t, dt)
118        if E.ndim == 2:
119            if self.mp:  # PARALLELIZED:
120                # Note: the relative overhead for parallelization decreases
121                # as the ratio dtout/dt increases.
122                # But the overhead is already negligible with a ratio of 4.
123                import dapper.tools.multiproc as multiproc
124
125                with multiproc.Pool(self.mp) as pool:
126                    E = pool.map(lambda x: self.step_1(x, t=t, dt=dt), E)
127                E = np.array(E)
128            else:  # NON-PARALLELIZED:
129                for n, x in enumerate(E):
130                    E[n] = self.step_1(x, t, dt)
131            return E

Vector and 2D-array (ens) input, with multiproc for ens case.

NX1 = 2
NY1 = 2
res = 7
nx = 129
ny = 129
shape = (129, 129)
order = 'F'
def py2f(x): 
159def py2f(x):
160    return x.reshape(shape, order=order)
def f2py(X): 
163def f2py(X):
164    return X.flatten(order=order)
def square(x): 
169def square(x):
170    return x.reshape(shape[::-1])
def ind2sub(ind): 
173def ind2sub(ind):
174    return np.unravel_index(ind, shape[::-1])
def gen_sample(model, nSamples, SpinUp, Spacing): 
180def gen_sample(model, nSamples, SpinUp, Spacing):
181    simulator = modelling.with_recursion(model.step, prog="Simulating")
182    K = SpinUp + nSamples * Spacing
183    Nx = np.prod(shape)  # total state length
184    sample = simulator(np.zeros(Nx), K, 0.0, model.prms["dtout"])
185    return sample[SpinUp::Spacing]
sample_filename = PosixPath('/home/runner/dpr_data/samples/QG_samples.npz')
cm = <matplotlib.colors.ListedColormap object>
center = 8320
def LP_setup(jj): 
206def LP_setup(jj):
207    return [
208        (1, LP.spatial2d(square, ind2sub, jj, cm)),
209        (0, LP.spectral_errors),
210        (0, LP.sliding_marginals(dims=center + np.arange(4))),
211    ]