initial steps adding py_driver_2d model

based on test_problem

only iage TracerModule implemented
    advection, vert_mix, and horiz_mix processes not vetted with multiple tracers

bldepth variation in time for vert_mix currently turned off
    adaptive time-stepping use many steps when bldepth varies
    perhaps this will improve when #54 is addressed

comp_fcn uses scipy.integrate.solve_ivp to integrate ODEs
    will eventually migrate to PETSc

apply_precond_jacobian not functional yet
    will do this after comp_fcn migration to PETSC

On branch py_driver_2d
Changes to be committed:
	new file:   input/py_driver_2d/model_params.cfg
	new file:   input/py_driver_2d/newton_krylov.cfg
	new file:   input/py_driver_2d/tracer_module_defs.yaml
	new file:   src/py_driver_2d/__init__.py
	new file:   src/py_driver_2d/advection.py
	new file:   src/py_driver_2d/horiz_mix.py
	new file:   src/py_driver_2d/iage.py
	new file:   src/py_driver_2d/model_process.py
	new file:   src/py_driver_2d/model_state.py
	new file:   src/py_driver_2d/setup_solver.py
	new file:   src/py_driver_2d/tracer_module_state.py
	new file:   src/py_driver_2d/vert_mix.py

input/py_driver_2d/model_params.cfg

@@ -0,0 +1,29 @@
+[modelinfo]
+
+# depth axis definition
+depth_axisname=depth
+depth_units=m
+depth_edge_start=0.0
+depth_edge_end=4000.0
+
+# depth_nlevs=125
+# depth_delta_ratio_max=11.8
+
+# depth_nlevs=80
+# depth_delta_ratio_max=9.0
+
+depth_nlevs=40
+depth_delta_ratio_max=19.0
+
+# ypos axis definition
+ypos_axisname=ypos
+ypos_units=m
+ypos_edge_start=0.0
+ypos_edge_end=50.0e5
+ypos_delta_ratio_max=1.0
+
+# ypos_nlevs=160
+
+# ypos_nlevs=80
+
+ypos_nlevs=40

input/py_driver_2d/newton_krylov.cfg

@@ -0,0 +1,65 @@
+# configuration file of defaults for model=py_driver_2d
+
+[DEFAULT]
+
+# name of model
+model_name=py_driver_2d
+
+# should logfile output avoid user/cfg specific content
+# to make logging output reproducible
+logging_reproducible=False
+
+# directory where files are stored
+workdir=%(HOME)s/py_driver_2d_work
+
+# cfg vars that are allowed to have no value
+no_value_allowed=cfg_fname_out,region_mask_fname,region_mask_varname
+
+[solverinfo]
+
+# name of file that cfg contents are written to
+cfg_out_fname=%(workdir)s/newton_krylov.cfg.out
+
+# name of file that logging entries are written to
+logging_fname=%(workdir)s/newton_krylov.log
+
+# level of logging entries to be written (e.g., INFO or DEBUG)
+logging_level=INFO
+
+# relative tolerance for Newton convergence
+newton_rel_tol=1.0e-8
+
+# maximum Newton iteration
+newton_max_iter=3
+
+# perform a fixed-point iteration at the end of a Newton iteration
+# this is only appropriate for fixed-point problems
+post_newton_fp_iter=1
+
+[modelinfo]
+
+# should solver exit after each comp_fcn invocation and reinvoke solver
+reinvoke=True
+
+# name of script for invoking nk_driver.py
+invoker_script_fname=%(workdir)s/nk_driver.sh
+
+# name of file and var in that file for weight for
+# mean and dot_prod operations
+# model implementation assumes that axes info is in grid_weight_fname
+grid_weight_fname=%(workdir)s/grid_vars.nc
+grid_weight_varname=grid_weight
+
+# name and file and var in that file for region_mask
+# either unset indicates that region_mask is all ones
+region_mask_fname
+region_mask_varname
+
+# names of tracer modules that solver is being applied to
+tracer_module_names=iage
+
+# name of file data with tracer module definitions
+tracer_module_defs_fname=%(repo_root)s/input/%(model_name)s/tracer_module_defs.yaml
+
+# name of file containing initial iterate
+init_iterate_fname=%(workdir)s/gen_init_iterate/init_iterate.nc

