Minor updates.

mechanisms.py

@@ -1,8 +1,5 @@
 """
 Mechanisms to be implemented in a `Model` instance.
-
-TODO:
-    - Generalise stochastic forcing and improve docstrings.
 """
 
 import pyfftw
@@ -13,26 +10,6 @@
 fftw = pyfftw.interfaces.numpy_fft
 
 
-class SpectralFilter:
-    """Exponential spectral filter for applying highly scale-selective but
-    non-physical dissipation at small scales.
-    """
-    solution_mode = 'exact'
-
-    def __init__(self, model):
-        filterfac = 23.6
-        cphi = 0.65 * np.pi
-        wvx = np.sqrt(
-            (model.kx * model.dx) ** 2. + (model.k_y * model.dx) ** 2.)
-        exp_filter = np.exp(-filterfac * (wvx - cphi) ** 4.)
-        exp_filter[wvx <= cphi] = 1.
-        self.exp_filter = exp_filter
-        self.model = model
-
-    def __call__(self):
-        self.model.zk *= self.exp_filter
-
-
 class Advection:
     """Advection implemented without dealiasing.
     """
@@ -127,9 +104,9 @@ def __call__(self):
 
 class Diffusion:
     """Diffusion with `order` to be specified. Order refers to the power of the
-    Laplacian. `order=1.` gives standard Newtonian viscosity; `order>2.` gives
+    Laplacian. `order=1.` gives standard Newtonian viscosity; `order>1.` gives
     hyperviscosity; `order=0.` gives linear drag; `order=-1.` gives
-    large-scale friction, etc.
+    large-scale/hyper friction, etc.
     """
     solution_mode = 'exact'
 
@@ -138,33 +115,36 @@ def __init__(self, model, order=1., coefficient=None):
         self.order = order
         self.coefficient = coefficient
         if self.order >= 0.:
