From 98eadc51b24cc569ff667a686627e7bd2510374f Mon Sep 17 00:00:00 2001 From: Martin Brolly Date: Fri, 19 Apr 2024 15:54:14 +0100 Subject: [PATCH] Minor updates. --- mechanisms.py | 62 +++++++++++------------ model.py | 6 +-- plots.py | 115 +++++++++++++++++++++++++++++++++++++++++- spatial_statistics.py | 42 ++++++++++++--- 4 files changed, 182 insertions(+), 43 deletions(-) diff --git a/mechanisms.py b/mechanisms.py index 173ca06..93dcdce 100644 --- a/mechanisms.py +++ b/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 diff --git a/model.py b/model.py index 2f0bb89..55dea17 100644 --- a/model.py +++ b/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() diff --git a/plots.py b/plots.py index 6007625..4f19445 100644 --- a/plots.py +++ b/plots.py @@ -4,10 +4,13 @@ - 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() @@ -15,7 +18,7 @@ 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,6 +29,84 @@ 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. @@ -33,9 +114,13 @@ def plot_isotropic_energy_spectrum(model, filename='figures/tmp_E.png', 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() diff --git a/spatial_statistics.py b/spatial_statistics.py index fe4ada9..6b2cc70 100644 --- a/spatial_statistics.py +++ b/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 = []