input/py_driver_2d/tracer_module_defs.yaml

@@ -0,0 +1,15 @@
+# tracer module definitions for model=py_driver_2d
+
+tracer_module_defs:
+    iage:  # module name
+        tracers:
+            iage:
+                attrs: {long_name: "ideal age", units: "years"}
+                init_iterate_val_depths: [55.0, 200.0]
+                init_iterate_vals: [0.0, 2.0]
+
+precond_matrix_defs:
+    base:
+        hist_to_precond_varnames:
+            - 'vert_mixing_coeff:mean'
+            - 'vert_mixing_coeff:log_mean'

src/py_driver_2d/advection.py

@@ -0,0 +1,99 @@
+"""functions related to advection"""
+
+import numpy as np
+
+from .model_process import ModelProcess
+
+
+class Advection(ModelProcess):
+    """class related to advection"""
+
+    def __init__(self, depth, ypos):
+
+        Advection.depth = depth
+        Advection.ypos = ypos
+
+        self.gen_vel_field(depth, ypos)
+
+        self._tend_work_y = np.zeros((len(self.depth), len(self.ypos) + 1))
+        self._tend_work_z = np.zeros((len(self.depth) + 1, len(self.ypos)))
+
+        super().__init__(depth, ypos)
+
+    @staticmethod
+    def gen_vel_field(depth, ypos):
+        """generate streamfunction and velocity field"""
+
+        depth_norm = (depth.edges - depth.edges.min()) / (
+            depth.edges.max() - depth.edges.min()
+        )
+        stretch = 2.0
+        depth_norm = stretch * depth_norm / (1 + (stretch - 1) * depth_norm)
+        depth_fcn = (27.0 / 4.0) * depth_norm * (1.0 - depth_norm) ** 2
+
+        ypos_norm = (ypos.edges - ypos.edges.min()) / (
+            ypos.edges.max() - ypos.edges.min()
+        )
+        ypos_fcn = 4.0 * ypos_norm * (1.0 - ypos_norm)
+
+        stream = np.outer(depth_fcn, ypos_fcn)
+
+        # normalize so that max vvel ~ 0.1 m/s
+        vvel = (stream[1:, :] - stream[:-1, :]) * depth.delta_r[:, np.newaxis]
+        stream = stream * 0.1 / abs(vvel).max()
+
+        vvel = (stream[1:, :] - stream[:-1, :]) * depth.delta_r[:, np.newaxis]
+        wvel = (stream[:, 1:] - stream[:, :-1]) * ypos.delta_r
+
+        Advection.stream = stream
+        Advection.vvel = vvel
+        Advection.wvel = wvel
+
+    def comp_tend(self, time, tracer_vals):
+        """single tracer tendency from advection"""
+        self._tend_work_y[:, 1:-1] = 0.5 * (tracer_vals[:, 1:] + tracer_vals[:, :-1])
+        self._tend_work_y *= self.vvel
+
+        res = (self._tend_work_y[:, :-1] - self._tend_work_y[:, 1:]) * self.ypos.delta_r
+
+        self._tend_work_z[1:-1, :] = 0.5 * (tracer_vals[1:, :] + tracer_vals[:-1, :])
+        self._tend_work_z *= self.wvel
+
+        res += (
+            self._tend_work_z[1:, :] - self._tend_work_z[:-1, :]
+        ) * self.depth.delta_r[:, np.newaxis]
+
+        return res
+
+    def get_hist_vars_metadata(self):
+        """return dict of process-specific history variable metadata"""
+
+        depth_name = self.depth.axisname
+        depth_edges_name = self.depth.dump_names["edges"]
+        ypos_name = self.ypos.axisname
+        ypos_edges_name = self.ypos.dump_names["edges"]
+
+        hist_vars_metadata = {}
+        hist_vars_metadata["stream"] = {
+            "dimensions": (depth_edges_name, ypos_edges_name),
+            "attrs": {"long_name": "velocity streamfunction", "units": "m^2 / s"},
+        }
+        hist_vars_metadata["vvel"] = {
+            "dimensions": (depth_name, ypos_edges_name),
+            "attrs": {"long_name": "velocity in ypos direction", "units": "m / s"},
+        }
+        hist_vars_metadata["wvel"] = {
+            "dimensions": (depth_edges_name, ypos_name),
+            "attrs": {"long_name": "velocity in depth direction", "units": "m / s"},
+        }
+
+        return hist_vars_metadata
+
+    def hist_write(self, sol, fptr_hist):
+        """write processs-specific history variables"""
+
+        fptr_hist.variables["stream"][:] = self.stream
+        fptr_hist.variables["vvel"][:] = self.vvel
+        fptr_hist.variables["wvel"][:] = self.wvel
+
+        fptr_hist.sync()

