Add model equations and configuration to crash course notebook (#7)

* Add initial code for u tendency discretization

* Further updates for numba routines

* Add constants file, oops

* In progress of model scripting

* Finish code updates, minus w tendency and final model test

* bring in w tendency

* results of some quick tests...runs but blows up

* Add written-out crash course with equations, pending c-grid figure and error analysis holoviews

* Add summary, c-grid figure, and update author list

README.md

@@ -25,13 +25,19 @@ Numerical models are widely used, but gaining expertise in how they work has oft
  - [JT Thielen](https://github.com/jthielen)
  - [Sam Gardner](https://github.com/wx4stg)
  - [Roger Riggin](https://github.com/riotrogerriot)
+ - [Justin Spotts](https://github.com/jrspotts)
+ - [Mathieu R](https://github.com/shenronUber)
 
 ### Contributors
 
 <a href="https://github.com/ProjectPythia/cookbook-template/graphs/contributors">
   <img src="https://contrib.rocks/image?repo=ProjectPythia/cookbook-template" />
 </a>
 
+Addition contributions to discussions and decisions for this notebook by:
+
+- [Rachel Smith](https://github.com/rachaellsmith)
+
 ## Resources
 
 This cookbook would not be possible without the vast collection of academic texts and prior work in atmospheric modeling. The key resources used in building this notebook include:

notebooks/config/grid_staggering.ipynb

@@ -295,7 +295,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.8"
+   "version": "3.11.6"
   },
   "nbdime-conflicts": {
    "local_diff": [

notebooks/intro/constants.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+R_d = 8.314462618 / 28.96546e-3 # R / Md, (J/(mol*K)) / (kg/mol) = J/(kg*K)
+c_p = 1.4 * R_d / (1.4 - 1) # J/(kg*K)
+c_v = c_p - R_d
+gravity = 9.80665 # m/(s^(2))
+p0 = 1.0e5

notebooks/intro/driver.py

@@ -0,0 +1,218 @@
+import numpy as np
+import xarray as xr
+
+from util import *
+
+metadata_attrs = {
+    'u': {
+        'units': 'm/s'
+    },
+    'w': {
+        'units': 'm/s'
+    },
+    'theta_p': {
+        'units': 'K'
+    },
+    'pi': {
+        'units': 'dimensionless'
+    },
+    'x': {
+        'units': 'm'
+    },
+    'x_stag': {
+        'units': 'm'
+    },
+    'z': {
+        'units': 'm'
+    },
+    'z_stag': {
+        'units': 'm'
+    },
+    't': {
+        'units': 's'
+    }
+}
+
+class ModelDriver:
+
+    coords = {}
+    prognostic_arrays = {}
+    base_state_arrays = {}
+    diagnostic_arrays = {}
+    params = {}
+
+    def __init__(self, nx, nz, dx, dz, dt, **kwargs):
+        # Set parameters
+        self.nx = nx
+        self.nz = nz
+        self.dx = dx
+        self.dz = dz
+        self.dt = dt
+        for k, v in kwargs.items():
+            if k.endswith('_tendency'):
+                setattr(self, k, v)
+            else:
+                self.params[k] = v
+        self.dtype = dtype = getattr(self, 'dtype', np.float32)
+        self.t_count = 0
+
+        # Define arrays
+        self.coords['x'] = np.arange(self.nx) * self.dx - self.nx * self.dx / 2
+        self.coords['x_stag'] = np.arange(self.nx + 1) * self.dx - (self.nx + 1) * self.dx / 2
+        self.coords['z'] = np.arange(self.nz) * self.dz
+        self.coords['z_stag'] = np.arange(self.nz + 1) * self.dz - self.dz / 2
+        self.prognostic_arrays['u'] = np.zeros((3, nz, nx + 1), dtype=dtype)
+        self.prognostic_arrays['w'] = np.zeros((3, nz + 1, nx), dtype=dtype)
+        self.prognostic_arrays['theta_p'] = np.zeros((3, nz, nx), dtype=dtype)
+        self.prognostic_arrays['pi'] = np.zeros((3, nz, nx), dtype=dtype)
+        self.active_prognostic_variables = ['u', 'w', 'theta_p', 'pi']
+        self.base_state_arrays['theta_base'] = np.zeros(nz, dtype=dtype)
+        self.base_state_arrays['PI_base'] = np.zeros(nz, dtype=dtype)
+        self.base_state_arrays['rho_base'] = np.zeros(nz, dtype=dtype)
+        ## Todo do we need others??
+
+    def initialize_isentropic_base_state(self, theta, pressure_surface):
+        # Set uniform potential temperature
+        self.base_state_arrays['theta_base'] = np.full(
+            self.base_state_arrays['theta_base'].shape, theta, dtype=self.dtype
+        )
+        # Calculate pi based on hydrostatic balance (from surface)
+        self.base_state_arrays['PI_base'] = nondimensional_pressure_hydrostatic(
+            self.base_state_arrays['theta_base'],
+            self.coords['z'],
+            pressure_surface
+        )
+        # Calculate density from theta and pi
+        self.base_state_arrays['rho_base'] = density_from_ideal_gas_law(
+            self.base_state_arrays['theta_base'],
+            self.base_state_arrays['PI_base']
+        )
+
+    def initialize_warm_bubble(self, amplitude, x_radius, z_radius, z_center):
+        if np.min(self.base_state_arrays['theta_base']) <= 0.:
+            raise ValueError("Base state theta must be initialized as positive definite")
+
+        # Create thermal bubble (2D)
+        theta_p, pi = create_thermal_bubble(
+            amplitude, self.coords['x'], self.coords['z'], x_radius, z_radius, 0.0, z_center, 
+            self.base_state_arrays['theta_base']
+        )
+        # Ensure boundary conditions, and add time stacking (future, current, past)
+        self.prognostic_arrays['theta_p'] = np.stack([apply_periodic_lateral_zerograd_vertical(theta_p)] * 3)
+        self.prognostic_arrays['pi'] = np.stack([apply_periodic_lateral_zerograd_vertical(pi)] * 3)
+
+    def prep_new_timestep(self):
+        for var in self.active_prognostic_variables:
+            # Future-current to current-past
+            self.prognostic_arrays[var][0:2] = self.prognostic_arrays[var][1:3]
+
+    def take_first_timestep(self):
+        # check for needed parameters and methods
+        if not 'c_s_sqr' in self.params:
+            raise ValueError("Must set squared speed of sound prior to first timestep")
+        if not (
+            getattr(self, 'u_tendency')
+            and getattr(self, 'w_tendency')
+            and getattr(self, 'theta_p_tendency')
+            and getattr(self, 'pi_tendency')
+        ):
+            raise ValueError("Must set tendency equations prior to first timestep")
+
+        # Increment
+        self.t_count = 1
+
+        # Integrate forward-in-time
+        self.prognostic_arrays['u'][2] = (
+            self.prognostic_arrays['u'][1]
+            + self.dt * apply_periodic_lateral_zerograd_vertical(self.u_tendency(
+                self.prognostic_arrays['u'][1], self.prognostic_arrays['w'][1],
+                self.prognostic_arrays['pi'][1], self.base_state_arrays['theta_base'], self.dx, self.dz
+            ))
+        )
+        self.prognostic_arrays['w'][2] = (
+            self.prognostic_arrays['w'][1]
+            + self.dt * apply_periodic_lateral_zerow_vertical(self.w_tendency(
+                self.prognostic_arrays['u'][1], self.prognostic_arrays['w'][1],
+                self.prognostic_arrays['pi'][1], self.prognostic_arrays['theta_p'][1],
+                self.base_state_arrays['theta_base'], self.dx, self.dz
+            ))
+        )
+        self.prognostic_arrays['theta_p'][2] = (
+            self.prognostic_arrays['theta_p'][1]
+            + self.dt * apply_periodic_lateral_zerograd_vertical(self.theta_p_tendency(
+                self.prognostic_arrays['u'][1], self.prognostic_arrays['w'][1],
+                self.prognostic_arrays['theta_p'][1], self.base_state_arrays['theta_base'], self.dx, self.dz
+            ))
+        )
+        self.prognostic_arrays['pi'][2] = (
+            self.prognostic_arrays['pi'][1]
+            + self.dt * apply_periodic_lateral_zerograd_vertical(self.pi_tendency(
+                self.prognostic_arrays['u'][1], self.prognostic_arrays['w'][1],
+                self.prognostic_arrays['pi'][1], self.base_state_arrays['theta_base'], 
+                self.base_state_arrays['rho_base'], self.params['c_s_sqr'], self.dx, self.dz
+            ))
+        )
+
+        self.prep_new_timestep()
+
+    def take_single_timestep(self):
+        # Check if initialized
+        if self.t_count == 0:
+            raise RuntimeError("Must run initial timestep!")
+        self.t_count += 1
+
+        # Integrate leapfrog
+        self.prognostic_arrays['u'][2] = (
+            self.prognostic_arrays['u'][0]
+            + 2 * self.dt * apply_periodic_lateral_zerograd_vertical(self.u_tendency(
+                self.prognostic_arrays['u'][1], self.prognostic_arrays['w'][1],
+                self.prognostic_arrays['pi'][1], self.base_state_arrays['theta_base'], self.dx, self.dz
+            ))
+        )
+        self.prognostic_arrays['w'][2] = (
+            self.prognostic_arrays['w'][0]
+            + 2 * self.dt * apply_periodic_lateral_zerow_vertical(self.w_tendency(
+                self.prognostic_arrays['u'][1], self.prognostic_arrays['w'][1],
+                self.prognostic_arrays['pi'][1], self.prognostic_arrays['theta_p'][1],
+                self.base_state_arrays['theta_base'], self.dx, self.dz
+            ))
+        )
+        self.prognostic_arrays['theta_p'][2] = (
+            self.prognostic_arrays['theta_p'][0]
+            + 2 * self.dt * apply_periodic_lateral_zerograd_vertical(self.theta_p_tendency(
+                self.prognostic_arrays['u'][1], self.prognostic_arrays['w'][1],
+                self.prognostic_arrays['theta_p'][1], self.base_state_arrays['theta_base'], self.dx, self.dz
+            ))
+        )
+        self.prognostic_arrays['pi'][2] = (
+            self.prognostic_arrays['pi'][0]
+            + 2 * self.dt * apply_periodic_lateral_zerograd_vertical(self.pi_tendency(
+                self.prognostic_arrays['u'][1], self.prognostic_arrays['w'][1],
+                self.prognostic_arrays['pi'][1], self.base_state_arrays['theta_base'], 
+                self.base_state_arrays['rho_base'], self.params['c_s_sqr'], self.dx, self.dz
+            ))
+        )
+
+        self.prep_new_timestep()
+
+    def integrate(self, n_steps):
+        for _ in range(n_steps):
+            self.take_single_timestep()
+
+    def current_state(self):
+        """Export the prognostic variables, with coordinates, at current time."""
+        data_vars = {}
+        for var in self.active_prognostic_variables:
+            if var == 'u':
+                dims = ('t', 'z', 'x_stag')
+            elif var == 'w':
+                dims = ('t', 'z_stag', 'x')
+            else:
+                dims = ('t', 'z', 'x')
+            data_vars[var] = xr.Variable(dims, self.prognostic_arrays[var][1:2].copy(), metadata_attrs[var])
+        data_vars['x'] = xr.Variable('x', self.coords['x'], metadata_attrs['x'])
+        data_vars['x_stag'] = xr.Variable('x_stag', self.coords['x_stag'], metadata_attrs['x_stag'])
+        data_vars['z'] = xr.Variable('z', self.coords['z'], metadata_attrs['z'])
+        data_vars['z_stag'] = xr.Variable('z_stag', self.coords['z_stag'], metadata_attrs['z_stag'])
+        data_vars['t'] = xr.Variable('t', [self.t_count * self.dt], metadata_attrs['t'])
+        return xr.Dataset(data_vars)

notebooks/intro/util.py

@@ -0,0 +1,70 @@
+# Helper functions for dry model
+import numpy as np
+import numba
+
+from constants import *
+
+@numba.njit()
+def nondimensional_pressure_hydrostatic(theta, z, pressure_surface):
+    pi = np.zeros_like(theta)
+    # Start at (w level) surface as given
+    pi_sfc = (pressure_surface / p0)**(R_d / c_p)
+    # Go down, half level, for u level, using above-surface theta
+    pi[0] = pi_sfc + gravity * (z[1] - z[0]) / (2 * c_p * theta[1])
+    # Now, integrate upward over full u levels
+    for i in range(1, pi.shape[0]):
+        theta_at_level = 0.5 * (theta[i] + theta[i - 1])
+        pi[i] = pi[i - 1] - gravity * (z[i] - z[i - 1]) / (c_p * theta_at_level)
+    return pi
+
+@numba.njit()
+def density_from_ideal_gas_law(theta, pi):
+    return p0 * pi ** (c_v / R_d) / (R_d * theta)
+
+@numba.njit()
+def create_thermal_bubble(amplitude, x, z, x_radius, z_radius, x_center, z_center, theta_base):
+    # Coordinates in 2d
+    xx = np.broadcast_to(x[None, :], (z.shape[0], x.shape[0]))
+    zz = np.broadcast_to(z[:, None], (z.shape[0], x.shape[0]))
+    rad = np.sqrt(((zz - z_center) / z_radius)**2 + ((xx - x_center) / x_radius)**2)
+    # Create thermal bubble
+    theta_p = np.zeros_like(xx)
+    for k in range(rad.shape[0]):
+        for i in range(rad.shape[1]):
+            if rad[k, i] <= 1.0:
+                theta_p[k, i] = 0.5 * amplitude * (np.cos(np.pi * rad[k, i]) + 1.0)
+    # Create balanced pi, integrating downward from assumed zero at topmost level
+    pi = np.zeros_like(theta_p)
+    for k in range(rad.shape[0] - 2, -1, -1):
+        for i in range(rad.shape[1]):
+            integrand_trapz = 0.5 * (
+                theta_p[k + 1, i] / theta_base[k + 1]**2
+                + theta_p[k, i] / theta_base[k]**2
+            )
+            pi[k, i] = pi[k + 1, i] - gravity * (z[k + 1] - z[k]) / c_p * integrand_trapz
+    # Return results
+    return theta_p, pi
+
+@numba.njit()
+def apply_periodic_lateral_zerograd_vertical(a):
+    # Bottom and top (no gradient)
+    for i in range(0, a.shape[1]):
+        a[0, i] = a[1, i]
+        a[a.shape[0] - 1, i] = a[a.shape[0] - 2, i]
+    # Left and right (mirrored)
+    for k in range(1, a.shape[0] - 1):
+        a[k, 0] = a[k, a.shape[1] - 2]
+        a[k, a.shape[1] - 1] = a[k, 1]
+    return a
+
+@numba.njit()
+def apply_periodic_lateral_zerow_vertical(a):
+    # Bottom and top (fixed zero)
+    for i in range(0, a.shape[1]):
+        a[0, i] = a[1, i] = 0
+        a[a.shape[0] - 1, i] = a[a.shape[0] - 2, i] = 0
+    # Left and right (mirrored)
+    for k in range(1, a.shape[0] - 1):
+        a[k, 0] = a[k, a.shape[1] - 2]
+        a[k, a.shape[1] - 1] = a[k, 1]
+    return a