merge py_driver_2d (#63)

* 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

* add remap_linear_interpolant method to SpatialAxis class

On branch py_driver_2d
Changes to be committed:
	modified:   src/spatial_axis.py
	modified:   tests/test_spatial_axis.py

* py_driver_2d vert_mix update:
    replace linear interpolation of piecewise linear mixing coefficients profile with conservative remapping
    modest decrease in runtime, due to longer timesteps

    enable bldepth variation in time

On branch py_driver_2d
Changes to be committed:
	modified:   src/py_driver_2d/iage.py
	modified:   src/py_driver_2d/vert_mix.py

* py_driver_2d update: rename dtracer_vals_dt to tracer_tend_vals

On branch py_driver_2d
Changes to be committed:
	modified:   src/py_driver_2d/iage.py
	modified:   src/py_driver_2d/tracer_module_state.py

* py_driver_2d update: add comp_jacobian methods and pass to solve_ivp

On branch py_driver_2d
Changes to be committed:
	modified:   src/py_driver_2d/advection.py
	modified:   src/py_driver_2d/horiz_mix.py
	modified:   src/py_driver_2d/iage.py
	modified:   src/py_driver_2d/model_process.py
	modified:   src/py_driver_2d/model_state.py
	modified:   src/py_driver_2d/tracer_module_state.py
	modified:   src/py_driver_2d/vert_mix.py

* py_driver_2d performance improvements
    use numpy arrays to construct sparse Jacobians, precomputing nnz
    cache Advection and HorizMix Jacobians
    cache remap_linear_interpolant results in VertMix.mixing_coeff

On branch py_driver_2d
Changes to be committed:
	modified:   src/py_driver_2d/advection.py
	modified:   src/py_driver_2d/horiz_mix.py
	modified:   src/py_driver_2d/model_state.py
	modified:   src/py_driver_2d/tracer_module_state.py
	modified:   src/py_driver_2d/vert_mix.py

* py_driver_2d coupled to overall solver

implement apply_precond_jacobian
    precond jacobian is time mean of instantaneous jacobians

fix typo in np.einsum call in TracerModuleState:stats_vars_vals

enable --persist option for py_driver_2d

tighten tolerances in solve_ivp call to 1.0e-6

loosen newton_rel_tol in newton_krylov.cfg to 1.0e-5
    doesn't make sense to be more accurate than solve_ivp tolerances

On branch py_driver_2d
Changes to be committed:
	modified:   input/py_driver_2d/newton_krylov.cfg
	modified:   input/py_driver_2d/tracer_module_defs.yaml
	modified:   src/py_driver_2d/iage.py
	modified:   src/py_driver_2d/model_state.py
	modified:   src/py_driver_2d/tracer_module_state.py
	modified:   src/share.py

* py_driver_2d updates
    adjust vert_mix paramters to increase iage local max
    add more layers in ypos direction

On branch py_driver_2d
Changes to be committed:
	modified:   input/py_driver_2d/model_params.cfg
	modified:   src/py_driver_2d/vert_mix.py

* py_driver_2d: add support for multiple tracers

add 2nd iage tracer with weaker restoring in surface layer

construct jacobian sparsity matrix from jacobian at initial time
    instead of recomputing pattern

On branch py_driver_2d
Changes to be committed:
	modified:   input/py_driver_2d/tracer_module_defs.yaml
	modified:   src/py_driver_2d/advection.py
	modified:   src/py_driver_2d/horiz_mix.py
	modified:   src/py_driver_2d/iage.py
	modified:   src/py_driver_2d/model_process.py
	modified:   src/py_driver_2d/model_state.py
	modified:   src/py_driver_2d/tracer_module_state.py
	modified:   src/py_driver_2d/vert_mix.py
	modified:   src/test_problem/model_state.py

* py_driver_2d: rework apply_precond_jacobian for iage

On branch py_driver_2d
Changes to be committed:
	modified:   src/py_driver_2d/iage.py

* py_driver_2d: add phosphorus tracer module

add {tracer}_depth_ypos_int to history output

rm unused comp_jac_sparsity from class ModelProcess

On branch py_driver_2d
Changes to be committed:
	modified:   input/py_driver_2d/newton_krylov.cfg
	modified:   input/py_driver_2d/tracer_module_defs.yaml
	modified:   src/py_driver_2d/model_process.py
	modified:   src/py_driver_2d/model_state.py
	new file:   src/py_driver_2d/phosphorus.py
	modified:   src/py_driver_2d/tracer_module_state.py

* py_driver_2d: add preformed tracer module

On branch py_driver_2d
Changes to be committed:
	modified:   input/py_driver_2d/model_params.cfg
	modified:   input/py_driver_2d/newton_krylov.cfg
	new file:   input/py_driver_2d/po4_surf.nc
	modified:   input/py_driver_2d/tracer_module_defs.yaml
	new file:   src/py_driver_2d/preformed.py

* py_driver_2d: correct iage restoring rate formula

mv formula to iage.__init__, to avoid it being in more than 1 place

On branch py_driver_2d
Changes to be committed:
	modified:   src/py_driver_2d/iage.py

* update for black: rm trailing comma to satisfy black 20.8b1

On branch py_driver_2d
Changes to be committed:
	modified:   src/py_driver_2d/phosphorus.py

* py_driver_2d: generalize parameter specification for phosphorus tracers

place parameters in dict in phosphorus class

enable specification of parameters in cfg modelinfo

introduce function for evaluation of basic arithmetic operations in a string
    enables phosphorus parameters in cfg file to be expressions

On branch py_driver_2d
Changes to be committed:
	modified:   src/py_driver_2d/phosphorus.py
	modified:   src/utils.py

* py_driver_2d: generalize preformed.gen_surf_restore_fcn and mv to utils

On branch py_driver_2d
Changes to be committed:
	modified:   src/py_driver_2d/preformed.py
	modified:   src/utils.py

* py_driver_2d: modify parameters to increase regenerated po4

On branch py_driver_2d
Changes to be committed:
	modified:   input/py_driver_2d/po4_surf.nc
	modified:   src/py_driver_2d/phosphorus.py
	modified:   src/py_driver_2d/vert_mix.py

* py_driver_2d: refactor __init__ in ModelProcess class and subclasses

On branch py_driver_2d
Changes to be committed:
	modified:   src/py_driver_2d/advection.py
	modified:   src/py_driver_2d/horiz_mix.py
	modified:   src/py_driver_2d/model_process.py
	modified:   src/py_driver_2d/vert_mix.py

input/py_driver_2d/model_params.cfg

@@ -0,0 +1,32 @@
+[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=150
+
+# ypos_nlevs=100
+
+ypos_nlevs=50
+
+surf_restore_fname=%(repo_root)s/input/%(model_name)s/po4_surf.nc
+surf_restore_varname=po4

input/py_driver_2d/newton_krylov.cfg

@@ -0,0 +1,67 @@
+# 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-5
+
+# 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
+# tracer_module_names=phosphorus
+tracer_module_names=preformed
+
+# 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,52 @@
+# 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]
+            iage_slow_rest:
+                attrs:
+                    long_name: "ideal age, slower surface restoring"
+                    units: "years"
+                init_iterate_val_depths: [55.0, 200.0]
+                init_iterate_vals: [0.0, 2.0]
+    phosphorus:  # module name
+        tracers:
+            po4:
+                attrs:
+                    long_name: "phosphate"
+                    units: "mmol / m^3"
+                init_iterate_val_depths: [1.3e2, 2.6e2]
+                init_iterate_vals: [5.5e-3, 4.1e+0]
+                precond_matrix: phosphorus
+            dop:
+                attrs:
+                    long_name: "dissolved organic phosphorus"
+                    units: "mmol / m^3"
+                init_iterate_val_depths: [9.5e1, 1.4e2]
+                init_iterate_vals: [7.1e-2, 1.5e-4]
+            pop:
+                attrs:
+                    long_name: "particulate organic phosphorus"
+                    units: "mmol / m^3"
+                init_iterate_val_depths: [1.7e2, 2.5e2]
+                init_iterate_vals: [1.8e-2, 7.9e-4]
+    preformed:  # module name
+        tracers:
+            po4_pf:
+                attrs:
+                    long_name: "preformed phosphate"
+                    units: "mmol / m^3"
+                init_iterate_val_depths: [0.0]
+                init_iterate_vals: [0.0]
+
+precond_matrix_defs:
+    base:
+        hist_to_precond_varnames:
+            - "time"
+    phosphorus:
+        hist_to_precond_varnames:
+            - "po4"

src/py_driver_2d/advection.py

@@ -0,0 +1,181 @@
+"""functions related to advection"""
+
+import numpy as np
+from scipy import sparse
+
+from .model_process import ModelProcess
+
+
+class Advection(ModelProcess):
+    """class related to advection"""
+
+    def __init__(self, depth, ypos):
+
+        super().__init__(depth, 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)))
+
+        Advection.jacobian_cache = None
+
+    @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"""
+
+        shape = tracer_vals.shape
+        res = np.empty(shape)
+
+        for tracer_ind in range(shape[0]):
+            self._tend_work_y[:, 1:-1] = 0.5 * (
+                tracer_vals[tracer_ind, :, 1:] + tracer_vals[tracer_ind, :, :-1]
+            )
+            self._tend_work_y *= self.vvel
+
+            res[tracer_ind, :] = self.ypos.delta_r * (
+                self._tend_work_y[:, :-1] - self._tend_work_y[:, 1:]
+            )
+
+            self._tend_work_z[1:-1, :] = 0.5 * (
+                tracer_vals[tracer_ind, 1:, :] + tracer_vals[tracer_ind, :-1, :]
+            )
+            self._tend_work_z *= self.wvel
+
+            res[tracer_ind, :] += self.depth.delta_r[:, np.newaxis] * (
+                self._tend_work_z[1:, :] - self._tend_work_z[:-1, :]
+            )
+
+        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()
+
+    def comp_jacobian(self, time, tracer_cnt):
+        """
+        compute jacobian of tracer tendencies from advection
+        jacobian units are 1 / s
+        """
+        if Advection.jacobian_cache is None:
+            # First construct matrix with sparsity pattern for a single tracer.
+            # Then concatenate using sparse.block_diag.
+            depth_n = len(self.depth)
+            ypos_n = len(self.ypos)
+            nnz = 5 * (depth_n - 2) * (ypos_n - 2)
+            nnz += 4 * 2 * (depth_n - 2) + 4 * 2 * (ypos_n - 2) + 3 * 4
+            data = np.empty(nnz)
+            row_ind = np.empty(nnz, int)
+            col_ind = np.empty(nnz, int)
+            mat_ind = 0
+            for depth_i in range(depth_n):
+                for ypos_i in range(ypos_n):
+                    cell_ind = ypos_i + ypos_n * depth_i
+                    tmp_sum = 0.0
+                    # cell shallower
+                    if depth_i > 0:
+                        tmp = -0.5 * self.wvel[depth_i, ypos_i]
+                        tmp *= self.depth.delta_r[depth_i]
+                        tmp_sum += tmp
+                        data[mat_ind] = tmp
+                        row_ind[mat_ind] = cell_ind
+                        col_ind[mat_ind] = cell_ind - ypos_n
+                        mat_ind += 1
+                    # cell to the south
+                    if ypos_i > 0:
+                        tmp = 0.5 * self.vvel[depth_i, ypos_i]
+                        tmp *= self.ypos.delta_r[ypos_i]
+                        tmp_sum += tmp
+                        data[mat_ind] = tmp
+                        row_ind[mat_ind] = cell_ind
+                        col_ind[mat_ind] = cell_ind - 1
+                        mat_ind += 1
+                    # cell to the north
+                    if ypos_i < ypos_n - 1:
+                        tmp = -0.5 * self.vvel[depth_i, ypos_i + 1]
+                        tmp *= self.ypos.delta_r[ypos_i]
+                        tmp_sum += tmp
+                        data[mat_ind] = tmp
+                        row_ind[mat_ind] = cell_ind
+                        col_ind[mat_ind] = cell_ind + 1
+                        mat_ind += 1
+                    # cell deeper
+                    if depth_i < depth_n - 1:
+                        tmp = 0.5 * self.wvel[depth_i + 1, ypos_i]
+                        tmp *= self.depth.delta_r[depth_i]
+                        tmp_sum += tmp
+                        data[mat_ind] = tmp
+                        row_ind[mat_ind] = cell_ind
+                        col_ind[mat_ind] = cell_ind + ypos_n
+                        mat_ind += 1
+                    # cell itself
+                    data[mat_ind] = tmp_sum
+                    row_ind[mat_ind] = cell_ind
+                    col_ind[mat_ind] = cell_ind
+                    mat_ind += 1
+            if mat_ind != nnz:
+                msg = "mat_ind = %d, nnz = %d" % (mat_ind, nnz)
+                raise RuntimeError(msg)
+            dof = ypos_n * depth_n
+            Advection.jacobian_cache = sparse.csr_matrix(
+                (data, (row_ind, col_ind)), shape=(dof, dof)
+            )
+
+        return sparse.block_diag(tracer_cnt * [Advection.jacobian_cache], "csr")