src/py_driver_2d/horiz_mix.py

@@ -0,0 +1,81 @@
+"""functions related to horizontal mixing"""
+
+import numpy as np
+
+from .advection import Advection
+from .model_process import ModelProcess
+
+
+class HorizMix(ModelProcess):
+    """class related to horizontal mixing"""
+
+    def __init__(self, depth, ypos):
+
+        HorizMix.depth = depth
+        HorizMix.ypos = ypos
+
+        self._tend_work = np.zeros((len(depth), len(ypos) + 1))
+
+        # compute (time-invariant) mixing coefficients
+        self._mixing_coeff = self._comp_mixing_coeff(mixing_coeff_const=1000.0)
+
+        super().__init__(depth, ypos)
+
+    def _comp_mixing_coeff(self, mixing_coeff_const):
+        """
+        compute mixing_coeff values
+        includes 1/dypos term
+        """
+
+        res = np.full((len(self.depth), len(self.ypos) - 1), mixing_coeff_const)
+
+        # increase mixing_coeff to keep grid Peclet number <= 2.0
+        peclet_p5 = (
+            (0.5 / mixing_coeff_const)
+            * self.ypos.delta_mid[:]
+            * abs(Advection.vvel[:, 1:-1])
+        )
+        res *= np.where(peclet_p5 > 1.0, peclet_p5, 1.0)
+
+        return res * self.ypos.delta_mid_r
+
+    def comp_tend(self, time, tracer_vals):
+        """single tracer tendency from mixing, with zero flux boundary conditions"""
+
+        # compute horiz_mix fluxes
+        self._tend_work[:, 1:-1] = self._mixing_coeff * (
+            tracer_vals[:, 1:] - tracer_vals[:, :-1]
+        )
+
+        return (self._tend_work[:, 1:] - self._tend_work[:, :-1]) * self.ypos.delta_r
+
+    def get_hist_vars_metadata(self):
+        """return dict of process-specific history variable metadata"""
+
+        depth_name = self.depth.axisname
+        ypos_edges_name = self.ypos.dump_names["edges"]
+
+        hist_vars_metadata = {}
+        hist_vars_metadata["horiz_mixing_coeff"] = {
+            "dimensions": (depth_name, ypos_edges_name),
+            "attrs": {"long_name": "horizontal mixing coefficient", "units": "m^2 / s"},
+        }
+
+        return hist_vars_metadata
+
+    def hist_write(self, sol, fptr_hist):
+        """write processs-specific history variables"""
+
+        # write tracer module independent vars
+        fptr_hist.variables["horiz_mixing_coeff"][:, 1:-1] = (
+            self._mixing_coeff * self.ypos.delta_mid
+        )
+        # kludge to avoid missing values
+        fptr_hist.variables["horiz_mixing_coeff"][:, 0] = fptr_hist.variables[
+            "horiz_mixing_coeff"
+        ][:, 1]
+        fptr_hist.variables["horiz_mixing_coeff"][:, -1] = fptr_hist.variables[
+            "horiz_mixing_coeff"
+        ][:, -2]
+
+        fptr_hist.sync()