-            self.n_order_visc = np.exp(
+            self.multiplier = np.exp(
                 -self.coefficient * self.model.timestepper.dt
                 * self.model.wv2 ** self.order)
         else:
-            self.n_order_visc = np.exp(
+            self.multiplier = np.exp(
                 -self.coefficient * self.model.timestepper.dt
                 * self.model.wv2i ** (-self.order))
 
     def __call__(self):
-        self.model.zk *= self.n_order_visc
+        self.model.zk *= self.multiplier
 
 
 class Beta:
     """Beta plane.
     """
-    solution_mode = 'approximate'
+    solution_mode = 'exact'
 
     def __init__(self, model, beta):
         self.model = model
         self.beta = beta
+        self.solution_multiplier = np.exp(self.model.timestepper.dt * self.beta
+                                          * self.model.ikx * self.model.wv2i)
 
     def __call__(self):
-        self.model.rhs -= self.beta * (self.model.ikx * self.model.psik)
+        self.model.zk *= self.solution_multiplier
 
 
 class StochasticRingForcing:
-    """White-in-time stochastic forcing. Details to follow.
+    """White-in-time stochastic forcing concentrated on a band of wavenumbers.
+    Energy is input at a mean rate which is uniform across forced wavenumbers.
     """
     solution_mode = 'discrete'
 
@@ -185,3 +165,23 @@ def __call__(self):
                              + 1j * self.rng.normal(size=self.model.wv.size),
                              self.model.wv.shape) * self.fk_vars ** 0.5
         self.model.zk += self.fk * self.model.timestepper.dt ** 0.5
+
+
+class SpectralFilter:
+    """Exponential spectral filter for applying highly scale-selective but
+    non-physical dissipation at small scales.
+    """
+    solution_mode = 'exact'
+
+    def __init__(self, model):
+        filterfac = 23.6
+        cphi = 0.65 * np.pi
+        wvx = np.sqrt(
+            (model.kx * model.dx) ** 2. + (model.ky * model.dx) ** 2.)
+        exp_filter = np.exp(-filterfac * (wvx - cphi) ** 4.)
+        exp_filter[wvx <= cphi] = 1.
+        self.exp_filter = exp_filter
+        self.model = model
+
+    def __call__(self):
+        self.model.zk *= self.exp_filter

model.py

@@ -64,17 +64,17 @@ def _evolve_one_step(self):
         """
         self.psik = -self.wv2i * self.zk
         approximate_mode_mechanisms = False
-        for mechanism in self.mechanisms:
+        for _, mechanism in self.mechanisms.items():
             if mechanism.solution_mode == 'approximate':
                 mechanism()
                 approximate_mode_mechanisms = True
         if not approximate_mode_mechanisms:
             self.rhs = 0.
         self.timestepper.step()
-        for mechanism in self.mechanisms:
+        for _, mechanism in self.mechanisms.items():
             if mechanism.solution_mode == 'discrete':
                 mechanism()
-        for mechanism in self.mechanisms:
+        for _, mechanism in self.mechanisms.items():
             if mechanism.solution_mode == 'exact':
                 mechanism()
 

plots.py

@@ -4,18 +4,21 @@
     - Add a function for producing videos.
 """
 
+import pyfftw
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
 import spatial_statistics
+from scipy.interpolate import RegularGridInterpolator
+fftw = pyfftw.interfaces.numpy_fft
 plt.ioff()
 
 
 def plot_vorticity_field(model, halfrange=None, filename='figures/tmp_z.png'):
     """Plot the vorticity field.
     """
     fig, ax = plt.subplots()
-    ax.pcolormesh(model.x, model.y, model.z.T,
+    ax.pcolormesh(model.x, model.y, model.z,
                   norm=mpl.colors.CenteredNorm(halfrange=halfrange),
                   cmap='RdBu')
     ax.set_xlim(0., 2. * np.pi)
@@ -26,16 +29,98 @@ def plot_vorticity_field(model, halfrange=None, filename='figures/tmp_z.png'):
     plt.close()
 
 
+def plot_vorticity_field_upscale(model, halfrange=None, upscale_factor=4,
+                                 filename='figures/tmp_z.png'):
+    """Plot the vorticity field.
+    """
+    x_l = model.x[0]
+    y_l = model.y[:, 0]
+    x_lp = np.concatenate(
+        (x_l[0:1] - (x_l[1] - x_l[0]), x_l, x_l[-1:] + (x_l[1] - x_l[0])))
+    y_lp = np.concatenate(
+        (y_l[0:1] - (y_l[1] - y_l[0]), y_l, y_l[-1:] + (y_l[1] - y_l[0])))
+    z_lp = np.pad(model.z, pad_width=1, mode='wrap')
+    interp = RegularGridInterpolator((x_lp, y_lp), z_lp, method='linear')
+    L = 2 * np.pi
+    n_x = upscale_factor * model.n_x
+    x, y = np.meshgrid(
+        L * np.arange(0.5, n_x) / n_x,
+        L * np.arange(0.5, n_x) / n_x)
+    z = interp(
+        np.concatenate((x.reshape((-1, 1)), y.reshape((-1, 1))), axis=1)
+        ).reshape(x.shape)
+    fig, ax = plt.subplots()
+    ax.pcolormesh(x, y, z,
+                  norm=mpl.colors.CenteredNorm(halfrange=halfrange),
+                  cmap='RdBu')
+    ax.set_xlim(0., 2. * np.pi)
+    ax.set_ylim(0., 2. * np.pi)
+    ax.set_aspect('equal')
+    fig.tight_layout()
+    plt.savefig(filename, dpi=576)
+    plt.close()
+
+
+def plot_vorticity_field_upscalef(model, halfrange=None, upscale_factor=4,
+                                  filename='figures/tmp_z.png'):
+    """Plot the vorticity field.
+    """
+    pad = upscale_factor
+    m_x = int(pad * model.n_x)
+    m_k = m_x // 2 + 1
+    padder = np.ones(m_x, dtype=bool)
+    padder[int(model.n_x / 2):
+           int(model.n_x * (pad - 0.5)):] = False
+    zk_padded = np.zeros((m_x, m_k), dtype='complex128')
+    zk_padded[padder, :model.n_kx] = (
+        model.zk)[:, :] * pad ** 2
+    z_up = fftw.irfft2(zk_padded)
+
+    L = 2 * np.pi
+    n_x = upscale_factor * model.n_x
+    x, y = np.meshgrid(
+        L * np.arange(0.5, n_x) / n_x,
+        L * np.arange(0.5, n_x) / n_x)
+    fig, ax = plt.subplots()
+    ax.pcolormesh(x, y, z_up,
+                  norm=mpl.colors.CenteredNorm(halfrange=halfrange),
+                  cmap='RdBu')
+    ax.set_xlim(0., 2. * np.pi)
+    ax.set_ylim(0., 2. * np.pi)
+    ax.set_aspect('equal')
+    fig.tight_layout()
+    plt.savefig(filename, dpi=576)
+    plt.close()
+
+
+def plot_speed_field(model, filename='figures/tmp_speed.png'):
+    """Plot the vorticity field.
+    """
+    fig, ax = plt.subplots()
+    ax.pcolormesh(model.x, model.y, (model.u ** 2 + model.v ** 2) ** 0.5,
+                  vmin=0., cmap='Greys_r')
+    ax.set_xlim(0., 2. * np.pi)
+    ax.set_ylim(0., 2. * np.pi)
+    ax.set_aspect('equal')
+    fig.tight_layout()
+    plt.savefig(filename, dpi=576)
+    plt.close()
+
+
 def plot_isotropic_energy_spectrum(model, filename='figures/tmp_E.png',
                                    ymin=None):
     """Plot the isotropic energy spectrum.
     """
     kr, spec_iso = spatial_statistics.isotropic_energy_spectrum(model)
     fig, ax = plt.subplots()
     ax.loglog(kr, spec_iso, 'k')
+    ax.loglog(kr, kr ** -(5 / 3), 'g--')
+    ax.loglog(kr, kr ** -3, 'b--')
+    ax.loglog(kr, kr ** -2, 'r--')
     ax.set_ylim(ymin, None)
     ax.set_xlabel(r"$k$")
     ax.set_ylabel(r"$E(k)$")
+    ax.grid(True)
     fig.tight_layout()
     plt.savefig(filename, dpi=576)
     plt.close()
@@ -48,9 +133,35 @@ def plot_isotropic_enstrophy_spectrum(model, filename='figures/tmp_Z.png',
     kr, spec_iso = spatial_statistics.isotropic_enstrophy_spectrum(model)
     fig, ax = plt.subplots()
     ax.loglog(kr, spec_iso, 'k')
+    ax.loglog(kr, kr ** +(1 / 3), 'g--')
+    ax.loglog(kr, kr ** -1, 'b--')
+    ax.loglog(kr, kr ** 0., 'r--')
     ax.set_ylim(ymin, None)
     ax.set_xlabel(r"$k$")
-    ax.set_ylabel(r"$E(k)$")
+    ax.set_ylabel(r"$Z(k)$")
+    ax.grid(True)
+    fig.tight_layout()
+    plt.savefig(filename, dpi=576)
+    plt.close()
+
+
+def plot_time_series(t, quantity, ymin=None, ylabel=None,
+                     filename='figures/tmp_E'):
+    fig, ax = plt.subplots()
+    ax.plot(t, quantity, 'k')
+    ax.set_ylim(ymin, None)
+    ax.set_xlabel(r"$t$")
+    ax.set_ylabel(ylabel)
+    fig.tight_layout()
+    plt.savefig(filename, dpi=576)
+    plt.close()
+
+
+def plot_zonally_averaged_velocity(model, filename='figures/tmp_ubar.png'):
+    fig, ax = plt.subplots()
+    ax.plot(model.u.mean(axis=1), model.y, 'k')
+    ax.set_xlabel(r"$u(y)$")
+    ax.set_ylabel(r"$y$")
     fig.tight_layout()
     plt.savefig(filename, dpi=576)
     plt.close()

spatial_statistics.py

@@ -1,8 +1,5 @@
 """
 A module for computing spatial statistics of a `Model` instance.
-
-TODO:
-    - Add function for computing a set of statistics for a full dataset.
 """
 
 import numpy as np
@@ -18,10 +15,10 @@ def spectral_variance(model, fk):
     fk.
     """
 
-    var_dens = 2. * np.abs(fk) ** 2 / model.n_x ** 4
+    var_dens = 2. * np.abs(fk) ** 2
     var_dens[..., 0] /= 2
     var_dens[..., -1] /= 2
-    return var_dens.sum(axis=(-1, -2))
+    return var_dens.sum(axis=(-1, -2)) / model.n_x ** 4
 
 
 def energy(model):
@@ -30,8 +27,14 @@ def energy(model):
     return 0.5 * spectral_variance(model, model.wv * model.psik)
 
 
+def energy_via_vorticity(model):
+    """Calculate mean energy per unit area using zk.
+    """
+    return 0.5 * spectral_variance(model, model.zk * model.wv2i ** 0.5)
+
+
 def enstrophy(model):
-    """Calculate mean enstrophy per unit area using psik.
+    """Calculate mean enstrophy per unit area using zk.
     """
     return 0.5 * spectral_variance(model, model.zk)
 
@@ -103,8 +106,33 @@ def eddy_turnover_time(model):
     return 2 * np.pi / np.sqrt(enstrophy(model))
 
 
+def velocity_kurtosis(model):
+    """Calculate the kurtosis of the velocity field (u component only).
+    """
+    return np.mean(model.u ** 4) / np.var(model.u) ** 2
+
+
+def vorticity_kurtosis(model):
+    """Calculate the vorticity of the vorticity field.
+    """
+    return np.mean(model.z ** 4) / np.var(model.z) ** 2
+
+
+def reynolds_number(model):
+    """Calculate Reynold's number.
+    """
+    return np.sqrt(np.mean(model.u ** 2 + model.v ** 2)) / (
+        energy_centroid(model) * model.mechanisms['viscosity'].coefficient)
+
+
+def energy_dissipation_due_to_viscosity(model):
+    """Calculate rate of energy dissipation per unit time due to viscosity.
+    """
+    return 2 * model.mechanisms['viscosity'].coefficient * enstrophy(model)
+
+
 def time_series(data_dir, n_x, twrite, n_snapshots):
-    """Calculate and save time series of some specified statistics for data"""
+    """Calculate and save time series of some specified statistics."""
     m = Model(n_x)
     E = []
     Z = []