src/py_driver_2d/iage.py

@@ -0,0 +1,59 @@
+"""iage subclass of py_driver_2d's TracerModuleState"""
+
+import numpy as np
+from scipy.linalg import solve_banded
+
+from .tracer_module_state import TracerModuleState
+
+
+class iage(TracerModuleState):  # pylint: disable=invalid-name
+    """iage tracer module specifics for TracerModuleState"""
+
+    def comp_tend(self, time, tracer_vals_flat, processes):
+        """
+        compute tendency for iage
+        tendency units are tr_units / s
+        """
+        print(time / (365.0 * 86400.0))
+        shape = (len(self.depth), len(self.ypos))
+        dtracer_vals_dt = super().comp_tend(time, tracer_vals_flat, processes)
+
+        # restore surface layer to zero at rate of 24 / day over 10 m
+        rate = 24.0 / 86400.0 * self.depth.delta[0] / 10.0
+        tracer_vals_2d = tracer_vals_flat.reshape(shape)
+        dtracer_vals_dt_2d = dtracer_vals_dt.reshape(shape)
+        dtracer_vals_dt_2d[0, :] -= rate * tracer_vals_2d[0, :]
+
+        # age 1/year
+        dtracer_vals_dt[:] += 1.0 / (365.0 * 86400.0)
+
+        return dtracer_vals_dt
+
+    def apply_precond_jacobian(self, time_range, res_tms, mca):
+        """
+        apply preconditioner of jacobian of iage fcn
+
+        time_range: length-2 sequence with start and end times of model
+        res_tms: TracerModuleState object where results are stored
+        mca: (vertical) mixing coefficient, time average
+        """
+
+        self_vals = self.get_tracer_vals_all()[0, :]
+        rhs_vals = (1.0 / (time_range[1] - time_range[0])) * self_vals
+
+        l_and_u = (1, 1)
+        matrix_diagonals = np.zeros((3, len(self.depth)))
+        # d tend[k] / d tracer[k-1]
+        matrix_diagonals[0, 1:] = mca * self.depth.delta_mid_r * self.depth.delta_r[:-1]
+        # d tend[k] / d tracer[k]
+        matrix_diagonals[1, :-1] -= (
+            mca * self.depth.delta_mid_r * self.depth.delta_r[:-1]
+        )
+        matrix_diagonals[1, 1:] -= mca * self.depth.delta_mid_r * self.depth.delta_r[1:]
+        matrix_diagonals[1, 0] -= 240.0 * self.depth.delta_r[0]
+        # d tend[k] / d tracer[k+1]
+        matrix_diagonals[2, :-1] = mca * self.depth.delta_mid_r * self.depth.delta_r[1:]
+
+        res_vals = solve_banded(l_and_u, matrix_diagonals, rhs_vals)
+
+        res_tms.set_tracer_vals_all(res_vals - self_vals)

src/py_driver_2d/model_process.py

@@ -0,0 +1,23 @@
+"""generic model process class"""
+
+
+class ModelProcess:
+    """generic model process class"""
+
+    def __init__(self, depth, ypos):
+        pass
+
+    def comp_tend(self, time, tracer_vals):
+        """tracer tendency from process"""
+        raise NotImplementedError("Method must be implemented in derived class")
+
+    @staticmethod
+    def get_hist_vars_metadata():
+        """
+        return dict of process-specific history variable metadata
+        return empty dict, for subclasses that do not provide this method
+        """
+        return {}
+
+    def hist_write(self, sol, fptr_hist):
+        """write processs-specific history variables"""