polaris conda environment, spack environment, compilers and system modules

Requirement: Shared steps are run before steps that depend on their output.

Requirement: Working directory step paths are easily discoverable by users.

Requirement: The output of shared steps may be used by multiple tasks.

Requirement: tasks do not rely on outputs from steps in other tasks

Implementation: Shared steps are run before steps that depend on their output.

Implementation: Working directory step paths are easily discoverable by users.

Implementation: The output of shared steps may be used by multiple tasks.

Implementation: tasks do not rely on outputs from steps in other tasks

Testing And Validation: Shared steps are run before steps that depend on their output.

Testing And Validation: Shared steps are not daughters of a task

Testing And Validation: Working directory structure is intuitive.

Testing And Validation: Working directory step paths are easily discoverable by users.

Testing And Validation: The output of shared steps may be used by multiple tasks.

Testing And Validation: tasks do not rely on outputs from steps in other tasks

Testing and Validation: name-of-topic-here (same as Requirement)

get_cell_width() (polaris.mesh.IcosahedralMeshStep static method)

get_subdivisions() (polaris.mesh.IcosahedralMeshStep static method)

build_cell_width_lat_lon() (polaris.mesh.QuasiUniformSphericalMeshStep method)

make_jigsaw_mesh() (polaris.mesh.QuasiUniformSphericalMeshStep method)

save_and_plot_cell_width() (polaris.mesh.spherical.SphericalBaseStep method)

ConvergenceAnalysis (class in polaris.ocean.convergence.analysis)

__init__() (polaris.ocean.convergence.analysis.ConvergenceAnalysis method)

compute_error() (polaris.ocean.convergence.analysis.ConvergenceAnalysis method)

exact_solution() (polaris.ocean.convergence.analysis.ConvergenceAnalysis method)

run() (polaris.ocean.convergence.analysis.ConvergenceAnalysis method)

setup() (polaris.ocean.convergence.analysis.ConvergenceAnalysis method)

__init__() (polaris.ocean.convergence.forward.ConvergenceForward method)

SphericalConvergenceForward (class in polaris.ocean.convergence.spherical)

setup_ssh_adjustment_steps() (polaris.ocean.ice_shelf.IceShelfTask method)

compute_cell_count() (polaris.ocean.ice_shelf.SshForward method)

dynamic_model_config() (polaris.ocean.ice_shelf.SshForward method)

add_spherical_base_mesh_step() (in module polaris.ocean.mesh.spherical)

map_from_native_model_vars() (polaris.ocean.model.OceanIOStep method)

map_to_native_model_vars() (polaris.ocean.model.OceanIOStep method)

compute_cell_count() (polaris.ocean.model.OceanModelStep method)

constrain_resources() (polaris.ocean.model.OceanModelStep method)

add_baroclinic_channel_tasks() (in module polaris.ocean.tasks.baroclinic_channel)

__init__() (polaris.ocean.tasks.baroclinic_channel.decomp.Decomp method)

Default (class in polaris.ocean.tasks.baroclinic_channel.default)

__init__() (polaris.ocean.tasks.baroclinic_channel.default.Default method)

Forward (class in polaris.ocean.tasks.baroclinic_channel.forward)

__init__() (polaris.ocean.tasks.baroclinic_channel.forward.Forward method)

__init__() (polaris.ocean.tasks.baroclinic_channel.init.Init method)

setup() (polaris.ocean.tasks.baroclinic_channel.init.Init method)

Restart (class in polaris.ocean.tasks.baroclinic_channel.restart)

__init__() (polaris.ocean.tasks.baroclinic_channel.restart.Restart method)

RestartStep (class in polaris.ocean.tasks.baroclinic_channel.restart.restart_step)

__init__() (polaris.ocean.tasks.baroclinic_channel.rpe.Rpe method)

configure() (polaris.ocean.tasks.baroclinic_channel.rpe.Rpe method)

Analysis (class in polaris.ocean.tasks.baroclinic_channel.rpe.analysis)

__init__() (polaris.ocean.tasks.baroclinic_channel.rpe.analysis.Analysis method)

run() (polaris.ocean.tasks.baroclinic_channel.rpe.analysis.Analysis method)

setup() (polaris.ocean.tasks.baroclinic_channel.rpe.analysis.Analysis method)

Threads (class in polaris.ocean.tasks.baroclinic_channel.threads)

__init__() (polaris.ocean.tasks.baroclinic_channel.threads.Threads method)

Validate (class in polaris.ocean.tasks.baroclinic_channel.validate)

__init__() (polaris.ocean.tasks.baroclinic_channel.validate.Validate method)

run() (polaris.ocean.tasks.baroclinic_channel.validate.Validate method)

__init__() (polaris.ocean.tasks.baroclinic_channel.viz.Viz method)

__init__() (polaris.ocean.tasks.barotropic_gyre.BarotropicGyre method)

add_barotropic_gyre_tasks() (in module polaris.ocean.tasks.barotropic_gyre)

Analysis (class in polaris.ocean.tasks.barotropic_gyre.analysis)

__init__() (polaris.ocean.tasks.barotropic_gyre.analysis.Analysis method)

compute_error() (polaris.ocean.tasks.barotropic_gyre.analysis.Analysis method)

exact_solution() (polaris.ocean.tasks.barotropic_gyre.analysis.Analysis method)

run() (polaris.ocean.tasks.barotropic_gyre.analysis.Analysis method)

__init__() (polaris.ocean.tasks.barotropic_gyre.forward.Forward method)

compute_max_time_step() (in module polaris.ocean.tasks.barotropic_gyre.forward)

__init__() (polaris.ocean.tasks.barotropic_gyre.init.Init method)

add_cosine_bell_tasks() (in module polaris.ocean.tasks.cosine_bell)

__init__() (polaris.ocean.tasks.cosine_bell.analysis.Analysis method)

exact_solution() (polaris.ocean.tasks.cosine_bell.analysis.Analysis method)

__init__() (polaris.ocean.tasks.cosine_bell.forward.Forward method)

compute_cell_count() (polaris.ocean.tasks.cosine_bell.forward.Forward method)

dynamic_model_config() (polaris.ocean.tasks.cosine_bell.forward.Forward method)

configure() (polaris.ocean.tasks.geostrophic.Geostrophic method)

add_geostrophic_tasks() (in module polaris.ocean.tasks.geostrophic)

__init__() (polaris.ocean.tasks.geostrophic.analysis.Analysis method)

exact_solution() (polaris.ocean.tasks.geostrophic.analysis.Analysis method)

get_output_field() (polaris.ocean.tasks.geostrophic.analysis.Analysis method)

__init__() (polaris.ocean.tasks.geostrophic.forward.Forward method)

add_ice_shelf_2d_tasks() (in module polaris.ocean.tasks.ice_shelf_2d)

__init__() (polaris.ocean.tasks.ice_shelf_2d.default.Default method)

__init__() (polaris.ocean.tasks.ice_shelf_2d.forward.Forward method)

compute_cell_count() (polaris.ocean.tasks.ice_shelf_2d.forward.Forward method)

dynamic_model_config() (polaris.ocean.tasks.ice_shelf_2d.forward.Forward method)

__init__() (polaris.ocean.tasks.ice_shelf_2d.validate.Validate method)

run() (polaris.ocean.tasks.ice_shelf_2d.validate.Validate method)

InertialGravityWave (class in polaris.ocean.tasks.inertial_gravity_wave)

add_inertial_gravity_wave_tasks() (in module polaris.ocean.tasks.inertial_gravity_wave)

Analysis (class in polaris.ocean.tasks.inertial_gravity_wave.analysis)

__init__() (polaris.ocean.tasks.inertial_gravity_wave.analysis.Analysis method)

ExactSolution (class in polaris.ocean.tasks.inertial_gravity_wave.exact_solution)

Forward (class in polaris.ocean.tasks.inertial_gravity_wave.forward)

__init__() (polaris.ocean.tasks.inertial_gravity_wave.forward.Forward method)

__init__() (polaris.ocean.tasks.inertial_gravity_wave.init.Init method)

run() (polaris.ocean.tasks.inertial_gravity_wave.init.Init method)

__init__() (polaris.ocean.tasks.inertial_gravity_wave.viz.Viz method)

run() (polaris.ocean.tasks.inertial_gravity_wave.viz.Viz method)

add_internal_wave_tasks() (in module polaris.ocean.tasks.internal_wave)

__init__() (polaris.ocean.tasks.internal_wave.default.Default method)

__init__() (polaris.ocean.tasks.internal_wave.forward.Forward method)

compute_cell_count() (polaris.ocean.tasks.internal_wave.forward.Forward method)

Analysis (class in polaris.ocean.tasks.internal_wave.rpe.analysis)

__init__() (polaris.ocean.tasks.internal_wave.rpe.analysis.Analysis method)

run() (polaris.ocean.tasks.internal_wave.rpe.analysis.Analysis method)

ManufacturedSolution (class in polaris.ocean.tasks.manufactured_solution)

add_manufactured_solution_tasks() (in module polaris.ocean.tasks.manufactured_solution)

Analysis (class in polaris.ocean.tasks.manufactured_solution.analysis)

__init__() (polaris.ocean.tasks.manufactured_solution.analysis.Analysis method)

ExactSolution (class in polaris.ocean.tasks.manufactured_solution.exact_solution)

Forward (class in polaris.ocean.tasks.manufactured_solution.forward)

__init__() (polaris.ocean.tasks.manufactured_solution.forward.Forward method)

__init__() (polaris.ocean.tasks.manufactured_solution.init.Init method)

run() (polaris.ocean.tasks.manufactured_solution.init.Init method)

setup() (polaris.ocean.tasks.manufactured_solution.init.Init method)

__init__() (polaris.ocean.tasks.manufactured_solution.viz.Viz method)

run() (polaris.ocean.tasks.manufactured_solution.viz.Viz method)

add_single_column_tasks() (in module polaris.ocean.tasks.single_column)

__init__() (polaris.ocean.tasks.single_column.cvmix.CVMix method)

__init__() (polaris.ocean.tasks.single_column.forward.Forward method)

__init__() (polaris.ocean.tasks.single_column.ideal_age.IdealAge method)

__init__() (polaris.ocean.tasks.sphere_transport.SphereTransport method)

configure() (polaris.ocean.tasks.sphere_transport.SphereTransport method)

add_sphere_transport_tasks() (in module polaris.ocean.tasks.sphere_transport)

Analysis (class in polaris.ocean.tasks.sphere_transport.analysis)

__init__() (polaris.ocean.tasks.sphere_transport.analysis.Analysis method)

FilamentAnalysis (class in polaris.ocean.tasks.sphere_transport.filament_analysis)

__init__() (polaris.ocean.tasks.sphere_transport.forward.Forward method)

__init__() (polaris.ocean.tasks.sphere_transport.init.Init method)

MixingAnalysis (class in polaris.ocean.tasks.sphere_transport.mixing_analysis)

__init__() (polaris.ocean.tasks.sphere_transport.viz.Viz method)

alter_bottom_depth() (in module polaris.ocean.vertical.partial_cells)

init_sigma_vertical_coord() (in module polaris.ocean.vertical.sigma)

update_sigma_layer_thickness() (in module polaris.ocean.vertical.sigma)

compute_min_max_level_cell() (in module polaris.ocean.vertical.zlevel)

compute_z_level_layer_thickness() (in module polaris.ocean.vertical.zlevel)

compute_z_level_resting_thickness() (in module polaris.ocean.vertical.zlevel)

init_z_level_vertical_coord() (in module polaris.ocean.vertical.zlevel)

update_z_level_layer_thickness() (in module polaris.ocean.vertical.zlevel)

init_z_star_vertical_coord() (in module polaris.ocean.vertical.zstar)

update_z_star_layer_thickness() (in module polaris.ocean.vertical.zstar)

find_planar_transect_cells_and_weights() (in module polaris.ocean.viz.transect.horiz)

find_spherical_transect_cells_and_weights() (in module polaris.ocean.viz.transect.horiz)

make_triangle_tree() (in module polaris.ocean.viz.transect.horiz)

mesh_to_triangles() (in module polaris.ocean.viz.transect.horiz)

find_transect_levels_and_weights() (in module polaris.ocean.viz.transect.vert)

interp_mpas_to_transect_cells() (in module polaris.ocean.viz.transect.vert)

interp_mpas_to_transect_nodes() (in module polaris.ocean.viz.transect.vert)

add_single_column_tasks() (in module polaris.seaice.tasks.single_column)

ExactRestart (class in polaris.seaice.tasks.single_column.exact_restart)

Validate (class in polaris.seaice.tasks.single_column.exact_restart.validate)

__init__() (polaris.seaice.tasks.single_column.forward.Forward method)

StandardPhysics (class in polaris.seaice.tasks.single_column.standard_physics)

Viz (class in polaris.seaice.tasks.single_column.standard_physics.viz)

__init__() (polaris.seaice.tasks.single_column.standard_physics.viz.Viz method)

run() (polaris.seaice.tasks.single_column.standard_physics.viz.Viz method)

Extension internationalization (i18n) and localization (l10n) using i18n API

sphinx.ext.imgconverter – A reference image converter using Imagemagick

sphinx.ext.autosectionlabel – Allow referencing sections by their title

sphinx.ext.napoleon – Support for NumPy and Google style docstrings

Creating artists specifically for adding to the legend (aka. Proxy artists)

boxplot and bxp vert parameter, and rcParams["boxplot.vertical"]

MEP15: Fix axis autoscaling when limits are specified for one axis only

f-strings support = for self-documenting expressions and debugging

ZERO-CLAUSE BSD LICENSE FOR CODE IN THE PYTHON 3.13.1 DOCUMENTATION

email.message.Message: Representing an email message using the compat32 API

What is the most efficient way to concatenate many strings together?

Why does a_tuple[i] += [‘item’] raise an exception when the addition works?

What does the slash(/) in the parameter list of a function mean?

Why am I getting an UnboundLocalError when the variable has a value?

PEP 3118: New memoryview implementation and buffer protocol documentation

PEP 476: Enabling certificate verification by default for stdlib http clients

PEP 471 - os.scandir() function – a better and faster directory iterator

PEP 655: Marking individual TypedDict items as required or not-required

Why must ‘self’ be used explicitly in method definitions and calls?

argparse — Parser for command-line options, arguments and subcommands

functools — Higher-order functions and operations on callable objects

importlib.resources – Package resource reading, opening and access

multiprocessing.shared_memory — Shared memory for direct access across processes

sysconfig — Provide access to Python’s configuration information

Gumbel (LogWeibull, Fisher-Tippetts, Type I Extreme Value) Distribution

Laplace (Double Exponential, Bilateral Exponential) Distribution

Chi-square test of independence of variables in a contingency table

Multivariate data interpolation on a regular grid  (RegularGridInterpolator)

How do xarray data structures differ from those found in pandas?

Non-standard calendars and dates outside the nanosecond-precision range

import importlib.resources as imp_res
import json

class Component:
    """
    The base class for housing all the tasks for a given component, such as
    ocean, landice, or seaice

    Attributes
    ----------
    name : str
        the name of the component

    tasks : dict
        A dictionary of tasks in the component with the subdirectories of the
        tasks in the component as keys

    steps : dict
        A dictionary of steps in the component with the subdirectories of the
        steps in the component as keys

    cached_files : dict
        A dictionary that maps from output file names in steps within tasks to
        cached files in the ``polaris_cache`` database for the component. These
        file mappings are read in from ``cached_files.json`` in the component.
    """

    def __init__(self, name):
        """
        Create a new container for the tasks for a given component

        Parameters
        ----------
        name : str
            the name of the component
        """
        self.name = name

        # tasks are added with add_task()
        self.tasks = dict()
        # steps are added with add_step()
        self.steps = dict()
        # configs are added with add_config()
        self.configs = dict()

        self.cached_files = dict()
        self._read_cached_files()

    def add_task(self, task):
        """
        Add a task to the component

        task : polaris.Task
            The task to add
        """
        if task.subdir in self.tasks:
            raise ValueError(f'A task has already been added with path '
                             f'{task.subdir}')
        self.tasks[task.subdir] = task

    def add_step(self, step):
        """
        Add a step to the component

        step : polaris.Step
            The step to add
        """
        if step.subdir in self.steps and self.steps[step.subdir] != step:
            raise ValueError(f'A different step has already been added with '
                             f'path {step.subdir}')
        self.steps[step.subdir] = step

    def remove_step(self, step):
        """
        Remove the given step from this component

            The step to add if adding by Step object, not subdirectory
        """
        if step.subdir not in self.steps:
            raise ValueError(f'step {step.name} at {step.subdir} not in the '
                             f'{self.name} component')
        self.steps.pop(step.subdir)

    def add_config(self, config):
        """
        Add a shared config to the component

        config : polaris.config.PolarisConfigParser
            The config to add
        """
        if config.filepath is None:
            raise ValueError('The filepath attribute of a config must be set '
                             'before it can be added to a component')
        if config.filepath in self.configs and \
                self.configs[config.filepath] != config:
            raise ValueError(f'A different shared config has already been '
                             f'added at {config.filepath}')
        self.configs[config.filepath] = config

    def configure(self, config):
        """
        Configure the component

            config options to modify
        """
        pass

    def _read_cached_files(self):
        """ Read in the dictionary of cached files from cached_files.json """

        package = f'polaris.{self.name}'
        filename = 'cached_files.json'
        try:
            pkg_file = imp_res.files(package).joinpath(filename)
            with pkg_file.open('r') as data_file:
                self.cached_files = json.load(data_file)
        except FileNotFoundError:
            # no cached files for this core
            pass

import os
import shutil
from collections import OrderedDict
from typing import Dict, List, Union

import numpy as np
import xarray as xr
from lxml import etree
from mpas_tools.logging import check_call

import polaris.namelist
import polaris.streams
from polaris.step import Step
from polaris.yaml import PolarisYaml, yaml_to_mpas_streams

class ModelStep(Step):
    """
    namelist : str
        The name of the namelist file

    streams : str
        The name of the streams file

    yaml : str
        The name of the yaml file

    config_models : list of str
        If config options are available for multiple models, a list of valid
        models from which config options should be taken.  For example, for
        MPAS-Ocean this would be ``['ocean', 'mpas-ocean']`` and for Omega it
        is ``['ocean', 'omega']``, since both models share the generic
        ``ocean`` config options.

    update_pio : bool
        Whether to modify the namelist so the number of PIO tasks and the
        stride between them is consistent with the number of nodes and
        cores (one PIO task per node).

    make_graph : bool
        Whether to make a graph file from the given MPAS mesh file.  If
        ``True``, ``mesh_filename`` must be given.

    mesh_filename : str
        The name of an MPAS mesh file to use to make the graph file

    partition_graph : bool
        Whether to partition the domain for the requested number of cores.
        If so, the partitioning executable is taken from the ``partition``
        option of the ``[executables]`` config section.

    graph_filename : str
        The name of the graph file to partition

    model_config_data : list
        a list used internally to keep track of updates to the default
        model config options from calls to
        :py:meth:`polaris.ModelStep.add_yaml_file`
        :py:meth:`polaris.ModelStep.add_namelist_file`
        and :py:meth:`polaris.ModelStep.add_model_config_options`

    streams_data : list
        streams from calls to :py:meth:`polaris.ModelStep.add_streams_file`

    make_yaml : bool
        Whether to create a yaml file with model config options and streams
        instead of MPAS namelist and streams files

    streams_section : str
        The name of the streams section in yaml files
    """
    def __init__(self, component, name, subdir=None, indir=None, ntasks=None,
                 min_tasks=None, openmp_threads=None, max_memory=None,
                 cached=False, namelist=None, streams=None, yaml=None,
                 update_pio=True, make_graph=False, mesh_filename=None,
                 partition_graph=True, graph_filename='graph.info',
                 make_yaml=False):
        """
        Make a step for running the model

        component : polaris.Component
            the component that the step belongs to

        name : str
            The name of the step

        subdir : str, optional
            the subdirectory for the step.  If neither this nor ``indir``
             are provided, the directory is the ``name``

        indir : str, optional
            the directory the step is in, to which ``name`` will be appended

        ntasks : int, optional
            the target number of tasks the step would ideally use.  If too
            few cores are available on the system to accommodate the number of
            tasks and the number of cores per task, the step will run on
            fewer tasks as long as as this is not below ``min_tasks``

        min_tasks : int, optional
            the number of tasks the step requires.  If the system has too
            few cores to accommodate the number of tasks and cores per task,
            the step will fail

        openmp_threads : int, optional
            the number of OpenMP threads to use

        max_memory : int, optional
            the amount of memory that the step is allowed to use in MB.
            This is currently just a placeholder for later use with task
            parallelism

        cached : bool, optional
            Whether to get all of the outputs for the step from the database of
            cached outputs for this component

        namelist : str, optional
            The name of the namelist file, default is ``namelist.<component>``

        streams : str, optional
            The name of the streams file, default is ``streams.<component>``

        update_pio : bool, optional
            Whether to modify the namelist so the number of PIO tasks and the
            stride between them is consistent with the number of nodes and
            cores (one PIO task per node).

        make_graph : bool, optional
            Whether to make a graph file from the given MPAS mesh file.  If
            ``True``, ``mesh_filename`` must be given.

        mesh_filename : str, optional
            The name of an MPAS mesh file to use to make the graph file

        partition_graph : bool, optional
            Whether to partition the domain for the requested number of cores.
            If so, the partitioning executable is taken from the ``partition``
            option of the ``[executables]`` config section.

        graph_filename : str, optional
            The name of the graph file to partition

        make_yaml : bool, optional
            Whether to create a yaml file with model config options and streams
            instead of MPAS namelist and streams files
        """
        super().__init__(component=component, name=name, subdir=subdir,
                         indir=indir, cpus_per_task=openmp_threads,
                         min_cpus_per_task=openmp_threads, ntasks=ntasks,
                         min_tasks=min_tasks, openmp_threads=openmp_threads,
                         max_memory=max_memory, cached=cached)

        if namelist is None:
            namelist = f'namelist.{component.name}'

        if streams is None:
            streams = f'streams.{component.name}'

        if yaml is None:
            yaml = f'{component.name}.yaml'

        self.namelist = namelist
        self.streams = streams
        self.yaml = yaml
        self.config_models: Union[List[None], List[str]] = [None]
        self.update_pio = update_pio
        self.make_graph = make_graph
        self.mesh_filename = mesh_filename
        self.partition_graph = partition_graph
        self.graph_filename = graph_filename

        self.make_yaml = make_yaml
        self.streams_section = 'streams'

        self.add_input_file(filename='<<<model>>>')

        self.model_config_data = list()
        self.streams_data = list()

        # used during generating yaml, namelist and streams files
        self._yaml = None
        self._namelist = None
        self._streams_tree = None

    def setup(self):
        """ Setup the command-line arguments """
        config = self.config
        component_path = config.get('executables', 'component')
        model_basename = os.path.basename(component_path)
        if self.make_yaml:
            self.args = [[f'./{model_basename}']]
        else:
            self.args = [[f'./{model_basename}',
                          '-n', self.namelist,
                          '-s', self.streams]]

    def set_model_resources(self, ntasks=None, min_tasks=None,
                            openmp_threads=None, max_memory=None):
        """
        Update the resources for the step.  This can be done within init,
        ``setup()`` or ``runtime_setup()`` for the step that this step
        belongs to, or init, ``configure()`` or ``run()`` for the task
        that this step belongs to.

            the number of tasks the step would ideally use.  If too few
            cores are available on the system to accommodate the number of

        """
        self.set_resources(cpus_per_task=openmp_threads,
                           min_cpus_per_task=openmp_threads, ntasks=ntasks,
                           min_tasks=min_tasks, openmp_threads=openmp_threads,
                           max_memory=max_memory)

    def add_model_config_options(self, options, config_model=None):
        """
        Add the replacement model config options to be parsed when generating
        a namelist or yaml file if and when the step gets set up.  The config
        values should be in the appropriate python data type: bool, float,
        int or str.

        options : dict
            A dictionary of options and value to replace model config options
            with new values

        config_model : str, optional
            If config options are available for multiple models, the model that
            the config options are from
        """
        self.model_config_data.append(dict(options=options,
                                           config_model=config_model))

    def add_yaml_file(self, package, yaml, template_replacements=None):
        """
        Add a file with updates to yaml config options to the step to be parsed
        when generating a complete yaml file if and when the step gets set
        up.

        package : Package
            The package name or module object that contains ``yaml``

        yaml : str
            The name of the yaml file with replacement config options to read

        template_replacements : dict, optional
            A dictionary of replacements, in which case ``yaml`` must be a
            Jinja2 template to be rendered with these replacements
        """
        self.model_config_data.append(dict(package=package, yaml=yaml,
                                           replacements=template_replacements))
        self.streams_data.append(dict(package=package, yaml=yaml,
                                      replacements=template_replacements))

    def map_yaml_options(self, options, config_model):
        """
        A mapping between model config options from different models.  This
        method should be overridden for situations in which yaml config
        options have diverged in name or structure from their counterparts in
        another model (e.g. when translating from MPAS-Ocean namelist options
        to Omega config options)

            A dictionary of yaml options and value to use as replacements for
            existing values

        config_model : str or None

        Returns
        -------
            A revised dictionary of yaml options and value to use as
            replacements for existing values
        """
        return options

    def map_yaml_configs(self, configs, config_model):
        """

        configs : dict
            A nested dictionary of yaml sections, options and value to use as

        Returns
        -------
            A revised nested dictionary of yaml sections, options and value to
            use as replacements for existing values
        """
        return configs

    def map_yaml_streams(self, streams, config_model):
        """
        A mapping between model streams from different models.  This method
        should be overridden for situations in which yaml streams have diverged
        in name or structure from their counterparts in another model (e.g.
        when translating from MPAS-Ocean streams to Omega IOStreams)

        streams : dict
            A nested dictionary of streams data

            If streams are available for multiple models, the model that the
            streams are from

        Returns
        -------
            A revised nested dictionary of streams data
        """
        return streams

    def map_yaml_to_namelist(self, options):
        """
        A mapping from yaml model config options to namelist options.  This
        options have diverged in name or structure from their namelist
        counterparts (e.g. when translating from Omega yaml to MPAS-Ocean
        namelists)

        Returns
        -------
            A nested dictionary of namelist sections, options and value to use
            as replacements for existing values
        """
        namelist = dict()
        # flatten the dictionary, since MPAS is expecting a flat set of options
        # and values
        for section_or_option in options:
            if isinstance(options[section_or_option], (dict, OrderedDict)):
                section = section_or_option
                for option in options[section]:
                    namelist[option] = options[section][option]
            else:
                option = section_or_option
                namelist[option] = options[option]

        for option in namelist:
            value = namelist[option]
            if isinstance(value, bool):
                if value:
                    namelist[option] = '.true.'
                else:
                    namelist[option] = '.false.'
            elif isinstance(value, str):
                # extra set of quotes
                namelist[option] = f"'{value}'"
            elif isinstance(value, float):
                namelist[option] = f'{value:g}'
            else:
                namelist[option] = f'{value}'

        return namelist

    def add_namelist_file(self, package, namelist):
        """
        Add a file with updates to namelist options to the step to be parsed
        when generating a complete namelist file if and when the step gets set
        up.

            The package name or module object that contains ``namelist``

        namelist : str
            The name of the namelist replacements file to read from
        """
        self.model_config_data.append(dict(package=package, namelist=namelist))

    def add_streams_file(self, package, streams, template_replacements=None):
        """
        Add a streams file to the step to be parsed when generating a complete
        streams file if and when the step gets set up.

            The package name or module object that contains the streams file

        streams : str
            The name of the streams file to read from

            A dictionary of replacements, in which case ``streams`` must be a
        """
        self.streams_data.append(
            dict(package=package, streams=streams,
                 replacements=template_replacements))

    def dynamic_model_config(self, at_setup):
        """
        Add model config options, namelist, streams and yaml files using config
        options or template replacements that need to be set both during step
        setup and at runtime

        at_setup : bool
            Whether this method is being run during setup of the step, as
            opposed to at runtime
        """
        if self.update_pio and not at_setup:
            self.update_namelist_pio()

    def runtime_setup(self):
        """
        Update PIO namelist options, make graph file, and partition graph file
        (if any of these are requested)
        """
        quiet = False
        self._read_model_config()

        # start fresh for dynamic model configuration
        self.dynamic_model_config(at_setup=False)

            self._process_yaml(quiet=quiet, remove_unrequested_streams=False)
        else:
            self._process_namelists(quiet=quiet)
            self._process_streams(quiet=quiet, remove_unrequested=False)

        self._write_model_config()

        if self.make_graph:
            make_graph_file(mesh_filename=self.mesh_filename,
                            graph_filename=self.graph_filename)

        if self.partition_graph:
            self.partition(graph_file=self.graph_filename)

    def process_inputs_and_outputs(self):
        """
        Process the model as an input, then call the parent class' version

        Also generates namelist and streams files
        """
        for entry in self.input_data:
            filename = entry['filename']

            if filename == '<<<model>>>':
                filename, target = ModelStep._process_model(
                    self.config, self.base_work_dir)

                entry['filename'] = filename
                entry['target'] = target

        super().process_inputs_and_outputs()

        if self.cached:
            return

        self.dynamic_model_config(at_setup=True)

        quiet = True
        self._create_model_config()

            self._process_yaml(quiet=quiet, remove_unrequested_streams=True)
        else:
            self._process_streams(quiet=quiet, remove_unrequested=True)

    def update_namelist_pio(self):
        """
        Modify the namelist so the number of PIO tasks and the stride between
        them consistent with the number of nodes and cores (one PIO task per
        node).
        """

        cores = self.ntasks * self.cpus_per_task

        cores_per_node = config.getint('parallel', 'cores_per_node')

        # update PIO tasks based on the machine settings and the available
        # number or cores
        pio_num_iotasks = int(np.ceil(cores / cores_per_node))
        pio_stride = self.ntasks // pio_num_iotasks
        if pio_stride > cores_per_node:
            raise ValueError(f'Not enough nodes for the number of cores.  '
                             f'cores: {cores}, cores per node: '
                             f'{cores_per_node}')

        replacements = {'config_pio_num_iotasks': pio_num_iotasks,
                        'config_pio_stride': pio_stride}

        self.add_model_config_options(options=replacements)

    def partition(self, graph_file='graph.info'):
        """
        Partition the domain for the requested number of tasks

        graph_file : str, optional

        """
        ntasks = self.ntasks
        if ntasks > 1:
            executable = self.config.get('parallel', 'partition_executable')
            args = [executable, graph_file, f'{ntasks}']
            check_call(args, self.logger)

    @staticmethod
    def _process_model(config, base_work_dir):
        """ Process the E3SM component model as an input to the step """
        model = config.get('executables', 'component')
        filename = os.path.basename(model)
        copy_executable = config.getboolean('setup', 'copy_executable')
        if copy_executable:
            # make a copy of the model executable, then link to that
            component_subdir = os.path.basename(
                config.get('paths', 'component_path'))
            component_workdir = os.path.join(base_work_dir, component_subdir)
            target = os.path.join(component_workdir, filename)
            try:
                os.makedirs(component_workdir)
            except FileExistsError:
                pass

            try:
                shutil.copy(model, target)
                pass
        else:
            target = os.path.abspath(model)
        return filename, target

    def _create_model_config(self):
        """
        At setup, create either the yaml or the namelist/streams file for the
        step
        """
            defaults_filename = config.get('model_config', 'defaults')
            self._yaml = PolarisYaml.read(defaults_filename,
                                          streams_section=self.streams_section)
        else:
            defaults_filename = config.get('namelists', 'forward')
            self._namelist = polaris.namelist.ingest(defaults_filename)

            defaults_filename = config.get('streams', 'forward')
            self._streams_tree = etree.parse(defaults_filename)

    def _read_model_config(self):
        """
        At runtime, read either the yaml or the namelist/streams file for the
        step
        """
            filename = os.path.join(self.work_dir, self.yaml)
            self._yaml = PolarisYaml.read(filename,
        else:
            filename = os.path.join(self.work_dir, self.namelist)
            self._namelist = polaris.namelist.ingest(filename)

            filename = os.path.join(self.work_dir, self.streams)
            self._streams_tree = etree.parse(filename)

    def _write_model_config(self):
        """
        At setup or runtime, write either the yaml or the namelist/streams
        file for the step
        """
        step_work_dir = self.work_dir
            yaml_filename = f'{step_work_dir}/{self.yaml}'
            if self._yaml is None:
                raise ValueError('Trying to write a yaml file but no yaml '
                                 'object was created.')
            self._yaml.write(yaml_filename)
        else:
            namelist_filename = f'{step_work_dir}/{self.namelist}'
            if self._namelist is None:
                raise ValueError('Trying to write a namelist file but no '
                                 'namelist object was created.')
            polaris.namelist.write(self._namelist, namelist_filename)

            streams_filename = f'{step_work_dir}/{self.streams}'
            if self._streams_tree is None:
                raise ValueError('Trying to write a streams file but no '
                                 'streams XML tree was created.')
            polaris.streams.write(self._streams_tree, streams_filename)
        # set these back to None because we don't need to keep them around
        # and the streams tree can't be pickled

    def _process_namelists(self, quiet):
        """
        Processes changes to a namelist file from the files and dictionaries
        in the step's ``model_config_data``.
        """
        if not self.model_config_data:
            return

        replacements = dict()

        for entry in self.model_config_data:
            if 'namelist' in entry:
                options = polaris.namelist.parse_replacements(
                    entry['package'], entry['namelist'])
                replacements.update(options)
            for config_model in self.config_models:
                if 'options' in entry and \
                        entry['config_model'] == config_model:
                    # this is a dictionary of replacement model config options
                    options = entry['options']
                    options = self.map_yaml_options(options=entry['options'],
                                                    config_model=config_model)
                    options = self.map_yaml_to_namelist(options)
                    replacements.update(options)
                if 'yaml' in entry:
                    yaml = PolarisYaml.read(
                        filename=entry['yaml'], package=entry['package'],
                        replacements=entry['replacements'], model=config_model,
                        streams_section=self.streams_section)

                    configs = self.map_yaml_configs(configs=yaml.configs,
                    configs = self.map_yaml_to_namelist(configs)
                    replacements.update(configs)

        if not quiet:
            print(f'Warning: replacing namelist options in {self.namelist}')
            for key, value in replacements.items():
                print(f'{key} = {value}')

        if self._streams_tree is None:
            raise ValueError('Trying to update a namelist object but it was '
                             'never created.')
        self._namelist = polaris.namelist.replace(self._namelist, replacements)

    def _process_streams(self, quiet, remove_unrequested):
        """
        Processes changes to a streams file from the files and dictionaries
        in the step's ``streams_data``.
        """
        if not self.streams_data:
            return

        # generate the streams file
        tree = None

        processed_registry_filename = None
        if config.has_section('registry') and \
                config.has_option('registry', 'processed'):
            processed_registry_filename = config.get('registry', 'processed')

        for entry in self.streams_data:
            package = entry['package']
            replacements = entry['replacements']
            if 'streams' in entry:
                streams_filename = entry['streams']
                if not quiet:
                    print(f'{package} {streams_filename}')

                tree = polaris.streams.read(
                    package=package, streams_filename=streams_filename,
                    replacements=replacements, tree=tree)
                    tree = self._process_yaml_streams(
                        entry['yaml'], package, replacements, config_model,
                        processed_registry_filename, tree, quiet)

            if not quiet and replacements is not None:
                for key, value in replacements.items():
                    print(f'{key} = {value}')

        if tree is None:
            # nothing to add
            return

            raise ValueError('Trying to update a streams XML tree but it was '

        defaults = next(self._streams_tree.iter('streams'))
        streams = next(tree.iter('streams'))

        for stream in streams:
            polaris.streams.update_defaults(stream, defaults)

        if remove_unrequested:
            # during setup, we remove any streams that aren't requested but
            # at runtime we don't want to do this because we would lose any
            # streams added only during setup.
            for default in defaults:
                found = False
                for stream in streams:
                    if stream.attrib['name'] == default.attrib['name']:
                        found = True
                        break
                if not found:
                    defaults.remove(default)

    def _process_yaml_streams(self, yaml_filename, package, replacements,
                              config_model, processed_registry_filename,
                              tree, quiet):
            print(f'{package} {yaml_filename}')

        yaml = PolarisYaml.read(filename=yaml_filename,
                                package=package,
                                replacements=replacements,
                                model=config_model,
                                streams_section=self.streams_section)
        assert processed_registry_filename is not None
        new_tree = yaml_to_mpas_streams(
            processed_registry_filename, yaml)
        tree = polaris.streams.update_tree(tree, new_tree)
        return tree

    def _process_yaml(self, quiet, remove_unrequested_streams):
        """
        Processes changes to a yaml file from the files and dictionaries
        """
            return

        if self._yaml is None:
            raise ValueError('Trying to update a yaml object but it was '

            print(f'Warning: replacing yaml options in {self.yaml}')

        streams: Dict[str, Dict[str, Union[str, float, int, List[str]]]] = {}

                raise ValueError('Cannot generate a yaml config from an MPAS '
                                 'namelist file.')

                    self._yaml.update(options=options, quiet=quiet)

                    new_streams = self.map_yaml_streams(
                        streams=yaml.streams, config_model=config_model)
                    self._update_yaml_streams(streams, new_streams,
                                              quiet=quiet,
                                              remove_unrequested=False)
                    self._yaml.update(configs=configs, quiet=quiet)
        self._update_yaml_streams(
            self._yaml.streams, streams, quiet=quiet,
            remove_unrequested=remove_unrequested_streams)

    def _update_yaml_streams(streams, new_streams, quiet, remove_unrequested):
        """
        Update yaml streams, optionally removing any streams that aren't in
        new_streams
        """

        for stream_name, new_stream in new_streams.items():
            if stream_name in streams:
                streams[stream_name].update(new_stream)
                    print(f'  updating: {stream_name}')
            else:
                    print(f'  adding:   {stream_name}')
                streams[stream_name] = new_stream

            # during setup, we remove any default streams that aren't requested
            # but at runtime we don't want to do this because we would lose any
            for stream_name in list(streams.keys()):
                if stream_name not in new_streams:
                    if not quiet:
                        print(f'  dropping: {stream_name}')
                    streams.pop(stream_name)

def make_graph_file(mesh_filename, graph_filename='graph.info',
                    weight_field=None):
    """
    Make a graph file from the MPAS mesh for use in the Metis graph
    partitioning software

    Parameters
     mesh_filename : str
        The name of the input MPAS mesh file

    graph_filename : str, optional
        The name of the output graph file

    weight_field : str
        The name of a variable in the MPAS mesh file to use as a field of
        weights
    """

    with xr.open_dataset(mesh_filename) as ds:

        nCells = ds.sizes['nCells']

        nEdgesOnCell = ds.nEdgesOnCell.values
        cellsOnCell = ds.cellsOnCell.values - 1
        if weight_field is not None:
            if weight_field in ds:
                raise ValueError(f'weight_field {weight_field} not found in '
                                 f'{mesh_filename}')
            weights = ds[weight_field].values
        else:
            weights = None

    nEdges = 0
    for i in range(nCells):
        for j in range(nEdgesOnCell[i]):
            if cellsOnCell[i][j] != -1:
                nEdges = nEdges + 1

    nEdges = nEdges // 2

    with open(graph_filename, 'w+') as graph:
        if weights is None:
            graph.write(f'{nCells} {nEdges}\n')

            for i in range(nCells):
                for j in range(0, nEdgesOnCell[i]):
                    if cellsOnCell[i][j] >= 0:
                        graph.write(f'{cellsOnCell[i][j] + 1} ')
                graph.write('\n')
        else:
            graph.write(f'{nCells} {nEdges} 010\n')

                graph.write(f'{int(weights[i])} ')

import logging
import os

from mache import MachineInfo

from polaris.config import PolarisConfigParser
from polaris.io import download, symlink, update_permissions
from polaris.validate import compare_variables

class Step:
    """
    The base class for a step of a tasks, such as setting up a mesh,
    creating an initial condition, or running the component forward in time.
    The step is the smallest unit of work in polaris that can be run on its
    own by a user, though users will typically run full tasks or suites.

    Below, the terms "input" and "output" refer to inputs and outputs to the
    step itself, not necessarily the MPAS model.  In fact, the MPAS model
    itself is often an input to the step.

    name : str
        the name of the step

    component : polaris.Component
        The component the step belongs to

    subdir : str
        the subdirectory for the step

    path : str
        the path within the base work directory of the step, made up of
        ``component``, the task's ``subdir`` and the step's ``subdir``

    cpus_per_task : int, optional
        the number of cores per task the step would ideally use.  If
        fewer cores per node are available on the system, the step will
        run on all available cores as long as this is not below
        ``min_cpus_per_task``

    min_cpus_per_task : int, optional
        the number of cores per task the step requires.  If the system
        has fewer than this number of cores per node, the step will fail

    ntasks : int, optional
        the number of tasks the step would ideally use.  If too few
        cores are available on the system to accommodate the number of
        tasks and the number of cores per task, the step will run on
        fewer tasks as long as as this is not below ``min_tasks``

    min_tasks : int, optional
        the number of tasks the step requires.  If the system has too
        few cores to accommodate the number of tasks and cores per task,
        the step will fail

    openmp_threads : int
        the number of OpenMP threads to use

    max_memory : int
        the amount of memory that the step is allowed to use in MB.
        This is currently just a placeholder for later use with task
        parallelism

    input_data : list of dict
        a list of dict used to define input files typically to be
        downloaded to a database and/or symlinked in the work directory

    inputs : list of str
        a list of absolute paths of input files produced from ``input_data`` as
        part of setting up the step.  These input files must all exist at run
        time or the step will raise an exception

    outputs : list of str
        a list of absolute paths of output files produced by this step (or
        cached) and available as inputs to other tasks and steps.  These
        files must exist after the task has run or an exception will be raised

    dependencies : dict of polaris.Step
        A dictionary of steps that this step depends on (i.e. it can't run
        until they have finished). Dependencies are used when the names of
        the files produced by the dependency aren't known at setup (e.g.
        because they depend on config options or data read in from files).
        Under other circumstances, it is sufficient to indicate that an output
        file from another step is an input of this step to establish a
        dependency.

    has_shared_config : bool
        Whether this step uses a shared config file.

    is_dependency : bool
        Whether this step is the dependency of one or more other steps.

        The tasks this step is used in

    config : polaris.config.PolarisConfigParser
        Configuration options for this step, possibly shared with other tasks
        and steps

    machine_info : mache.MachineInfo
        Information about E3SM supported machines

    config_filename : str
        The filename or symlink within the step where ``config`` is written to
        during setup and read from during run

    work_dir : str
        The step's work directory, defined during setup as the combination
        of ``base_work_dir`` and ``path``

    base_work_dir : str
        The base work directory

    baseline_dir : str
        Location of the same task within the baseline work directory,
        for use in comparing variables and timers

    validate_vars : dict of list
        A list of variables for each output file for which a baseline
        comparison should be performed if a baseline run has been provided. The
        baseline validation is performed after the step has run.

    logger : logging.Logger
        A logger for output from the step

    log_filename : str
        At run time, the name of a log file where output/errors from the step
        are being logged, or ``None`` if output is to stdout/stderr

    cached : bool
        Whether to get all of the outputs for the step from the database of
        cached outputs for this component

    run_as_subprocess : bool
        Whether to run this step as a subprocess, rather than just running
        it directly from the task.  It is useful to run a step as a
        subprocess if there is not a good way to redirect output to a log
        file (e.g. if the step calls external code that, in turn, calls
        additional subprocesses).

    args : {list of list of str, None}
        A list of lists of command-line arguments to call in parallel.  Each
        inner list represents a single command.  All commands must use the
        same parallel resources.
    """

    def __init__(self, component, name, subdir=None, indir=None,
                 cpus_per_task=1, min_cpus_per_task=1, ntasks=1, min_tasks=1,
                 openmp_threads=1, max_memory=None, cached=False,
                 run_as_subprocess=False):
        """
        Create a new task

            The component the step belongs to

        name : str
            the name of the task

        cpus_per_task : int, optional
            the number of cores per task the step would ideally use.  If
            fewer cores per node are available on the system, the step will
            run on all available cores as long as this is not below
            ``min_cpus_per_task``

        min_cpus_per_task : int, optional
            the number of cores per task the step requires.  If the system
            has fewer than this number of cores per node, the step will fail

        openmp_threads : int

        run_as_subprocess : bool
            Whether to run this step as a subprocess, rather than just running
            it directly from the task.  It is useful to run a step as a
            subprocess if there is not a good way to redirect output to a log
            file (e.g. if the step calls external code that, in turn, calls
            additional subprocesses).
        """
        self.component = component
        if subdir is not None:
            self.subdir = subdir
        elif indir is not None:
            self.subdir = os.path.join(indir, name)
        else:
            self.subdir = name

        self.cpus_per_task = cpus_per_task
        self.min_cpus_per_task = min_cpus_per_task
        self.ntasks = ntasks
        self.min_tasks = min_tasks
        self.openmp_threads = openmp_threads
        self.max_memory = max_memory

        self.path = os.path.join(self.component.name, self.subdir)

        self.run_as_subprocess = run_as_subprocess

        self.has_shared_config = False

        self.config = PolarisConfigParser()
        self.config_filename = ""

        # child steps (or tasks) will add to these
        self.input_data = list()
        self.inputs = list()
        self.outputs = list()
        self.args = None
        self.dependencies = dict()
        self.is_dependency = False

        # these will be set later during setup, dummy placeholders for now
        self.machine_info = MachineInfo(machine='default')
        self.work_dir = ""
        self.base_work_dir = ""
        # may be set during setup if there is a baseline for comparison
        self.baseline_dir = None
        self.validate_vars = dict()
        self.setup_complete = False

        # these will be set before running the step, dummy placeholders for now
        self.logger = logging.getLogger('dummy')
        self.log_filename = None

        # output caching
        self.cached = cached

    def set_resources(self, cpus_per_task=None, min_cpus_per_task=None,
                      ntasks=None, min_tasks=None, openmp_threads=None,
                      max_memory=None):
        """
        Update the resources for the subtask.  This can be done within init,
        ``setup()`` or ``runtime_setup()`` of this step, or init,
        ``configure()`` or ``run()`` for the task that this step belongs
        to.

        """
        if cpus_per_task is not None:
            self.cpus_per_task = cpus_per_task
        if min_cpus_per_task is not None:
            self.min_cpus_per_task = min_cpus_per_task
        if ntasks is not None:
            self.ntasks = ntasks
        if min_tasks is not None:
            self.min_tasks = min_tasks
        if openmp_threads is not None:
            self.openmp_threads = openmp_threads
        if max_memory is not None:
            self.max_memory = max_memory

    def constrain_resources(self, available_resources):
        """
        Constrain ``cpus_per_task`` and ``ntasks`` based on the number of
        cores available to this step

        available_resources : dict
            A dictionary containing available resources (cores, tasks, nodes
            and cores_per_node)
        """
        mpi_allowed = available_resources['mpi_allowed']
        if not mpi_allowed and self.ntasks > 1:
            raise ValueError(
                'You are trying to run an MPI job on a login node.\n'
                'Please switch to a compute node.')

        available_cores = available_resources['cores']
        cores_per_node = available_resources['cores_per_node']
        self.cpus_per_task = min(self.cpus_per_task,
                                 min(available_cores, cores_per_node))
        if self.cpus_per_task < self.min_cpus_per_task:
                f'Available cpus_per_task ({self.cpus_per_task}) is below the '
                f'minimum of {self.min_cpus_per_task} for step {self.name}')

        available_tasks = available_cores // self.cpus_per_task
        self.ntasks = min(self.ntasks, available_tasks)

        if self.ntasks < self.min_tasks:
                f'Available number of MPI tasks ({self.ntasks}) is below the '
                f'minimum of {self.min_tasks} for step {self.name}')

        """
        Set up the task in the work directory, including downloading any
        dependencies.  The step should override this function to perform setup
        operations such as generating namelist and streams files, adding inputs
        and outputs.
        """
        pass

        """
        Update attributes of the step at runtime before calling the ``run()``
        method.  The most common reason to override this method is to
        determine the number of cores and threads to run with.  It may also
        be useful for performing small (python) tasks such as creating
        graph and partition files before running a parallel command. When
        running with parallel tasks in the future, this method will be called
        for each step in serial before steps are run in task parallel.
        """
        pass

    def run(self):
        """
        Run the step.  Every child class must override this method to perform
        the main work.
        """
        pass

    def add_input_file(self, filename=None, target=None, database=None,
                       database_component=None, url=None, work_dir_target=None,
                       package=None, copy=False):
        """
        Add an input file to the step (but not necessarily to the MPAS model).
        The file can be local, a symlink to a file that will be created in
        another step, a symlink to a file in one of the databases for files
        cached after download, and/or come from a specified URL.

        filename : str, optional
            The relative path of the input file within the step's work
            directory. The default is the file name (without the path) of
            ``target``.

        target : str, optional
            A file that will be the target of a symlink to ``filename``.  If
            ``database`` is not specified, this should be an absolute path or a
            relative path from the step's work directory.  If ``database`` is
            specified, this is a relative path within the database and the name
            of the remote file to download.

        database : str, optional
            The name of a database for caching local files.  This will be a
            subdirectory of the local cache directory for this core.  If
            ``url`` is not provided, the URL for downloading the file will be
            determined by combining the base URL of the data server, the
            relative path for the core, ``database`` and ``target``.

        database_component : {'ocean', 'seaice', 'landice', None}, optional
            The prefix of a database root that the database resides in.  By
            default this is the component name for this step.  The suffix
            ``_database_root`` is appended to this string to determine the
            config option where the database root resides on the local machine.

        url : str, optional
            The URL (including file name) for downloading the file.  This
            option should be set if the file is not in a database on the data
            server. The ``filename``, ``target`` and ``database`` are not
            added to URL even if they are provided.

        work_dir_target : str, optional
            Same as ``target`` but with a path relative to the base work
            directory.  This is useful if it is not easy to determine the
            relative path between the step's work directory and the target.

        package : str or package, optional
            A package within ``polaris`` from which the file should be linked

        copy : bool, optional
            Whether to make a copy of the file, rather than a symlink
        """
        if filename is None:
            if target is None:
                raise ValueError('At least one of local_name and target are '
                                 'required.')
            filename = os.path.basename(target)

        self.input_data.append(dict(filename=filename, target=target,
                                    database=database,
                                    database_component=database_component,
                                    url=url, work_dir_target=work_dir_target,
                                    package=package, copy=copy))

    def add_output_file(self, filename, validate_vars=None):
        """
        Add the output file that must be produced by this step and may be made
        available as an input to steps, perhaps in other tasks.  This file
        must exist after the task has run or an exception will be raised.

        Optionally, a list of variables can be provided for validation against
        a baseline (if one is provided), once the step has been run.

        filename : str
            The relative path of the output file within the step's work
            directory

        validate_vars : list, optional
            A list of variable names to compare with a baseline (if one is
            provided)
        """
        self.outputs.append(filename)
        if validate_vars is not None:
            self.validate_vars[filename] = validate_vars

    def add_dependency(self, step, name=None):
        """
        Add `step` as a dependency of this step (i.e. this step can't run
        until the dependency has finished). A dependency should be used when
        the names of the files produced by the dependency aren't known at setup
        (e.g. because they depend on config options or data read in from
        files). Under other circumstances, a dependency can be established by
        indicating that an output (added with the ``add_output_file()`` method)
        from another step is an input (added with ``add_input_file()``) of this
        step .

            The step that is a dependency

        name : str, optional
            The name of the step used to access it in the ``dependencies``
            dictionary.  By default, it is ``step.name`` but another name may
            be required if 2 dependencies have the same ``step.name``.
        """
        if name is None:
            name = step.name
        if name in self.dependencies:
            raise ValueError('Adding a dependency that is already in '
                             'dependencies.')
        self.dependencies[name] = step
        step.is_dependency = True
        step.add_output_file('step_after_run.pickle')
        filename = f'dependencies/{name}_after_run.pickle'
        target = f'{step.path}/step_after_run.pickle'
        self.add_input_file(filename=filename, work_dir_target=target)

    def validate_baselines(self):
        """
        Compare variables between output files in this step and in the same
        step from a baseline run if one was provided.

        Returns
        -------
        compared : bool
            Whether a baseline comparison was performed

        success : bool
            Whether the outputs were successfully validated against a baseline
        """
        if self.work_dir is None:
            raise ValueError('The work directory must be set before the step '
                             'outputs can be validated against baselines.')
        compared = False
        success = True
        if self.baseline_dir is not None:
            for filename, variables in self.validate_vars.items():
                logger = self.logger
                filename = str(filename)

                this_filename = os.path.join(self.work_dir, filename)
                baseline_filename = os.path.join(self.baseline_dir, filename)
                result = compare_variables(
                    variables, this_filename, baseline_filename, logger=logger)
                success = success and result
                compared = True
        return compared, success

    def set_shared_config(self, config, link=None):
        """
        Replace the step's config parser with the shared config parser

            A shared config parser whose ``filepath`` attribute must have been
            set

        link : str, optional
            A link to the shared config file to go in the step's work
            directory. If not provided, the config file itself must be in
            the step's work directory
        """
        self.has_shared_config = True
        self._set_config(config=config, link=link)

        """
        Process the inputs to and outputs from a step added with
        :py:meth:`polaris.Step.add_input_file` and
        :py:meth:`polaris.Step.add_output_file`.  This includes downloading
        files, making symlinks, and converting relative paths to absolute
        paths.
       """
        component = self.component.name
        step_dir = self.work_dir
        base_work_dir = self.base_work_dir

        # process the outputs first because cached outputs will add more inputs
            # forget about the inputs -- we won't used them, but we will add
            # the cached outputs as inputs
            self.input_data = list()
            for output in self.outputs:
                filename = os.path.join(self.path, output)
                if filename not in self.component.cached_files:
                    raise ValueError(f'The file {filename} has not been added '
                                     f'to the cache database')
                target = self.component.cached_files[filename]
                self.add_input_file(
                    filename=output,
                    target=target,
                    database='polaris_cache')

        inputs = []
        databases_with_downloads = set()
            input_file, database_subdirs = Step._process_input(
                entry, config, base_work_dir, component, step_dir)
            if database_subdirs is not None:
                databases_with_downloads.update(database_subdirs)
            inputs.append(input_file)
        self.inputs = inputs

        if len(databases_with_downloads) > 0 and \
                config.has_option('e3sm_unified', 'group'):
            group = config.get('e3sm_unified', 'group')
            update_permissions(databases_with_downloads, group)

        # inputs are already absolute paths, convert outputs to absolute paths
        self.outputs = [os.path.abspath(os.path.join(step_dir, filename)) for
                        filename in self.outputs]

    def _set_config(self, config, link=None):
        """

            set

        """
        self.component.add_config(config)

        self.config = config
        if link is None:
            directory, basename = os.path.split(config.filepath)
            if directory != self.subdir:
                raise ValueError('No link parameter was provided but the '
                                 'config file is not in this step\'s work '
                                 'directory.')
            self.config_filename = basename
        else:
            self.config_filename = link

    def _process_input(entry, config, base_work_dir, component, step_dir):
        database_subdirs = None
        filename = entry['filename']
        target = entry['target']
        database = entry['database']
        database_component = entry['database_component']
        url = entry['url']
        work_dir_target = entry['work_dir_target']
        package = entry['package']
        copy = entry['copy']

        if package is not None:
                target = filename
            target = str(imp_res.as_file(
                imp_res.files(package).joinpath(target)))

        if work_dir_target is not None:
            target = os.path.join(base_work_dir, work_dir_target)

        if target is not None:
            download_target = target
        else:
            download_target = filename

        download_path = None

        if database is not None:
            # we're downloading a file to a cache of a database (if it's
            # not already there.
            if database_component is None:
                database_component = component

            if url is None:
                base_url = config.get('download', 'server_base_url')
                url = f'{base_url}/{database_component}/{database}/{target}'

            database_root = config.get('paths', 'database_root')
            download_path = os.path.join(database_root, database_component,
                                         database, download_target)
            if not os.path.exists(download_path):
                database_subdirs = {
                    database_root,
                    os.path.join(database_root, database_component),
                    os.path.join(database_root, database_component, database)
                }
        elif url is not None:
            download_path = download_target

        if url is not None:
            download_target = download(url, download_path, config)
            if target is not None:
                # this is the absolute path that we presumably want
                target = download_target

            filepath = os.path.join(step_dir, filename)
            dirname = os.path.dirname(filepath)
            if copy:
                shutil.copy(target, filepath)
            else:
                try:
                    os.makedirs(dirname)
                except FileExistsError:
                    pass
                symlink(target, filepath)
            input_file = os.path.join(dirname, target)
        else:
            input_file = os.path.join(step_dir, filename)

        input_file = os.path.abspath(input_file)

        return input_file, database_subdirs

import os

class Task:
    """
    The base class for tasks---such as a decomposition, threading or
    restart test---that are made up of one or more steps

    name : str
        the name of the task

        The component the task belongs to

        A dictionary of steps in the task with step names as keys

    step_symlinks : dict
        A dictionary of symlink paths within the test case's work directory
        to shared steps outside the test case

    steps_to_run : list
        A list of the steps to run when ``run()`` gets called.  This list
        includes all steps by default but can be replaced with a list of only
        those tasks that should run by default if some steps are optional and
        should be run manually by the user.

        the subdirectory for the task

    path : str
        the path within the base work directory of the task, made up of
        ``component`` and the task's ``subdir``

        Configuration options for this task, possibly shared with other tasks
        and steps

        The filename or symlink within the task where ``config`` is written to

        The task's work directory, defined during setup as the combination

    stdout_logger : logging.Logger
        A logger for output from the task that goes to stdout regardless
        of whether ``logger`` is a log file or stdout

        A logger for output from the task

        At run time, the name of a log file where output/errors from the task

    new_step_log_file : bool
        Whether to create a new log file for each step or to log output to a
        common log file for the whole task.  The latter is used when
        running the task as part of a suite
    """

    def __init__(self, component, name, subdir=None, indir=None):
        """

            the component that this task belongs to

        name : str

            the subdirectory for the task.  If neither this nor ``indir``

            the directory the task is in, to which ``name`` will be appended
        """
        else:

        self.config_filename = ''
        self.config.tasks.add(self)

        # steps will be added by calling add_step()
        self.step_symlinks = dict()
        self.steps_to_run = list()

        # these will be set during setup, dummy values for now
        self.work_dir = ''
        self.base_work_dir = ''

        # these will be set when running the task, dummy values for now
        self.new_step_log_file = True
        self.stdout_logger = None

    def configure(self):
        """
        Modify the configuration options for this task. Tasks should
        override this method if they want to add config options specific to
        the task, e.g. from a config file stored in the task's python
        package.  If a task overrides this method, it should assume that
        the ``<self.name>.cfg`` file in its package has already been added
        to the config options prior to calling ``configure()``.
        """
        pass

    def add_step(self, step=None, subdir=None, symlink=None,
                 run_by_default=True):
        """
        Add a step to the task and component (if not already present)

        step : polaris.Step, optional

            The subdirectory of the step within the component if wish to add
            the step by path, and it has already been added to the component

        symlink : str, optional
            A location for a symlink to the step, relative to the test case's
            work directory. This is typically used for a shared step that lives
            outside of the test case

        run_by_default : bool, optional
            Whether to add this step to the list of steps to run when the
            ``run()`` method gets called.  If ``run_by_default=False``, users
            would need to run this step manually.
        """
        if step is None and subdir is None:
            raise ValueError('One of step or subdir must be provided.')
        if step is not None and subdir is not None:
            raise ValueError('Only one of step or subdir should be provided.')

        component = self.component

            if subdir not in self.component.steps:
                raise ValueError(f'Could not find {subdir} in the steps in '
                                 f'the component.  Add the step to the '
                                 f'component first, then to the task.')
            step = component.steps[subdir]

        if step.name in self.steps:
            raise ValueError(f'A step has already been added to this task '
                             f'with name {step.name}')

        # add the step to the component (if it's not already there)
        component.add_step(step)

        self.steps[step.name] = step
        step.tasks[self.subdir] = self
        if symlink:
            self.step_symlinks[step.name] = symlink
        if run_by_default:
            self.steps_to_run.append(step.name)

        """
        Remove the given step from this task and the component

        """
        if step.name not in self.steps:
            raise ValueError(f'step {step.name} not in this task {self.name}')

        self.steps.pop(step.name)
        step.tasks.pop(self.subdir)
        if not step.tasks:
            # no tasks are using this step
            self.component.remove_step(step)
        if step.name in self.step_symlinks:
            self.step_symlinks.pop(step.name)
        if step.name in self.steps_to_run:
            self.steps_to_run.remove(step.name)

        """
        Replace the task's config parser with the shared config parser

            set

            A link to the shared config file to go in the task's work
            the task's work directory
        """

        config.tasks.add(self)
                                 'config file is not in this task\'s work '
        else:

import os
from typing import Union

from mpas_tools.config import MpasConfigParser

class PolarisConfigParser(MpasConfigParser):
    """
    A "meta" config parser that keeps a dictionary of config parsers and their
    sources to combine when needed.  The custom config parser allows provenance
    of the source of different config options and allows the "user" config
    options to always take precedence over other config options (even if they
    are added later).

    This version simply overrides the ``combine()`` method to also ensure that
    certain paths are absolute, rather than relative.

    filepath : str
        A filepath within the component's work directory where this config
        will be written out

    tasks : set of polaris.Task
        A list of tasks that use this config
    """

    def __init__(self, filepath=None):
        """
        Make a new (empty) config parser

        filepath : str, optional
            A filepath within the component's work directory where this config
            will be written out
        """
        super().__init__()
        self.filepath: Union[str, None] = filepath
        self.tasks = set()

        """
        A method that can be overridden to add config options during polaris
        setup
        """
        pass

    def combine(self, raw=False):
        """
        Combine the config files into one

        raw : bool, optional
            Whether to combine config "raw" config options, rather than using
            extended interpolation
        """
        super().combine(raw=raw)
        self._ensure_absolute_paths()

    def _ensure_absolute_paths(self):
        """
        make sure all paths in the paths, namelists, streams, and executables
        sections are absolute paths
        """
        config = self.combined
        for section in ['paths', 'namelists', 'streams', 'executables']:
            if not config.has_section(section):
                continue
            for option, value in config.items(section):
                # not safe to make paths that start with other config options
                # into absolute paths
                if not value.startswith('$'):
                    value = os.path.abspath(value)
                    config.set(section, option, value)

import grp
import os
import stat
import tempfile
from urllib.parse import urlparse

import progressbar
import requests

def download(url, dest_path, config, exceptions=True):  # noqa: C901
    """
    Download a file from a URL to the given path or path name

    url : str
        The URL (including file name) to download

    dest_path : str
        The path (including file name) where the downloaded file should be
        saved

        Configuration options used to find custom paths if ``dest_path`` is
        a config option

    exceptions : bool, optional
        Whether to raise exceptions when the download fails

    Returns
    -------
        The resulting file name if the download was successful, or None if not
    """

    in_file_name = os.path.basename(urlparse(url).path)
    dest_path = os.path.abspath(dest_path)
    out_file_name = os.path.basename(dest_path)

    do_download = config.getboolean('download', 'download')
    check_size = config.getboolean('download', 'check_size')
    verify = config.getboolean('download', 'verify')

    if not do_download:
        if not os.path.exists(dest_path):
            raise OSError(f'File not found and downloading is disabled: '
                          f'{dest_path}')
        return dest_path

    if not check_size and os.path.exists(dest_path):

    session = requests.Session()
    if not verify:
        session.verify = False

    # dest_path contains full path, so we need to make the relevant
    # subdirectories if they do not exist already
    directory = os.path.dirname(dest_path)
    try:
        os.makedirs(directory)
    except OSError:
        pass

    try:
        response = session.get(url, stream=True)
        total_size = response.headers.get('content-length')
    except requests.exceptions.RequestException:
        if exceptions:
            raise
        else:
            print(f'  {url} could not be reached!')
            return None

    try:
        response.raise_for_status()
    except requests.exceptions.HTTPError as e:
            raise
        else:
            print(f'ERROR while downloading {in_file_name}:')
            print(e)

    if total_size is None:
        # no content length header
            dest_dir = os.path.dirname(dest_path)
            with open(dest_path, 'wb') as f:
                print(f'Downloading {in_file_name}\n'
                      f'  to {dest_dir}...')
                try:
                    f.write(response.content)
                except requests.exceptions.RequestException:
                    if exceptions:
                        raise
                    else:
                        print(f'  {in_file_name} failed!')
                        return None
                else:
                    print('  {in_file_name} done.')
    else:
        # we can do the download in chunks and use a progress bar, yay!

        total_size_int = int(total_size)
        if os.path.exists(dest_path) and \
                total_size_int == os.path.getsize(dest_path):
            # we already have the file, so just return
            return dest_path

        if out_file_name == in_file_name:
            file_names = in_file_name
        else:
            file_names = f'{in_file_name} as {out_file_name}'
        dest_dir = os.path.dirname(dest_path)
        print(f'Downloading {file_names} ({_sizeof_fmt(total_size_int)})\n'
              f'  to {dest_dir}')
        widgets = [progressbar.Percentage(), ' ', progressbar.Bar(),
                   ' ', progressbar.ETA()]
        bar = progressbar.ProgressBar(widgets=widgets,
                                      max_value=total_size_int).start()
        size = 0
        with open(dest_path, 'wb') as f:
            try:
                for data in response.iter_content(chunk_size=4096):
                    size += len(data)
                    f.write(data)
                    bar.update(size)
                bar.finish()
            except requests.exceptions.RequestException:
                if exceptions:
                    raise
                else:
                    print(f'  {in_file_name} failed!')
                    return None
            else:
                print(f'  {in_file_name} done.')
    return dest_path

def symlink(target, link_name, overwrite=True):
    """
    From https://stackoverflow.com/a/55742015/7728169
    Create a symbolic link named link_name pointing to target.
    If link_name exists then FileExistsError is raised, unless overwrite=True.
    When trying to overwrite a directory, IsADirectoryError is raised.

    target : str
        The file path to link to

    link_name : str
        The name of the new link

    overwrite : bool, optional
        Whether to replace an existing link if one already exists
    """

    # make the directory that the link is in if it doesn't exist
    directory = os.path.dirname(os.path.abspath(link_name))
    try:
    except FileExistsError:
        pass

    if not overwrite:
        os.symlink(target, link_name)
        return

    # os.replace() may fail if files are on different filesystems
    link_dir = os.path.dirname(link_name)

    # Create link to target with temporary file_name
    while True:
        temp_link_name = tempfile.mktemp(dir=link_dir)

        # os.* functions mimic as closely as possible system functions
        # The POSIX symlink() returns EEXIST if link_name already exists
        # https://pubs.opengroup.org/onlinepubs/9699919799/functions/symlink.html
        try:
            os.symlink(target, temp_link_name)
            break
        except FileExistsError:
            pass

    # Replace link_name with temp_link_name
    try:
        # Preempt os.replace on a directory with a nicer message
        if not os.path.islink(link_name) and os.path.isdir(link_name):
            raise IsADirectoryError(
                f"Cannot symlink over existing directory: '{link_name}'")
        os.replace(temp_link_name, link_name)
    except Exception:
        if os.path.islink(temp_link_name):
            os.remove(temp_link_name)
        raise

def update_permissions(directories, group, show_progressbar=True):
    """
    Fix permissions on the databases where files were downloaded so
    everyone in the group can read/write to them

    directories : list
        directories to change permissions on

    group : str
        The name of the group to set permissions to

    show_progressbar : bool, optional
        Whether to show a progress bar as permissions are updated
    """
    new_uid = os.getuid()
    new_gid = grp.getgrnam(group).gr_gid

    write_perm = (stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP |
                  stat.S_IWGRP | stat.S_IROTH)
    exec_perm = (stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR |
                 stat.S_IRGRP | stat.S_IWGRP | stat.S_IXGRP |
                 stat.S_IROTH | stat.S_IXOTH)

    mask = stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO

    if show_progressbar:
        print('changing permissions on downloaded files')

    # first the base directories that don't seem to be included in
    # os.walk()
    for directory in directories:
        root = None
        _set_dir_perms(root, directory, mask, exec_perm, new_uid, new_gid)

    files_and_dirs = _walk_dirs(directories)

                                      maxval=len(files_and_dirs)).start()
    else:
        bar = None
    progress = 0
    for base in directories:
        for root, dirs, files in os.walk(base):
            for directory in dirs:
                progress += 1
                if show_progressbar:
                    bar.update(progress)

                _set_dir_perms(root, directory, mask, exec_perm, new_uid,
                               new_gid)

            for file_name in files:
                file_name = os.path.join(root, file_name)
                _set_file_perms(root, file_name, mask, exec_perm, write_perm,
                                new_uid, new_gid)

        bar.finish()
        print('  done.')

def _walk_dirs(directories):
    """
    Walk through directories to find files to change
    """
    files_and_dirs = []
            files_and_dirs.extend(dirs)
            files_and_dirs.extend(files)
    return files_and_dirs

def _set_dir_perms(root, directory, mask, exec_perm, new_uid, new_gid):
    """
    Set permissions for a directory
    """
    if root is not None:
        directory = os.path.join(root, directory)

    try:
        dir_stat = os.stat(directory)
        return

    if dir_stat.st_uid != new_uid:
        # current user doesn't own this dir so let's move on
        return

    perm = dir_stat.st_mode & mask

    if perm == exec_perm and dir_stat.st_gid == new_gid:
        return

    try:
        os.chown(directory, new_uid, new_gid)
        os.chmod(directory, exec_perm)
        return

def _set_file_perms(root, file_name, mask, exec_perm, write_perm, new_uid,
                    new_gid):
    """
    """
    file_name = os.path.join(root, file_name)
    try:
        file_stat = os.stat(file_name)
        return

    if file_stat.st_uid != new_uid:
        # current user doesn't own this file so let's move on
        return

    perm = file_stat.st_mode & mask

    if perm & stat.S_IXUSR:
        # executable, so make sure others can execute it
        new_perm = exec_perm
    else:
        new_perm = write_perm

    if perm == new_perm and file_stat.st_gid == new_gid:
        return

    try:
        os.chown(file_name, new_uid, new_gid)
        os.chmod(file_name, new_perm)
        return

# From https://stackoverflow.com/a/1094933/7728169
def _sizeof_fmt(num, suffix='B'):
    """
    Covert a number of bytes to a human-readable file size
    """
    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
        if abs(num) < 1024.0:
            return f"{num:3.1f}{unit}{suffix}"
        num /= 1024.0
    return f"{num:.1f}{'Yi'}{suffix}"

import os

from jinja2 import Template

def write_job_script(config, machine, target_cores, min_cores, work_dir,
                     suite=''):
    """

        Configuration options for this test case, a combination of user configs
        and the defaults for the machine and component

    machine : {str, None}
        The name of the machine

    target_cores : int
        The target number of cores for the job to use

    min_cores : int
        The minimum number of cores for the job to use

        The work directory where the job script should be written

    suite : str, optional
        The name of the suite
    """

    if config.has_option('parallel', 'account'):
        account = config.get('parallel', 'account')
    else:
        account = ''

    cores_per_node = config.getint('parallel', 'cores_per_node')

    # as a rule of thumb, let's do the geometric mean between min and target
    cores = np.sqrt(target_cores * min_cores)
    nodes = int(np.ceil(cores / cores_per_node))

    partition, qos, constraint, gpus_per_node, wall_time = get_slurm_options(
        config, machine, nodes)

    job_name = config.get('job', 'job_name')
    if job_name == '<<<default>>>':
        if suite == '':
            job_name = 'polaris'
        else:
            job_name = f'polaris_{suite}'

    template = Template(imp_res.files('polaris.job').joinpath(
        'job_script.template').read_text())

    text = template.render(job_name=job_name, account=account,
                           nodes=f'{nodes}', wall_time=wall_time, qos=qos,
                           partition=partition, constraint=constraint,
                           gpus_per_node=gpus_per_node, suite=suite)
    text = clean_up_whitespace(text)
    if suite == '':
        script_filename = 'job_script.sh'
    else:
        script_filename = f'job_script.{suite}.sh'
    script_filename = os.path.join(work_dir, script_filename)
    with open(script_filename, 'w') as handle:
        handle.write(text)

def get_slurm_options(config, machine, nodes):
    """
    Get Slurm options

        Config options

    machine : str
        Name of the machine

    nodes : int
        Number of nodes

    Returns
    -------
    partition : str
        Slurm partition

    qos : str
        Slurm quality of service

    constraint : str
        Slurm constraint

    gpus_per_node : str
        The numer of GPUs per node (if any)

    wall_time : str
        Slurm wall time
    """

    partition = config.get('job', 'partition')
    if partition == '<<<default>>>':
        if machine == 'anvil':
            # choose the partition based on the number of nodes
            if nodes <= 5:
                partition = 'acme-small'
            elif nodes <= 60:
                partition = 'acme-medium'
            else:
                partition = 'acme-large'
        elif config.has_option('parallel', 'partitions'):
            # get the first, which is the default
            partition = config.getlist('parallel', 'partitions')[0]
        else:
            partition = ''

    qos = config.get('job', 'qos')
    if qos == '<<<default>>>':
        if config.has_option('parallel', 'qos'):
            qos = config.getlist('parallel', 'qos')[0]
        else:
            qos = ''

    constraint = config.get('job', 'constraint')
    if constraint == '<<<default>>>':
        if config.has_option('parallel', 'constraints'):
            constraint = config.getlist('parallel', 'constraints')[0]
        else:
            constraint = ''

    if config.has_option('parallel', 'gpus_per_node'):
        gpus_per_node = config.get('parallel', 'gpus_per_node')
    else:
        gpus_per_node = ''

    wall_time = config.get('job', 'wall_time')

    return partition, qos, constraint, gpus_per_node, wall_time

def clean_up_whitespace(text):
    """
    Clean up whitespace after jinja templating

    text : str
        Text to clean up

    Returns
    -------
    text : str
        Text with extra blank lines removed
    """
    prev_line = None
    lines = text.split('\n')
    trimmed = list()
    # remove extra blank lines
    for line in lines:
        if line != '' or prev_line != '':
            trimmed.append(line)
            prev_line = line

    line = ''
    lines = list()
    # remove blank lines between comments
    for next_line in trimmed:
        if line != '' or not next_line.startswith('#'):
            lines.append(line)
        line = next_line

    # add the last line that we missed and an extra blank line
    lines.extend([trimmed[-1], ''])
    text = '\n'.join(lines)
    return text

import argparse
import os
import re
import sys

from polaris.components import get_components

def list_cases(task_expr=None, number=None, verbose=False):
    """
    List the available tasks

    task_expr : str, optional
        A regular expression for a task path name to search for

    number : int, optional
        The number of the task to list

    verbose : bool, optional
        Whether to print details of each task or just the subdirectories.
        When applied to suites, verbose will list the tasks in the suite.
    """
    components = get_components()

    if number is None:
        print('Tasks:')

    tasks = []
    for component in components:
        for task in component.tasks.values():
            tasks.append(task)

    for task_number, task in enumerate(tasks):
        print_number = False
        print_task = False
        if number is not None:
            if number == task_number:
                print_task = True
        elif task_expr is None or re.match(task_expr, task.path):
            print_task = True
            print_number = True

        if print_task:
            number_string = f'{task_number:d}: '.rjust(6)
            if print_number:
                prefix = number_string
            else:
                prefix = ''
            if verbose:
                lines = list()
                to_print = {'path': task.path,
                            'name': task.name,
                            'component': task.component.name,
                            'subdir': task.subdir}
                for key in to_print:
                    key_string = f'{key}: '.ljust(15)
                    lines.append(f'{prefix}{key_string}{to_print[key]}')
                    if print_number:
                        prefix = '      '
                lines.append(f'{prefix}steps:')
                longest = 0
                for step in task.steps.values():
                    longest = max(longest, len(step.name))
                    step_name = f'{step.name}: '.ljust(longest + 2)
                    lines.append(f'{prefix} - {step_name}{step.path}')
                lines.append('')
                print_string = '\n'.join(lines)
            else:
                print_string = f'{prefix}{task.path}'

            print(print_string)

def list_machines():
    machine_configs = sorted(
        [resource.name for resource in
         imp_res.files('polaris.machines').iterdir() if
         resource.is_file()])

    print('Machines:')
    for config in machine_configs:
        if config.endswith('.cfg'):
            print(f'   {os.path.splitext(config)[0]}')

def list_suites(components=None, verbose=False):
    if components is None:
        components = [component.name for component in get_components()]
    print('Suites:')
        package = f'polaris.{component}.suites'
        try:
            suites = sorted([resource.name for resource in
                             imp_res.files(package).iterdir() if
                             resource.is_file()])
            continue
        for suite in sorted(suites):
            print_header_if_cached = True
            if suite.endswith('.txt'):
                print(f'  -c {component} -t {os.path.splitext(suite)[0]}')
                if verbose:
                    text = imp_res.files(package).joinpath(suite).read_text()
                    tasks = list()
                    for task in text.split('\n'):
                        task = task.strip()
                        if task == 'cached':
                            print('\t  Cached steps: all')
                        elif task.startswith('cached'):
                            if print_header_if_cached:
                                print('\t  Cached steps:')
                            # don't print the header again if there are
                            # multiple lines of cached steps
                            print_header_if_cached = False
                            print(task.replace('cached: ', '\t    '))
                        elif (len(task) > 0 and task not in tasks and
                              not task.startswith('#')):
                            print(f'\t* {task}')
                            print_header_if_cached = True

def main():
    parser = argparse.ArgumentParser(
        description='List the available tasks or machines',
        prog='polaris list')
    parser.add_argument('-t', '--task_expr', dest='task_expr',
                        help='A regular expression for a task path name to '
                             'search for.',
                        metavar='TASK')
    parser.add_argument('-n', '--number', dest='number', type=int,
                        help='The number of the task to list.')
    parser.add_argument('--machines', dest='machines', action='store_true',
                        help='List supported machines (instead of task '
                             'cases).')
    parser.add_argument('--suites', dest='suites', action='store_true',
                        help='List suites (instead of tasks).')
    parser.add_argument('-v', '--verbose', dest='verbose', action='store_true',
                        help='List details of each task, not just the '
                             'path.  When applied to suites, verbose lists '
                             'the tasks contained in each suite.')
    args = parser.parse_args(sys.argv[2:])
    if args.machines:
        list_machines()
    elif args.suites:
        list_suites(verbose=args.verbose)
    else:
        list_cases(task_expr=args.task_expr, number=args.number,
                   verbose=args.verbose)

import cartopy
import cartopy.crs as ccrs
import jigsawpy
import matplotlib.pyplot as plt
import xarray
import xarray.plot
from jigsawpy.savejig import savejig
from mpas_tools.cime.constants import constants
from mpas_tools.mesh.creation.jigsaw_to_netcdf import jigsaw_to_netcdf
from mpas_tools.ocean.inject_meshDensity import inject_spherical_meshDensity
from mpas_tools.viz.colormaps import register_sci_viz_colormaps
from mpas_tools.viz.paraview_extractor import extract_vtk

from polaris import Step
from polaris.model_step import make_graph_file

class SphericalBaseStep(Step):
    """
    A base class for steps that create a JIGSAW spherical mesh

    opts : jigsawpy.jigsaw_jig_t
        JIGSAW options for creating the mesh
    """
    def __init__(self, component, name, subdir):
        """
        Create a new step

        name : str
            the name of the step

        subdir : {str, None}
            the subdirectory for the step.  The default is ``name``
        """
        super().__init__(component=component, name=name, subdir=subdir)

        # setup files for JIGSAW
        self.opts = jigsawpy.jigsaw_jig_t()

    def save_and_plot_cell_width(self, lon, lat, cell_width):
        """
        Save the cell width field on a lon/lat grid to
        ``self.cell_width_filename`` and plot

        lon : numpy.ndarray
            longitude in degrees (length n and between -180 and 180)

        lat : numpy.ndarray
            longitude in degrees (length m and between -90 and 90)

        cell_width : numpy.ndarray
            m x n array of cell width in km
        """
        section = self.config['spherical_mesh']
        da = xarray.DataArray(
            cell_width, dims=['lat', 'lon'], coords={'lat': lat, 'lon': lon},
            name='cellWidth')
        cell_width_filename = section.get('cell_width_filename')
        da.to_netcdf(cell_width_filename)

        if section.getboolean('plot_cell_width'):
            self._plot_cell_width(cell_width)

        """
        Add output files
        """
        for option in ['jigsaw_mesh_filename', 'mpas_mesh_filename',
                       'cell_width_filename']:
            filename = config.get('spherical_mesh', option)
            self.add_output_file(filename=filename)
        self.add_output_file(filename='graph.info')

        """
        Finish up the step. Run this after the subclass's run() method
        """
        logger = self.logger

        section = config['spherical_mesh']
        earth_radius = constants['SHR_CONST_REARTH']
        jigsaw_mesh_filename = section.get('jigsaw_mesh_filename')
        mpas_mesh_filename = section.get('mpas_mesh_filename')

        logger.info('Convert triangles from jigsaw format to netcdf')
        jigsaw_to_netcdf(msh_filename=jigsaw_mesh_filename,
                         output_name='mesh_triangles.nc', on_sphere=True,
                         sphere_radius=earth_radius)

        logger.info('Convert from triangles to MPAS mesh')

        args = ['MpasMeshConverter.x',
                'mesh_triangles.nc',
                mpas_mesh_filename]
        check_call(args=args, logger=logger)

        if section.getboolean('add_mesh_density'):
            logger.info('Add meshDensity into the mesh file')
            ds = xarray.open_dataset('cellWidthVsLatLon.nc')
            inject_spherical_meshDensity(
                ds.cellWidth.values, ds.lon.values, ds.lat.values,
                mesh_filename=mpas_mesh_filename)

        if section.getboolean('convert_to_vtk'):
            vtk_dir = section.get('vtk_dir')
            # only use progress bars if we're not writing to a log file
            use_progress_bar = self.log_filename is None
            lat_lon = section.getboolean('vtk_lat_lon')

            logger.info('Create vtk file for visualization')
            extract_vtk(ignore_time=True, lonlat=lat_lon,
                        dimension_list=['maxEdges='],
                        variable_list=['allOnCells'],
                        filename_pattern=mpas_mesh_filename,
                        out_dir=vtk_dir, use_progress_bar=use_progress_bar)

        make_graph_file(mesh_filename=mpas_mesh_filename,
                        graph_filename='graph.info')

    def _plot_cell_width(self, cell_width):
        """
        Plot a lat/lon map of cell widths (mesh resolution)

        """
        cmap = config.get('spherical_mesh', 'cell_width_colormap')
        image_filename = config.get('spherical_mesh',
                                    'cell_width_image_filename')
        register_sci_viz_colormaps()
        fig = plt.figure(figsize=[16.0, 8.0])
        ax = plt.axes(projection=ccrs.PlateCarree())
        ax.set_global()
        im = ax.imshow(cell_width, origin='lower',
                       transform=ccrs.PlateCarree(),
                       extent=[-180, 180, -90, 90], cmap=cmap, zorder=0)
        ax.add_feature(cartopy.feature.LAND, edgecolor='black', zorder=1)
        gl = ax.gridlines(
            crs=ccrs.PlateCarree(),
            draw_labels=True,
            linewidth=1,
            color='gray',
            alpha=0.5,
            linestyle='-', zorder=2)
        gl.top_labels = False
        gl.right_labels = False
        min_width = np.amin(cell_width)
        max_width = np.amax(cell_width)
        plt.title(
            f'Grid cell size, km, min: {min_width:.1f} max: {max_width:.1f}')
        plt.colorbar(im, shrink=.60)
        fig.canvas.draw()
        plt.tight_layout()
        plt.savefig(image_filename, bbox_inches='tight')
        plt.close()

class QuasiUniformSphericalMeshStep(SphericalBaseStep):
    """
    A step for creating a quasi-uniform JIGSAW mesh with a constant approximate
    cell width.  Subclasses can override the ``build_cell_width_lat_lon()``
    method to define a lon/lat map of cell width.  They can override the
    ``make_jigsaw_mesh()`` method to change how the jigsaw mesh is generated.

    cell_width : float
        The approximate cell width in km of the mesh if constant resolution
    """

    def __init__(self, component, name='base_mesh', subdir=None,
                 cell_width=None):
        """

        subdir : {str, None}, optional
            the subdirectory for the step

        cell_width : float, optional
            The approximate cell width in km of the mesh if constant resolution
        """
        self.cell_width = cell_width

        # build mesh via JIGSAW!
        self.opts.hfun_scal = 'absolute'
        self.opts.hfun_hmax = float('inf')
        self.opts.hfun_hmin = 0.0
        # 2-dim. simplexes
        self.opts.mesh_dims = 2
        self.opts.optm_qlim = 0.9375
        self.opts.verbosity = 1

        """
        Add JIGSAW options based on config options
        """
        self.opts.mesh_file = section.get('jigsaw_mesh_filename')
        self.opts.geom_file = section.get('jigsaw_geom_filename')
        self.opts.jcfg_file = section.get('jigsaw_jcfg_filename')
        self.opts.hfun_file = section.get('jigsaw_hfun_filename')
        super().setup()

        """
        Run this step of the task
        """
        cell_width, lon, lat = self.build_cell_width_lat_lon()
        self.save_and_plot_cell_width(lon, lat, cell_width)

        self.make_jigsaw_mesh(lon, lat, cell_width)

        # do the rest of the step (converting to the MPAS base mesh)
        super().run()

    def build_cell_width_lat_lon(self):
        """
        A function for creating cell width array for this mesh on a regular
        latitude-longitude grid.  The default method use constant resolution
        by setting the ``cell_width`` config option in the ``[spherical_mesh]``
        section. Subclasses should override this function if a more complex
        cell-width map is needed.

        Returns
        -------

        """

        cell_width = self.cell_width
        if cell_width is None:
            raise ValueError('The cell width was not set.')

        logger.info(f'  cell width: {cell_width} km')

        # save the constant approximate resolution on a 10 degree grid
        dlon = 10.
        dlat = dlon
        nlon = int(360. / dlon) + 1
        nlat = int(180. / dlat) + 1
        lon = np.linspace(-180., 180., nlon)
        lat = np.linspace(-90., 90., nlat)
        cell_width = cell_width * np.ones((nlat, nlon))
        return cell_width, lon, lat

    def make_jigsaw_mesh(self, lon, lat, cell_width):
        """
        Build the JIGSAW mesh.  A subclass can override this method to build
        the mesh in a different way.

        """
        opts = self.opts

        # save HFUN data to file
        hmat = jigsawpy.jigsaw_msh_t()
        hmat.mshID = 'ELLIPSOID-GRID'
        hmat.xgrid = np.radians(lon)
        hmat.ygrid = np.radians(lat)
        hmat.value = cell_width
        jigsawpy.savemsh(opts.hfun_file, hmat)

        # define JIGSAW geometry
        geom = jigsawpy.jigsaw_msh_t()
        geom.mshID = 'ELLIPSOID-MESH'
        geom.radii = earth_radius * 1e-3 * np.ones(3, float)
        jigsawpy.savemsh(opts.geom_file, geom)

        savejig(opts.jcfg_file, opts)
        check_call(['jigsaw', opts.jcfg_file], logger=logger)

class IcosahedralMeshStep(SphericalBaseStep):
    """
    A step for creating an icosahedral JIGSAW mesh

    subdivisions : int
        The number of subdivisions of the icosahedral mesh to perform
    """

                 cell_width=None, subdivisions=None):
        """

        subdivisions : int, optional
            The number of subdivisions of the icosahedral mesh to perform
        """

        # run as a subprocess so output goes to a log file
        self.run_as_subprocess = True

        self.subdivisions = subdivisions

        self.opts.hfun_hmax = 1.

        self.opts.optm_iter = 512
        self.opts.optm_qtol = 1e-6

        """
        """

        """
        """
        logger.info('Generate JIGSAW icosahedral mesh')

        subdivisions, cell_width, lon, lat = \
            self.build_subdivisions_cell_width_lat_lon()

        self.make_jigsaw_mesh(subdivisions)

    def make_jigsaw_mesh(self, subdivisions):
        """
        Make the JIGSAW mesh.  A subclass can override this method to build

        subdivisions : int
            The number of subdivisions of the icosahedron
        """

        geom.mshID = 'ellipsoid-mesh'

        icos = jigsawpy.jigsaw_msh_t()
        jigsawpy.cmd.icosahedron(opts, subdivisions, icos)

    def build_subdivisions_cell_width_lat_lon(self):
        """
        latitude-longitude grid.

        Returns
        -------
            The number of subdivisions of the icosahedron to make

        """

        icosahedral_method = config.get('spherical_mesh', 'icosahedral_method')
        if icosahedral_method == 'subdivisions':
            subdivisions = self.subdivisions
            if subdivisions is None:
                raise ValueError('The number of subdivisions was not set.')
        elif icosahedral_method == 'cell_width':
            cell_width = self.cell_width
            if cell_width is None:
                raise ValueError('The cell width was not set.')
            subdivisions = self.get_subdivisions(cell_width)
        else:
            raise ValueError(f'"icosahedral_method" must be either '
                             f'"cell_width" or the "subdivisions".  Got '
                             f'{icosahedral_method}.')

        cell_width = self.get_cell_width(subdivisions)

        logger.info(f'  subdivisions: {subdivisions}')

        # save the constant approximate resolution on a 1 degree grid
        nlon = 361
        nlat = 181
        return subdivisions, cell_width, lon, lat

    def get_subdivisions(cell_width):
        """
        Find the number of subdivisions of an icosahedron to achieve a
        resolution as close as possible to ``cell_width``.

        cell_width : float
            The approximate size in km of each cell in the resulting mesh.

        Returns
        -------
        """
        earth_area = 4 * np.pi * earth_radius ** 2

        # Using Euclidean, not spherical area, so not accurate for large cell
        # widths
        triangle_area = np.sqrt(3) / 4. * (cell_width * 1e3)**2
        triangle_count = earth_area / triangle_area
        subdivisions = 0.5 * np.log2(triangle_count / 20)
        subdivisions = max(0, int(np.round(subdivisions)))

        return subdivisions

    def get_cell_width(subdivisions):
        """
        Get the approximate cell width for an icosahedral mesh given either a
        number of subdivisions of the icosahedron.  As a rule of thumb:

        ==============  =================
         subdivisions    cell width (km)
        5               240
        6               120
        7               60
        8               30
        9               15
        10              7.5
        11              3.8
        12              1.9
        13              0.94

        Returns
        -------
            The approximate size in km of each cell
        """

        # compute and save the cell widths for later use in computing mesh
        # density
        triangle_count = 20 * 2**(2 * subdivisions)
        triangle_area = earth_area / triangle_count
        cell_width = 1e-3 * np.sqrt(triangle_area * 4. / np.sqrt(3))
        return cell_width

from typing import Dict

def parse_replacements(package, namelist):
    """
    Parse the replacement namelist options from the given file

    package : Package
        The package name or module object that contains ``namelist``

        The name of the namelist replacements file to read from

    Returns
    -------
    replacements : dict
        A dictionary of replacement namelist options
    """

    lines = imp_res.files(package).joinpath(namelist).read_text().split('\n')
    replacements = dict()
        if '=' in line:
            opt, val = line.split('=')
            replacements[opt.strip()] = val.strip()

    return replacements

def ingest(defaults_filename):
    """ Read the defaults file """
    with open(defaults_filename, 'r') as f:
        lines = f.readlines()

    namelist: Dict[str, Dict[str, str]] = dict()
    record = None
        if '&' in line:
            record = line.strip('&').strip('\n').strip()
            namelist[record] = dict()
        elif '=' in line:
            if record is not None:
                opt, val = line.strip('\n').split('=')
                namelist[record][opt.strip()] = val.strip()

    return namelist

def replace(namelist, replacements):
    """ Replace entries in the namelist using the replacements dict """
    new = dict(namelist)
    is_not_replaced = [True for _ in replacements.keys()]
    for record in new:
        for idx, key in enumerate(replacements):
            if key in new[record]:
                new[record][key] = replacements[key]
                is_not_replaced[idx] = False
    for idx, key in enumerate(replacements):
        if is_not_replaced[idx]:
            print(f'Warning: {key} is not in the namelist and replacements '
                  'will not be used')

    return new

def write(namelist, filename):
    """ Write the namelist out """

    with open(filename, 'w') as f:
        for record in namelist:
            f.write(f'&{record}\n')
            rec = namelist[record]
            for key in rec:
                f.write(f'    {key.strip()} = {rec[key].strip()}\n')
            f.write('/\n')

import multiprocessing
import os
import subprocess
import warnings

def get_available_parallel_resources(config):
    """
    Get the number of total cores and nodes available for running steps

        Configuration options for the test case

    Returns
    -------
    available_resources : dict
        A dictionary containing available resources (cores, tasks, nodes
        and cores_per_node)
    """

    parallel_system = config.get('parallel', 'system')
    if parallel_system == 'slurm' and 'SLURM_JOB_ID' not in os.environ:
        parallel_system = 'login'

    if parallel_system == 'slurm':
        job_id = os.environ['SLURM_JOB_ID']
        node = os.environ['SLURMD_NODENAME']
        args = ['sinfo', '--noheader', '--node', node, '-o', '%C']
        # get allocated, idle, other and total cores
        aiot = _get_subprocess_str(args).split('/')

        # we can only use the allocated cores
        cores_per_node = int(aiot[0])

        if cores_per_node == 0:
            # hmm, no allocated cores so I guess we'll go with total cores
            cores_per_node = int(aiot[3])

        args = ['sinfo', '--noheader', '--node', node, '-o', '%Z']
        slurm_threads_per_core = _get_subprocess_int(args)

        if config.has_option('parallel', 'threads_per_core'):
            threads_per_core = config.getint('parallel', 'threads_per_core')
            # correct for this in allocated_cores, which might not match
            cores_per_node = ((cores_per_node * threads_per_core) //
                              slurm_threads_per_core)
        else:
            threads_per_core = slurm_threads_per_core
        args = ['squeue', '--noheader', '-j', job_id, '-o', '%D']
        nodes = _get_subprocess_int(args)
        cores = cores_per_node * nodes
        mpi_allowed = True
    elif parallel_system == 'login':
        cores = min(multiprocessing.cpu_count(),
                    config.getint('parallel', 'login_cores'))
        cores_per_node = cores
        nodes = 1
        mpi_allowed = False
    elif parallel_system == 'single_node':
        cores = multiprocessing.cpu_count()
        if config.has_option('parallel', 'cores_per_node'):
            cores = min(cores, config.getint('parallel', 'cores_per_node'))
        nodes = 1
    else:
        raise ValueError(f'Unexpected parallel system: {parallel_system}')

    available_resources = dict(
        cores=cores,
        nodes=nodes,
        cores_per_node=cores_per_node,
        mpi_allowed=mpi_allowed
    )

        available_resources['gpus_per_node'] = \
            config.getint('parallel', 'gpus_per_node')

    return available_resources

def set_cores_per_node(config, cores_per_node):
    """
    If the system has Slurm, find out the ``cpus_per_node`` and set the config
    option accordingly.

        Configuration options
    """
        old_cores_per_node = config.getint('parallel', 'cores_per_node')
        config.set('parallel', 'cores_per_node', f'{cores_per_node}')
        if old_cores_per_node != cores_per_node:
            warnings.warn(f'Slurm found {cores_per_node} cpus per node but '
                          f'config from mache was {old_cores_per_node}')
        if not config.has_option('parallel', 'cores_per_node'):
            config.set('parallel', 'cores_per_node', f'{cores_per_node}')

def run_command(args, cpus_per_task, ntasks, openmp_threads, config, logger):
    """
    Run a subprocess with the given command-line arguments and resources

    args : list of str
        The command-line arguments to run in parallel

    cpus_per_task : int
        the number of cores per task the process would ideally use.  If
        fewer cores per node are available on the system, the substep will

    ntasks : int
        the number of tasks the process would ideally use.  If too few
        tasks and the number of cores per task, the substep will run on

    config : configparser.ConfigParser

    """
    env = dict(os.environ)

    env['OMP_NUM_THREADS'] = f'{openmp_threads}'
    if openmp_threads > 1:
        logger.info(f'Running with {openmp_threads} OpenMP threads')

    command_line_args = get_parallel_command(args, cpus_per_task, ntasks,
                                             config)
    check_call(command_line_args, logger, env=env)

def get_parallel_command(args, cpus_per_task, ntasks, config):
    """

    Returns
    -------
    command_line_args : list
        The full parallel command
    """
    parallel_executable = config.get('parallel', 'parallel_executable')

    # split the parallel executable into constituents in case it includes flags
    command_line_args = parallel_executable.split(' ')
        cores = ntasks * cpus_per_task
        nodes = int(np.ceil(cores / cores_per_node))
        command_line_args.extend(['-c', f'{cpus_per_task}',
                                  '-N', f'{nodes}',
                                  '-n', f'{ntasks}'])
        command_line_args.extend(['-n', f'{ntasks}'])
    else:

    command_line_args.extend(args)
    return command_line_args

def _get_subprocess_str(args):
    value = subprocess.check_output(args)
    value_str = value.decode('utf-8').strip('\n')
    return value_str

def _get_subprocess_int(args):
    value_int = int(_get_subprocess_str(args))
    return value_int

import os.path

import pyproj
from pyremap import (
    LatLon2DGridDescriptor,
    LatLonGridDescriptor,
    MpasCellMeshDescriptor,
    MpasEdgeMeshDescriptor,
    MpasVertexMeshDescriptor,
    PointCollectionDescriptor,
    ProjectionGridDescriptor,
    Remapper,
    get_lat_lon_descriptor,
)

class MappingFileStep(Step):
    """
    A step for creating a mapping file between grids

    src_grid_info : dict
        Information about the source grid

    dst_grid_info : dict
        Information about the destination grid

    method : {'bilinear', 'neareststod', 'conserve'}
        The method of interpolation used

    map_filename : str or None
        The name of the output mapping file

    expand_distance : float or None
        The distance in meters over which to expand destination grid cells for
         smoothing

    expand_factor : float or None
        The factor by which to expand destination grid cells for smoothing

    src_mesh_filename : str
        The name of the SCRIP file for the source mesh

    dst_mesh_filename : str
        The name of the SCRIP file for the destination mesh
    """

                 min_tasks=None, map_filename=None, method='bilinear'):
        """

        name : str

            the target number of MPI tasks the step would ideally use

            the number of MPI tasks the step requires

        map_filename : str, optional
            The name of the output mapping file,
            ``map_{source_type}_{dest_type}_{method}.nc`` by default

        method : {'bilinear', 'neareststod', 'conserve'}, optional
            The method of interpolation used
        """
                         indir=indir, ntasks=ntasks, min_tasks=min_tasks)
        self.src_grid_info = dict()
        self.dst_grid_info = dict()
        self.map_filename = map_filename
        self.method = method
        self.expand_distance = None
        self.expand_factor = None
        self.src_mesh_filename = 'src_mesh.nc'
        self.dst_mesh_filename = 'dst_mesh.nc'

    def src_from_mpas(self, filename, mesh_name, mesh_type='cell'):
        """
        Set the source grid from an MPAS mesh file

            A file containing the MPAS mesh

        mesh_name : str
            The name of the MPAS mesh

        mesh_type : {'cell', 'edge', 'vertex'}, optional
            Which type of MPAS mesh
        """
        src = dict()
        src['type'] = 'mpas'
        src['filename'] = filename
        src['name'] = mesh_name
        src['mpas_mesh_type'] = mesh_type
        self.src_grid_info = src

    def dst_from_mpas(self, filename, mesh_name, mesh_type='cell'):
        """
        Set the destination grid from an MPAS mesh file

        """
        dst = dict()
        dst['type'] = 'mpas'
        dst['filename'] = filename
        dst['name'] = mesh_name
        dst['mpas_mesh_type'] = mesh_type
        self.dst_grid_info = dst

    def src_from_lon_lat(self, filename, mesh_name=None, lon_var='lon',
                         lat_var='lat'):
        """
        Set the source grid from a file with a longitude-latitude grid.  The
        latitude and longitude variables can be 1D or 2D.

            A file containing the latitude-longitude grid

        mesh_name : str, optional
            The name of the lon-lat grid (defaults to resolution and units,
            something like "0.5x0.5degree")

        lon_var : str, optional
            The name of the longitude coordinate in the file

        lat_var : str, optional
            The name of the latitude coordinate in the file
        """
        src['type'] = 'lon-lat'
        src['lon'] = lon_var
        src['lat'] = lat_var
        if mesh_name is not None:
            src['name'] = mesh_name

    def dst_from_lon_lat(self, filename, mesh_name=None, lon_var='lon',
        """
        Set the destination grid from a file with a longitude-latitude grid.
        The latitude and longitude variables can be 1D or 2D.

        """
        dst['type'] = 'lon-lat'
        dst['lon'] = lon_var
        dst['lat'] = lat_var
            dst['name'] = mesh_name

    def dst_global_lon_lat(self, dlon, dlat, lon_min=-180., mesh_name=None):
        """

        dlon : float
            The longitude resolution in degrees

        dlat : float
            The latitude resolution in degrees

        lon_min : float, optional
            The longitude for the left-hand edge of the global grid in degrees

        """

        dst['dlon'] = dlon
        dst['dlat'] = dlat
        dst['lon_min'] = lon_min

    def src_from_proj(self, filename, mesh_name, x_var='x', y_var='y',
                      proj_attr=None, proj_str=None):
        """
        Set the source grid from a file with a projection grid.

            A file containing the projection grid

            The name of the projection grid

        x_var : str, optional
            The name of the x coordinate in the file

        y_var : str, optional
            The name of the y coordinate in the file

        proj_attr : str, optional
            The name of a global attribute in the file containing the proj
            string for the projection

        proj_str : str, optional
            A proj string defining the projection, ignored if ``proj_attr``
            is provided
        """
        src['type'] = 'proj'
        src['x'] = x_var
        src['y'] = y_var
        if proj_attr is not None:
            src['proj_attr'] = proj_attr
        elif proj_str is not None:
            src['proj_str'] = proj_str
        else:
            raise ValueError('Must provide one of "proj_attr" or "proj_str".')

    def dst_from_proj(self, filename, mesh_name, x_var='x', y_var='y',
        """
        Set the destination grid from a file with a projection grid.

        """
        dst['type'] = 'proj'
        dst['x'] = x_var
        dst['y'] = y_var
            dst['proj_attr'] = proj_attr
            dst['proj_str'] = proj_str
        else:

    def dst_from_points(self, filename, mesh_name, lon_var='lon',
                        lat_var='lat'):
        """
        Set the destination grid from a file with a collection of points.

            The name of the point collection

        """
        dst['type'] = 'points'

    def get_remapper(self):
        """
        Get the remapper for this step.  After the step has been run, it can
        be used to remap data between the source and destination grids by
        calling its ``remap_file()`` or ``remap()`` methods.

        Returns
        -------
        remapper : pyremap.Remapper
            The remapper between the source and destination grids
        """
        src = self.src_grid_info
        dst = self.dst_grid_info

        if 'type' not in src:
            raise ValueError('None of the "src_from_*()" methods were called')

        if 'type' not in dst:
            raise ValueError('None of the "dst_from_*()" methods were called')

        # to absolute paths for when the remapper is used in another step
        for info in [src, dst]:
            if 'filename' in info:
                info['filename'] = os.path.abspath(os.path.join(
                    self.work_dir, info['filename']))

        in_descriptor = _get_descriptor(src)
        out_descriptor = _get_descriptor(dst)

        if self.map_filename is None:
            map_tool = self.config.get('mapping', 'map_tool')
            prefixes = {
                'esmf': 'esmf',
                'moab': 'mbtr'
            }
            suffixes = {
                'conserve': 'aave',
                'bilinear': 'bilin',
                'neareststod': 'neareststod'
            }
            suffix = f'{prefixes[map_tool]}{suffixes[self.method]}'

            self.map_filename = \
                f'map_{in_descriptor.meshName}_to_{out_descriptor.meshName}' \
                f'_{suffix}.nc'

        self.map_filename = os.path.abspath(os.path.join(
            self.work_dir, self.map_filename))

        remapper = Remapper(in_descriptor, out_descriptor, self.map_filename)
        return remapper

        """
        Create a remapper and set the command-line arguments
        """
        remapper = self.get_remapper()
        map_tool = self.config.get('mapping', 'map_tool')
        _check_remapper(remapper, self.method, map_tool=map_tool)

        src_descriptor = remapper.sourceDescriptor
        src_descriptor.to_scrip(self.src_mesh_filename)

        dst_descriptor = remapper.destinationDescriptor
        dst_descriptor.to_scrip(self.dst_mesh_filename,
                                expandDist=self.expand_distance,
                                expandFactor=self.expand_factor)

        if map_tool == 'esmf':
            self.args = _esmf_build_map_args(remapper, self.method,
                                             src_descriptor,
                                             self.src_mesh_filename,
                                             dst_descriptor,
                                             self.dst_mesh_filename)
        elif map_tool == 'moab':
            self.args = _moab_build_map_args(remapper, self.method,

def _check_remapper(remapper, method, map_tool):
    """
    Check for inconsistencies in the remapper
    """
    if map_tool not in ['moab', 'esmf']:
        raise ValueError(f'Unexpected map_tool {map_tool}. Valid '
                         f'values are "esmf" or "moab".')

    if isinstance(remapper.destinationDescriptor,
                  PointCollectionDescriptor) and \
            method not in ['bilinear', 'neareststod']:
        raise ValueError(f'method {method} not supported for destination '
                         f'grid of type PointCollectionDescriptor.')

    if map_tool == 'moab' and method == 'neareststod':
        raise ValueError('method neareststod not supported by mbtempest.')

def _esmf_build_map_args(remapper, method, src_descriptor, src_mesh_filename,
                         dst_descriptor, dst_mesh_filename):
    """
    Get command-line arguments for making a mapping file with
    ESMF_RegridWeightGen
    """

    args = ['ESMF_RegridWeightGen',
            '--source', src_mesh_filename,
            '--destination', dst_mesh_filename,
            '--weight', remapper.mappingFileName,
            '--method', method,
            '--netcdf4']

    if src_descriptor.regional:
        args.append('--src_regional')

    if dst_descriptor.regional:
        args.append('--dst_regional')

    if src_descriptor.regional or dst_descriptor.regional:
        args.append('--ignore_unmapped')

    return [args]

def _moab_build_map_args(remapper, method, src_descriptor, src_mesh_filename,
    """
    Get command-line arguments for making a mapping file with mbtempest
    """
    fvmethod = {
        'conserve': 'none',
        'bilinear': 'bilin'}

    map_filename = remapper.mappingFileName
    intx_filename = \
        f'moab_intx_{src_descriptor.meshName}_to_{dst_descriptor.meshName}.h5m'

    intx_args = [
        'mbtempest',
        '--type', '5',
        '--load', src_mesh_filename,
        '--load', dst_mesh_filename,
        '--intx', intx_filename
    ]

        intx_args.append('--rrmgrids')

    map_args = [
        '--intx', intx_filename,
        '--weights',
        '--method', 'fv',
        '--file', map_filename,
        '--order', '1',
        '--fvmethod', fvmethod[method]
    ]

    if method == 'conserve' and (src_descriptor.regional or
                                 dst_descriptor.regional):
        map_args.append('--rrmgrids')

    return [intx_args, map_args]

def _get_descriptor(info):
    """
    Get a mesh descriptor from the mesh info
    """
    grid_type = info['type']
    if grid_type == 'mpas':
        descriptor = _get_mpas_descriptor(info)
    elif grid_type == 'lon-lat':
        descriptor = _get_lon_lat_descriptor(info)
    elif grid_type == 'proj':
        descriptor = _get_proj_descriptor(info)
    elif grid_type == 'points':
        descriptor = _get_points_descriptor(info)
    else:
        raise ValueError(f'Unexpected grid type {grid_type}')

    # for compatibility with mbtempest
    descriptor.format = 'NETCDF3_64BIT_DATA'
    return descriptor

def _get_mpas_descriptor(info):
    """
    Get an MPAS mesh descriptor from the given info
    """
    mesh_type = info['mpas_mesh_type']
    filename = info['filename']
    mesh_name = info['name']

    if mesh_type == 'cell':
        descriptor = MpasCellMeshDescriptor(fileName=filename,
                                            meshName=mesh_name)
    elif mesh_type == 'vertex':
        descriptor = MpasVertexMeshDescriptor(fileName=filename,
                                              meshName=mesh_name)
    elif mesh_type == 'edge':
        descriptor = MpasEdgeMeshDescriptor(fileName=filename,
    else:
        raise ValueError(f'Unexpected MPAS mesh type {mesh_type}')

def _get_lon_lat_descriptor(info):
    """
    Get a lon-lat descriptor from the given info
    """

    if 'dlat' in info and 'dlon' in info:
        lon_min = info['lon_min']
        lon_max = lon_min + 360.
        descriptor = get_lat_lon_descriptor(dLon=info['dlon'],
                                            dLat=info['dlat'],
                                            lonMin=lon_min,
                                            lonMax=lon_max)
    else:
        filename = info['filename']
        lon = info['lon']
        lat = info['lat']
        with xr.open_dataset(filename) as ds:
            lon_lat_1d = len(ds[lon].dims) == 1 and len(ds[lat].dims) == 1
            lon_lat_2d = len(ds[lon].dims) == 2 and len(ds[lat].dims) == 2
            if not lon_lat_1d and not lon_lat_2d:
                raise ValueError(f'longitude and latitude coordinates {lon} '
                                 f'and {lat} have unexpected sizes '
                                 f'{len(ds[lon].dims)} and '
                                 f'{len(ds[lat].dims)}.')

        if lon_lat_1d:
            descriptor = LatLonGridDescriptor.read(fileName=filename,
                                                   lonVarName=lon,
                                                   latVarName=lat)
        else:
            descriptor = LatLon2DGridDescriptor.read(fileName=filename,
                                                     lonVarName=lon,
                                                     latVarName=lat)

    if 'name' in info:
        descriptor.meshName = info['name']

def _get_proj_descriptor(info):
    """
    Get a ProjectionGridDescriptor from the given info
    """
    grid_name = info['name']
    x = info['x']
    y = info['y']
    if 'proj_attr' in info:
            proj_str = ds.attrs[info['proj_attr']]
    else:
        proj_str = info['proj_str']

    proj = pyproj.Proj(proj_str)

    descriptor = ProjectionGridDescriptor.read(projection=proj,
                                               fileName=filename,
                                               meshName=grid_name,
                                               xVarName=x,
                                               yVarName=y)

def _get_points_descriptor(info):
    """
    Get a PointCollectionDescriptor from the given info
    """
    collection_name = info['name']
    lon_var = info['lon']
    lat_var = info['lat']
    with xr.open_dataset(filename) as ds:
        lon = ds[lon_var].value
        lat = ds[lat_var].values
        unit_attr = lon.attrs['units'].lower()
        if 'deg' in unit_attr:
            units = 'degrees'
        elif 'rad' in unit_attr:
            units = 'radians'
        else:
            raise ValueError(f'Unexpected longitude unit unit {unit_attr}')

    descriptor = PointCollectionDescriptor(lons=lon, lats=lat,
                                           collectionName=collection_name,
                                           units=units)

import os
import pickle

def unpickle_suite(suite_name):
    """
    Unpickle the suite

    suite_name : str
        The name of the test suite

    Returns
    -------
    test_suite : dict
        The test suite
    """
    # Allow a suite name to either include or not the .pickle suffix
    if suite_name.endswith('.pickle'):
        # code below assumes no suffix, so remove it
        suite_name = suite_name[:-len('.pickle')]
    # Now open the the suite's pickle file
    if not os.path.exists(f'{suite_name}.pickle'):
        raise ValueError(f'The suite "{suite_name}" does not appear to have '
                         f'been set up here.')
    with open(f'{suite_name}.pickle', 'rb') as handle:
        test_suite = pickle.load(handle)
    return test_suite

def setup_config(base_work_dir, component_name, config_filepath):
    """
    Set up the config object from the config file

        The base work directory where suites and tasks are being set up

    component_name : str
        The component the config belongs to

    config_filepath : str
        The path to the config file within the component's work directory

    Returns
    -------
        The config object
    """
    config_filename = os.path.join(base_work_dir, component_name,
                                   config_filepath)
    config = PolarisConfigParser()
    config.add_from_file(config_filename)
    config.filepath = config_filepath
    return config

def load_dependencies(step):
    """
    Load each dependency from its pickle file to pick up changes that may have
    happened since it ran

    step : polaris.step.Step
        The step object
    """
    for name, old_dependency in step.dependencies.items():
        if old_dependency.cached:
            continue

        pickle_filename = os.path.join(old_dependency.work_dir,
                                       'step_after_run.pickle')
        if not os.path.exists(pickle_filename):
            raise ValueError(f'The dependency {name} of '
                             f'step {step.path} was not run.')

        with open(pickle_filename, 'rb') as handle:
            dependency = pickle.load(handle)
            step.dependencies[name] = dependency

def complete_step_run(step):
    """
    Write a file to indicate that the step has completed. If this step is a
    dependency of other steps, pickle the step after it has run so its
    dependencies will pick up any changes in its attributes

    """
    with open('polaris_step_complete.log', 'w') as log_file:
        log_file.write(f'{step.path} finished successfully.')
    if step.is_dependency:
        # pickle the test case and step for use at runtime
        pickle_filename = os.path.join(step.work_dir, 'step_after_run.pickle')
        with open(pickle_filename, 'wb') as handle:
            pickle.dump(step, handle,
                        protocol=pickle.HIGHEST_PROTOCOL)

import glob
import os
import sys
import time
from datetime import timedelta

import mpas_tools.io
from mpas_tools.logging import LoggingContext, check_call

from polaris import Task
from polaris.logging import log_function_call, log_method_call
from polaris.parallel import (
    get_available_parallel_resources,
    run_command,
    set_cores_per_node,
)
from polaris.run import (
    complete_step_run,
    load_dependencies,
    setup_config,
    unpickle_suite,
)

# ANSI fail text: https://stackoverflow.com/a/287944/7728169
start_fail = '\033[91m'
start_pass = '\033[92m'
start_time_color = '\033[94m'
end_color = '\033[0m'
pass_str = f'{start_pass}PASS{end_color}'
success_str = f'{start_pass}SUCCESS{end_color}'
fail_str = f'{start_fail}FAIL{end_color}'
error_str = f'{start_fail}ERROR{end_color}'

def run_tasks(suite_name, quiet=False, is_task=False, steps_to_run=None,
              steps_to_skip=None):
    """
    Run the given suite or task

    quiet : bool, optional
        Whether step names are not included in the output as the suite
        progresses

    is_task : bool
        Whether this is a task instead of a full suite

    steps_to_run : list of str, optional
        A list of the steps to run if this is a task, not a full suite.
        The default behavior is to run the default steps unless they are in
        ``steps_to_skip``

    steps_to_skip : list of str, optional
        A list of steps not to run if this is a task, not a full suite.
        Typically, these are steps to remove from the defaults
    """

    suite = unpickle_suite(suite_name)

    # get the config file for the first task in the suite
    task = next(iter(suite['tasks'].values()))
    component = task.component
    common_config = setup_config(task.base_work_dir,
                                 component.name,
                                 f'{component.name}.cfg')
    available_resources = get_available_parallel_resources(common_config)

    # start logging to stdout/stderr
    with LoggingContext(suite_name) as stdout_logger:

        os.environ['PYTHONUNBUFFERED'] = '1'

        if not is_task:
            try:
                os.makedirs('case_outputs')
            except OSError:
                pass

        failures = 0
        cwd = os.getcwd()
        suite_start = time.time()
        task_times = dict()
        result_strs = dict()
        for task_name in suite['tasks']:
            stdout_logger.info(f'{task_name}')

            task = suite['tasks'][task_name]

            if is_task:
                log_filename = None
                task_logger = stdout_logger
            else:
                task_prefix = task.path.replace('/', '_')
                log_filename = f'{cwd}/case_outputs/{task_prefix}.log'
                task_logger = None

            result_str, success, task_time = _log_and_run_task(
                task, stdout_logger, task_logger, quiet, log_filename,
                is_task, steps_to_run, steps_to_skip,
                available_resources)
            result_strs[task_name] = result_str
            if not success:
                failures += 1
            task_times[task_name] = task_time

        suite_time = time.time() - suite_start

        os.chdir(cwd)
        _log_task_runtimes(stdout_logger, task_times, result_strs, suite_time,
                           failures)

def run_single_step(step_is_subprocess=False, quiet=False):
    """
    Used by the framework to run a step when ``polaris serial`` gets called in
    the step's work directory

    step_is_subprocess : bool, optional
        Whether the step is being run as a subprocess of a task or suite

    """
    with open('step.pickle', 'rb') as handle:
        step = pickle.load(handle)
    task = Task(component=step.component, name='dummy_task')
    task.add_step(step)
    task.new_step_log_file = False

    # This prevents infinite loop of subprocesses
    if step_is_subprocess:
        step.run_as_subprocess = False

    config = setup_config(step.base_work_dir,
                          step.component.name,
                          step.config.filepath)
    task.config = config
    available_resources = get_available_parallel_resources(config)
    set_cores_per_node(task.config, available_resources['cores_per_node'])

    mpas_tools.io.default_format = config.get('io', 'format')
    mpas_tools.io.default_engine = config.get('io', 'engine')

    logger_name = step.path.replace('/', '_')
    with LoggingContext(name=logger_name) as stdout_logger:
        task.logger = stdout_logger
        if quiet:
            task.stdout_logger = None
        else:
            task.stdout_logger = stdout_logger
            log_function_call(function=_run_task, logger=stdout_logger)
            stdout_logger.info('')
            stdout_logger.info(f'Running step: {step.name}')
        _run_task(task, available_resources)

        description='Run a suite, task or step',
        prog='polaris serial')
    parser.add_argument("suite", nargs='?',
                        help="The name of a suite to run. Can exclude "
                        "or include the .pickle filename suffix.")
    parser.add_argument("--steps", dest="steps", nargs='+',
                        help="The steps of a task to run.")
    parser.add_argument("--skip_steps", dest="skip_steps", nargs='+',
                        help="The steps of a task not to run, see "
                             "steps_to_run in the config file for defaults.")
    parser.add_argument("-q", "--quiet", dest="quiet", action="store_true",
                        help="If set, step names are not included in the "
                             "output as the suite progresses.  Has no "
                             "effect when running tasks or steps on "
                             "their own.")
    parser.add_argument("--step_is_subprocess", dest="step_is_subprocess",
                        action="store_true",
                        help="Used internally by polaris to indicate that "
                             "a step is being run as a subprocess.")

    if args.suite is not None:
        # Running a specified suite from the base work directory
        run_tasks(args.suite, quiet=args.quiet)
    elif os.path.exists('task.pickle'):
        # Running a task inside of its work directory
        run_tasks(suite_name='task', quiet=args.quiet, is_task=True,
                  steps_to_run=args.steps, steps_to_skip=args.skip_steps)
    elif os.path.exists('step.pickle'):
        # Running a step inside of its work directory
        run_single_step(step_is_subprocess=args.step_is_subprocess,
                        quiet=args.quiet)
    else:
        pickles = glob.glob('*.pickle')
        if len(pickles) == 1:
            # Running an unspecified suite from the base work directory
            suite = os.path.splitext(os.path.basename(pickles[0]))[0]
            run_tasks(suite, quiet=args.quiet)
        elif len(pickles) == 0:
            raise OSError('No pickle files were found. Are you sure this is '
                          'a polaris suite, task or step work directory?')
        else:
            raise ValueError('More than one suite was found. Please specify '
                             'which to run: polaris serial <suite>')

def _update_steps_to_run(task_name, steps_to_run, steps_to_skip, config,
                         steps):
    """
    Update the steps to run
    """
    if steps_to_run is None:
        step_str = config.get(task_name, 'steps_to_run').replace(',', ' ')
        steps_to_run = step_str.split()

    for step in steps_to_run:
        if step not in steps:
                f'A step "{step}" was requested but is not one of the steps '
                f'in this task:'
                f'\n{list(steps)}')

    if steps_to_skip is not None:
        for step in steps_to_skip:
            if step not in steps:
                raise ValueError(
                    f'A step "{step}" was flagged not to run but is not one '
                    f'of the steps in this task:'
                    f'\n{list(steps)}')

        steps_to_run = [step for step in steps_to_run if step not in
                        steps_to_skip]

    return steps_to_run

def _log_task_runtimes(stdout_logger, task_times, result_strs, suite_time,
                       failures):
    """
    Log the runtimes for the task(s)
    """
    stdout_logger.info('Task Runtimes:')
    for task_name, task_time in task_times.items():
        task_time_str = str(timedelta(seconds=round(task_time)))
        stdout_logger.info(f'{task_time_str} '
                           f'{result_strs[task_name]} {task_name}')
    suite_time_str = str(timedelta(seconds=round(suite_time)))
    stdout_logger.info(f'Total runtime: {suite_time_str}')

    if failures == 0:
        stdout_logger.info('PASS: All passed successfully!')
    else:
        if failures == 1:
            message = '1 task'
        else:
            message = f'{failures} tasks'
        stdout_logger.error(f'FAIL: {message} failed, see above.')
        sys.exit(1)

def _print_to_stdout(task, message):
    """
    Write out a message to stdout if we're not running a single step
    """
    if task.stdout_logger is not None:
        task.stdout_logger.info(message)
        if task.logger != task.stdout_logger:
            # also write it to the log file
            task.logger.info(message)

def _log_and_run_task(task, stdout_logger, task_logger, quiet,
                      log_filename, is_task, steps_to_run,
                      steps_to_skip, available_resources):

    task_name = task.path.replace('/', '_')
    with LoggingContext(task_name, logger=task_logger,
                        log_filename=log_filename) as task_logger:
        if quiet:
            # just log the step names and any failure messages to the
            # log file
            task.stdout_logger = task_logger
        else:
            # log steps to stdout
        task.logger = task_logger
        task.log_filename = log_filename

        # If we are running a task on its own, we want a log file per step
        # If we are running within a suite, we want a log file per task, with
        #   output from each of its steps
        task.new_step_log_file = is_task

        os.chdir(task.work_dir)

        config = setup_config(task.base_work_dir,
                              task.component.name,
                              task.config.filepath)
        task.config = config
        set_cores_per_node(task.config,
                           available_resources['cores_per_node'])

        mpas_tools.io.default_format = config.get('io', 'format')
        mpas_tools.io.default_engine = config.get('io', 'engine')

        task.steps_to_run = _update_steps_to_run(
            task.name, steps_to_run, steps_to_skip, config, task.steps)

        task_start = time.time()

        log_function_call(function=_run_task, logger=task_logger)
        task_logger.info('')
        task_list = ', '.join(task.steps_to_run)
        task_logger.info(f'Running steps: {task_list}')
        try:
            baselines_passed = _run_task(task, available_resources)
            run_status = success_str
            task_pass = True
        except Exception:
            run_status = error_str
            task_pass = False
            task_logger.exception('Exception raised while running '
                                  'the steps of the task')

        status = f'  task execution:   {run_status}'
        if task_pass:
            stdout_logger.info(status)
            if baselines_passed is None:
                result_str = pass_str
                success = True
            else:
                if baselines_passed:
                    baseline_str = pass_str
                    result_str = pass_str
                    success = True
                else:
                    baseline_str = fail_str
                    result_str = fail_str
                    success = False
                status = f'  baseline comp.:   {baseline_str}'
                stdout_logger.info(status)

        else:
            stdout_logger.error(status)
            if not is_task:
                stdout_logger.error(f'  see: case_outputs/{task_name}.log')
            result_str = fail_str
            success = False

        task_time = time.time() - task_start

        stdout_logger.info(f'  task runtime:     '
                           f'{start_time_color}{task_time_str}{end_color}')

        return result_str, success, task_time

def _run_task(task, available_resources):
    """
    Run each step of the task
    """
    logger = task.logger
    cwd = os.getcwd()
    baselines_passed = None
    for step_name in task.steps_to_run:
        step = task.steps[step_name]
        complete_filename = os.path.join(step.work_dir,
                                         'polaris_step_complete.log')

        _print_to_stdout(task, f'  * step: {step_name}')

        if os.path.exists(complete_filename):
            _print_to_stdout(task, '          already completed')
            continue
        if step.cached:
            _print_to_stdout(task, '          cached')
            continue

        step_start = time.time()

        step.config = setup_config(step.base_work_dir,
                                   step.component.name,
                                   step.config.filepath)
        if task.log_filename is not None:
            step_log_filename = task.log_filename
        else:
            step_log_filename = None

        try:
            if step.run_as_subprocess:
                _run_step_as_subprocess(
                    logger, step, task.new_step_log_file)
            else:
                _run_step(task, step, task.new_step_log_file,
                          available_resources, step_log_filename)
            _print_to_stdout(task,
                             f'          execution:        {error_str}')
            raise
        _print_to_stdout(task,
                         f'          execution:        {success_str}')
        step_time = time.time() - step_start
        step_time_str = str(timedelta(seconds=round(step_time)))

        compared, status = step.validate_baselines()
        if compared:
            if status:
                baseline_str = pass_str
            else:
                baseline_str = fail_str
                             f'          baseline comp.:   {baseline_str}')
                baselines_passed = status
            elif not status:
                baselines_passed = False

                         f'          runtime:          '
                         f'{start_time_color}{step_time_str}{end_color}')

    return baselines_passed

def _run_step(task, step, new_log_file, available_resources,
              step_log_filename):
    """
    Run the requested step
    """

    missing_files = list()
    for input_file in step.inputs:
        if not os.path.exists(input_file):
            missing_files.append(input_file)

    if len(missing_files) > 0:
        raise OSError(
            f'input file(s) missing in step {step.name} in '
            f'{step.component.name}/{step.subdir}: {missing_files}')

    load_dependencies(step)

    # each logger needs a unique name
    if new_log_file:
        # we want to create new log file and point the step to that name
        new_log_filename = f'{cwd}/{step.name}.log'
        step_log_filename = new_log_filename
        step_logger = None
    else:
        # either we don't want a log file at all or there is an existing one
        # to use.  Either way, we don't want a new log filename and we want
        # to use the existing logger.  The step log filename will be whatever
        # is passed as a parameter
        step_logger = logger
        new_log_filename = None

    step.log_filename = step_log_filename

    with LoggingContext(name=logger_name, logger=step_logger,
                        log_filename=new_log_filename) as step_logger:
        step.logger = step_logger
        os.chdir(step.work_dir)

        step_logger.info('')
        log_method_call(method=step.constrain_resources, logger=step_logger)
        step.constrain_resources(available_resources)

        # runtime_setup() will perform small tasks that require knowing the
        # resources of the task before the step runs (such as creating
        # graph partitions)
        log_method_call(method=step.runtime_setup, logger=step_logger)
        step.runtime_setup()

        if step.args is not None:
            step_logger.info('\nBypassing step\'s run() method and running '
                             'with command line args\n')
            for args in step.args:
                log_function_call(function=run_command, logger=step_logger)
                step_logger.info('')
                run_command(args, step.cpus_per_task, step.ntasks,
                            step.openmp_threads, step.config, step.logger)
        else:
            step_logger.info('')
            log_method_call(method=step.run, logger=step_logger)
            step.run()

    complete_step_run(step)

    for output_file in step.outputs:
        if not os.path.exists(output_file):
            missing_files.append(output_file)

        # We want to indicate that the step failed by removing the pickle
        try:
            os.remove('step_after_run.pickle')
            pass
            f'output file(s) missing in step {step.name} in '

def _run_step_as_subprocess(logger, step, new_log_file):
    """
    Run the requested step as a subprocess
    """
        log_filename = f'{cwd}/{step.name}.log'
    else:
        log_filename = None

    step.log_filename = log_filename

                        log_filename=log_filename) as step_logger:

        step_args = ['polaris', 'serial', '--step_is_subprocess']
        check_call(step_args, step_logger)

import os
import sys
from typing import Dict, List

from polaris import Task, provenance
from polaris.io import symlink
from polaris.job import write_job_script
from polaris.machines import discover_machine

def setup_tasks(work_dir, task_list=None, numbers=None, config_file=None,
                machine=None, baseline_dir=None, component_path=None,
                suite_name='custom', cached=None, copy_executable=False,
                clean=False, model=None):
    """
    Set up one or more tasks

        A directory that will serve as the base for creating task
        directories

    task_list : list of str, optional
        Relative paths for a tasks to set up

    numbers : list of str, optional
        Task numbers to setup, as listed from ``polaris list``, optionally with
        a suffix ``c`` to indicate that all steps in that task should be
        cached

    config_file : {str, None}, optional
        Configuration file with custom options for setting up and running tasks

    machine : str, optional
        The name of one of the machines with defined config options, which can
        be listed with ``polaris list --machines``

    baseline_dir : str, optional
        Location of baselines that can be compared to

    component_path : str, optional
        The relative or absolute path to the location where the model and
        default namelists have been built

    suite_name : str, optional
        The name of the suite if tasks are being set up through a suite or
        ``'custom'`` if not

    cached : list of list of str, optional
        For each task in ``tasks``, which steps (if any) should be cached,
        or a list with "_all" as the first entry if all steps in the task
        should be cached

    copy_executable : bool, optional
        Whether to copy the model executable to the work directory

    clean : bool, optional
        Whether to delete the contents of the base work directory before
        setting up tasks

    model : str, optional
        The model to run

    Returns
    -------
    tasks : dict of polaris.Task
        A dictionary of tasks, with the relative path in the work
        directory as keys
    """
    machine = __get_machine_and_check_params(machine, config_file, task_list,
                                             numbers, cached)

    if work_dir is None:
        print('Warning: no base work directory was provided so setting up in '
              'the current directory.')
        work_dir = os.getcwd()
    work_dir = os.path.abspath(work_dir)

    all_tasks = dict()
            all_tasks[task.path] = task

    tasks: Dict[str, Task] = dict()
    cached_steps: Dict[str, List[str]] = dict()
    _add_tasks_by_number(numbers, all_tasks, tasks, cached_steps)
    _add_tasks_by_name(task_list, all_tasks, cached, tasks, cached_steps)

    # get the component of the first task.  We'll ensure that all tasks are
    # for this component
    first_path = next(iter(tasks))
    component = tasks[first_path].component

    basic_config = _get_basic_config(config_file, machine, component_path,
                                     component, model)

    provenance.write(work_dir, tasks, config=basic_config)

    _expand_and_mark_cached_steps(tasks, cached_steps)

    if clean:
        print('')
        print('Cleaning task and step work directories:')
        _clean_tasks_and_steps(tasks, work_dir)
        print('')

    _setup_configs(component, tasks, work_dir, config_file, machine,
                   component_path, copy_executable, model)

    print('Setting up tasks:')
    for path, task in tasks.items():
        setup_task(path, task, machine, work_dir,
                   baseline_dir, cached_steps=cached_steps[path], model=model)

    _check_dependencies(tasks)

    suite = {'name': suite_name,
             'tasks': tasks,
             'work_dir': work_dir}

    # pickle the task or step dictionary for use at runtime
    pickle_file = os.path.join(suite['work_dir'],
                               f'{suite_name}.pickle')
    with open(pickle_file, 'wb') as handle:
        pickle.dump(suite, handle, protocol=pickle.HIGHEST_PROTOCOL)

    _symlink_load_script(work_dir)

    max_cores, max_of_min_cores = _get_required_cores(tasks)

    print(f'target cores: {max_cores}')
    print(f'minimum cores: {max_of_min_cores}')

    if machine is not None:
        write_job_script(basic_config, machine, max_cores, max_of_min_cores,
                         work_dir, suite=suite_name)

    return tasks

def setup_task(path, task, machine, work_dir, baseline_dir, cached_steps,
               model):
    """

    path : str
        Relative path for a tasks to set up

    task : polaris.Task
        A task to set up

        A directory that will serve as the base for creating task directories

    cached_steps : list of str
        Which steps (if any) should be cached, identified by a list of
        subdirectories in the component

    """

    print(f'  {path}')

    task_dir = os.path.join(work_dir, path)
    try:
        os.makedirs(task_dir)
        pass
    task.work_dir = task_dir
    task.base_work_dir = work_dir

    # add the baseline directory for this task
    if baseline_dir is not None:
        task.baseline_dir = os.path.join(baseline_dir, path)

    if len(cached_steps) > 0:
        print_steps = ' '.join(cached_steps)
        print(f'    steps with cached outputs: {print_steps}')

    # iterate over steps
    for step in task.steps.values():
        _setup_step(task, step, work_dir, baseline_dir, task_dir)

    # wait until we've set up all the steps before pickling because steps may
    # need other steps to be set up
        if step.setup_complete:
            # this is a shared step that has already been set up
            continue

        # pickle the task and step for use at runtime
        pickle_filename = os.path.join(step.work_dir, 'step.pickle')

        _symlink_load_script(step.work_dir)

        if machine is not None:
            cores = step.cpus_per_task * step.ntasks
            min_cores = step.min_cpus_per_task * step.min_tasks
            write_job_script(step.config, machine, cores, min_cores,
                             step.work_dir)
        step.setup_complete = True

    # pickle the task and step for use at runtime
    pickle_filename = os.path.join(task.work_dir, 'task.pickle')
    with open(pickle_filename, 'wb') as handle:
        suite = {'name': 'task',
                 'tasks': {task.path: task},
                 'work_dir': task.work_dir}

    _symlink_load_script(task_dir)

        max_cores, max_of_min_cores = _get_required_cores({path: task})
        write_job_script(task.config, machine, max_cores, max_of_min_cores,
                         task_dir)

        description='Set up one or more tasks', prog='polaris setup')
    parser.add_argument("-t", "--tasks", nargs='+', dest="tasks",
                        help="Relative path for a task(s) to set up.",
                        metavar="PATH")
    parser.add_argument("-n", "--task_number", nargs='+', dest="task_num",
                        type=str,
                        help="Task number(s) to setup, as listed from "
                             "'polaris list'. Can be a space-separated "
                             "list of task numbers.  A suffix 'c' indicates "
                             "that all steps in the task should use cached "
                             "outputs.", metavar="NUM")
    parser.add_argument("-f", "--config_file", dest="config_file",
                        help="Configuration file for task setup.",
                        metavar="FILE")
    parser.add_argument("-m", "--machine", dest="machine",
                        help="The name of the machine for loading machine-"
                             "related config options.", metavar="MACH")
    parser.add_argument("-w", "--work_dir", dest="work_dir", required=True,
                        help="A base directory for setting up tasks.",
    parser.add_argument("-b", "--baseline_dir", dest="baseline_dir",
                        help="Location of baselines that can be compared to.",
    parser.add_argument("-p", "--component_path", dest="component_path",
                        help="The path where the component executable and "
                             "default namelists have been built.",
    parser.add_argument("--suite_name", dest="suite_name", default="custom",
                        help="The name to use for the 'custom' suite "
                             "containing all setup tasks.",
                        metavar="SUITE")
    parser.add_argument("--cached", dest="cached", nargs='+',
                        help="A list of steps in a single task supplied with "
                             "--tasks or --task_number that should use cached "
                             "outputs, or '_all' if all steps should be "
                             "cached.",
                        metavar="STEP")
    parser.add_argument("--copy_executable", dest="copy_executable",
                        help="If the model executable should be copied to the "
                             "work directory.")
    parser.add_argument("--clean", dest="clean", action="store_true",
                        help="If the base work directory should be deleted "
                             "before setting up the tasks.")
    parser.add_argument("--model", dest="model",
                        help="The model to run (one of 'mpas-ocean', 'omega', "
                             "or 'mpas-seaice')")

    cached = None
    if args.cached is not None:
        if args.tasks is not None and len(args.tasks) != 1:
            raise ValueError('You can only cache steps for one task at at '
                             'time.')
        if args.task_num is not None and len(args.task_num) != 1:
                             'time.')
        # cached is a list of lists
        cached = [args.cached]

    setup_tasks(task_list=args.tasks, numbers=args.task_num,
                config_file=args.config_file, machine=args.machine,
                work_dir=args.work_dir, baseline_dir=args.baseline_dir,
                component_path=args.component_path, suite_name=args.suite_name,
                cached=cached, copy_executable=args.copy_executable,
                clean=args.clean, model=args.model)

def _expand_and_mark_cached_steps(tasks, cached_steps):
    """
    Mark any steps that will be cached.  If any task asked for a step to
    be cached, it will be cached for all tasks that share the step.
    """
        cached_names = cached_steps[path]
        if len(cached_names) > 0 and cached_names[0] == '_all':
            cached_steps[path] = list(task.steps.keys())

        for step_name in cached_steps[path]:
            task.steps[step_name].cached = True

def _setup_configs(component, tasks, work_dir, config_file, machine,
                   component_path, copy_executable, model):
    """ Set up config parsers for this component """

    common_config = _get_basic_config(config_file, machine, component_path,
                                      component, model)
    if copy_executable:
        common_config.set('setup', 'copy_executable', 'True')

    if 'POLARIS_BRANCH' in os.environ:
        polaris_branch = os.environ['POLARIS_BRANCH']
        common_config.set('paths', 'polaris_branch', polaris_branch)
    else:
        common_config.set('paths', 'polaris_branch', os.getcwd())

    initial_configs = _add_task_configs(component, tasks, common_config)

    # okay, we're finally ready to configure all the tasks and add configs
    # to the "owned" steps
    configs = _configure_tasks_and_add_step_configs(tasks, component,
                                                    initial_configs,
                                                    common_config)

    _write_configs(common_config, configs, component.name, work_dir)

    _symlink_configs(tasks, component.name, work_dir)

def _add_task_configs(component, tasks, common_config):
    """
    Add config parsers for tasks and steps that don't already have shared ones
    """

    # get a list of shared steps and add config files for tasks to the
    # component
    configs = dict()
    for task in tasks.values():
        if task.config.filepath is None:
            task.config_filename = f'{task.name}.cfg'
            task.config.filepath = os.path.join(task.subdir,
                                                task.config_filename)
        component.add_config(task.config)
        configs[task.config.filepath] = task.config

    # now go through all the configs and prepend the common config options,
    # then run the setup() method for each in case there is some customization
    for config in configs.values():
        config.prepend(common_config)
        config.setup()

    return configs

def _configure_tasks_and_add_step_configs(tasks, component, initial_configs,
                                          common_config):
    """
    Call the configure() method for each task and add configs to "owned" steps
    """

    for config in initial_configs.values():
        for task in config.tasks:
            task.configure()
            config.set(section=f'{task.name}',
                       option='steps_to_run',
                       value=' '.join(task.steps_to_run),
                       comment=f'A list of steps to include when running the '
                               f'{task.name} task')

    # add configs to steps after calling task.configure() on all tasks in case
    # new steps were added
    new_configs = dict()
        for step in task.steps.values():
            if step.has_shared_config:
                configs[step.config.filepath] = step.config
                if step.config.filepath is None:
                    step.config_filename = f'{step.name}.cfg'
                    step.config.filepath = os.path.join(step.subdir,
                                                        step.config_filename)
                if step.config.filepath not in initial_configs:
                    new_configs[step.config.filepath] = step.config
                    component.add_config(step.config)
            else:
                step._set_config(task.config, link=task.config_filename)

    for config in new_configs.values():

def _write_configs(common_config, configs, component_name, work_dir):
    """ Write out all the config files """

    # add the common config at the component level
    common_config.filepath = f'{component_name}.cfg'
    configs[common_config.filepath] = common_config

    # finally, write out the config files
    component_work_dir = os.path.join(work_dir, component_name)
        config_filepath = os.path.join(component_work_dir, config.filepath)
        config_dir = os.path.dirname(config_filepath)
        try:
            os.makedirs(config_dir)
            pass
        with open(config_filepath, 'w') as f:
            config.write(f)

def _symlink_configs(tasks, component_name, work_dir):
    """ Symlink config files for requested tasks and steps """

    symlinks = dict()
        config = task.config
        link_path = os.path.join(component_work_dir, task.subdir,
                                 task.config_filename)
        if not os.path.exists(link_path) and link_path not in symlinks:
            symlinks[link_path] = config_filepath

            config = step.config
            config_filepath = os.path.join(component_work_dir, config.filepath)
            link_path = os.path.join(component_work_dir, step.subdir,
                                     step.config_filename)
            if not os.path.exists(link_path) and link_path not in symlinks:
                symlinks[link_path] = config_filepath

        for link_path, config_filepath in symlinks.items():
            link_dir = os.path.dirname(link_path)
            try:
                os.makedirs(link_dir)
                pass
            symlink(config_filepath, link_path)

def _clean_tasks_and_steps(tasks, base_work_dir):
    """
    Remove contents of task and step work directories to start fresh
    """
    print(f'{base_work_dir}:')
        task_work_dir = os.path.join(base_work_dir, path)
        try:
            shutil.rmtree(task_work_dir)
            print(f' {path}')
            pass

            step_work_dir = os.path.join(base_work_dir, step.path)
            try:
                shutil.rmtree(step_work_dir)
                print(f'  {step.path}')
            except FileNotFoundError:
                pass

def _get_required_cores(tasks):
    """ Get the maximum number of target cores and the max of min cores """

    max_cores = 0
    max_of_min_cores = 0
        for step_name in task.steps_to_run:
            step = task.steps[step_name]
            if step.ntasks is None:
                    f'The number of tasks (ntasks) was never set for '
                    f'{task.path} step {step_name}')
            if step.cpus_per_task is None:
                    f'The number of CPUs per task (cpus_per_task) was never '
                    f'set for {task.path} step {step_name}')
            max_cores = max(max_cores, cores)
            max_of_min_cores = max(max_of_min_cores, min_cores)

    return max_cores, max_of_min_cores

def __get_machine_and_check_params(machine, config_file, tasks, numbers,
                                   cached):
    if machine is None and 'POLARIS_MACHINE' in os.environ:
        machine = os.environ['POLARIS_MACHINE']

    if machine is None:
        machine = discover_machine()

    if config_file is None and machine is None:
        raise ValueError('At least one of config_file and machine is needed.')

    if config_file is not None and not os.path.exists(config_file):
        raise FileNotFoundError(
            f'The user config file wasn\'t found: {config_file}')

    if tasks is None and numbers is None:
        raise ValueError('At least one of tasks or numbers is needed.')

    if cached is not None:
        if tasks is None:
            warnings.warn('Ignoring "cached" argument because "tasks" was '
                          'not provided')
        elif len(cached) != len(tasks):
            raise ValueError('A list of cached steps must be provided for '
                             'each task in "tasks"')

    return machine

def _get_basic_config(config_file, machine, component_path, component, model):
    """
    Get a base config parser for the machine and component but not a specific
    task
    """

    if config_file is not None:
        config.add_user_config(config_file)

    # set the model from the command line if provided
    if model is not None:
        config.set(component.name, 'model', model, user=True)

    # start with default polaris config options
    config.add_from_package('polaris', 'default.cfg')

    # add the E3SM config options from mache
        config.add_from_package('mache.machines', f'{machine}.cfg',
                                exception=False)

    # add the polaris machine config file
        machine = 'default'
    config.add_from_package('polaris.machines', f'{machine}.cfg')

        config.set('paths', 'polaris_branch', polaris_branch)
    else:
        config.set('paths', 'polaris_branch', os.getcwd())

    # add the config options for the component
    config.add_from_package(f'polaris.{component.name}',
                            f'{component.name}.cfg')

    component.configure(config)

    # set the component_path path from the command line if provided
    if component_path is not None:
        component_path = os.path.abspath(component_path)
        config.set('paths', 'component_path', component_path, user=True)

def _add_tasks_by_number(numbers, all_tasks, tasks, cached_steps):
    if numbers is not None:
        keys = list(all_tasks)
        for number in numbers:
            cache_all = False
            if number.endswith('c'):
                cache_all = True
                number = int(number[:-1])
            else:
                number = int(number)

            if number >= len(keys):
                raise ValueError(f'task number {number} is out of range.  '
                                 f'There are only {len(keys)} tasks.')
            path = keys[number]
            if cache_all:
                cached_steps[path] = ['_all']
            else:
                cached_steps[path] = list()
            tasks[path] = all_tasks[path]

def _add_tasks_by_name(task_list, all_tasks, cached, tasks, cached_steps):
    if task_list is not None:
        for index, path in enumerate(task_list):
            if path not in all_tasks:
                raise ValueError(f'Task with path {path} is not in '
                                 f'tasks')
            if cached is not None:
                cached_steps[path] = cached[index]
            else:

def _setup_step(task, step, work_dir, baseline_dir, task_dir):
    """ Set up a step in a task """

    # make the step directory if it doesn't exist
    step_dir = os.path.join(work_dir, step.path)

    if step.name in task.step_symlinks:
        symlink(step_dir,
                os.path.join(task_dir, task.step_symlinks[step.name]))

    if step.setup_complete:
        # this is a shared step that has already been set up
        return
    try:
        os.makedirs(step_dir)
        pass

    step.work_dir = step_dir
    step.base_work_dir = work_dir

    # set up the step
    step.setup()

    # add the baseline directory for this step
        step.baseline_dir = os.path.join(baseline_dir, step.path)

    # process input, output, namelist and streams files
    step.process_inputs_and_outputs()

def _symlink_load_script(work_dir):
    """ make a symlink to the script for loading the polaris conda env. """
    if 'LOAD_POLARIS_ENV' in os.environ:
        script_filename = os.environ['LOAD_POLARIS_ENV']
        symlink(script_filename,
                os.path.join(work_dir, 'load_polaris_env.sh'))

def _check_dependencies(tasks):
            for name, dependency in step.dependencies.items():
                if dependency.work_dir == '':
                    raise ValueError(f'The dependency {name} of '
                                     f'{task.path} step {step.name} was '
                                     f'not set up.')

from copy import deepcopy

def read(package, streams_filename, tree=None, replacements=None):
    """
    Parse the given streams file

        The package name or module object that contains the streams file

    streams_filename : str
        The name of the streams file to read from

    tree : lxml.etree, optional
        An existing set of streams to add to or modify

    replacements : dict, optional
        A dictionary of replacements, in which case ``streams_filename`` is
        assumed to be a Jinja2 template to be rendered with these replacements

    Returns
    -------
    tree : lxml.etree
        A tree of XML data describing MPAS i/o streams with the content from
        the given streams file
    """
    text = imp_res.files(package).joinpath(streams_filename).read_text()
    if replacements is not None:
        template = Template(text)
        text = template.render(**replacements)

    new_tree = etree.fromstring(text)

    tree = update_tree(tree, new_tree)

    return tree

def write(streams, out_filename):
    """ write the streams XML data to the file """

    with open(out_filename, 'w') as stream_file:

        stream_file.write('<streams>\n')

        # Write out all immutable streams first
        for stream in streams.findall('immutable_stream'):
            stream_name = stream.attrib['name']

            stream_file.write('\n')
            stream_file.write(f'<immutable_stream name="{stream_name}"')
            # Process all attributes on the stream
            for attr, val in stream.attrib.items():
                if attr.strip() != 'name':
                    stream_file.write(f'\n                  {attr}="{val}"')

            stream_file.write('/>\n')

        # Write out all immutable streams
        for stream in streams.findall('stream'):

            stream_file.write(f'<stream name="{stream_name}"')

            # Process all attributes
                    stream_file.write(f'\n        {attr}="{val}"')

            stream_file.write('>\n\n')

            # Write out all contents of the stream
            for tag in ['stream', 'var_struct', 'var_array', 'var']:
                for child in stream.findall(tag):
                    child_name = child.attrib['name']
                    if tag == 'stream' and child_name == stream_name:
                        # don't include the stream itself
                        continue
                    if 'packages' in child.attrib.keys():
                        package_name = child.attrib['packages']
                        entry = f'    <{tag} name="{child_name}" ' \
                                f'packages="{package_name}"/>\n'
                    else:
                        entry = f'    <{tag} name="{child_name}"/>\n'
                    stream_file.write(entry)

            stream_file.write('</stream>\n')

        stream_file.write('\n')
        stream_file.write('</streams>\n')

def update_defaults(new_child, defaults):
    """
    Update a stream or its children (sub-stream, var, etc.) starting from the
    defaults or add it if it's new.
    """
    if 'name' not in new_child.attrib:
        return

    name = new_child.attrib['name']
    found = False
    for child in defaults:
        if child.attrib['name'] == name:
            found = True
            if child.tag != new_child.tag:
                raise ValueError(f'Trying to update stream "{name}" with '
                                 f'inconsistent tags {child.tag} vs. '
                                 f'{new_child.tag}.')

            # copy the attributes
            for attr, value in new_child.attrib.items():
                child.attrib[attr] = value

            if len(new_child) > 0:
                # we don't want default grandchildren
                for grandchild in child:
                    child.remove(grandchild)

            # copy or add the grandchildren's contents
            for new_grandchild in new_child:
                update_defaults(new_grandchild, child)

    if not found:
        # add a deep copy of the element
        defaults.append(deepcopy(new_child))

def update_tree(tree, new_tree):
    """

    new_tree : lxml.etree
        A new set of streams to add or modify

    Returns
    -------
    """

    if tree is None:
        tree = new_tree
    else:
        new_streams = next(new_tree.iter('streams'))

        for new_stream in new_streams:
            _update_element(new_stream, streams)

def _update_element(new_child, elements):
    """
    add the new child/grandchildren or add/update attributes if they exist
    """
        return

    for child in elements:

                _update_element(new_grandchild, child)

        elements.append(deepcopy(new_child))

import configparser

import cmocean  # noqa: F401
import matplotlib.colors as cols
import mosaic
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from pyremap.descriptor.utility import interp_extrap_corner

from polaris.viz.style import use_mplstyle

def plot_global_mpas_field(mesh_filename, da, out_filename, config,
                           colormap_section, title=None,
                           plot_land=True, colorbar_label='',
                           central_longitude=0., figsize=(8, 4.5), dpi=200,
                           patch_edge_color=None, antialiased=False):
    """
    Plots a data set as a longitude-latitude map

        A filename containing the MPAS mesh

    da : xarray.DataArray
        The horizontal field to plot

    out_filename : str
        The image file name to be written

        The config options to use for colormap settings

    colormap_section : str
        The name of a section in the config options. Options must include:

        colormap_name
            The name of the colormap

        norm_type
            The norm: {'linear', 'log'}

    title : str, optional
        The subtitle of the plot

    plot_land : bool
        Whether to plot continents over the data

    colorbar_label : str, optional
        Label on the colorbar

    central_longitude : float, optional
        The longitude of the center of the plot

    figsize : tuple, optional
        The size of the figure in inches

    dpi : int, optional
        Dots per inch for the output plot

    patch_edge_color : str, optional
        The color of patch edges (if not the same as the face)

    antialiased : bool, optional
        Whether each patch should be drawn with antialiasing. Default is False.
        Using antialiasing is desirable for plots with a limited number of
        patches, but it can cause artifact patterns for high resolution meshes
        so use with caution.
    """

    use_mplstyle()

    transform = cartopy.crs.Geodetic()
    projection = cartopy.crs.PlateCarree(central_longitude=central_longitude)

    mesh_ds = xr.open_dataset(mesh_filename)
    descriptor = mosaic.Descriptor(
        mesh_ds, projection=projection, transform=transform, use_latlon=True
    )

    fig, ax = plt.subplots(figsize=figsize,
                           constrained_layout=True,
                           subplot_kw=dict(projection=projection))

    if title is not None:
        fig.suptitle(title, y=0.935)

    colormap, norm, ticks = _setup_colormap(config, colormap_section)

    pcolor_kwargs = dict(cmap=colormap, norm=norm, zorder=1, edgecolors='face',
                         linewidths=0.2, antialiased=antialiased)

    if patch_edge_color is not None:
        pcolor_kwargs['edgecolors'] = patch_edge_color

    gl = ax.gridlines(
        color='gray', linestyle=':', zorder=5, draw_labels=True, linewidth=0.5
    )
    gl.right_labels = False
    gl.top_labels = False

    pc = mosaic.polypcolor(ax, descriptor, da, **pcolor_kwargs)

    if plot_land:
        _add_land_lakes_coastline(ax)

    cbar = fig.colorbar(
        pc, ax=ax, label=colorbar_label, extend='both', shrink=0.6
    )

    if ticks is not None:
        cbar.set_ticks(ticks)
        cbar.set_ticklabels([f'{tick}' for tick in ticks])

    fig.savefig(out_filename, dpi=dpi, bbox_inches='tight', pad_inches=0.1)

def plot_global_lat_lon_field(lon, lat, data_array, out_filename, config,
                              colormap_section, title=None, plot_land=True,
                              colorbar_label=None):
    """

    lon : numpy.ndarray
        1D longitude coordinate

    lat : numpy.ndarray
        1D latitude coordinate

    data_array : numpy.ndarray
        2D data array to plot

        norm_type
            The norm: {'symlog', 'log', 'linear'}

        norm_args
            A dict of arguments to pass to the norm

        It may also include:

        colorbar_ticks
            An array of values where ticks should be placed

    """

    nlat, nlon = data_array.shape
    if lon.shape[0] == nlon:
        lon_corner = interp_extrap_corner(lon)
    elif lon.shape[0] == nlon + 1:
        lon_corner = lon
    else:
        raise ValueError(f'Unexpected length of lon {lon.shape[0]}. Should '
                         f'be either {nlon} or {nlon + 1}')

    if lat.shape[0] == nlat:
        lat_corner = interp_extrap_corner(lat)
    elif lat.shape[0] == nlat + 1:
        lat_corner = lat
    else:
        raise ValueError(f'Unexpected length of lat {lat.shape[0]}. Should '
                         f'be either {nlat} or {nlat + 1}')

    figsize = (8, 4.5)
    dpi = 200
    fig = plt.figure(figsize=figsize, dpi=dpi)

    subplots = [111]
    ref_projection = cartopy.crs.PlateCarree()
    central_longitude = 0.5 * (lon_corner[0] + lon_corner[-1])

    extent = [lon_corner[0], lon_corner[-1], lat_corner[0], lat_corner[-1]]

    ax = plt.subplot(subplots[0], projection=projection)

    ax.set_extent(extent, crs=ref_projection)

    gl = ax.gridlines(crs=ref_projection, color='gray', linestyle=':',
                      zorder=5, draw_labels=True, linewidth=0.5)

    plotHandle = ax.pcolormesh(lon_corner, lat_corner, data_array,
                               cmap=colormap, norm=norm,
                               transform=ref_projection, zorder=1)

    cax = inset_axes(ax, width='3%', height='60%', loc='center right',
                     bbox_to_anchor=(0.08, 0., 1, 1),
                     bbox_transform=ax.transAxes, borderpad=0)

    cbar = plt.colorbar(plotHandle, cax=cax, extend='both')
    cbar.set_label(colorbar_label)

    plt.savefig(out_filename, dpi='figure', bbox_inches='tight',
                pad_inches=0.2)

    plt.close()

def _setup_colormap(config, colormap_section):
    """
    Set up a colormap from the registry

        Configuration options for the test case, including a section for
        the colormap

        norm_type

        norm_args

    Returns
    -------
    colormap : str
        the name of the new colormap

    norm : matplotlib.colors.Normalize
        a matplotlib norm object used to normalize the colormap

    ticks : list of float
        is an array of values where ticks should be placed
    """

    colormap = plt.get_cmap(config.get(colormap_section, 'colormap_name'))

    norm_type = config.get(colormap_section, 'norm_type')

    kwargs = config.getexpression(colormap_section, 'norm_args')

    if norm_type == 'symlog':
        norm = cols.SymLogNorm(**kwargs)
    elif norm_type == 'log':
        norm = cols.LogNorm(**kwargs)
    elif norm_type == 'linear':
        norm = cols.Normalize(**kwargs)
    else:
        raise ValueError(f'Unsupported norm type {norm_type} in section '
                         f'{colormap_section}')

    try:
        ticks = config.getexpression(colormap_section, 'colorbar_ticks',
                                     use_numpyfunc=True)
    except configparser.NoOptionError:
        ticks = None

    return colormap, norm, ticks

def _add_land_lakes_coastline(ax, ice_shelves=True):
    land_color = cartopy.feature.COLORS['land']
    water_color = cartopy.feature.COLORS['water']

    land_50m = cartopy.feature.NaturalEarthFeature(
        'physical', 'land', '50m', edgecolor='k',
        facecolor=land_color, linewidth=0.5)
    lakes_50m = cartopy.feature.NaturalEarthFeature(
        'physical', 'lakes', '50m', edgecolor='k',
        facecolor=water_color, linewidth=0.5)
    ax.add_feature(land_50m, zorder=2)
    if ice_shelves:
        ice_50m = cartopy.feature.NaturalEarthFeature(
            'physical', 'antarctic_ice_shelves_polys', '50m',
            edgecolor='k', facecolor=land_color, linewidth=0.5)
        ax.add_feature(ice_50m, zorder=3)
    ax.add_feature(lakes_50m, zorder=4)

from ruamel.yaml import YAML

class PolarisYaml:
    """
    A class for reading writing and combining config files in yaml format
    (e.g. as used in Omega).

    configs : dict
        Nested dictionaries containing config options

    streams : dict
        Nested dictionaries containing data about streams

        The name of the streams section

    model : str
        The name of the E3SM component
    """

    def __init__(self):
        """
        Create a yaml config object
        """
        self.streams = dict()
        self.model = None

    @classmethod
    def read(cls, filename, package=None, replacements=None, model=None,
             streams_section='streams'):
        """
        Add config options from a yaml file

            A template yaml config file

        package : str, optional
            The name of a package the filename is found in

        replacements : dict, optional
            A dictionary of replacements, which, if provided, is used to
            replace Jinja variables and the yaml file is assumed do be a Jinja
            template

        model : str, optional
            The name of the model to parse if the yaml file might have multiple
            models

        streams_section : str, optional
            The name of the streams section

        Returns
        -------
        yaml : polaris.yaml.PolarisYaml
            A yaml object read in from the given file (and optionally package)
        """
        # read the text from a file (possibly in a package)
            text = imp_res.files(package).joinpath(filename).read_text()
        else:
            with open(filename, 'r') as infile:
                text = infile.read()

        # if this is a jinja template, render the template with the
        # replacements
        if replacements is not None:
            template = Template(text)
            text = template.render(**replacements)

        yaml = cls()
        yaml.streams_section = streams_section
        yaml_data = YAML(typ='rt')
        configs = yaml_data.load(text)

        if model is None:
            keys = list(configs)
            if len(keys) > 1:
                    f'Config yaml file contains unexpected sections: \n '
                    f'{keys[1:]}')
            model = keys[0]

        yaml.model = model
        yaml.streams = {}
        if model in configs:
            configs = configs[model]
            if streams_section in configs:
                yaml.streams = configs[streams_section]
                configs = dict(configs)
                configs.pop(streams_section)
        else:
            configs = {}

        yaml.configs = configs
        return yaml

    def update(self, configs=None, options=None, quiet=True):
        """
        Add config options from a dictionary

        configs : dict, optional
            A nested dictionary of config sections, options and values

        options : dict, optional
            A flat dictionary of options and values

        quiet : bool, optional
            Whether or not to print the updated config options as they are
            replaced
        """
        if configs is not None:
            if self.model in configs:
                # we want one layer deeper
                configs = configs[self.model]
            _update_section(configs, self.configs, quiet)

        if options is not None:
            _update_options(options, self.configs, quiet)

    def write(self, filename):
        """
        Write config options to a yaml file

            A yaml config file
        """
        yaml = YAML(typ='rt')
        configs = dict(self.configs)
        if self.streams:
            configs[self.streams_section] = self.streams

        model_configs = dict()
        model_configs[self.model] = configs

        with open(filename, 'w') as outfile:
            yaml.dump(model_configs, outfile)

def mpas_namelist_and_streams_to_yaml(model, namelist_template=None,
                                      namelist=None,
                                      streams=None):
    """
    Add config options from a yaml file

    model : str
        The name of the model

    namelist_template : str
        An MPAS namelist template file

    namelist : str, optional
        An MPAS namelist file

    streams : str, optional
        An MPAS streams file

    Returns
    -------
    yaml : polaris.yaml.PolarisYaml
        A yaml object with the namelists and streams
    """
    yaml = PolarisYaml()
    yaml.model = model
    if namelist is not None:
        yaml.configs = _read_namelist(namelist_template, namelist)
    if streams is not None:
        yaml.streams = _streams_xml_to_dict(streams)

    return yaml

def main_mpas_to_yaml():
        description='Convert a namelist and/or streams file to yaml')
    parser.add_argument("-n", "--namelist", dest="namelist",
                        required=False,
                        help="MPAS namelist file")
    parser.add_argument("-s", "--streams", dest="streams",
                        help="MPAS streams file")
    parser.add_argument("-t", "--namelist_template", dest="namelist_template",
                        help="MPAS namelist template file (with all namelist "
                             "options). For MPAS-Ocean, this will typically be"
                             " ${PATH_TO_MPASO}/default_inputs/"
                             "namelist.ocean.forward")
    parser.add_argument("-y", "--yaml", dest="yaml",
                        required=True,
                        help="Output yaml file")
    parser.add_argument("-m", "--model", dest="model", default='mpas-ocean',
                        help="Model name for the yaml")

    args = parser.parse_args()

    yaml = mpas_namelist_and_streams_to_yaml(
        model=args.model,
        namelist_template=args.namelist_template,
        namelist=args.namelist,
        streams=args.streams)

    yaml.write(args.yaml)

def yaml_to_mpas_streams(processed_registry_filename, yaml):
    """

    processed_registry_filename : str
        The processed registry file, used to determine the types of variables
        each steam (since the yaml format doesn't supply that information).

    Returns
    -------
        the streams in the yaml file
    """
    with open(processed_registry_filename, 'r') as reg_file:
        registry_string = reg_file.read()
    registry_string = registry_string.lstrip('\n')
    registry = etree.fromstring(registry_string)

    root = etree.Element('streams')
    for stream in yaml.streams:
        # find out if stream or immutable_stream
        tag = _get_stream_tag(registry, stream)
        attrs = dict(yaml.streams[stream])
        contents = None
        if 'contents' in attrs:
            contents = attrs.pop('contents')
        attrs['name'] = stream
        child = etree.SubElement(root, tag, attrib=attrs)
        if contents is not None:
            for var in contents:
                # find out what type it has
                tag = _get_var_tag(registry, var)
                etree.SubElement(child, tag, attrib=dict(name=var))

    tree = etree.ElementTree(element=root)

def _update_section(src, dst, quiet, print_section=None):
    """
    Recursively update config options in a section of a config from a source
    section to the associate destination dictionary
    """
    for name in src:
        if isinstance(src[name], (dict, OrderedDict)):
            if print_section is not None:
                print_subsection = f'{print_section}: {name}'
            else:
                print_subsection = name
            if name not in dst:
                    f'Attempting to modify config options to a '
                    f'nonexistent config\n'
                    f'(sub)section: {print_subsection}')
            # this is a subsection
            src_sub = src[name]
            dst_sub = dst[name]
            _update_section(src_sub, dst_sub, quiet, print_subsection)
        else:
                    f'Attempting to modify a nonexistent config '
                    f'options: {print_section}: {name}')
            if not quiet:
                print(f'  {print_section}: {name} = {src[name]}')
            dst[name] = src[name]

def _update_options(src, dst, quiet):
    """
    Update config options by searching in the destination nested dictionary
    """
        success = _update_option(name, src[name], dst, quiet)
        if not success:
                f'Attempting to modify a nonexistent config '
                f'options: {name}')

def _update_option(option, value, dst, quiet, print_section=None):
    """
    Recursively attempt to find and replace the value of the
    given option
    """
    for name in dst:
        if isinstance(dst[name], (dict, OrderedDict)):
            else:
            success = _update_option(option, value, dst[name], quiet,
                                     print_subsection)
            if success:
                return True
        elif name == option:
            dst[name] = value
                print(f'  {print_section}: {name} = {value}')
            return True
    return False

def _read_namelist(namelist_template, namelist_filename):
    record_map = _read_namelist_template(namelist_template)

    with open(namelist_filename, 'r') as f:

    namelist: Dict[str, Dict[str, int | float | bool | str]] = dict()
        _, opt, value = _read_namelist_line(line)
        if opt is not None and value is not None:
            record = record_map[opt]
            if record not in namelist:
                namelist[record] = dict()
            namelist[record][opt] = value

def _read_namelist_template(namelist_template):
    with open(namelist_template, 'r') as f:

    record_map: Dict[str, str] = dict()
        new_record, opt, _ = _read_namelist_line(line)
        if new_record is not None:
            record = new_record
        elif opt is not None and record is not None:
            record_map[opt] = record

    return record_map

def _read_namelist_line(line):
    opt = None
    value: int | float | bool | str | None = None
    if '&' in line:
        record = line.strip('&').strip('\n').strip()
    elif '=' in line:
        opt, val = line.strip('\n').split('=')
        opt = opt.strip()
        str_value = \
            val.strip().strip('\"').strip('\'').strip()
        try:
            value = int(str_value)
        except ValueError:
            try:
                value = float(str_value)
            except ValueError:
                if str_value in ['true', '.true.']:
                    value = True
                elif str_value in ['false', '.false.']:
                    value = False
        if value is None:
            value = str_value

    return record, opt, value

def _streams_xml_to_dict(streams_filename):
    """ Convert a streams XML file to nested dictionaries """
    streams: Dict[str, Dict[str, str | list]] = dict()
    tree = etree.parse(streams_filename)
    xml_streams = next(tree.iter('streams'))
    for child in xml_streams:
        if child.tag not in ['stream', 'immutable_stream']:
            raise ValueError(f'Unexpected tag {child.tag} instead of stream or'
                             f'immutable stream')
        stream_name = child.attrib['name']
        streams[stream_name] = dict()
        for attr, value in child.attrib.items():
            if attr != 'name':
                streams[stream_name][attr] = value

        contents = list()
        for grandchild in child:
            if grandchild.tag == 'var':
                contents.append(grandchild.attrib['name'])
            elif grandchild.tag == 'var_struct':
            elif grandchild.tag == 'var_array':
            elif grandchild.tag == 'stream':
            else:
                raise ValueError(f'Unexpected tag {grandchild.tag}')
        if len(contents) > 0:
            streams[stream_name]['contents'] = contents

    return streams

def _get_stream_tag(registry, stream):
    """ Get the xml tag, 'stream' or 'immutable_stream' for a given stream """
    streams = next(next(registry.iter('registry')).iter('streams'))
    # if we don't find the stream, it can't be an immutable stream
    tag = 'stream'
    for child in streams:
        if child.tag == 'stream' and child.attrib['name'] == stream:
            if 'immutable' in child.attrib and \
                    child.attrib['immutable'] == 'true':
                tag = 'immutable_stream'
            break
    return tag

def _get_var_tag(registry, variable):
    """
    Get the xml tag -- 'stream', 'var_struct', 'var_array' or 'var' -- for a
    variable
    """
    tag = None
        if child.tag == 'stream' and child.attrib['name'] == variable:
            return 'stream'

    tree = next(registry.iter('registry'))
    for child in tree:
        if child.tag == 'var_struct':
            if child.attrib['name'] == variable:
                return 'var_struct'

            for grandchild in child:
                if grandchild.tag in ['var_struct', 'var_array', 'var'] and \
                        grandchild.attrib['name'] == variable:
                    return grandchild.tag
                if grandchild.tag in ['var_struct', 'var_array']:
                    for greatgrand in grandchild:
                        if greatgrand.tag in ['var_array', 'var'] and \
                                greatgrand.attrib['name'] == variable:
                            return greatgrand.tag

    if tag is None:
        raise ValueError(f'Could not find {variable} in preprocessed registry')
    return tag

import os

import pandas as pd

from polaris.mpas import area_for_field, time_since_start
from polaris.ocean.convergence import (
    get_resolution_for_task,
    get_timestep_for_task,
)
from polaris.ocean.model import OceanIOStep
from polaris.viz import use_mplstyle

class ConvergenceAnalysis(OceanIOStep):
    """
    A step for analyzing the output from convergence tests

    dependencies_dict : dict of dict of polaris.Steps
        The dependencies of this step must be given as separate keys in the
        dict:

            mesh : dict of polaris.Steps
                Keys of the dict correspond to `resolutions`
                Values of the dict are polaris.Steps, which must have the
                attribute `path`, the path to `base_mesh.nc` of that
                resolution
            init : dict of polaris.Steps
                attribute `path`, the path to `initial_state.nc` of that
            forward : dict of polaris.Steps
                attribute `path`, the path to `forward.nc` of that

    convergence_vars: list of dict
        The attributes for each variable for which to analyze the convergence
        rate. Each dict must contain the following keys:

            name : str
                The name of the variable as given in the output netcdf file

            title : str
                The name of the variable to use in the plot title

            zidx : int
                The z-index to use for variables that have an nVertLevels
                dimension, which should be None for variables that don't

    refinement : str
        Refinement type. One of 'space', 'time' or 'both' indicating both
        space and time
    """
    def __init__(self, component, subdir, dependencies, convergence_vars,
                 refinement='both'):
        """
        Create the step

        subdir : str
            The subdirectory that the step resides in

        dependencies : dict of dict of polaris.Steps
            The dependencies of this step must be given as separate keys in the
            dict:

                mesh : dict of polaris.Steps
                    Keys of the dict correspond to `refinement_factors`
                    Values of the dict are polaris.Steps, which must have the
                    attribute `path`, the path to `base_mesh.nc` of that
                    resolution
                init : dict of polaris.Steps
                    attribute `path`, the path to `initial_state.nc` of that
                forward : dict of polaris.Steps
                    attribute `path`, the path to `forward.nc` of that

        convergence_vars: list of dict
            The convergence attributes for each variable. Each dict must
            contain the following keys:

                name : str
                    The name of the variable as given in the output netcdf file

                title : str
                    The name of the variable to use in the plot title

                zidx : int
                    The z-index to use for variables that have an nVertLevels
                    dimension, which should be None for variables that don't

        refinement : str, optional
            Refinement type. One of 'space', 'time' or 'both' indicating both
            space and time
        """
        super().__init__(component=component, name='analysis', subdir=subdir)
        self.dependencies_dict = dependencies
        self.convergence_vars = convergence_vars
        self.refinement = refinement

        for var in convergence_vars:
            self.add_output_file(f'convergence_{var["name"]}.png')

        """
        Add input files based on resolutions, which may have been changed by
        user config options
        """
        dependencies = self.dependencies_dict
        if self.refinement == 'time':
            option = 'refinement_factors_time'
        else:
            option = 'refinement_factors_space'
        refinement_factors = config.getlist('convergence', option,
                                            dtype=float)
        for refinement_factor in refinement_factors:
            base_mesh = dependencies['mesh'][refinement_factor]
            init = dependencies['init'][refinement_factor]
            forward = dependencies['forward'][refinement_factor]
            self.add_input_file(
                filename=f'mesh_r{refinement_factor:02g}.nc',
                work_dir_target=f'{base_mesh.path}/base_mesh.nc')
                filename=f'init_r{refinement_factor:02g}.nc',
                work_dir_target=f'{init.path}/initial_state.nc')
                filename=f'output_r{refinement_factor:02g}.nc',
                work_dir_target=f'{forward.path}/output.nc')

        """
        Run this step of the test case
        """
        plt.switch_backend('Agg')
        convergence_vars = self.convergence_vars
            self.plot_convergence(
                variable_name=var["name"],
                title=var["title"],
                zidx=var["zidx"])

    def plot_convergence(self, variable_name, title, zidx):
        """
        Compute the error norm for each resolution and produce a convergence
        plot

        variable_name : str
            The name of the variable as given in the output netcdf file

        title : str
            The name of the variable to use in the plot title

        zidx : int
            The z-index to use for variables that have an nVertLevels
            dimension, which should be None for variables that don't
        """
        conv_thresh, error_type = self.convergence_parameters(
            field_name=variable_name)

        error = []
            header = 'time step'
        else:
            header = 'resolution'
        resolutions = list()
        timesteps = list()
            timestep, _ = get_timestep_for_task(
                config, refinement_factor, refinement=self.refinement)
            timesteps.append(timestep)
            resolution = get_resolution_for_task(
            resolutions.append(resolution)
            error_res = self.compute_error(
                refinement_factor=refinement_factor,
                variable_name=variable_name,
                zidx=zidx,
                error_type=error_type)
            error.append(error_res)

            refinement_array = np.array(timesteps)
            x_label = 'Time (s)'
        else:
            refinement_array = np.array(resolutions)
            x_label = 'Horizontal resolution (km)'
        error_array = np.array(error)
        filename = f'convergence_{variable_name}.csv'
        data = np.stack((refinement_array, error_array), axis=1)
        df = pd.DataFrame(data, columns=[header, error_type])
        df.to_csv(f'convergence_{variable_name}.csv', index=False)

        convergence_failed = False
        poly = np.polyfit(np.log10(refinement_array), np.log10(error_array), 1)
        convergence = poly[0]
        conv_round = np.round(convergence, 3)

        fit = refinement_array ** poly[0] * 10 ** poly[1]

        order1 = 0.5 * error_array[-1] * \
            (refinement_array / refinement_array[-1])
        order2 = 0.5 * error_array[-1] * \
            (refinement_array / refinement_array[-1])**2

        use_mplstyle()
        fig = plt.figure()

        error_dict = {'l2': 'L2 norm', 'inf': 'L-infinity norm'}
        error_title = error_dict[error_type]

        ax = fig.add_subplot(111)
        ax.loglog(resolutions, order1, '--k', label='first order',
                  alpha=0.3)
        ax.loglog(resolutions, order2, 'k', label='second order',
        ax.loglog(refinement_array, fit, 'k',
                  label=f'linear fit (order={conv_round})')
        ax.loglog(refinement_array, error_array, 'o', label='numerical')

            baseline_filename = os.path.join(self.baseline_dir, filename)
            if os.path.exists(baseline_filename):
                data = pd.read_csv(baseline_filename)
                if error_type not in data.keys():
                    raise ValueError(
                        f'{error_type} not available for baseline')
                else:
                    refinement_array = data.resolution.to_numpy()
                    error_array = data[error_type].to_numpy()
                    poly = np.polyfit(
                        np.log10(refinement_array), np.log10(error_array), 1)
                    base_convergence = poly[0]
                    conv_round = np.round(base_convergence, 3)

                    fit = refinement_array ** poly[0] * 10 ** poly[1]
                    ax.loglog(refinement_array, error_array, 'o',
                              color='#ff7f0e', label='baseline')
                    ax.loglog(refinement_array, fit, color='#ff7f0e',
                              label=f'linear fit, baseline '
                                    f'(order={conv_round})')
        ax.set_xlabel(x_label)
        ax.set_ylabel(f'{error_title}')
        ax.set_title(f'Error Convergence of {title}')
        ax.legend(loc='lower left')
        ax.invert_xaxis()
        fig.savefig(f'convergence_{variable_name}.png',
                    bbox_inches='tight', pad_inches=0.1)

        logger.info(f'Order of convergence for {title}: '
                    f'{conv_round}')

        if conv_round < conv_thresh:
            logger.error(f'Error: order of convergence for {title}\n'
                         f'  {conv_round} < min tolerance '
                         f'{conv_thresh}')
            convergence_failed = True

        if convergence_failed:
            raise ValueError('Convergence rate below minimum tolerance.')

    def compute_error(self, refinement_factor, variable_name, zidx=None,
                      error_type='l2'):
        """
        Compute the error for a given resolution

        refinement_factor : float
            The factor by which step is refined in space, time or both

            The variable name of which to evaluate error with respect to the
            exact solution

        zidx : int, optional
            The z index to use to slice the field given by variable name

        error_type: {'l2', 'inf'}, optional
            The type of error to compute

        Returns
        -------
        error : float
            The error of the variable given by variable_name
        """
        norm_type = {'l2': None, 'inf': np.inf}
        ds_mesh = self.open_model_dataset(f'mesh_r{refinement_factor:02g}.nc')
        section = config['convergence']
        eval_time = section.getfloat('convergence_eval_time')
        s_per_hour = 3600.0

        field_exact = self.exact_solution(refinement_factor, variable_name,
                                          time=eval_time * s_per_hour,
                                          zidx=zidx)
        field_mpas = self.get_output_field(refinement_factor, variable_name,
                                           time=eval_time * s_per_hour,
                                           zidx=zidx)
        diff = field_exact - field_mpas

        # Only the L2 norm is area-weighted
        if error_type == 'l2':
            area = area_for_field(ds_mesh, diff)
            field_exact = field_exact * area
            diff = diff * area
        error = np.linalg.norm(diff, ord=norm_type[error_type])

        # Normalize the error norm by the vector norm of the exact solution
        den = np.linalg.norm(field_exact, ord=norm_type[error_type])
        error = np.divide(error, den)

        return error

    def exact_solution(self, refinement_factor, field_name, time, zidx=None):
        """
        Get the exact solution

        field_name : str
            The name of the variable of which we evaluate convergence
            For the default method, we use the same convergence rate for all
            fields

        time : float
            The time at which to evaluate the exact solution in seconds.
            For the default method, we always use the initial state.

            The z-index for the vertical level at which to evaluate the exact
            solution

        Returns
        -------
        solution : xarray.DataArray
            The exact solution as derived from the initial condition
        """

        ds_init = self.open_model_dataset(f'init_r{refinement_factor:02g}.nc')
        ds_init = ds_init.isel(Time=0)
        if zidx is not None:
            ds_init = ds_init.isel(nVertLevels=zidx)

        return ds_init[field_name]

    def get_output_field(self, refinement_factor, field_name, time, zidx=None):
        """
        Get the model output field at the given time and z index

        e str
            The mesh name which is the prefix for the output file

            The time at which to evaluate the exact solution in seconds

            The z-index for the vertical level to take the field from

        Returns
        -------
        field_mpas : xarray.DataArray
            model output field
        """
        ds_out = self.open_model_dataset(f'output_r{refinement_factor:02g}.nc',
                                         decode_times=False)

        model = config.get('ocean', 'model')
        if model == 'mpas-o':
            dt = time_since_start(ds_out.xtime.values)
        else:
            # time is seconds since the start of the simulation in Omega
            dt = ds_out.Time.values

        tidx = np.argmin(np.abs(dt - time))

        ds_out = ds_out.isel(Time=tidx)
        field_mpas = ds_out[field_name]
            field_mpas = field_mpas.isel(nVertLevels=zidx)
        return field_mpas

    def convergence_parameters(self, field_name=None):
        """
        Get convergence parameters

            fields

        Returns
        -------
        conv_thresh : float
            The minimum convergence rate

        error_type : {'l2', 'inf'}, str
            The error norm to compute
        """
        conv_thresh = section.getfloat('convergence_thresh')
        error_type = section.get('error_type')
        return conv_thresh, error_type

from polaris.ocean.convergence import get_timestep_for_task
from polaris.ocean.model import OceanModelStep, get_time_interval_string

class ConvergenceForward(OceanModelStep):
    """
    A step for performing forward ocean component runs as part of a spherical
    convergence test

        The package name or module object that contains the YAML file

    yaml_filename : str
        A YAML file that is a Jinja2 template with model config options

    refinement_factor : float
        Refinement factor use to scale space, time or both depending on
        refinement option
    """

    def __init__(self, component, name, subdir, refinement_factor, mesh, init,
                 package, yaml_filename='forward.yaml',
                 options=None, graph_target=None, output_filename='output.nc',
                 validate_vars=None, refinement='both'):
        """

        name : str

            The subdirectory for the step

            Refinement factor use to scale space, time or both depending on
            refinement option

            The package name or module object that contains the YAML file

        yaml_filename : str, optional
            A YAML file that is a Jinja2 template with model config options

            A nested dictionary of options and value for each ``config_model``
            to replace model config options with new values

        graph_target : str, optional
            The graph file name (relative to the base work directory).
            If none, it will be created.

        output_filename : str, optional
            The output file that will be written out at the end of the forward
            run

        validate_vars : list of str, optional
            A list of variables to validate against a baseline if requested

        """
                         openmp_threads=1, graph_target=graph_target)

        self.refinement_factor = refinement_factor
        self.package = package
        self.yaml_filename = yaml_filename

        # make sure output is double precision
        self.add_yaml_file('polaris.ocean.config', 'output.yaml')

            for config_model in options:
                self.add_model_config_options(options=options[config_model],
                                              config_model=config_model)

        self.add_input_file(
            filename='init.nc',
            work_dir_target=f'{init.path}/initial_state.nc')

        self.add_output_file(filename=output_filename,
                             validate_vars=validate_vars)

    def compute_cell_count(self):
        """
        Compute the approximate number of cells in the mesh, used to constrain
        resources

        Returns
        -------
        cell_count : int or None
            The approximate number of cells in the mesh
        """
        raise ValueError('compute_cell_count method must be overridden by '
                         'spherical or planar method')

        """
        Set the model time step from config options at setup and runtime

        """
        super().dynamic_model_config(at_setup=at_setup)

        refinement_factor = self.refinement_factor

        vert_levels = config.getfloat('vertical_grid', 'vert_levels')
        if not at_setup and vert_levels == 1:
            self.add_yaml_file('polaris.ocean.config', 'single_layer.yaml')

        timestep, btr_timestep = get_timestep_for_task(
            config, refinement_factor, refinement=self.refinement)
        dt_str = get_time_interval_string(seconds=timestep)
        btr_dt_str = get_time_interval_string(seconds=btr_timestep)

        s_per_hour = 3600.
        section = config['convergence_forward']
        run_duration = section.getfloat('run_duration')
        run_duration_str = get_time_interval_string(
            seconds=run_duration * s_per_hour)

        output_interval = section.getfloat('output_interval')
        output_interval_str = get_time_interval_string(
            seconds=output_interval * s_per_hour)

        # For Omega, we want the output interval as a number of seconds
        output_freq = int(output_interval * s_per_hour)

        time_integrator = section.get('time_integrator')
        time_integrator_map = dict([('RK4', 'RungeKutta4')])
        if model == 'omega':
            if time_integrator in time_integrator_map.keys():
                time_integrator = time_integrator_map[time_integrator]
            else:
                print('Warning: mapping from time integrator '
                      f'{time_integrator} to omega not found, '
                      'retaining name given in config')

        replacements = dict(
            time_integrator=time_integrator,
            dt=dt_str,
            btr_dt=btr_dt_str,
            run_duration=run_duration_str,
            output_interval=output_interval_str,
            output_freq=f'{output_freq}'
        )

        self.add_yaml_file(self.package, self.yaml_filename,
                           template_replacements=replacements)

from polaris.ocean.convergence import get_resolution_for_task
from polaris.ocean.convergence.forward import ConvergenceForward

class SphericalConvergenceForward(ConvergenceForward):
    """

    resolution : float
        The resolution of the (uniform) mesh in km

    """

        """
        resources

        Returns
        -------
        """
        # use a heuristic based on QU30 (65275 cells) and QU240 (10383 cells)
        resolution = get_resolution_for_task(
            self.config, self.refinement_factor,
            refinement=self.refinement)
        cell_count = 6e8 / resolution**2
        return cell_count

from polaris import Step, Task
from polaris.ocean.ice_shelf.ssh_adjustment import SshAdjustment
from polaris.ocean.ice_shelf.ssh_forward import SshForward

class IceShelfTask(Task):

    """
    A class for tasks with domains containing ice shelves

    sshdir : string
        The directory to put the ssh_adjustment steps in

    component : polaris.ocean.Ocean
        The ocean component that this task belongs to

    min_resolution : float
        The resolution of the test case in km
    """
    def __init__(self, component, min_resolution, name, subdir, sshdir=None):
        """
        Construct ice shelf task

        component : polaris.ocean.Ocean
            The ocean component that this task belongs to

        min_resolution : float
            The resolution of the test case in km

        name : string

        subdir : string

        sshdir : string, optional
            The directory to put the ssh_adjustment steps in. If None,
            defaults to subdir.
        """
        if sshdir is None:
            sshdir = subdir
        self.sshdir = sshdir
        self.min_resolution = min_resolution

    def setup_ssh_adjustment_steps(self, mesh_filename, graph_target,
                                   init_filename, config, config_filename,
                                   ForwardStep, package=None,
                                   yaml_filename='ssh_forward.yaml',
                                   yaml_replacements=None):

        """
        Setup ssh_forward and ssh_adjustment steps for all iterations

            The configuration for this task

        mesh_filename : str
            the mesh filename (relative to the base work directory)

        graph_target: str
            the graph filename (relative to the base work directory)

        init_filename : str
            the initial condition filename (relative to the base work
            directory)

        config_filename : str
            the configuration filename

        ForwardStep : polaris.ocean.ice_shelf.ssh_forward.SshForward
            the step class used to create ssh_forward steps

            from which ssh_forward steps will derive their configuration

            the yaml filename used for ssh_forward steps

        yaml_replacements : Dict, optional
            key, string combinations for templated replacements in the yaml
            file

        Returns
        -------
        shared_step : polaris.Step
            the final ssh_adjustment step that produces the input to the next
            forward step
        """
        min_resolution = self.min_resolution
        indir = self.sshdir

        num_iterations = config.getint('ssh_adjustment', 'iterations')

        # for the first iteration, ssh_forward step is run from the initial
        # state
        current_init_filename = init_filename
        for iteration in range(num_iterations):
            name = f'ssh_forward_{iteration}'
            subdir = f'{indir}/ssh_adjustment/{name}'
            if subdir in component.steps:
                shared_step = component.steps[subdir]
            else:
                ssh_forward = ForwardStep(
                    component=component, min_resolution=min_resolution,
                    indir=indir, graph_target=graph_target,
                    init_filename=current_init_filename, name=name,
                    package=package, yaml_filename=yaml_filename,
                    yaml_replacements=yaml_replacements)
                ssh_forward.set_shared_config(config, link=config_filename)
                shared_step = ssh_forward
            if indir == self.subdir:
                symlink = None
            else:
                symlink = f'ssh_adjustment/{name}'
            self.add_step(shared_step, symlink=symlink)

            name = f'ssh_adjust_{iteration}'
            else:
                ssh_adjust = SshAdjustment(
                    component=component, indir=indir, name=name,
                    mesh_filename=mesh_filename, forward=ssh_forward)
                ssh_adjust.set_shared_config(config, link=config_filename)
                shared_step = ssh_adjust
            else:
            # for the next iteration, ssh_forward is run from the adjusted
            # initial state (the output of ssh_adustment)
            current_init_filename = f'{shared_step.path}/output.nc'

        return shared_step

from mpas_tools.io import write_netcdf

from polaris.ocean.vertical import update_layer_thickness

class SshAdjustment(Step):
    """
    A step for iteratively adjusting the pressure from the weight of the ice
    shelf to match the sea-surface height as part of ice-shelf 2D test cases
    """
    def __init__(self, component, mesh_filename, forward, indir=None,
                 name='ssh_adjust'):
        """

        forward: polaris.Step
            the step that produced the state which will be adjusted

            the name of this step
        """
        super().__init__(component=component, name=name,
                         indir=f'{indir}/ssh_adjustment')

        self.add_input_file(filename='final.nc',
                            work_dir_target=f'{forward.path}/output.nc')
        self.add_input_file(filename='init.nc',
                            work_dir_target=f'{forward.path}/init.nc')
        self.add_input_file(filename='mesh.nc',
                            work_dir_target=mesh_filename)
        self.add_output_file(filename='output.nc')

    # no setup() is needed

        """
        Adjust the sea surface height or land-ice pressure to be dynamically
        consistent with one another.
        """
        adjust_variable = config.get('ssh_adjustment', 'adjust_variable')
        mask_variable = config.get('ssh_adjustment', 'mask_variable')
        mesh_filename = 'mesh.nc'
        init_filename = 'init.nc'
        final_filename = 'final.nc'
        out_filename = 'output.nc'
        ds_mesh = xr.open_dataset(mesh_filename)
        ds_init = xr.open_dataset(init_filename)
        ds_final = xr.open_dataset(final_filename)
        ds_out = ds_init.copy()

        if adjust_variable not in ['ssh', 'landIcePressure']:
            raise ValueError(f"Unknown variable to modify: {adjust_variable}")
        if mask_variable not in ds_init.keys():
            raise ValueError(f"Mask variable {mask_variable} is not contained "
                             f"in {init_filename}")

        logger.info("   * Updating SSH or land-ice pressure")

        # keep the data set with Time for output
        # and generate these time slices
        ds_final = ds_final.isel(Time=-1)

        on_a_sphere = ds_out.attrs['on_a_sphere'].lower() == 'yes'

        if 'minLevelCell' in ds_final:
            min_level_cell = ds_final.minLevelCell - 1
        else:
            min_level_cell = ds_mesh.minLevelCell - 1

        init_ssh = ds_init.ssh
        final_ssh = ds_final.ssh
        top_density = ds_final.density.isel(nVertLevels=min_level_cell)

        mask = ds_init[mask_variable]
        delta_ssh = mask * (final_ssh - init_ssh)

        # then, modify the SSH or land-ice pressure
        if adjust_variable == 'ssh':
            final_ssh = final_ssh.expand_dims(dim='Time', axis=0)
            ds_out['ssh'] = final_ssh
            # also update the landIceDraft variable, which will be used to
            # compensate for the SSH due to land-ice pressure when
            # computing sea-surface tilt
            ds_out['landIceDraft'] = final_ssh
            # we also need to stretch layerThickness to be compatible with
            # the new SSH
            update_layer_thickness(config, ds_out)
            land_ice_pressure = ds_out.landIcePressure.values
        else:
            # Moving the SSH up or down by deltaSSH would change the
            # land-ice pressure by density(SSH)*g*deltaSSH. If deltaSSH is
            # positive (moving up), it means the land-ice pressure is too
            # small and if deltaSSH is negative (moving down), it means
            # land-ice pressure is too large, the sign of the second term
            # makes sense.
            gravity = constants['SHR_CONST_G']
            delta_land_ice_pressure = top_density * gravity * delta_ssh

            land_ice_pressure = np.maximum(
                0.0, ds_final.landIcePressure + delta_land_ice_pressure)

            ds_out['landIcePressure'] = \
                land_ice_pressure.expand_dims(dim='Time', axis=0)

            final_ssh = init_ssh

        write_netcdf(ds_out, out_filename)

        # Write the largest change in SSH and its lon/lat to a file
        with open('maxDeltaSSH.log', 'w') as log_file:

            mask = land_ice_pressure > 0.
            i_cell = np.abs(delta_ssh.where(mask)).argmax().values

            ds_cell = ds_final.isel(nCells=i_cell)
            ds_mesh = ds_mesh.isel(nCells=i_cell)
            delta_ssh_max = delta_ssh.isel(nCells=i_cell).values

            if on_a_sphere:
                coords = (f'lon/lat: '
                          f'{np.rad2deg(ds_cell.lonCell.values):f} '
                          f'{np.rad2deg(ds_cell.latCell.values):f}')
            else:
                coords = (f'x/y: {1e-3 * ds_mesh.xCell.values:f} '
                          f'{1e-3 * ds_mesh.yCell.values:f}')
            string = (f'deltaSSHMax: '
                      f'{delta_ssh_max:g}, {coords}')
            logger.info(f'     {string}')
            log_file.write(f'{string}\n')
            string = (f'ssh: {final_ssh.isel(nCells=i_cell).values:g}, '
                      f'land_ice_pressure: '
                      f'{land_ice_pressure.isel(nCells=i_cell).values:g}')

        logger.info("   - Complete\n")

class SshForward(OceanModelStep):
    """
    A step for performing forward ocean component runs as part of ssh
    adjustment.

        The minimum resolution in km used to determine the time step

    yaml_filename : str, optional
        the yaml filename used for ssh_forward steps

    yaml_replacements : Dict, optional
        key, string combinations for templated replacements in the yaml
        file
    """
    def __init__(self, component, min_resolution, init_filename,
                 graph_target, name='ssh_forward', subdir=None,
                 package=None, yaml_filename='ssh_forward.yaml',
                 yaml_replacements=None, indir=None, ntasks=None,
                 min_tasks=None, openmp_threads=1):
        """

            The minimum resolution in km used to determine the time step

            where ssh_forward steps will derive their configuration

            file

            the number of tasks the step would ideally use.  If fewer tasks
            are available on the system, the step will run on all available
            tasks as long as this is not below ``min_tasks``

            the number of tasks the step requires.  If the system has fewer
            than this number of tasks, the step will fail

            the number of OpenMP threads the step will use
        """
                         indir=f'{indir}/ssh_adjustment', ntasks=ntasks,
                         graph_target=graph_target)

        self.yaml_replacements = yaml_replacements

                            work_dir_target=init_filename)
        self.add_input_file(filename='graph.info',
                            work_dir_target=graph_target)

        self.add_output_file(
            filename='output.nc',
            validate_vars=['temperature', 'salinity', 'layerThickness',
                           'normalVelocity'])

        """
        resources

        Returns
        -------
        """

        """

        """
        super().dynamic_model_config(at_setup)

        section = config['ssh_adjustment']

        # dt is proportional to resolution

        if time_integrator == 'RK4':
            dt_per_km = section.getfloat('rk4_dt_per_km')
        else:
            dt_per_km = section.getfloat('split_dt_per_km')
        dt_str = get_time_interval_string(
            seconds=dt_per_km * self.min_resolution)

        # btr_dt is also proportional to resolution
        btr_dt_per_km = section.getfloat('btr_dt_per_km')
        btr_dt_str = get_time_interval_string(
            seconds=btr_dt_per_km * self.min_resolution)

        )

        self.add_yaml_file('polaris.ocean.ice_shelf',
                           'ssh_forward.yaml',
        if self.package is not None:
            self.add_yaml_file(package=self.package,
                               yaml=self.yaml_filename,
                               template_replacements=self.yaml_replacements)

from typing import Dict, Union

class OceanIOStep(Step):
    """
    A step that writes input and/or output files for Omega or MPAS-Ocean

    mpaso_to_omega_dim_map : dict
        A map from MPAS-Ocean dimension names to their Omega equivalents

    omega_to_mpaso_dim_map : dict
        A map from Omega dimension names to their MPAS-Ocean equivalents, the
        inverse of ``mpaso_to_omega_dim_map``

    mpaso_to_omega_var_map : dict
        A map from MPAS-Ocean variable names to their Omega equivalents

    omega_to_mpaso_var_map : dict
        A map from Omega variable names to their MPAS-Ocean equivalents, the
        inverse of ``mpaso_to_omega_var_map``
    """
    def __init__(self, component, name, **kwargs):
        """

        name : str

        kwargs
            keyword arguments passed to `polaris.Step()`
        """
        super().__init__(
            component=component, name=name, **kwargs)

        self.mpaso_to_omega_dim_map: Union[None, Dict[str, str]] = None
        self.omega_to_mpaso_dim_map: Union[None, Dict[str, str]] = None
        self.mpaso_to_omega_var_map: Union[None, Dict[str, str]] = None
        self.omega_to_mpaso_var_map: Union[None, Dict[str, str]] = None

        """
        If the ocean model is Omega, set up maps between Omega and MPAS-Ocean
        variable names
        """
            self._read_var_map()
        elif model != 'mpas-ocean':
            raise ValueError(f'Unexpected ocean model: {model}')

    def map_to_native_model_vars(self, ds):
        """
        If the model is Omega, rename dimensions and variables in a dataset
        from their MPAS-Ocean names to the Omega equivalent (appropriate for
        input datasets like an initial condition)

        ds : xarray.Dataset
            A dataset containing MPAS-Ocean variable names

        Returns
        -------
            The same dataset with variables renamed as appropriate for the
            ocean model being run
        """
            assert self.mpaso_to_omega_dim_map is not None
            rename = {k: v for k, v in self.mpaso_to_omega_dim_map.items()
                      if k in ds.dims}
            assert self.mpaso_to_omega_var_map is not None
            rename_vars = {k: v for k, v in self.mpaso_to_omega_var_map.items()
                           if k in ds}
            rename.update(rename_vars)
            ds = ds.rename(rename)
        return ds

    def write_model_dataset(self, ds, filename):
        """
        Write out the given dataset, mapping dimension and variable names from
        MPAS-Ocean to Omega names if appropriate

            The path for the NetCDF file to write
        """
        ds = self.map_to_native_model_vars(ds)
        write_netcdf(ds=ds, fileName=filename)

    def map_from_native_model_vars(self, ds):
        """
        from their Omega names to the MPAS-Ocean equivalent (appropriate for
        datasets that are output from the model)

            A dataset containing variable names native to either ocean model

        Returns
        -------
            The same dataset with variables named as expected in MPAS-Ocean
        """
            assert self.omega_to_mpaso_dim_map is not None
            rename = {k: v for k, v in self.omega_to_mpaso_dim_map.items()
            assert self.omega_to_mpaso_var_map is not None
            rename_vars = {k: v for k, v in self.omega_to_mpaso_var_map.items()
                           if k in ds}
        return ds

    def open_model_dataset(self, filename, **kwargs):
        """
        Open the given dataset, mapping variable and dimension names from Omega
        to MPAS-Ocean names if appropriate

            The path for the NetCDF file to open

        kwargs
            keyword arguments passed to `xarray.open_dataset()`

        Returns
        -------
            The dataset with variables named as expected in MPAS-Ocean
        """
        ds = xr.open_dataset(filename, **kwargs)
        ds = self.map_from_native_model_vars(ds)
        return ds

    def _read_var_map(self):
        """
        Read the map from MPAS-Ocean to Omega dimension and variable names
        """
        package = 'polaris.ocean.model'
        filename = 'mpaso_to_omega.yaml'
        text = imp_res.files(package).joinpath(filename).read_text()

        nested_dict = yaml_data.load(text)
        self.mpaso_to_omega_dim_map = nested_dict['dimensions']
        self.mpaso_to_omega_var_map = nested_dict['variables']
        assert self.mpaso_to_omega_dim_map is not None
        self.omega_to_mpaso_dim_map = {
            v: k for k, v in self.mpaso_to_omega_dim_map.items()}
        assert self.mpaso_to_omega_var_map is not None
        self.omega_to_mpaso_var_map = {
            v: k for k, v in self.mpaso_to_omega_var_map.items()}

from polaris.model_step import ModelStep

class OceanModelStep(ModelStep):
    """
    An Omega or MPAS-Ocean step

    dynamic_ntasks : bool
        Whether the target and minimum number of MPI tasks (``ntasks`` and
        ``min_tasks``) are computed dynamically from the number of cells
        in the mesh

    config_map : dict
        A nested dictionary that maps from MPAS-Ocean to Omega model config
        options

    graph_target : str
        The name of the graph partition file to link to (relative to the base
        working directory)
    """
                 cached=False, yaml=None, update_pio=True, make_graph=False,
                 mesh_filename=None, partition_graph=True, graph_target=None):
        """

        name : str

        """
        if graph_target is None:
            self.make_graph = True

            component=component, name=name, subdir=subdir, indir=indir,
            ntasks=ntasks, min_tasks=min_tasks, openmp_threads=openmp_threads,
            max_memory=max_memory, cached=cached, yaml=yaml,
            update_pio=update_pio, make_graph=make_graph,
            mesh_filename=mesh_filename, partition_graph=partition_graph,
            graph_filename='graph.info')

        self.dynamic_ntasks = (ntasks is None and min_tasks is None)

        self.config_map: Union[None, List[Dict[str, Dict[str, str]]]] = None
        self.graph_target = graph_target

        """
        Determine if we will make yaml files or namelists and streams files,
        then, determine the number of MPI tasks to use based on the estimated
        mesh size
        """
            self.make_yaml = True
            self.config_models = ['ocean', 'Omega']
            self.yaml = 'omega.yml'
            self.streams_section = 'IOStreams'
            self._read_config_map()
            self.partition_graph = False
        elif model == 'mpas-ocean':
            self.config_models = ['ocean', 'mpas-ocean']
            self.make_yaml = False
                filename='graph.info',
                work_dir_target=self.graph_target)
            self.streams_section = 'streams'
        else:

        self.dynamic_ntasks = (self.ntasks is None and self.min_tasks is None)

        if self.dynamic_ntasks:
            self._update_ntasks()

    def constrain_resources(self, available_cores):
        """
        Update the number of MPI tasks to use based on the estimated mesh size
        """
        super().constrain_resources(available_cores)

        """
        resources

        Returns
        -------
        """
        return None

        """
        A mapping between model config options from MPAS-Ocean to Omega

        Returns
        -------
        """
        if model == 'omega' and config_model == 'ocean':
            options = self._map_mpaso_to_omega_options(options)

        """
        A mapping between model sections and config options from MPAS-Ocean to
        Omega

        Returns
        -------
        """
            configs = self._map_mpaso_to_omega_configs(configs)

        """
        up.

        """
        raise ValueError('Input namelist files are not supported in '
                         'OceanModelStep')

        """

        """
        raise ValueError('Input streams files are not supported in '

    def _update_ntasks(self):
        """
        Update ``ntasks`` and ``min_tasks`` for the step based on the estimated
        mesh size
        """
        cell_count = self.compute_cell_count()
        if cell_count is None:
            raise ValueError('ntasks and min_tasks were not set explicitly '
                             'but they also cannot be computed because '
                             'compute_cell_count() does not appear to have '
                             'been overridden.')

        goal_cells_per_core = config.getfloat('ocean', 'goal_cells_per_core')
        max_cells_per_core = config.getfloat('ocean', 'max_cells_per_core')

        # machines (e.g. Perlmutter) seem to be happier with ntasks that
        # are multiples of 4
        # ideally, about 200 cells per core
        self.ntasks = max(1, 4 * round(cell_count / (4 * goal_cells_per_core)))
        # In a pinch, about 2000 cells per core
        self.min_tasks = max(1,
                             4 * round(cell_count / (4 * max_cells_per_core)))

    def _read_config_map(self):
        """
        Read the map from MPAS-Ocean to Omega config options
        """

        self.config_map = nested_dict['config']

    def _map_mpaso_to_omega_options(self, options):
        """
        Map MPAS-Ocean namelist options to Omega config options
        """

        out_options: Dict[str, str] = {}
        not_found = []
        for mpaso_option, mpaso_value in options.items():
            try:
                omega_option, omega_value = self._map_mpaso_to_omega_option(
                    option=mpaso_option, value=mpaso_value)
                out_options[omega_option] = omega_value
                not_found.append(mpaso_option)

        self._warn_not_found(not_found)

        return out_options

    def _map_mpaso_to_omega_option(self, option, value):
        """
        Map MPAS-Ocean namelist option to Omega equivalent
        """
        out_option = option
        found = False

        assert self.config_map is not None
        # traverse the map
        for entry in self.config_map:
            options_dict = entry['options']
            for mpaso_option, omega_option in options_dict.items():
                if option == mpaso_option:
                    found = True
                    out_option = omega_option
                    break

            if found:
                break

        if not found:
            raise ValueError(f'No mapping found for {option}')

        out_option, out_value = self._map_handle_not(out_option, value)

        return out_option, out_value

    def _map_mpaso_to_omega_configs(self, configs):
        """
        """
        out_configs: Dict[str, Dict[str, str]] = {}
        for section, options in configs.items():
            for option, mpaso_value in options.items():
                try:
                    omega_section, omega_option, omega_value = \
                        self._map_mpaso_to_omega_section_option(
                            section=section, option=option, value=mpaso_value)
                    if omega_section not in out_configs:
                        out_configs[omega_section] = {}
                    out_configs[omega_section][omega_option] = omega_value
                except ValueError:
                    not_found.append(f'{section}/{option}')

        return out_configs

    def _map_mpaso_to_omega_section_option(self, section, option, value):
        """
        Map MPAS-Ocean namelist section and option to Omega equivalent
        """
        out_section = section

        option_found = False
            section_dict = entry['section']
            try:
                omega_section = section_dict[section]
            except KeyError:
                continue
            else:
                options_dict = entry['options']
                option_found = False
                try:
                    omega_option = options_dict[option]
                except KeyError:
                    continue
                else:
                    option_found = True
                    out_section = omega_section
                    break

        if not option_found:
            raise ValueError(f'No mapping found for {section}/{option}')

        return out_section, out_option, out_value

    def _warn_not_found(not_found):
        """ Warn about options that were not found in the map """
        if len(not_found) == 0:
            return

        print('WARNING: No Omega mapping found for these MPASO options:')
        for string in not_found:
            print(f'    {string}')
        print()

    def _map_handle_not(option, value):
        """
        Handle negation of boolean value if the option starts with "not"
        """
        if option.startswith('not '):
            # a special case where we want the opposite of a boolean value
            option = option[4:]
            value = not value

        return option, value

from polaris.mesh.planar import compute_planar_hex_nx_ny
from polaris.ocean.model import OceanModelStep

class Forward(OceanModelStep):
    """
    A step for performing forward ocean component runs as part of baroclinic
    channel tasks.

        The resolution of the task in km

    dt : float
        The model time step in seconds

    btr_dt : float
        The model barotropic time step in seconds

    run_time_steps : int or None
        Number of time steps to run for
    """
    def __init__(self, component, resolution, name='forward', subdir=None,
                 indir=None, ntasks=None, min_tasks=None, openmp_threads=1,
                 nu=None, run_time_steps=None, graph_target='graph.info'):
        """

        resolution : km
            The resolution of the task in km

        name : str

        nu : float, optional
            the viscosity (if different from the default for baroclinic channel
            tests)

        run_time_steps : int, optional
            Number of time steps to run for

        """
        self.resolution = resolution
        self.run_time_steps = run_time_steps
                         indir=indir, ntasks=ntasks, min_tasks=min_tasks,
                         openmp_threads=openmp_threads,

        self.add_input_file(filename='initial_state.nc',
                            target='../../init/initial_state.nc')

        self.add_yaml_file('polaris.ocean.tasks.baroclinic_channel',
                           'forward.yaml')

        if nu is not None:
            # update the viscosity to the requested value *after* loading
            # forward.yaml
            self.add_model_config_options(options=dict(config_mom_del2=nu),
                                          config_model='mpas-ocean')

        self.dt = None
        self.btr_dt = None

        """
        resources

        Returns
        -------
        """
        section = self.config['baroclinic_channel']
        lx = section.getfloat('lx')
        ly = section.getfloat('ly')
        nx, ny = compute_planar_hex_nx_ny(lx, ly, self.resolution)
        cell_count = nx * ny

        """

        """

        options = dict()

        # dt is proportional to resolution: default 30 seconds per km
        dt_per_km = config.getfloat('baroclinic_channel', 'dt_per_km')
        dt = dt_per_km * self.resolution
        # https://stackoverflow.com/a/1384565/7728169
        options['config_dt'] = \
            time.strftime('%H:%M:%S', time.gmtime(dt))

        if self.run_time_steps is not None:
            # default run duration is a few time steps
            run_seconds = self.run_time_steps * dt
            options['config_run_duration'] = \
                time.strftime('%H:%M:%S', time.gmtime(run_seconds))
            options['config_stop_time'] = 'none'

        # btr_dt is also proportional to resolution: default 1.5 seconds per km
        btr_dt_per_km = config.getfloat('baroclinic_channel', 'btr_dt_per_km')
        btr_dt = btr_dt_per_km * self.resolution
        options['config_btr_dt'] = \
            time.strftime('%H:%M:%S', time.gmtime(btr_dt))

        self.dt = dt
        self.btr_dt = btr_dt

        self.add_model_config_options(options=options,
                                      config_model='mpas-ocean')

from mpas_tools.mesh.conversion import convert, cull
from mpas_tools.planar_hex import make_planar_hex_mesh

from polaris.mpas import cell_mask_to_edge_mask
from polaris.ocean.vertical import init_vertical_coord
from polaris.ocean.viz import compute_transect, plot_transect
from polaris.viz import plot_horiz_field

class Init(Step):
    """
    A step for creating a mesh and initial condition for baroclinic channel
    tasks

    """
    def __init__(self, component, resolution, indir):
        """

        resolution : float

        indir : str
        """
        super().__init__(component=component, name='init', indir=indir)

        for file in ['base_mesh.nc', 'culled_mesh.nc', 'culled_graph.info']:
            self.add_output_file(file)
        self.add_output_file('initial_state.nc',
                             validate_vars=['temperature', 'salinity',
                                            'layerThickness'])

        """
        """

        section = config['baroclinic_channel']
        resolution = self.resolution

        # these could be hard-coded as functions of specific supported
        # resolutions but it is preferable to make them algorithmic like here
        # for greater flexibility
        nx, ny = compute_planar_hex_nx_ny(lx, ly, resolution)
        dc = 1e3 * resolution

        ds_mesh = make_planar_hex_mesh(nx=nx, ny=ny, dc=dc,
                                       nonperiodic_x=False,
                                       nonperiodic_y=True)
        write_netcdf(ds_mesh, 'base_mesh.nc')

        ds_mesh = cull(ds_mesh, logger=logger)
        ds_mesh = convert(ds_mesh, graphInfoFileName='culled_graph.info',
                          logger=logger)
        write_netcdf(ds_mesh, 'culled_mesh.nc')

        use_distances = section.getboolean('use_distances')
        gradient_width_dist = section.getfloat('gradient_width_dist')
        gradient_width_frac = section.getfloat('gradient_width_frac')
        bottom_temperature = section.getfloat('bottom_temperature')
        surface_temperature = section.getfloat('surface_temperature')
        temperature_difference = section.getfloat('temperature_difference')
        salinity = section.getfloat('salinity')
        coriolis_parameter = section.getfloat('coriolis_parameter')

        ds = ds_mesh.copy()
        x_cell = ds.xCell
        y_cell = ds.yCell

        bottom_depth = config.getfloat('vertical_grid', 'bottom_depth')

        ds['bottomDepth'] = bottom_depth * xr.ones_like(x_cell)
        ds['ssh'] = xr.zeros_like(x_cell)

        init_vertical_coord(config, ds)

        x_min = x_cell.min().values
        x_max = x_cell.max().values
        y_min = y_cell.min().values
        y_max = y_cell.max().values

        y_mid = 0.5 * (y_min + y_max)
        x_perturb_min = x_min + 4.0 * (x_max - x_min) / 6.0
        x_perturb_max = x_min + 5.0 * (x_max - x_min) / 6.0

        if use_distances:
            perturb_width = gradient_width_dist
        else:
            perturb_width = (y_max - y_min) * gradient_width_frac

        y_offset = perturb_width * np.sin(
            6.0 * np.pi * (x_cell - x_min) / (x_max - x_min))

        temp_vert = (bottom_temperature +
                     (surface_temperature - bottom_temperature) *
                     ((ds.refZMid + bottom_depth) / bottom_depth))

        frac = xr.where(y_cell < y_mid - y_offset, 1., 0.)

        mask = np.logical_and(y_cell >= y_mid - y_offset,
                              y_cell < y_mid - y_offset + perturb_width)
        frac = xr.where(mask,
                        1. - (y_cell - (y_mid - y_offset)) / perturb_width,
                        frac)

        temperature = temp_vert - temperature_difference * frac
        temperature = temperature.transpose('nCells', 'nVertLevels')

        # Determine y_offset for 3rd crest in sin wave
        y_offset = 0.5 * perturb_width * np.sin(
            np.pi * (x_cell - x_perturb_min) / (x_perturb_max - x_perturb_min))

        mask = np.logical_and(
            np.logical_and(y_cell >= y_mid - y_offset - 0.5 * perturb_width,
                           y_cell <= y_mid - y_offset + 0.5 * perturb_width),
            np.logical_and(x_cell >= x_perturb_min,
                           x_cell <= x_perturb_max))

        temperature = (temperature +
                       mask * 0.3 * (1. - ((y_cell - (y_mid - y_offset)) /
                                           (0.5 * perturb_width))))

        temperature = temperature.expand_dims(dim='Time', axis=0)

        normal_velocity = xr.zeros_like(ds_mesh.xEdge)
        normal_velocity, _ = xr.broadcast(normal_velocity, ds.refBottomDepth)
        normal_velocity = normal_velocity.transpose('nEdges', 'nVertLevels')
        normal_velocity = normal_velocity.expand_dims(dim='Time', axis=0)

        ds['temperature'] = temperature
        ds['salinity'] = salinity * xr.ones_like(temperature)
        ds['normalVelocity'] = normal_velocity
        ds['fCell'] = coriolis_parameter * xr.ones_like(x_cell)
        ds['fEdge'] = coriolis_parameter * xr.ones_like(ds_mesh.xEdge)
        ds['fVertex'] = coriolis_parameter * xr.ones_like(ds_mesh.xVertex)

        ds.attrs['nx'] = nx
        ds.attrs['ny'] = ny
        ds.attrs['dc'] = dc

        write_netcdf(ds, 'initial_state.nc')

        cell_mask = ds.maxLevelCell >= 1
        edge_mask = cell_mask_to_edge_mask(ds, cell_mask)

        plot_horiz_field(ds_mesh, ds['normalVelocity'],
                         'initial_normal_velocity.png', cmap='cmo.balance',
                         show_patch_edges=True, field_mask=edge_mask)

        y_min = ds_mesh.yVertex.min().values
        y_max = ds_mesh.yVertex.max().values
        x_mid = ds_mesh.xCell.median().values

        y = xr.DataArray(data=np.linspace(y_min, y_max, 2), dims=('nPoints',))
        x = x_mid * xr.ones_like(y)

        ds_transect = compute_transect(
            x=x, y=y, ds_horiz_mesh=ds_mesh,
            layer_thickness=ds.layerThickness.isel(Time=0),
            bottom_depth=ds.bottomDepth, min_level_cell=ds.minLevelCell - 1,
            max_level_cell=ds.maxLevelCell - 1, spherical=False)

        field_name = 'temperature'
        vmin = ds[field_name].min().values
        vmax = ds[field_name].max().values
        plot_transect(ds_transect=ds_transect,
                      mpas_field=ds[field_name].isel(Time=0),
                      title=f'{field_name} at x={1e-3 * x_mid:.1f} km',
                      out_filename=f'initial_{field_name}_section.png',
                      vmin=vmin, vmax=vmax, cmap='cmo.thermal',
                      colorbar_label=r'$^\circ$C', color_start_and_end=True)

        plot_horiz_field(ds_mesh, ds['temperature'], 'initial_temperature.png',
                         vmin=vmin, vmax=vmax, cmap='cmo.thermal',
                         field_mask=cell_mask, transect_x=x, transect_y=y)

from polaris.ocean.tasks.baroclinic_channel.forward import Forward

class RestartStep(Forward):
    """
    A forward model step in the restart test case
    """
    def __init__(self, component, resolution, name, indir, init):
        """
        Create a new test case

        name : str
            the name of the test case

        indir : str

        init: polaris.ocean.tasks.baroclinic_channel.init.Init
            the initial state step
        """
        super().__init__(component=component, name=name, indir=indir, ntasks=4,
                         min_tasks=4, openmp_threads=1,
                         resolution=resolution,
                         graph_target=f'{init.path}/culled_graph.info')

        """

        """

        dt = self.dt
        if dt is None:
            raise ValueError('dt was not set in the Forward class as expected')

        if self.name == 'full_run':
            # 2 time steps without a restart
            do_restart = False
            start_time = 0.
            run_duration = 2. * dt
            output_interval = 2. * dt
        elif self.name == 'restart_run':
            # 1 time step from the restart at 1 time step
            do_restart = True
            start_time = dt
            run_duration = dt
            output_interval = dt
        else:
            raise ValueError(f'Unexpected step name {self.name}')

        # to keep the time formatting from getting too complicated, we'll
        # assume 2 time steps is never more than a day
        start_time_str = time.strftime('0001-01-01_%H:%M:%S',
                                       time.gmtime(start_time))

        run_duration_str = time.strftime('0000-00-00_%H:%M:%S',
                                         time.gmtime(run_duration))

        output_interval_str = time.strftime('0000-00-00_%H:%M:%S',
                                            time.gmtime(output_interval))

        package = 'polaris.ocean.tasks.baroclinic_channel.restart'
            do_restart=do_restart,
            start_time=start_time_str,
        )

        self.add_yaml_file(package=package,
                           yaml='forward.yaml',

from polaris.ocean.tasks.baroclinic_channel.rpe.analysis import Analysis

class Rpe(Task):
    """
    The baroclinic channel reference potential energy (RPE) test case performs
    a 20-day integration of the model forward in time at 5 different values of
    the viscosity at the given resolution.

    """

    def __init__(self, component, resolution, indir, init, config):
        """
        Create the test case

        indir : str
            The directory the task is in, to which ``name`` will be appended

        init : polaris.ocean.tasks.baroclinic_channel.init.Init
            A shared step for creating the initial state

            A shared config parser
        """
        super().__init__(component=component, name='rpe', indir=indir)

        # this needs to be added before we can use the config options it
        # brings in to set up the steps
        self.set_shared_config(config, link='baroclinic_channel.cfg')
        self.add_step(init, symlink='init')
        self._add_rpe_and_analysis_steps()

        """
        Modify the configuration options for this test case.
        """
        super().configure()

    def _add_rpe_and_analysis_steps(self):
        """ Add the steps in the test case either at init or set-up """

        for step_name in list(self.steps.keys()):
            step = self.steps[step_name]
            if step_name.startswith('nu') or step_name == 'analysis':
                # remove previous RPE forward or analysis steps
                self.remove_step(step)

        nus = config.getlist('baroclinic_channel_rpe', 'viscosities',
                             dtype=float)
        for nu in nus:
            name = f'nu_{nu:g}'
            step = Forward(
                component=component, name=name, indir=self.subdir,
                ntasks=None, min_tasks=None, openmp_threads=1,
                resolution=resolution, nu=nu)

            step.add_yaml_file(
                'polaris.ocean.tasks.baroclinic_channel.rpe',
                'forward.yaml')
            self.add_step(step)

        self.add_step(
            Analysis(component=component, resolution=resolution, nus=nus,
                     indir=self.subdir))

from polaris.ocean.rpe import compute_rpe
from polaris.viz import plot_horiz_field, use_mplstyle

class Analysis(Step):
    """
    A step for plotting the results of a series of baroclinic channel RPE runs

    nus : list
        A list of viscosities
    """
    def __init__(self, component, indir, resolution, nus):
        """

        indir : str

        nus : list of float
            A list of viscosities
        """
        super().__init__(component=component, name='analysis', indir=indir)
        self.nus = nus

            filename='mesh.nc',
            target='../../init/culled_mesh.nc')

            target='../../init/initial_state.nc')

                filename=f'output_nu_{nu:g}.nc',
                target=f'../nu_{nu:g}/output.nc')

            filename=f'sections_baroclinic_channel_{resolution}.png')
        self.add_output_file(filename='rpe_t.png')
        self.add_output_file(filename='rpe.csv')

        """
        """
        output_filename = self.outputs[0]
        nus = self.nus
        section = self.config['baroclinic_channel_rpe']

        rpe = compute_rpe(mesh_filename=mesh_filename,
                          initial_state_filename=init_filename,
                          output_filenames=self.inputs[2:])

        sim_count = len(nus)
        time = section.getfloat('plot_time')
        min_temp = section.getfloat('min_temp')
        max_temp = section.getfloat('max_temp')

        ds = xr.open_dataset(f'output_nu_{nus[0]:g}.nc', decode_times=False)
        times = ds.daysSinceStartOfSim.values

        for i in range(sim_count):
            rpe_norm = np.divide((rpe[i, :] - rpe[i, 0]), rpe[i, 0])
            plt.plot(times, rpe_norm, label=f'$\\nu_h=${nus[i]}')
        plt.xlabel('Time, days')
        plt.ylabel('RPE-RPE(0)/RPE(0)')
        plt.legend()
        plt.savefig('rpe_t.png')
        plt.close(fig)

        fig, axes = plt.subplots(1, sim_count, figsize=(
            3 * sim_count, 5.0), constrained_layout=True)

        for row_index, nu in enumerate(nus):
            ax = axes[row_index]
            ds = xr.open_dataset(f'output_nu_{nu:g}.nc', decode_times=False)
            ds = ds.isel(nVertLevels=0)
            times = ds.daysSinceStartOfSim.values
            time_index = np.argmin(np.abs(times - time))

            cell_mask = ds_init.maxLevelCell >= 1
            plot_horiz_field(ds_mesh, ds['temperature'], ax=ax,
                             cmap='cmo.thermal', t_index=time_index,
                             vmin=min_temp, vmax=max_temp,
                             cmap_title='SST (C)', field_mask=cell_mask)
            ax.set_title(f'day {times[time_index]:g}, $\\nu_h=${nu:g}')

        plt.savefig(output_filename)

class Validate(Step):
    """
    A step for comparing outputs between steps in a baroclinic channel run

    step_subdirs : list of str
        The number of processors used in each run
    """
    def __init__(self, component, step_subdirs, indir):
        """

        step_subdirs : list of str
            The number of processors used in each run

        indir : str
       """
        super().__init__(component=component, name='validate', indir=indir)

        self.step_subdirs = step_subdirs

        for subdir in step_subdirs:
            self.add_input_file(filename=f'output_{subdir}.nc',
                                target=f'../{subdir}/output.nc')

        """
        Compare ``temperature``, ``salinity``, ``layerThickness`` and
        ``normalVelocity`` in the outputs of two previous steps with each other
        """
        step_subdirs = self.step_subdirs
        variables = ['temperature', 'salinity', 'layerThickness',
                     'normalVelocity']
        all_pass = compare_variables(variables=variables,
                                     filename1=self.inputs[0],
                                     filename2=self.inputs[1],
                                     logger=self.logger)
        if not all_pass:
            raise ValueError(f'Validation failed comparing outputs between '
                             f'{step_subdirs[0]} and {step_subdirs[1]}.')

class Viz(Step):
    """
    """
    def __init__(self, component, indir):
        """

        indir : str
        """
        super().__init__(component=component, name='viz', indir=indir)
            target='../forward/output.nc')

        """
        """
        ds_mesh = xr.load_dataset('mesh.nc')
        ds_init = xr.load_dataset('init.nc')
        ds = xr.load_dataset('output.nc')
        t_index = ds.sizes['Time'] - 1
        cell_mask = ds_init.maxLevelCell >= 1
        edge_mask = cell_mask_to_edge_mask(ds_init, cell_mask)
        max_velocity = np.max(np.abs(ds.normalVelocity.values))
                         'final_normalVelocity.png',
                         t_index=t_index,
                         vmin=-max_velocity, vmax=max_velocity,
                         cmap='cmo.balance', show_patch_edges=True,
                         field_mask=edge_mask)

            layer_thickness=ds.layerThickness.isel(Time=t_index),
            bottom_depth=ds_init.bottomDepth,
            min_level_cell=ds_init.minLevelCell - 1,
            max_level_cell=ds_init.maxLevelCell - 1,
            spherical=False)

        mpas_field = ds[field_name].isel(Time=t_index)
        plot_transect(ds_transect=ds_transect, mpas_field=mpas_field,
                      out_filename=f'final_{field_name}_section.png',

        plot_horiz_field(ds_mesh, ds['temperature'], 'final_temperature.png',
                         t_index=t_index, vmin=vmin, vmax=vmax,
                         cmap='cmo.thermal', field_mask=cell_mask,
                         transect_x=x, transect_y=y)

from polaris.ocean.tasks.barotropic_gyre.analysis import Analysis
from polaris.ocean.tasks.barotropic_gyre.forward import Forward
from polaris.ocean.tasks.barotropic_gyre.init import Init

def add_barotropic_gyre_tasks(component):
    """
    Add a task that defines a convergence test for inertial gravity waves

        the ocean component that the task will be added to
    """
    test_name = 'default'
    component.add_task(BarotropicGyre(component=component,
                                      test_name=test_name))

class BarotropicGyre(Task):
    """
    The convergence test case for inertial gravity waves
    """

    def __init__(self, component, test_name):
        """

        """
        group_name = 'barotropic_gyre'
        name = f'{group_name}_{test_name}'
        subdir = f'planar/{group_name}/{test_name}'

        config_filename = f'{group_name}.cfg'
        config.filepath = f'{subdir}/{config_filename}'
        config.add_from_package(f'polaris.ocean.tasks.{group_name}',
                                config_filename)
        self.set_shared_config(config, link=config_filename)

        init = Init(component=component, subdir=subdir)
        init.set_shared_config(config, link=config_filename)
        self.add_step(init)

        forward = Forward(component=component, indir=self.subdir, ntasks=None,
                          min_tasks=None, openmp_threads=1,
                          name='short_forward', run_time_steps=3,
                          graph_target=f'{init.path}/culled_graph.info')
        forward.set_shared_config(config, link=config_filename)
        self.add_step(forward)

                          name='long_forward',
        self.add_step(forward, run_by_default=False)

        analysis = Analysis(component=component, indir=self.subdir,
                            boundary_condition='free slip')
        analysis.set_shared_config(config, link=config_filename)
        self.add_step(analysis, run_by_default=False)

from math import pi

from matplotlib import colors as mcolors
from mpas_tools.ocean import compute_barotropic_streamfunction

from polaris.mpas import area_for_field

class Analysis(OceanIOStep):
    """
    A step for analysing the output from the barotropic gyre
    test case
    """

    def __init__(self, component, indir, boundary_condition='free slip'):
        """

        indir : str
        """
            target='../init/culled_mesh.nc')
            target='../init/init.nc')
            target='../long_forward/output.nc')
        self.boundary_condition = boundary_condition

        ds_mesh = xr.open_dataset('mesh.nc')
        ds_init = xr.open_dataset('init.nc')
        ds = xr.open_dataset('output.nc')

        field_mpas = compute_barotropic_streamfunction(
            ds_init.isel(Time=0), ds, prefix='', time_index=-1)
        field_exact = self.exact_solution(
            ds_mesh, self.config, loc='Vertex',
            boundary_condition=self.boundary_condition)

        ds['psi'] = field_mpas
        ds['psi_exact'] = field_exact
        ds['psi_error'] = field_mpas - field_exact

        error = self.compute_error(ds_mesh, ds, variable_name='psi',
                                   boundary_condition=self.boundary_condition)
        print(f'L2 error norm for {self.boundary_condition} bsf: {error:1.2e}')

        descriptor = mosaic.Descriptor(ds_mesh)

        pad = 20
        fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(12, 2))
        x0 = ds_mesh.xEdge.min().values
        y0 = ds_mesh.yEdge.min().values

        # offset coordinates
        descriptor.vertex_patches[..., 0] -= x0
        descriptor.vertex_patches[..., 1] -= y0
        # convert to km
        descriptor.vertex_patches *= 1.e-3

        eta0 = max(np.max(np.abs(field_exact.values)),
                   np.max(np.abs(field_mpas.values)))

        bounds = np.linspace(-eta0, eta0, 21)
        norm = mcolors.BoundaryNorm(bounds, cmocean.cm.amp.N)
        s = mosaic.polypcolor(axes[0], descriptor, field_mpas,
                              cmap='cmo.balance', norm=norm, antialiased=False)
        cbar = fig.colorbar(s, ax=axes[0])
        cbar.ax.set_title(r'$\psi$')
        s = mosaic.polypcolor(axes[1], descriptor, field_exact,
        cbar = fig.colorbar(s, ax=axes[1])

        eta0 = np.max(np.abs(field_mpas.values - field_exact.values))
        bounds = np.linspace(-eta0, eta0, 11)
        norm = mcolors.BoundaryNorm(bounds, cmocean.cm.balance.N)
        s = mosaic.polypcolor(axes[2], descriptor, field_mpas - field_exact,
        cbar = fig.colorbar(s, ax=axes[2])
        cbar.ax.set_title(r'$d\psi$')
        axes[0].set_title('Numerical solution', pad=pad)
        axes[0].set_ylabel('y (km)')
        axes[0].set_xlabel('x (km)')
        axes[1].set_title('Analytical solution', pad=pad)
        axes[1].set_xlabel('x (km)')
        axes[2].set_title('Error (Numerical - Analytical)', pad=pad)
        axes[2].set_xlabel('x (km)')

        xmin = descriptor.vertex_patches[..., 0].min()
        xmax = descriptor.vertex_patches[..., 0].max()
        ymin = descriptor.vertex_patches[..., 1].min()
        ymax = descriptor.vertex_patches[..., 1].max()
        for ax in axes:
            ax.set_xlim(xmin, xmax)
            ax.set_ylim(ymin, ymax)
            ax.set_aspect('equal')
        fig.savefig('comparison.png', bbox_inches='tight', pad_inches=0.1)

    def compute_error(self, ds_mesh, ds_out, variable_name, error_type='l2',
                      loc='Vertex', boundary_condition='free slip'):
        """

        ds_mesh: xarray.Dataset
            the mesh dataset

        ds_out: xarray.Dataset
            the output dataset

        Returns
        -------
        """

            ds_mesh, self.config, loc=loc,
        ds_out = ds_out.isel(Time=-1)
        field_mpas = ds_out[variable_name]

    def exact_solution(self, ds_mesh, config, loc='Cell',
                       boundary_condition='free slip'):
        """
        Exact solution to the sea surface height for the linearized Munk layer
        experiments.

        ds_mesh : xarray.Dataset
            Must contain the fields: f'x{loc}', f'y{loc}'
        """

        x = ds_mesh[f'x{loc}']
        x = x - ds_mesh.xEdge.min()
        y = ds_mesh[f'y{loc}']
        y = y - ds_mesh.yEdge.min()
        L_x = float(x.max() - x.min())
        L_y = float(y.max() - y.min())

        # df/dy where f is coriolis parameter
        beta = config.getfloat("barotropic_gyre", "beta")
        # Laplacian viscosity
        nu = config.getfloat("barotropic_gyre", "nu_2")

        # Compute some non-dimensional numbers
        delta_m = (nu / (beta * L_y**3.))**(1. / 3.)
        gamma = (np.sqrt(3.) * x) / (2. * delta_m * L_x)

        if boundary_condition == 'no slip':
            psi = (pi * np.sin(pi * y / L_y) *
                   (1. - (x / L_x) -
                    np.exp(-x / (2. * delta_m * L_x)) *
                    (np.cos(gamma) +
                     ((1. - 2 * delta_m) / np.sqrt(3.)) * np.sin(gamma)) +
                    delta_m * np.exp(((x / L_x) - 1) / delta_m)))

        elif boundary_condition == 'free slip':
            psi = (pi * (1. - x / L_x) * np.sin(pi * y / L_y) *
                   (1. -
                    (np.cos(gamma) + (1. / np.sqrt(3.)) * np.sin(gamma))))
        return psi

from math import ceil, floor

    """
    A step for performing forward ocean component runs as part of barotropic
    gyre tasks.

    dt : float

    """
    def __init__(self, component, name='forward', subdir=None,
        """

        name : str

        run_time_steps : int or None

        """

                            target='../init/init.nc')
        self.add_input_file(filename='forcing.nc',
                            target='../init/forcing.nc')

            validate_vars=['layerThickness', 'normalVelocity'])

        self.package = 'polaris.ocean.tasks.barotropic_gyre'
        self.yaml_filename = 'forward.yaml'

        """
        resources

        Returns
        -------
        """
        section = self.config['barotropic_gyre']
        resolution = section.getfloat('resolution')

        """

        """

        if model == 'mpas-ocean':

        rho_0 = config.getfloat("barotropic_gyre", "rho_0")

        dt_max = compute_max_time_step(config)
        dt = floor(dt_max / 3.)
        dt_str = get_time_interval_string(seconds=dt)

        options = {'config_dt': dt_str,
                   'config_density0': rho_0}

            run_duration = ceil(self.run_time_steps * dt)
            stop_time_str = time.strftime('0001-01-01_%H:%M:%S',
                                          time.gmtime(run_duration))
            output_interval_str = time.strftime('0000_%H:%M:%S',
                                                time.gmtime(run_duration))
        else:
            stop_time_str = time.strftime('0001-04-01_00:00:00')
            output_interval_str = time.strftime('0000-01-00_00:00:00')

            stop_time=stop_time_str,
            nu=f'{nu:02g}',
        )

def compute_max_time_step(config):
    """
    Compute the approximate maximum time step for stability

        Config options for test case

    Returns
    -------
    dt_max : float
        The approximate maximum time step for stability
    """
    u_max = 1.  # m/s
    stability_parameter_max = 0.25
    resolution = config.getfloat('barotropic_gyre', 'resolution')
    f_0 = config.getfloat("barotropic_gyre", "f_0")
    dt_max = min(stability_parameter_max * resolution * 1e3 /
                 (2 * u_max),
                 stability_parameter_max / f_0)
    return dt_max

from polaris.ocean.tasks.barotropic_gyre.forward import compute_max_time_step

    """
    tasks

    """
    def __init__(self, component, subdir):
        """

        """
        super().__init__(component=component, name='init', indir=subdir)

        for file in ['base_mesh.nc', 'culled_mesh.nc', 'culled_graph.info',
                     'forcing.nc']:
        self.add_output_file('init.nc', validate_vars=['layerThickness'])

        resolution = config.getfloat("barotropic_gyre", "resolution")
        # coriolis parameter
        # laplacian horizontal viscosity
        nu_2 = config.getfloat("barotropic_gyre", "nu_2")

        # calculate the boundary layer thickness for specified parameters
        m = (np.pi * 2) / np.sqrt(3) * (nu_2 / beta)**(1. / 3.)
        # ensure the boundary layer is at least 3 gridcells wide
        if m >= 3. * resolution * 1.e3:
            raise ValueError(f"Resolution {resolution} km is too coarse to "
                             "properly resolve the lateral boundary layer "
                             f"with anticipated width of {(m * 1e-3):03g} km")

        # check whether viscosity suitable for stability
        stability_parameter_max = 0.6
        dt = dt_max / 3.
        nu_2_max = stability_parameter_max * (resolution * 1.e3)**2. / \
            (8 * dt)
        if nu_2 > nu_2_max:
            raise ValueError(f'Laplacian viscosity cannot be set to {nu_2}; '
                             f'maximum value is {nu_2_max}')

        """Create the at rest inital condition for the barotropic gyre testcase
        """
        # domain parameters
        lx = config.getfloat("barotropic_gyre", "lx")
        ly = config.getfloat("barotropic_gyre", "ly")

        # convert cell spacing to meters
        dc = resolution * 1e3

        ds_mesh = make_planar_hex_mesh(
            nx=nx, ny=ny, dc=dc, nonperiodic_x=True, nonperiodic_y=True)

        # vertical coordinate parameters
        # coriolis parameters
        f0 = config.getfloat("barotropic_gyre", "f_0")
        # surface (wind) forcing parameters
        tau_0 = config.getfloat("barotropic_gyre", "tau_0")

        # create a copy of the culled mesh to place the IC's into

        # set the ssh initial condition to zero
        ds["ssh"] = xr.zeros_like(ds.xCell)
        ds['bottomDepth'] = bottom_depth * xr.ones_like(ds.xCell)

        # use polaris framework functions to initialize the vertical coordinate

        # set the coriolis values
        for loc in ["Cell", "Edge", "Vertex"]:
            ds[f"f{loc}"] = f0 + beta * ds[f"y{loc}"]

        # set the initial condition for normalVelocity
        normal_velocity, _ = xr.broadcast(
            xr.zeros_like(ds_mesh.xEdge), ds.refBottomDepth)

        # write the initial condition file
        write_netcdf(ds, 'init.nc')

        # set the wind stress forcing
        ds_forcing = xr.Dataset()
        # Convert from km to m
        ly = ly * 1e3
        wind_stress_zonal = -tau_0 * \
            np.cos(np.pi * (ds.yCell - ds.yCell.min()) / ly)
        wind_stress_meridional = xr.zeros_like(ds.xCell)
        ds_forcing["windStressZonal"] = \
            wind_stress_zonal.expand_dims(dim='Time', axis=0)
        ds_forcing["windStressMeridional"] = \
            wind_stress_meridional.expand_dims(dim='Time', axis=0)
        write_netcdf(ds_forcing, 'forcing.nc')

        plot_horiz_field(ds_mesh, ds_forcing['windStressZonal'],
                         'forcing_wind_stress_zonal.png', cmap='cmo.balance',
                         show_patch_edges=True, field_mask=cell_mask,
                         vmin=-0.1, vmax=0.1)

from math import ceil

)
from polaris.ocean.mesh.spherical import add_spherical_base_mesh_step
from polaris.ocean.tasks.cosine_bell.analysis import Analysis
from polaris.ocean.tasks.cosine_bell.forward import Forward
from polaris.ocean.tasks.cosine_bell.init import Init
from polaris.ocean.tasks.cosine_bell.viz import Viz

def add_cosine_bell_tasks(component):
    """
    Add tasks that define variants of the cosine bell test case

        the ocean component that the tasks will be added to
    """

    for icosahedral, prefix in [(True, 'icos'), (False, 'qu')]:

        filepath = f'spherical/{prefix}/cosine_bell/cosine_bell.cfg'
        config = PolarisConfigParser(filepath=filepath)
        config.add_from_package('polaris.ocean.convergence',
                                'convergence.cfg')
        config.add_from_package('polaris.ocean.convergence.spherical',
                                'spherical.cfg')
        config.add_from_package('polaris.ocean.tasks.cosine_bell',
                                'cosine_bell.cfg')

        for refinement in ['space', 'time', 'both']:
            for include_viz in [False, True]:
                component.add_task(CosineBell(component=component,
                                              config=config,
                                              icosahedral=icosahedral,
                                              include_viz=include_viz,
                                              refinement=refinement))

class CosineBell(Task):
    """
    A convergence test for the advection of a cosine-bell tracer

    icosahedral : bool
        Whether to use icosahedral, as opposed to less regular, JIGSAW meshes

    include_viz : bool
        Include VizMap and Viz steps for each resolution
    """
    def __init__(self, component, config, icosahedral, include_viz,
        """
        Create the convergence test

        icosahedral : bool
            Whether to use icosahedral, as opposed to less regular, JIGSAW
            meshes

        include_viz : bool
            Include VizMap and Viz steps for each resolution

        """
        if icosahedral:
            prefix = 'icos'
        else:
            prefix = 'qu'

        subdir = f'spherical/{prefix}/cosine_bell/convergence_{refinement}'
        name = f'{prefix}_cosine_bell_convergence_{refinement}'
        if include_viz:
            subdir = f'{subdir}/with_viz'
            name = f'{name}_with_viz'
        link = 'cosine_bell.cfg'
        self.icosahedral = icosahedral
        self.include_viz = include_viz

        self.set_shared_config(config, link=link)

        self._setup_steps(refinement=refinement)

        """
        Set config options for the test case
        """

        # set up the steps again in case a user has provided new resolutions
        self._setup_steps(self.refinement)

    def _setup_steps(self, refinement):
        """
        setup steps given resolutions

        refinement : str
           refinement type. One of 'space', 'time' or 'both'
        """
        icosahedral = self.icosahedral
        config_filename = self.config_filename

        else:

        if refinement == 'time':
        else:
        refinement_factors = config.getlist('spherical_convergence',
                                            f'{prefix}_{option}', dtype=str)
        refinement_factors = ', '.join(refinement_factors)
        config.set('convergence', option, value=refinement_factors)
        refinement_factors = config.getlist('convergence',
                                            option, dtype=float)

        base_resolution = config.getfloat('spherical_convergence',
                                          f'{prefix}_base_resolution')
        config.set('convergence', 'base_resolution',
                   value=f'{base_resolution:03g}')

        # start fresh with no steps
        for step in list(self.steps.values()):
            self.remove_step(step)

        analysis_dependencies: Dict[str, Dict[str, Step]] = (
            dict(mesh=dict(), init=dict(), forward=dict()))

        case_dir = f'spherical/{prefix}/cosine_bell'

                config, refinement_factor, refinement=refinement)

            base_mesh_step, mesh_name = add_spherical_base_mesh_step(
                component, resolution, icosahedral)
            analysis_dependencies['mesh'][refinement_factor] = base_mesh_step

            name = f'{prefix}_init_{mesh_name}'
            subdir = f'{case_dir}/init/{mesh_name}'
                init_step = component.steps[subdir]
            else:
                init_step = Init(component=component, name=name,
                                 subdir=subdir, base_mesh=base_mesh_step)
                init_step.set_shared_config(config, link=config_filename)
            analysis_dependencies['init'][refinement_factor] = init_step

            if resolution not in resolutions:
                self.add_step(base_mesh_step, symlink=f'base_mesh/{mesh_name}')
                self.add_step(init_step, symlink=f'init/{mesh_name}')
                resolutions.append(resolution)

            timestep = ceil(timestep)

            subdir = f'{case_dir}/forward/{mesh_name}_{timestep}s'
            symlink = f'forward/{mesh_name}_{timestep}s'
                forward_step = component.steps[subdir]
            else:
                name = f'{prefix}_forward_{mesh_name}_{timestep}s'
                forward_step = Forward(
                    component=component, name=name,
                    subdir=subdir,
                    refinement_factor=refinement_factor,
                    mesh=base_mesh_step,
                    init=init_step, refinement=refinement)
                forward_step.set_shared_config(
                    config, link=config_filename)
            self.add_step(forward_step, symlink=symlink)
            analysis_dependencies['forward'][refinement_factor] = forward_step

            if self.include_viz:
                name = f'{prefix}_viz_{mesh_name}_{timestep}s'
                subdir = f'{case_dir}/viz/{mesh_name}_{timestep}s'
                if subdir in component.steps:
                    viz_step = component.steps[subdir]
                else:
                    viz_step = Viz(component=component, name=name,
                                   subdir=subdir, base_mesh=base_mesh_step,
                                   init=init_step, forward=forward_step,
                                   mesh_name=mesh_name)
                    viz_step.set_shared_config(config, link=config_filename)
                self.add_step(viz_step)

        subdir = f'{case_dir}/convergence_{refinement}/analysis'
        if subdir in component.steps:
            step.resolutions = resolutions
            step.dependencies_dict = analysis_dependencies
        else:
            step = Analysis(component=component,
                            subdir=subdir,
                            dependencies=analysis_dependencies,
                            refinement=refinement)
            step.set_shared_config(config, link=config_filename)
        if self.include_viz:
            symlink = f'analysis_{refinement}'
            self.add_step(step, symlink=symlink)
        else:

from mpas_tools.transects import lon_lat_to_cartesian
from mpas_tools.vector import Vector

from polaris.ocean.convergence.analysis import ConvergenceAnalysis
from polaris.ocean.tasks.cosine_bell.init import cosine_bell

class Analysis(ConvergenceAnalysis):
    """
    A step for analyzing the output from the cosine bell test case
    """
    def __init__(self, component, subdir, dependencies, refinement='both'):
        """

            The dependencies of this step

            Whether to refine in space, time or both space and time
        """
        convergence_vars = [{'name': 'tracer1',
                             'title': 'tracer1',
                             'zidx': 0}]
        super().__init__(component=component, subdir=subdir,
                         dependencies=dependencies,
                         convergence_vars=convergence_vars,
                         refinement=refinement)

        """

            The factor by which to scale space, time or both

            fields

            solution

        Returns
        -------
            The exact solution with dimension nCells
        """

        if field_name != 'tracer1':
            print(f'Variable {field_name} not available as an analytic '
                  'solution for the cosine_bell test case')

        lat_center = config.getfloat('cosine_bell', 'lat_center')
        lon_center = config.getfloat('cosine_bell', 'lon_center')
        radius = config.getfloat('cosine_bell', 'radius')
        psi0 = config.getfloat('cosine_bell', 'psi0')
        vel_pd = config.getfloat('cosine_bell', 'vel_pd')

        ds_mesh = xr.open_dataset(f'mesh_r{refinement_factor:02g}.nc')
        sphere_radius = ds_mesh.sphere_radius

        latCell = ds_init.latCell.values
        lonCell = ds_init.lonCell.values

        # distance that the cosine bell center traveled in radians
        # based on equatorial velocity
        distance = 2.0 * np.pi * time / (3600.0 * vel_pd)

        # new location of blob center
        lon_new = lon_center + distance
        if lon_new > 2.0 * np.pi:
            lon_new -= 2.0 * np.pi

        x_center, y_center, z_center = lon_lat_to_cartesian(
            lon_new, lat_center, sphere_radius, degrees=False)
        x_cells, y_cells, z_cells = lon_lat_to_cartesian(
            lonCell, latCell, sphere_radius, degrees=False)
        xyz_center = Vector(x_center, y_center, z_center)
        xyz_cells = Vector(x_cells, y_cells, z_cells)
        ang_dist_from_center = xyz_cells.angular_distance(xyz_center)
        distance_from_center = ang_dist_from_center * sphere_radius

        bell_value = cosine_bell(psi0, distance_from_center, radius)
        tracer1 = np.where(distance_from_center < radius,
                           bell_value,
                           0.0)
        return xr.DataArray(data=tracer1, dims=('nCells',))

class Init(OceanIOStep):
    """
    A step for an initial condition for for the cosine bell test case
    """
    def __init__(self, component, name, subdir, base_mesh):
        """

        name : str

        base_mesh : polaris.Step
            The base mesh step
        """

            work_dir_target=f'{base_mesh.path}/base_mesh.nc')

        self.add_output_file(filename='initial_state.nc')

        """
        """

        section = config['cosine_bell']
        temperature = section.getfloat('temperature')
        lat_center = section.getfloat('lat_center')
        lon_center = section.getfloat('lon_center')
        radius = section.getfloat('radius')
        psi0 = section.getfloat('psi0')
        vel_pd = section.getfloat('vel_pd')

        section = config['vertical_grid']
        bottom_depth = section.getfloat('bottom_depth')

        angleEdge = ds_mesh.angleEdge
        latCell = ds_mesh.latCell
        latEdge = ds_mesh.latEdge
        lonCell = ds_mesh.lonCell

        ds['bottomDepth'] = bottom_depth * xr.ones_like(latCell)
        ds['ssh'] = xr.zeros_like(latCell)

        temperature_array = temperature * xr.ones_like(ds_mesh.latCell)
        temperature_array, _ = xr.broadcast(temperature_array, ds.refZMid)
        ds['temperature'] = temperature_array.expand_dims(dim='Time', axis=0)
        ds['salinity'] = salinity * xr.ones_like(ds.temperature)

            lon_center, lat_center, sphere_radius, degrees=False)
        debug_tracers = xr.where(distance_from_center < radius,
                                 bell_value,
                                 0.0)
        debug_tracers_array, _ = xr.broadcast(debug_tracers, ds.refZMid)
        ds['tracer1'] = debug_tracers_array.expand_dims(dim='Time', axis=0)
        ds['tracer2'] = ds.tracer1
        ds['tracer3'] = ds.tracer1

        # Initialize velocity
        velocity = (2.0 * np.pi * np.cos(angleEdge) * sphere_radius *
                    np.cos(latEdge) / (s_per_hour * vel_pd))
        velocity_array, _ = xr.broadcast(velocity, ds.refZMid)
        ds['normalVelocity'] = velocity_array.expand_dims(dim='Time', axis=0)

        ds['fCell'] = xr.zeros_like(ds_mesh.xCell)
        ds['fEdge'] = xr.zeros_like(ds_mesh.xEdge)
        ds['fVertex'] = xr.zeros_like(ds_mesh.xVertex)

        self.write_model_dataset(ds, 'initial_state.nc')

def cosine_bell(max_value, ri, r):
    """
    Compute values according to cosine bell function

    max_value : float
        Maximum value of the cosine bell function

    ri : np.ndarray of type float
        Distance from the center of the cosine bell in meters

    r : float
        Radius of the cosine bell in meters

    Returns
    -------
    f : np.ndarray of type float
        Cosine bell tracer values
    """
    return max_value / 2.0 * (1.0 + np.cos(np.pi * np.divide(ri, r)))

from polaris.viz import plot_global_mpas_field

class Viz(OceanIOStep):
    """
    A step for plotting fields from the cosine bell output

    mesh_name : str
        The name of the mesh
    """
    def __init__(self, component, name, subdir, base_mesh, init, forward,
                 mesh_name):
        """

        name : str

            The subdirectory in the test case's work directory for the step

        init : polaris.Step
            The init step

        forward : polaris.Step

            The name of the mesh
        """
            filename='initial_state.nc',
            work_dir_target=f'{forward.path}/output.nc')
        self.mesh_name = mesh_name
        self.add_output_file('init.png')
        self.add_output_file('final.png')

        """
        """
        mesh_name = self.mesh_name
        run_duration = config.getfloat('convergence_forward', 'run_duration')

        ds_init = self.open_model_dataset('initial_state.nc')
        da = ds_init['tracer1'].isel(Time=0, nVertLevels=0)

        plot_global_mpas_field(
            mesh_filename='mesh.nc', da=da, out_filename='init.png',
            config=config, colormap_section='cosine_bell_viz',
            title=f'{mesh_name} tracer at init', plot_land=False,
            central_longitude=180.)

        ds_out = self.open_model_dataset('output.nc')
        da = ds_out['tracer1'].isel(Time=-1, nVertLevels=0)

            mesh_filename='mesh.nc', da=da, out_filename='final.png',
            title=f'{mesh_name} tracer after {run_duration / 24.:g} days',
            plot_land=False, central_longitude=180.)

)
from polaris.ocean.tasks.geostrophic.analysis import Analysis
from polaris.ocean.tasks.geostrophic.forward import Forward
from polaris.ocean.tasks.geostrophic.init import Init
from polaris.ocean.tasks.geostrophic.viz import Viz

def add_geostrophic_tasks(component):
    """
    Add tasks that define variants of the geostrophic test

    """

        filepath = f'spherical/{prefix}/geostrophic/geostrophic.cfg'
        config.add_from_package('polaris.ocean.tasks.geostrophic',
                                'geostrophic.cfg')

                component.add_task(Geostrophic(component=component,
                                               config=config,
                                               icosahedral=icosahedral,
                                               include_viz=include_viz,
                                               refinement=refinement))

class Geostrophic(Task):
    """
    A convergence test for a configuration in geostrophic balance

    """
        """

            meshes

        """
        else:

        subdir = f'spherical/{prefix}/geostrophic/convergence_{refinement}'
        name = f'{prefix}_geostrophic_convergence_{refinement}'
        link = 'geostrophic.cfg'

        """
        """

        """

        """

        else:

        else:

        case_dir = f'spherical/{prefix}/geostrophic'

        analysis_dependencies: Dict[str, Dict[float, Step]] = (

            else:

            else:

                else:

        else:
        else:

from polaris.ocean.tasks.geostrophic.exact_solution import (
    compute_exact_solution,
)

    """
    A step for analyzing the output from the geostrophic test case
    """
        """

        """
        convergence_vars = [{'name': 'h',
                             'title': 'water-column thickness',
                             'zidx': None},
                            {'name': 'normalVelocity',
                             'title': 'normal velocity',

        """

            fields

            solution

        Returns
        -------
        solution: xarray.DataArray
        """

        if field_name not in ['h', 'normalVelocity']:
                  'solution for the geostrophic test case')

        section = config['geostrophic']
        alpha = section.getfloat('alpha')
        vel_period = section.getfloat('vel_period')
        gh_0 = section.getfloat('gh_0')

        mesh_filename = f'mesh_r{refinement_factor:02g}.nc'

        h, _, _, normalVelocity = compute_exact_solution(
            alpha, vel_period, gh_0, mesh_filename)

        if field_name == 'h':
            return h
        else:
            return normalVelocity

        """

        Returns
        -------
        """

            print(f'Variable {field_name} not available for analysis in the '
                  f'geostrophic test case')

        if field_name == 'normalVelocity':
            return super().get_output_field(
                field_name=field_name, time=time, zidx=zidx)
        else:
            ds_init = xr.open_dataset(f'init_r{refinement_factor:02g}.nc')
            bottom_depth = ds_init.bottomDepth
            ssh = super().get_output_field(
                field_name='ssh', time=time, zidx=None)
            h = ssh + bottom_depth
            return h

        """

            For cosine_bell, we use the same convergence rate for all fields
        Returns
        -------
        conv_thresh: float

            The maximum convergence rate
        """
        conv_thresh = config.getfloat('geostrophic',
                                      f'convergence_thresh_{field_name}')
        error_type = config.get('convergence', 'error_type')

)

    """
    A step for an initial condition for for the geostrophic test case
    """
        """

        name : str

        """

            filename='graph.info',
            work_dir_target=f'{base_mesh.path}/graph.info')

        self.add_output_file(filename='initial_state.nc',
                             validate_vars=['temperature',
                                            'salinity',
                                            'layerThickness',
                                            'normalVelocity'])

        """
        Run this step of the testcase
        """

        h, u_cell, v_cell, normalVelocity = compute_exact_solution(

        omega = 2 * np.pi / constants['SHR_CONST_SDAY']

        lonEdge = ds_mesh.lonEdge
        latVertex = ds_mesh.latVertex
        lonVertex = ds_mesh.lonVertex

        ds['ssh'] = -ds.bottomDepth + h

        salinity_array = salinity * xr.ones_like(temperature_array)

        normalVelocity, _ = xr.broadcast(normalVelocity, ds.refZMid)

        ds['salinity'] = salinity_array.expand_dims(dim='Time', axis=0)
        ds['normalVelocity'] = normalVelocity.expand_dims(dim='Time', axis=0)

        ds['fCell'] = _coriolis(lonCell, latCell, alpha, omega)
        ds['fEdge'] = _coriolis(lonEdge, latEdge, alpha, omega)
        ds['fVertex'] = _coriolis(lonVertex, latVertex, alpha, omega)

        # for visualization
        ds['velocityZonal'] = u_cell.broadcast_like(ds.temperature)
        ds['velocityMeridional'] = v_cell.broadcast_like(ds.temperature)

def _coriolis(lon, lat, alpha, omega):
    f = 2 * omega * (-np.cos(lon) * np.cos(lat) * np.sin(alpha) +
                     np.sin(lat) * np.cos(alpha))
    return f

    """

    """
        """

        name : str

        """

        """
        """
        run_duration = \
            config.getfloat('convergence_forward', 'run_duration') / 24.

        colormap_sections = dict(
            h='geostrophic_viz_h',
            u='geostrophic_viz_vel',
            v='geostrophic_viz_vel')

        ds_init = xr.open_dataset('initial_state.nc')
        bottom_depth = ds_init.bottomDepth
        ds_init = self._process_ds(ds_init, bottom_depth, time_index=0)
        ds_init.to_netcdf('remapped_init.nc')

        for var, colormap_section in colormap_sections.items():
            plot_global_mpas_field(
                mesh_filename='mesh.nc',
                da=ds_init[var],
                out_filename=f'init_{var}.png',
                config=config,
                colormap_section=colormap_section,
                colorbar_label=config.get(colormap_section, 'label'),
                title=f'{mesh_name} {var} at init',
                plot_land=False)

        ds_out = xr.open_dataset('output.nc')
        ds_out = self._process_ds(ds_out, bottom_depth, time_index=-1)
        ds_out.to_netcdf('remapped_final.nc')

                da=ds_out[var],
                out_filename=f'final_{var}.png',
                title=f'{mesh_name} {var} after {run_duration:g} days',

            h='geostrophic_viz_diff_h',
            u='geostrophic_viz_diff_vel',
            v='geostrophic_viz_diff_vel')

            diff = ds_out[var] - ds_init[var]
                da=diff,
                out_filename=f'diff_{var}.png',
                title=f'{mesh_name} {var} change after {run_duration:g} days',

    def _process_ds(ds, bottom_depth, time_index):
        ds_out = ds.isel(Time=time_index, nVertLevels=0)
        ds_out['h'] = ds_out.ssh + bottom_depth
        ds_out = ds_out.rename(dict(velocityZonal='u',
                                    velocityMeridional='v',))
        return ds_out

from math import floor

    """

    """
    def __init__(self, component, resolution, mesh, init,
                 name='forward', subdir=None, indir=None,
                 ntasks=None, min_tasks=None, openmp_threads=1,
                 do_restart=False, tidal_forcing=False):
        """

        mesh: polaris.Step
            The step used to produce the mesh

        init: polaris.Step
            The step used to produce the initial condition

        do_restart : bool, optional
            if true, the step is restart from another forward run

        tidal_forcing : bool, optional
            if true, apply tidal forcing in the forward step
        """
        if do_restart:
            name = 'restart'
                         graph_target=f'{mesh.path}/culled_graph.info')

        self.do_restart = do_restart
        self.tidal_forcing = tidal_forcing

                            work_dir_target=f'{init.path}/output.nc')

        """
        resources

        Returns
        -------
        """
        section = self.config['ice_shelf_2d']

        """

        """

        if self.tidal_forcing:
            section = config['ice_shelf_2d_default_tidal_forcing']
            run_duration = section.getfloat('forward_run_duration')
            run_duration = run_duration * constants['SHR_CONST_CDAY']
        else:
            section = config['ice_shelf_2d_default']
            run_duration = run_duration * 60.

        else:

            seconds=btr_dt_per_km * self.resolution)

        do_restart_str = 'false'
        if self.do_restart:
            do_restart_str = 'true'

        # For all forward runs we ensure that the time step is a factor of the
        # output interval because we might want to use the forward step output
        # for a restart test
        output_interval = dt * floor(run_duration / (2. * dt))
            run_duration_str = get_time_interval_string(
                seconds=output_interval)
            output_interval_str = get_time_interval_string(
        else:
                seconds=output_interval * 2.)

        start_time = '0001-01-01_00:00:00'
            start_time = f"{start_time.split('_')[0]}_" \
                         f"{run_duration_str.split('_')[1]}"

            land_ice_flux_mode = 'pressure_only'
        else:
            land_ice_flux_mode = 'standalone'

            do_restart=do_restart_str,
            start_time=start_time,
            land_ice_flux_mode=land_ice_flux_mode,
        )
        self.add_yaml_file('polaris.ocean.tasks.ice_shelf_2d',
                           'forward.yaml',
        # We only do this at set-up so that the user may change the
        # tidal forcing parameters
        if self.tidal_forcing and at_setup:
            self.add_yaml_file('polaris.ocean.tasks.ice_shelf_2d',
                               'tidal_forcing.yaml')
        if not self.tidal_forcing:
                               'global_stats.yaml',
                               template_replacements=replacements)

        if vert_levels == 1:

    """
    A step for creating a mesh and initial condition for ice shelf 2-d tasks

    """
        """

        indir : str
        """

        self.add_output_file('output.nc',

        """
        """

        section = config['ice_shelf_2d']

        # TODO change cfg options

        surface_salinity = section.getfloat('surface_salinity')
        bottom_salinity = section.getfloat('bottom_salinity')

        # points 1 and 2 are where angles on ice shelf are located.
        # point 3 is at the surface.
        # d variables are total water-column thickness below ice shelf
        y1 = section.getfloat('y1') * 1e3
        y2 = section.getfloat('y2') * 1e3
        y3 = section.getfloat('y3') * 1e3
        d1 = section.getfloat('y1_water_column_thickness')
        d2 = section.getfloat('y2_water_column_thickness')
        d3 = bottom_depth

        # Column thickness is a piecewise linear function
        column_thickness = xr.where(
            y_cell < y1,
            d1,
            d1 + (d2 - d1) * (y_cell - y1) / (y2 - y1))
            y_cell < y2,
            column_thickness,
            d2 + (d3 - d2) * (y_cell - y2) / (y3 - y2))
            y_cell < y3,
            d3)

        ds['ssh'] = -bottom_depth + column_thickness

        modify_mask = xr.where(y_cell < y3, 1, 0).expand_dims(
            dim='Time', axis=0)
        landIceFraction = modify_mask.astype(float)
        landIceMask = modify_mask.copy()
        landIceFloatingFraction = landIceFraction.copy()
        landIceFloatingMask = landIceMask.copy()

        ref_density = constants['SHR_CONST_RHOSW']
        landIceDraft = ds.ssh
        landIcePressure = _compute_land_ice_pressure_from_draft(
            land_ice_draft=landIceDraft, modify_mask=modify_mask,
            ref_density=ref_density)

        salinity = surface_salinity + ((bottom_salinity - surface_salinity) *
                                       (ds.zMid / (-bottom_depth)))
        salinity, _ = xr.broadcast(salinity, ds.layerThickness)
        ds['salinity'] = salinity.transpose('Time', 'nCells', 'nVertLevels')

        ds['temperature'] = temperature * xr.ones_like(ds.salinity)

        mask = xr.zeros_like(ds_mesh.yCell)
        mask = np.logical_and(ds.maxLevelCell > 0,
                              ds_mesh.yCell < y3).astype(float)
        ds[mask_variable] = mask

        ds['modifyLandIcePressureMask'] = modify_mask
        ds['landIceFraction'] = landIceFraction
        ds['landIceFloatingFraction'] = landIceFloatingFraction
        ds['landIceMask'] = landIceMask
        ds['landIceFloatingMask'] = landIceFloatingMask
        ds['landIcePressure'] = landIcePressure
        ds['landIceDraft'] = landIceDraft

        write_netcdf(ds, 'output.nc')

        # Generate the tidal forcing dataset whether it is used or not
        y_max = np.max(ds.yCell.values)
        ds_forcing['tidalInputMask'] = xr.where(
            y_cell > (y_max - 0.6 * dc), 1.0, 0.0)
        write_netcdf(ds_forcing, 'init_mode_forcing_data.nc')

def _compute_land_ice_pressure_from_draft(land_ice_draft, modify_mask,
                                          ref_density=None):
    """
    Compute the pressure from an overlying ice shelf from ice draft

    land_ice_draft : xarray.DataArray
        The ice draft (sea surface height)

    modify_mask : xarray.DataArray
        A mask that is 1 where ``landIcePressure`` can be deviate from 0

    ref_density : float, optional
        A reference density for seawater displaced by the ice shelf

    Returns
    -------
    land_ice_pressure : xarray.DataArray
        The pressure from the overlying land ice on the ocean
    """
    gravity = constants['SHR_CONST_G']
    if ref_density is None:
    land_ice_pressure = \
        modify_mask * np.maximum(-ref_density * gravity * land_ice_draft, 0.)
    return land_ice_pressure

    """
    A step for comparing outputs between steps in a ice shelf 2d run

        The steps to be compared (full run and restart run)
    """
        """

            The steps to be compared (full run and restart run)

        indir : str
            the directory the step is in, to which ``validate`` will be
            appended
       """

            self.add_input_file(filename=f'land_ice_fluxes_{subdir}.nc',
                                target=f'../{subdir}/land_ice_fluxes.nc')
            self.add_input_file(filename=f'frazil_{subdir}.nc',
                                target=f'../{subdir}/frazil.nc')

        """
        """
        output_variables = ['temperature', 'salinity', 'layerThickness',
                            'normalVelocity']
        land_ice_variables = ['ssh', 'landIcePressure', 'landIceDraft',
                              'landIceFraction',
                              'landIceMask', 'landIceFrictionVelocity',
                              'topDrag', 'topDragMagnitude',
                              'landIceFreshwaterFlux', 'landIceHeatFlux',
                              'heatFluxToLandIce',
                              'landIceBoundaryLayerTemperature',
                              'landIceBoundaryLayerSalinity',
                              'landIceHeatTransferVelocity',
                              'landIceSaltTransferVelocity',
                              'landIceInterfaceTemperature',
                              'landIceInterfaceSalinity',
                              'accumulatedLandIceMass',
                              'accumulatedLandIceHeat']
        frazil_variables = ['accumulatedFrazilIceMass',
                            'accumulatedFrazilIceSalinity',
                            'seaIceEnergy', 'frazilLayerThicknessTendency',
                            'frazilTemperatureTendency',
                            'frazilSalinityTendency',
                            'frazilSurfacePressure',
                            'accumulatedLandIceFrazilMass']
        output_pass = compare_variables(variables=output_variables,
                                        filename1=self.inputs[0],
                                        filename2=self.inputs[3],
                                        logger=self.logger)
        land_ice_pass = compare_variables(variables=land_ice_variables,
                                          filename1=self.inputs[1],
                                          filename2=self.inputs[4],
                                          logger=self.logger)
        frazil_pass = compare_variables(variables=frazil_variables,
                                        filename1=self.inputs[2],
                                        filename2=self.inputs[5],
        if not output_pass:
        if not land_ice_pass:
            raise ValueError(f'Validation failed comparing land ice outputs '
                             f'between {step_subdirs[0]} and '
                             f'{step_subdirs[1]}.')
        if not frazil_pass:
            raise ValueError(f'Validation failed comparing frazil ice outputs '

    """
    A step for visualizing a cross-section and horizontal planes through the
    ice shelf
    """
    def __init__(self, component, indir, mesh, init):
        """

            the directory the step is in, to which ``viz`` will be appended

        """
            work_dir_target=f'{mesh.path}/culled_mesh.nc')
            work_dir_target=f'{init.path}/init.nc')
            filename='adjusted_init.nc',
            target='../forward/init.nc')
            filename='land_ice_fluxes.nc',
            target='../forward/land_ice_fluxes.nc')

        """
        """
        ds_adj_init = xr.load_dataset('adjusted_init.nc')
        ds_ice = xr.load_dataset('land_ice_fluxes.nc')

        x_mid = ds_mesh.xCell.median()
        y_min = ds_mesh.yCell.min()
        y_max = ds_mesh.yCell.max()
        x = xr.DataArray(data=[x_mid, x_mid], dims=('nPoints',))
        y = xr.DataArray(data=[y_min, y_max], dims=('nPoints',))

        # Plot the time series of max velocity
        plt.figure(figsize=[12, 6], dpi=100)
        umax = np.amax(ds.velocityX[:, :, 0].values, axis=1)
        vmax = np.amax(ds.velocityY[:, :, 0].values, axis=1)
        t = ds.daysSinceStartOfSim.values
        time = pd.to_timedelta(t) / 1.e9
        plt.plot(time, umax, 'k', label='u')
        plt.plot(time, vmax, 'b', label='v')
        plt.xlabel('Time (s)')
        plt.ylabel('Velocity (m/s)')
        plt.savefig('velocity_max_t.png', dpi=200)

        ds_horiz = self._process_ds(ds_init, ds_ice, ds_adj_init,
                                    ds_init.bottomDepth,
                                    time_index=0)
        vmin_del_ssh = np.min(ds_horiz.delSsh.values)
        vmax_del_ssh = np.max(ds_horiz.delSsh.values)
        vmax_del_p = np.amax(ds_horiz.delLandIcePressure.values)
        vmin_temp = np.min(ds.temperature.values)
        vmax_temp = np.max(ds.temperature.values)
        vmin_salt = np.min(ds.salinity.values)
        vmax_salt = np.max(ds.salinity.values)
        vmax_uv = max(np.amax(ds.velocityX.values),
                      np.amax(ds.velocityY.values))

        time_index = 0  # Plot the initial time
            layer_thickness=ds_init.layerThickness.isel(Time=time_index),

        plot_transect(ds_transect,
                      mpas_field=ds_init.temperature.isel(Time=time_index),
                      out_filename='temperature_section_init.png',
                      title='temperature',
                      vmin=vmin_temp, vmax=vmax_temp,
                      colorbar_label=r'$^{\circ}$C', cmap='cmo.thermal')

                      mpas_field=ds_init.salinity.isel(Time=time_index),
                      out_filename='salinity_section_init.png',
                      title='salinity',
                      vmin=vmin_salt, vmax=vmax_salt,
                      colorbar_label=r'PSU', cmap='cmo.haline')

        # Plot water column thickness horizontal ds_init
        plot_horiz_field(ds_mesh, ds_horiz['columnThickness'],
                         'H_horiz_init.png', t_index=None,
                         field_mask=cell_mask)

        time_index = -1  # Plot the final state
        ds_horiz = self._process_ds(ds, ds_ice, ds_init,
                                    time_index=time_index)
        # Plot water column thickness horizontal
                         f'H_horiz_t{time_index}.png', t_index=None,
        plot_horiz_field(ds_mesh, ds_horiz['landIceFreshwaterFlux'],
                         f'melt_horiz_t{time_index}.png', t_index=None,
        if 'wettingVelocityFactor' in ds_horiz.keys():
            plot_horiz_field(ds_mesh, ds_horiz['wettingVelocityFactor'],
                             f'wet_horiz_t{time_index}.png', t_index=None,
                             z_index=None, field_mask=edge_mask,
                             vmin=0, vmax=1, cmap='cmo.ice')
        # Plot difference in ssh
        plot_horiz_field(ds_mesh, ds_horiz['delSsh'],
                         f'del_ssh_horiz_t{time_index}.png', t_index=None,
                         field_mask=cell_mask,
                         vmin=vmin_del_ssh, vmax=vmax_del_ssh)

        # Plot difference in land ice pressure
        plot_horiz_field(ds_mesh, ds_horiz['delLandIcePressure'],
                         f'del_land_ice_pressure_horiz_t{time_index}.png',
                         t_index=None, field_mask=cell_mask,
                         vmin=-vmax_del_p, vmax=vmax_del_p,
                         cmap='cmo.balance')

        # Plot transects
            layer_thickness=ds.layerThickness.isel(Time=time_index),

        plot_horiz_field(ds_mesh, ds['velocityX'],
                         f'u_surf_horiz_t{time_index}.png', t_index=time_index,
                         z_index=0, field_mask=cell_mask,
                         vmin=-vmax_uv, vmax=vmax_uv,
                         cmap_title=r'm/s', cmap='cmo.balance')
                         f'u_bot_horiz_t{time_index}.png', t_index=time_index,
                         z_index=-1, field_mask=cell_mask,
        plot_horiz_field(ds_mesh, ds['velocityY'],
                         f'v_surf_horiz_t{time_index}.png', t_index=time_index,
                         f'v_bot_horiz_t{time_index}.png', t_index=time_index,
                      mpas_field=ds.velocityX.isel(Time=time_index),
                      out_filename=f'u_section_t{time_index}.png',
                      title='x-velocity',
                      vmin=-vmax_uv, vmax=vmax_uv,
                      colorbar_label=r'm/s', cmap='cmo.balance')

                      mpas_field=ds.velocityY.isel(Time=time_index),
                      out_filename=f'v_section_t{time_index}.png',
                      title='y-velocity',

        plot_transect(
            ds_transect,
            mpas_field=ds.temperature.isel(Time=time_index),
            out_filename=f'temperature_section_t{time_index}.png',
            title='temperature',
            vmin=vmin_temp, vmax=vmax_temp,
            colorbar_label=r'$^{\circ}$C', cmap='cmo.thermal')

                      mpas_field=ds.salinity.isel(Time=time_index),
                      out_filename=f'salinity_section_t{time_index}.png',

    def _process_ds(ds, ds_ice, ds_init, bottom_depth, time_index):
        ds_out['columnThickness'] = ds_out.ssh + bottom_depth
        ds_out['landIceFreshwaterFlux'] = \
            ds_ice.landIceFreshwaterFlux.isel(Time=time_index)
        ds_out['delSsh'] = ds_out.ssh - ds_init.ssh.isel(Time=0)
        ds_out['delLandIcePressure'] = \
            ds_ice.landIcePressure.isel(Time=time_index) - \
            ds_init.landIcePressure.isel(Time=0)

)
from polaris.ocean.resolution import resolution_to_subdir
from polaris.ocean.tasks.inertial_gravity_wave.analysis import Analysis
from polaris.ocean.tasks.inertial_gravity_wave.forward import Forward
from polaris.ocean.tasks.inertial_gravity_wave.init import Init
from polaris.ocean.tasks.inertial_gravity_wave.viz import Viz

def add_inertial_gravity_wave_tasks(component):
    """

    """
    basedir = 'planar/inertial_gravity_wave'
    config_filename = 'inertial_gravity_wave.cfg'
    filepath = f'{basedir}/{config_filename}'
    config = PolarisConfigParser(filepath=filepath)
    config.add_from_package('polaris.ocean.convergence',
                            'convergence.cfg')
    config.add_from_package(
        'polaris.ocean.tasks.inertial_gravity_wave',
        config_filename)
    for refinement in ['space', 'time', 'both']:
        component.add_task(InertialGravityWave(component=component,

class InertialGravityWave(Task):
    """
    """

    def __init__(self, component, config, refinement='both'):
        """

        """
        name = f'inertial_gravity_wave_convergence_{refinement}'
        basedir = 'planar/inertial_gravity_wave'
        subdir = f'{basedir}/convergence_{refinement}'
        config_filename = 'inertial_gravity_wave.cfg'

        else:
        refinement_factors = self.config.getlist(
            'convergence', option, dtype=float)
                self.config, refinement_factor, refinement=refinement)
            mesh_name = resolution_to_subdir(resolution)

            subdir = f'{basedir}/init/{mesh_name}'
            symlink = f'init/{mesh_name}'
            else:
                init_step = Init(component=component, resolution=resolution,
                                 subdir=subdir)
                init_step.set_shared_config(self.config, link=config_filename)
                self.add_step(init_step, symlink=symlink)

            subdir = f'{basedir}/forward/{mesh_name}_{timestep}s'
            else:
                    component=component,
                    refinement=refinement,
                    name=f'forward_{mesh_name}_{timestep}s',
                    init=init_step)

            analysis_dependencies['mesh'][refinement_factor] = init_step

        self.add_step(Analysis(component=component,
                               subdir=f'{self.subdir}/analysis',
                               refinement=refinement,
                               dependencies=analysis_dependencies))
        self.add_step(Viz(component=component, resolutions=resolutions,
                          taskdir=self.subdir),
                      run_by_default=False)
        config.add_from_package(
            'polaris.ocean.tasks.inertial_gravity_wave',
            config_filename)

from polaris.ocean.tasks.inertial_gravity_wave.exact_solution import (
    ExactSolution,
)

    """
    A step for analysing the output from the inertial gravity wave
    test case

    resolutions : list of float
        The resolutions of the meshes that have been run
    """
        """

        resolutions : list of float
            The resolutions of the meshes that have been run

        """
        convergence_vars = [{'name': 'ssh',
                             'title': 'SSH',
                             'zidx': None}]

        """

            fields

            solution

        Returns
        -------
        """
        init = xr.open_dataset(f'init_r{refinement_factor:02g}.nc')
        exact = ExactSolution(init, self.config)
        if field_name != 'ssh':
            raise ValueError(f'{field_name} is not currently supported')
        return exact.ssh(time)

class ExactSolution():
    """
    Class to compute the exact solution for the inertial gravity wave
    test case

    angleEdge : xr.DataArray
        angle between edge normals and positive x direction

    xCell : xr.DataArray
        x coordinates of mesh cell centers

    yCell : xr.DataArray
        y coordinates of mesh cell centers

    xEdge: xr.DataArray
        x coordinates of mesh edges

    yEdge : xr.DataArray
        y coordinates of mesh edges

    f0 : float
        Coriolis parameter

    eta0 : float
        Amplitide of sea surface height

    kx : float
        Wave number in the x direction

    ky : float
        Wave number in the y direction

    omega : float
        Angular frequency
    """
    def __init__(self, ds, config):
        """
        Create a new exact solution object

        ds : xr.DataSet
            MPAS mesh information

            Config options for test case
        """
        self.angleEdge = ds.angleEdge
        self.xCell = ds.xCell
        self.yCell = ds.yCell
        self.xEdge = ds.xEdge
        self.yEdge = ds.yEdge

        section = config['inertial_gravity_wave']
        self.f0 = section.getfloat('coriolis_parameter')
        self.eta0 = section.getfloat('ssh_amplitude')
        npx = section.getfloat('n_wavelengths_x')
        npy = section.getfloat('n_wavelengths_y')

        self.g = constants['SHR_CONST_G']
        ly = np.sqrt(3.0) / 2.0 * lx
        self.kx = npx * 2.0 * np.pi / (lx * 1e3)
        self.ky = npy * 2.0 * np.pi / (ly * 1e3)
        self.omega = np.sqrt(self.f0**2 +
                             self.g * bottom_depth * (self.kx**2 + self.ky**2))

    def ssh(self, t):
        """
        Exact solution for sea surface height

        t : float
            time at which to evaluate exact solution

        Returns
        -------
        eta : xr.DataArray
            the exact sea surface height solution on cells at time t

        """
        eta = self.eta0 * np.cos(self.kx * self.xCell +
                                 self.ky * self.yCell -
                                 self.omega * t)

        return eta

    def normal_velocity(self, t):
        """
        Exact solution for normal velocity

        t : float

        Returns
        -------
        norm_vel : xr.DataArray
            the exact normal velocity solution on edges at time t
        """
        u = self.eta0 * (self.g / (self.omega**2.0 - self.f0**2.0) *
                         (self.omega * self.kx * np.cos(self.kx * self.xEdge +
                          self.ky * self.yEdge - self.omega * t) -
                          self.f0 * self.ky * np.sin(self.kx * self.xEdge +
                          self.ky * self.yEdge - self.omega * t)))

        v = self.eta0 * (self.g / (self.omega**2.0 - self.f0**2.0) *
                         (self.omega * self.ky * np.cos(self.kx * self.xEdge +
                          self.ky * self.yEdge - self.omega * t) +
                          self.f0 * self.kx * np.sin(self.kx * self.xEdge +

        norm_vel = u * np.cos(self.angleEdge) + v * np.sin(self.angleEdge)

        return norm_vel

class Forward(ConvergenceForward):
    """
    A step for performing forward ocean component runs as part of inertial
    gravity wave test cases.

        The factor by which to scale space, time or both

    """
    def __init__(self, component, name, refinement_factor, subdir,
                 init, refinement='both'):
        """

        name : str

            The subdirectory that the task belongs to

            The step which generates the mesh and initial condition

        """
        super().__init__(component=component,
                         name=name, subdir=subdir,
                         refinement=refinement,
                         refinement_factor=refinement_factor,
                         mesh=init, init=init,
                         package='polaris.ocean.tasks.inertial_gravity_wave',
                         yaml_filename='forward.yaml',
                         graph_target=f'{init.path}/culled_graph.info',
                         output_filename='output.nc',
                         validate_vars=['layerThickness', 'normalVelocity'])

        """
        resources

        Returns
        -------
        """
        section = self.config['inertial_gravity_wave']
            self.config, self.refinement_factor, refinement=self.refinement)

)

    """
    A step for creating a mesh and initial condition for the
    inertial gravity wave test cases

    """
    def __init__(self, component, resolution, subdir):
        """

        """
        mesh_name = resolution_to_subdir(resolution)
                         name=f'init_{mesh_name}',
                         subdir=subdir)
        for filename in ['culled_mesh.nc', 'initial_state.nc',
                         'culled_graph.info']:

        """
        """

        f0 = section.getfloat('coriolis_parameter')

                                       nonperiodic_y=False)

        ds['ssh'] = xr.zeros_like(ds_mesh.xCell)
        ds['bottomDepth'] = bottom_depth * xr.ones_like(ds_mesh.xCell)

        ds['fCell'] = f0 * xr.ones_like(ds_mesh.xCell)
        ds['fEdge'] = f0 * xr.ones_like(ds_mesh.xEdge)
        ds['fVertex'] = f0 * xr.ones_like(ds_mesh.xVertex)

        exact_solution = ExactSolution(ds, config)

        ssh = exact_solution.ssh(0.0)
        ssh = ssh.expand_dims(dim='Time', axis=0)
        ds['ssh'] = ssh

        layerThickness = ssh + bottom_depth
        layerThickness, _ = xr.broadcast(layerThickness, ds.refBottomDepth)
        layerThickness = layerThickness.transpose('Time', 'nCells',
                                                  'nVertLevels')
        ds['layerThickness'] = layerThickness

        normal_velocity = exact_solution.normal_velocity(0.0)

import datetime

)

    """
    A step for visualizing the output from the inertial gravity wave
    test case

    """
    def __init__(self, component, resolutions, taskdir):
        """

        taskdir : str
        """
        super().__init__(component=component, name='viz', indir=taskdir)
        self.resolutions = resolutions

        for resolution in resolutions:
                filename=f'mesh_{mesh_name}.nc',
                target=f'../init/{mesh_name}/culled_mesh.nc')
                filename=f'init_{mesh_name}.nc',
                target=f'../init/{mesh_name}/initial_state.nc')
                filename=f'output_{mesh_name}.nc',
                target=f'../forward/{mesh_name}/output.nc')

        self.add_output_file('comparison.png')

        """
        """
        resolutions = self.resolutions
        nres = len(resolutions)

        eta0 = section.getfloat('ssh_amplitude')

        fig, axes = plt.subplots(nrows=nres, ncols=3, figsize=(12, 2 * nres))
        rmse = []
        error_range = None
        for i, res in enumerate(resolutions):
            mesh_name = resolution_to_subdir(res)
            ds_mesh = xr.open_dataset(f'mesh_{mesh_name}.nc')
            ds_init = xr.open_dataset(f'init_{mesh_name}.nc')
            ds = xr.open_dataset(f'output_{mesh_name}.nc')
            exact = ExactSolution(ds_init, config)

            t0 = datetime.datetime.strptime(ds.xtime.values[0].decode(),
                                            '%Y-%m-%d_%H:%M:%S')
            tf = datetime.datetime.strptime(ds.xtime.values[-1].decode(),
            t = (tf - t0).total_seconds()
            ssh_model = ds.ssh.values[-1, :]
            rmse.append(np.sqrt(np.mean((ssh_model - exact.ssh(t).values)**2)))

            # Comparison plots
            ds['ssh_exact'] = exact.ssh(t)
            ds['ssh_error'] = ssh_model - exact.ssh(t)
            if error_range is None:
                error_range = np.max(np.abs(ds.ssh_error.values))

            descriptor = plot_horiz_field(
                ds_mesh, ds['ssh'], ax=axes[i, 0], cmap='cmo.balance',
                t_index=ds.sizes["Time"] - 1, vmin=-eta0, vmax=eta0,
                cmap_title="SSH (m)", field_mask=cell_mask)
            plot_horiz_field(ds_mesh, ds['ssh_exact'], ax=axes[i, 1],
                             cmap='cmo.balance',
                             vmin=-eta0, vmax=eta0, cmap_title="SSH (m)",
                             descriptor=descriptor)
            plot_horiz_field(ds_mesh, ds['ssh_error'], ax=axes[i, 2],
                             cmap_title=r"$\Delta$ SSH (m)",
                             vmin=-error_range, vmax=error_range,

        axes[0, 0].set_title('Numerical solution')
        axes[0, 1].set_title('Analytical solution')
        axes[0, 2].set_title('Error (Numerical - Analytical)')

        pad = 5
        for ax, res in zip(axes[:, 0], resolutions):
            ax.annotate(f'{res}km', xy=(0, 0.5),
                        xytext=(-ax.yaxis.labelpad - pad, 0),
                        xycoords=ax.yaxis.label, textcoords='offset points',
                        size='large', ha='right', va='center')

    """
    A step for performing forward MPAS-Ocean runs as part of internal wave
    test cases.

    """
    def __init__(self, component, init, name='forward', subdir=None,
                 nu=None, run_time_steps=None, vadv_method='standard'):
        """

        name : str

        init : polaris.ocean.tasks.internal_wave.init.Init

            are provided, the directory is the ``name``

            the viscosity (if different from the default for the test group)

        vadv_method : str, optional
            The vertical advection method, 'standard' or 'vlr'
        """

                            work_dir_target=f'{init.path}/initial_state.nc')

        self.add_yaml_file('polaris.ocean.tasks.internal_wave',

        vadv_dict = {'standard': 'flux-form',
                     'vlr': 'remap'}
        self.add_model_config_options({
            'config_vert_advection_method': f"{vadv_dict[vadv_method]}"},
            config_model='mpas-ocean')

        """
        resources

        Returns
        -------
        """
        section = self.config['internal_wave']

        """

        """

        resolution = config.getfloat('internal_wave', 'resolution')
        dt_per_km = config.getfloat('internal_wave', 'dt_per_km')
        dt = dt_per_km * resolution

    """
    A step for creating a mesh and initial condition for internal wave test
    cases
    """
        """

        indir : str
        """

        """
        """

        section = config['internal_wave']
        amplitude_width_dist = section.getfloat('amplitude_width_dist')
        amplitude_width_frac = section.getfloat('amplitude_width_frac')

        vert_levels = section.getint('vert_levels')

        ds['maxLevelCell'] = vert_levels * xr.ones_like(y_cell)
        ds['bottomDepth'] = bottom_depth * xr.ones_like(y_cell)
        ds['ssh'] = xr.zeros_like(y_cell)

            perturb_width = amplitude_width_dist
        else:
            perturb_width = (y_max - y_min) * amplitude_width_frac

        # Set stratified temperature

        depth_frac = xr.zeros_like(temp_vert)
        ref_bottom_depth = ds['refBottomDepth']
        for k in range(1, vert_levels):
            depth_frac[k] = (ref_bottom_depth[k - 1] /
                             ref_bottom_depth[vert_levels - 1])

        # If cell is in the southern half, outside the sin width, subtract
        # temperature difference
        frac = xr.where(
            np.abs(y_cell - y_mid) < perturb_width,
            np.cos(0.5 * np.pi * (y_cell - y_mid) / perturb_width) *
            np.sin(np.pi * depth_frac),
            0.)

        normal_velocity = xr.zeros_like(ds.xEdge)
        normal_velocity, _ = xr.broadcast(normal_velocity, ref_bottom_depth)

        ds['fCell'] = coriolis_parameter * xr.ones_like(y_cell)

from polaris.ocean.tasks.internal_wave.forward import Forward
from polaris.ocean.tasks.internal_wave.rpe.analysis import Analysis

    """
    The internal wave reference potential energy (RPE) test case performs

    init : polaris.ocean.tasks.baroclinic_channel.init.Init
        A shared step for creating the initial state
    """

    def __init__(self, component, indir, init, config, vadv_method='standard'):
        """

        indir : str

        """
        self.vadv_method = vadv_method
        super().__init__(component=component, name=f'{vadv_method}/rpe',
                         indir=indir)

        self.set_shared_config(config, link='internal_wave.cfg')
        self.init = init

        """
        """

        nus = config.getlist('internal_wave_rpe', 'viscosities',
                component=component, name=name, init=self.init,
                indir=self.subdir, ntasks=None, min_tasks=None,
                openmp_threads=1, nu=nu, vadv_method=self.vadv_method)

                'polaris.ocean.tasks.internal_wave.rpe',

            Analysis(component=component, nus=nus, indir=self.subdir))

    """
    A step for plotting the results of a series of internal wave RPE runs

    nus : list
    """
    def __init__(self, component, indir, nus):
        """

        indir : str

        """

            target='../../../init/culled_mesh.nc')

            target='../../../init/initial_state.nc')

            filename='sections_internal_wave.png')

        """
        """
        section = self.config['internal_wave_rpe']

        fig, axes = plt.subplots(1, sim_count, sharey=True,
                                 figsize=(3 * sim_count, 5.0),
                                 constrained_layout=True)
        fig.suptitle(f'day {time}')
            ds_transect = compute_transect(
                x=x, y=y, ds_horiz_mesh=ds_mesh,
                layer_thickness=ds.layerThickness.isel(Time=time_index),
                bottom_depth=ds_init.bottomDepth,
                min_level_cell=ds_init.minLevelCell - 1,
                max_level_cell=ds_init.maxLevelCell - 1,
                spherical=False)

            if row_index == len(nus) - 1:
                colorbar_label = r'$^{\circ}$C'
            else:
                colorbar_label = None
            plot_transect(ds_transect,
                          mpas_field=ds.temperature.isel(Time=time_index),
                          ax=ax,
                          title='temperature',
                          interface_color='grey',
                          vmin=min_temp, vmax=max_temp,
                          colorbar_label=colorbar_label, cmap='cmo.thermal')
            ax.set_title(f'$\\nu_h=${nu:g}')
            if row_index != 0:
                ax.set_ylabel(None)

    """
    A step for visualizing a cross-section through the internal wave
    """
        """

        test_case : polaris.TestCase
            The test case this step belongs to
        """

        """
        """
        ds_mesh = xr.load_dataset('init.nc')
        ds_init = ds_mesh

        vmax_v = np.max(np.abs(ds.vertVelocityTop.values))

        tidx = 0  # Plot the initial time
            layer_thickness=ds_init.layerThickness.isel(Time=tidx),
            bottom_depth=ds_mesh.bottomDepth,
            min_level_cell=ds_mesh.minLevelCell - 1,
            max_level_cell=ds_mesh.maxLevelCell - 1,

                      mpas_field=ds_init.temperature.isel(Time=tidx),
                      interface_color='grey',

        tidx = -1  # Plot the final time
            layer_thickness=ds.layerThickness.isel(Time=tidx),

                      mpas_field=ds.temperature.isel(Time=tidx),
                      out_filename='temperature_section_final.png',

        w_values = ds.vertVelocityTop.isel(Time=tidx).values[:, :-1]
        w = w_values * xr.ones_like(ds.temperature.isel(Time=tidx))
            mpas_field=w,
            out_filename='vertical_velocity_section_final.png',
            title='vertical velocity',
            vmin=-vmax_v, vmax=vmax_v,
            colorbar_label='m/s', cmap='cmo.balance')

)
from polaris.ocean.tasks.manufactured_solution.analysis import Analysis
from polaris.ocean.tasks.manufactured_solution.forward import Forward
from polaris.ocean.tasks.manufactured_solution.init import Init
from polaris.ocean.tasks.manufactured_solution.viz import Viz

def add_manufactured_solution_tasks(component):
    """
    Add a task that defines a convergence test that uses the method of
    manufactured solutions

    """
    basedir = 'planar/manufactured_solution'
    config_filename = 'manufactured_solution.cfg'
        'polaris.ocean.tasks.manufactured_solution',
        component.add_task(ManufacturedSolution(component=component,
                                                config=config,
                                                refinement=refinement))

class ManufacturedSolution(Task):
    """
    The convergence test case using the method of manufactured solutions
    """

        """

        """
        name = f'manufactured_solution_convergence_{refinement}'
        basedir = 'planar/manufactured_solution'
        config_filename = 'manufactured_solution.cfg'

        else:

            else:

            else:

                               dependencies=analysis_dependencies,
                               refinement=refinement))
        self.add_step(Viz(component=component,
                          dependencies=analysis_dependencies,
                          taskdir=self.subdir,
                          refinement=refinement),
            'polaris.ocean.tasks.manufactured_solution',

from polaris.ocean.tasks.manufactured_solution.exact_solution import (
)

    """
    A step for analysing the output from the manufactured solution
    test case
    """
        """

        """

        """

            fields

            solution

        Returns
        -------
        """
        init = self.open_model_dataset(f'init_r{refinement_factor:02g}.nc')
        exact = ExactSolution(self.config, init)

    """
    Class to compute the exact solution for the manufactured solution
    test case

    kx : float

    ky : float

    """

    def __init__(self, config, ds=None):
        """

        """
        section = config['manufactured_solution']
        self.g = 9.80616
        self.lambda_x = (lx * 1e3) / npx
        self.lambda_y = (ly * 1e3) / npy
        self.kx = 2.0 * np.pi / self.lambda_x
        self.ky = 2.0 * np.pi / self.lambda_y
        self.omega = np.sqrt(self.g * bottom_depth *
                             (self.kx**2 + self.ky**2))

        if ds is not None:
            self.angleEdge = ds.angleEdge
            self.xCell = ds.xCell
            self.yCell = ds.yCell
            self.xEdge = ds.xEdge
            self.yEdge = ds.yEdge

        """

        t : float

        Returns
        -------

        """

        eta = self.eta0 * np.sin(self.kx * self.xCell +

        return eta

        """

        t : float

        Returns
        -------
        normalvelocity : xr.DataArray
        """
        u = self.eta0 * np.cos(self.kx * self.xEdge +
                               self.ky * self.yEdge -
                               self.omega * t)
        v = u
        return u * np.cos(self.angleEdge) + v * np.sin(self.angleEdge)

)

    """
    A step for performing forward ocean component runs as part of manufactured
    solution test cases.

    """
        """

        name : str

        """
                         mesh=init, init=init, refinement=refinement,
                         package='polaris.ocean.tasks.manufactured_solution',

        """
        TEMP: symlink initial condition to name hard-coded in Omega
        """
        # TODO: remove as soon as Omega no longer hard-codes this file
            self.add_input_file(filename='OmegaMesh.nc', target='init.nc')

        """
        resources

        Returns
        -------
        """
        # no file to read from, so we'll compute it based on config options
        section = self.config['manufactured_solution']

        """

        """

        exact_solution = ExactSolution(self.config)
        options = {'config_manufactured_solution_amplitude':
                   float(exact_solution.eta0),
                   'config_manufactured_solution_wavelength_x':
                   float(exact_solution.lambda_x),
                   'config_manufactured_solution_wavelength_y':
                   float(exact_solution.lambda_y)}
        self.add_model_config_options(options, config_model='mpas-ocean')

)

    """
    manufactured solution test cases

    """
        """

        """

        output_filenames = ['culled_mesh.nc', 'initial_state.nc']
        model = self.config.get('ocean', 'model')
            output_filenames.append('culled_graph.info')
        for filename in output_filenames:

        """
        """

        self.write_model_dataset(ds_mesh, 'base_mesh.nc')

        self.write_model_dataset(ds_mesh, 'culled_mesh.nc')

        ds['fCell'] = coriolis_parameter * xr.ones_like(ds_mesh.xCell)

        # Evaluate the exact solution at time=0
        exact_solution = ExactSolution(config, ds)

        layer_thickness = ssh + bottom_depth
        layer_thickness, _ = xr.broadcast(layer_thickness, ds.refBottomDepth)
        layer_thickness = layer_thickness.transpose('Time', 'nCells',
                                                    'nVertLevels')
        ds['layerThickness'] = layer_thickness

)
)

    """
    A step for visualizing the output from the manufactured solution
    test case

        dict:

                Keys of the dict correspond to `refinement_factors`

    """
    def __init__(self, component, dependencies, taskdir, refinement='both'):
        """

            dict:

        """

        """
        """

        else:

        """
        """
        else:

        nres = len(refinement_factors)

        rmse = []

        for i, refinement_factor in enumerate(refinement_factors):
            ds_mesh = self.open_model_dataset(
                f'mesh_r{refinement_factor:02g}.nc')
            ds_init = self.open_model_dataset(
                f'init_r{refinement_factor:02g}.nc')
            ds = self.open_model_dataset(
                f'output_r{refinement_factor:02g}.nc', decode_times=False)
            exact = ExactSolution(config, ds_init)

            if model == 'mpas-o':
                t0 = datetime.datetime.strptime(ds.xtime.values[0].decode(),
                                                '%Y-%m-%d_%H:%M:%S')
                tf = datetime.datetime.strptime(ds.xtime.values[-1].decode(),
                t = (tf - t0).total_seconds()

            else:
                # time is seconds since the start of the simulation in Omega
                t = ds.Time[-1].values

            ssh_model = ds.ssh.isel(Time=-1)
            if 'nVertLevels' in ssh_model.dims:
                # Omega v0 uses stacked shallow water where ssh has nVertLevels
                ssh_model = ssh_model.isel(nVertLevels=0)
            rmse.append(np.sqrt(np.mean(
                (ssh_model.values - exact.ssh(t).values)**2)))

                ds, ds_mesh, 'ssh', ax=axes[i, 0], cmap='cmo.balance',
                cmap_title="SSH", field_mask=cell_mask)
                             vmin=-eta0, vmax=eta0, cmap_title="SSH",
                             cmap='cmo.balance', cmap_title="dSSH",

        pad = 5
        for ax, refinement_factor in zip(axes[:, 0], refinement_factors):

            ax.annotate(f'{resolution}km\n{timestep}s', xy=(0, 0.5),

    """
    A step for creating a mesh and initial condition for single column
    test cases
    """
    def __init__(self, component, indir, ideal_age=False):
        """

        indir : str
            The subdirectory that the task belongs to, that this step will
            go into a subdirectory of

        ideal_age : bool, optional
            Whether the initial condition should include the ideal age tracer
        """
        self.ideal_age = ideal_age
                     'initial_state.nc', 'forcing.nc']:

        """
        """
        section = config['single_column']
        ideal_age = self.ideal_age
        nx = section.getint('nx')
        ny = section.getint('ny')

        x_cell = ds_mesh.xCell

        surface_temperature = section.getfloat(
            'surface_temperature')
        temperature_gradient_mixed_layer = section.getfloat(
            'temperature_gradient_mixed_layer')
        temperature_difference_across_mixed_layer = section.getfloat(
            'temperature_difference_across_mixed_layer')
        temperature_gradient_interior = section.getfloat(
            'temperature_gradient_interior')
        mixed_layer_depth_temperature = section.getfloat(
            'mixed_layer_depth_temperature')
        surface_salinity = section.getfloat(
            'surface_salinity')
        salinity_gradient_mixed_layer = section.getfloat(
            'salinity_gradient_mixed_layer')
        salinity_difference_across_mixed_layer = section.getfloat(
            'salinity_difference_across_mixed_layer')
        salinity_gradient_interior = section.getfloat(
            'salinity_gradient_interior')
        mixed_layer_depth_salinity = section.getfloat(
            'mixed_layer_depth_salinity')
        coriolis_parameter = section.getfloat(
            'coriolis_parameter')

        z_mid = ds.refZMid

        temperature_at_mixed_layer_depth = (
            surface_temperature + temperature_difference_across_mixed_layer)
        temperature_vert = xr.where(
            z_mid > -mixed_layer_depth_temperature,
            surface_temperature + temperature_gradient_mixed_layer * z_mid,
            temperature_at_mixed_layer_depth +
            temperature_gradient_interior *
            (z_mid + mixed_layer_depth_temperature))
        temperature_vert[0] = surface_temperature
        temperature, _ = xr.broadcast(temperature_vert, x_cell)

        salinity_at_mixed_layer_depth = (
            surface_salinity + salinity_difference_across_mixed_layer)
        salinity_vert = xr.where(
            z_mid > -mixed_layer_depth_salinity,
            surface_salinity + salinity_gradient_mixed_layer * z_mid,
            salinity_at_mixed_layer_depth +
            salinity_gradient_interior *
            (z_mid + mixed_layer_depth_salinity))
        salinity_vert[0] = surface_salinity
        salinity, _ = xr.broadcast(salinity_vert, x_cell)
        salinity = salinity.transpose('nCells', 'nVertLevels')
        salinity = salinity.expand_dims(dim='Time', axis=0)

        if ideal_age:
            ds['idealAgeTracers'] = xr.zeros_like(x_cell)

        ds['salinity'] = salinity

        # create forcing stream
        forcing_array = xr.ones_like(temperature)
        forcing_array_surface = xr.ones_like(ds.bottomDepth)
        forcing_array_surface = forcing_array_surface.expand_dims(
        section = config['single_column_forcing']
        temperature_piston_velocity = section.getfloat(
            'temperature_piston_velocity')
        salinity_piston_velocity = section.getfloat(
            'salinity_piston_velocity')
        temperature_surface_restoring_value = section.getfloat(
            'temperature_surface_restoring_value')
        salinity_surface_restoring_value = section.getfloat(
            'salinity_surface_restoring_value')
        temperature_interior_restoring_rate = section.getfloat(
            'temperature_interior_restoring_rate')
        salinity_interior_restoring_rate = section.getfloat(
            'salinity_interior_restoring_rate')

        latent_heat_flux = section.getfloat('latent_heat_flux')
        sensible_heat_flux = section.getfloat('sensible_heat_flux')
        shortwave_heat_flux = section.getfloat('shortwave_heat_flux')
        evaporation_flux = section.getfloat('evaporation_flux')
        rain_flux = section.getfloat('rain_flux')
        wind_stress_zonal = section.getfloat('wind_stress_zonal')
        wind_stress_meridional = section.getfloat('wind_stress_meridional')

        ds_forcing['temperaturePistonVelocity'] = \
            temperature_piston_velocity * forcing_array_surface
        ds_forcing['salinityPistonVelocity'] = \
            salinity_piston_velocity * forcing_array_surface
        ds_forcing['temperatureSurfaceRestoringValue'] = \
            temperature_surface_restoring_value * forcing_array_surface
        ds_forcing['salinitySurfaceRestoringValue'] = \
            salinity_surface_restoring_value * forcing_array_surface
        ds_forcing['temperatureInteriorRestoringRate'] = \
            temperature_interior_restoring_rate * forcing_array
        ds_forcing['salinityInteriorRestoringRate'] = \
            salinity_interior_restoring_rate * forcing_array
        ds_forcing['temperatureInteriorRestoringValue'] = temperature
        ds_forcing['salinityInteriorRestoringValue'] = salinity
        ds_forcing['windStressZonal'] = \
            wind_stress_zonal * forcing_array_surface
        ds_forcing['windStressMeridional'] = \
            wind_stress_meridional * forcing_array_surface
        ds_forcing['latentHeatFlux'] = latent_heat_flux * forcing_array_surface
        ds_forcing['sensibleHeatFlux'] = \
            sensible_heat_flux * forcing_array_surface
        ds_forcing['shortWaveHeatFlux'] = \
            shortwave_heat_flux * forcing_array_surface
        ds_forcing['evaporationFlux'] = \
            evaporation_flux * forcing_array_surface
        ds_forcing['rainFlux'] = rain_flux * forcing_array_surface

    """
    A step for plotting the results of a single-column test
    """
        """

        indir : str

        """
            target='../init/initial_state.nc')

        """
        """
        t = ds.daysSinceStartOfSim[t_index]
        t_days = t.values.astype('timedelta64[D]')
        title = f'final time = {t_days / np.timedelta64(1, "D")} days'
        fields = {'temperature': 'degC',
                  'salinity': 'PSU'}
            fields['iAge'] = 'seconds'
        z_mid = ds['zMid'].mean(dim='nCells')
        z_mid_init = z_mid.isel(Time=0)
        z_mid_final = z_mid.isel(Time=t_index)
        for field_name, field_units in fields.items():
            if field_name not in ds.keys():
                raise ValueError(f'{field_name} not present in output.nc')
            var = ds[field_name].mean(dim='nCells')
            var_init = var.isel(Time=0)
            var_final = var.isel(Time=t_index)
            plt.figure(figsize=(3, 5))
            ax = plt.subplot(111)
            ax.plot(var_init, z_mid_init, '--k', label='initial')
            ax.plot(var_final, z_mid_final, '-k', label='final')
            ax.set_xlabel(f'{field_name} ({field_units})')
            ax.set_ylabel('z (m)')
            ax.legend()
            plt.title(title)
            plt.tight_layout(pad=0.5)
            plt.savefig(f'{field_name}.png')
            plt.close()

)
from polaris.ocean.tasks.sphere_transport.analysis import Analysis
from polaris.ocean.tasks.sphere_transport.filament_analysis import (
    FilamentAnalysis,
)
from polaris.ocean.tasks.sphere_transport.forward import Forward
from polaris.ocean.tasks.sphere_transport.init import Init
from polaris.ocean.tasks.sphere_transport.mixing_analysis import MixingAnalysis
from polaris.ocean.tasks.sphere_transport.viz import Viz

def add_sphere_transport_tasks(component):
    """
    Add tasks that define variants of sphere transport test cases:
    nondivergent_2d, divergent_2d, correlated_tracers_2d, rotation_2d

    """

        for case_name in ['rotation_2d', 'nondivergent_2d', 'divergent_2d',
                          'correlated_tracers_2d']:
            filepath = f'spherical/{prefix}/{case_name}/{case_name}.cfg'
            config = PolarisConfigParser(filepath=filepath)
            config.add_from_package('polaris.ocean.convergence',
                                    'convergence.cfg')
            config.add_from_package('polaris.ocean.convergence.spherical',
                                    'spherical.cfg')
            package = 'polaris.ocean.tasks.sphere_transport'
            config.add_from_package(package, 'sphere_transport.cfg')
            config.add_from_package(package, f'{case_name}.cfg')
                component.add_task(SphereTransport(
                    component=component, case_name=case_name,
                    icosahedral=icosahedral, config=config,
                    include_viz=include_viz))

class SphereTransport(Task):
    """
    A test case for testing properties of tracer advection

        A list of mesh resolutions

    """
    def __init__(self, component, config, case_name, icosahedral, include_viz,
        """
        Create test case for creating a global MPAS-Ocean mesh

        case_name: string
            The name of the case which determines what variant of the
            configuration to use

            meshes

        """
        else:

        subdir = f'spherical/{prefix}/{case_name}'
        name = f'{prefix}_{case_name}'
            link = f'{case_name}.cfg'
        else:
            # config options live in the task already so no need for a symlink
            link = None
        self.resolutions = list()
        self.case_name = case_name

        # add the steps with default resolutions so they can be listed

        """
        """

    def _setup_steps(self, refinement):  # noqa: C901
        """ setup steps given resolutions """
        case_name = self.case_name

        else:

        else:

        sph_trans_dir = f'spherical/{prefix}/{case_name}'

        for idx, refinement_factor in enumerate(refinement_factors):

            subdir = f'{sph_trans_dir}/init/{mesh_name}'
                symlink = f'init/{mesh_name}'
            else:
            else:
                                 subdir=subdir, base_mesh=base_mesh_step,
                                 case_name=case_name)
            self.add_step(base_mesh_step, symlink=f'base_mesh/{mesh_name}')
            self.add_step(init_step, symlink=symlink)

            name = f'{prefix}_forward_{mesh_name}_{timestep}s'
            subdir = f'{sph_trans_dir}/forward/{mesh_name}_{timestep}s'
                symlink = f'forward/{mesh_name}_{timestep}s'
            else:
            else:
                    base_mesh=base_mesh_step,
                    init=init_step,
                    case_name=case_name,
                    refinement=refinement)
                forward_step.set_shared_config(config, link=config_filename)

                with_viz_dir = f'{sph_trans_dir}/with_viz'
                subdir = f'{with_viz_dir}/viz/{mesh_name}_{timestep}s'
                step = Viz(component=component, name=name,
                           subdir=subdir, base_mesh=base_mesh_step,
                           init=init_step, forward=forward_step,
                           mesh_name=mesh_name)
                step.set_shared_config(config, link=config_filename)
                self.add_step(step)

        subdir = f'{sph_trans_dir}/analysis/{refinement}'
        else:
            symlink = None
        else:
                            subdir=subdir, case_name=case_name,
        self.add_step(step, symlink=symlink)

        if case_name == 'correlated_tracers_2d':
            subdir = f'{sph_trans_dir}/mixing_analysis'
                symlink = 'mixing_analysis'
            else:
                step = component.steps[subdir]
            else:
                step = MixingAnalysis(component=component,
                                      icosahedral=icosahedral, subdir=subdir,
                                      refinement_factors=refinement_factors,
                                      case_name=case_name,
                                      dependencies=analysis_dependencies,
                                      refinement=refinement)

        if case_name == 'nondivergent_2d':
            subdir = f'{sph_trans_dir}/filament_analysis'
                symlink = 'filament_analysis'
            else:
            else:
                step = FilamentAnalysis(component=component,
                                        refinement_factors=refinement_factors,
                                        icosahedral=icosahedral, subdir=subdir,
                                        case_name=case_name,
                                        dependencies=analysis_dependencies,
                                        refinement=refinement)

    """
    A step for analyzing the output from sphere transport test cases

    case_name : str
        The name of the test case
    """
    def __init__(self, component, subdir, case_name,
                 dependencies, refinement='both'):
        """

        case_name: str
            The name of the test case

        """
                             'zidx': 1},
                            {'name': 'tracer2',
                             'title': 'tracer2',
                            {'name': 'tracer3',
                             'title': 'tracer3',
                             'zidx': 1}]
        # Note: there is no need to overwrite the default method exact_solution
        # which uses the initial condition

        """

        Returns
        -------

        """
        section = config[self.case_name]
        conv_thresh = section.getfloat(f'convergence_thresh_{field_name}')

from polaris.mpas import time_index_from_xtime

class FilamentAnalysis(Step):
    """

        Whether to use icosahedral, as opposed to less regular, JIGSAW
        meshes

    """
    def __init__(self, component, refinement_factors, icosahedral, subdir,
                 case_name, dependencies, refinement='both'):
        """

            meshes

        """
        super().__init__(component=component, name='filament_analysis',
        self.refinement_factors = refinement_factors

        self.add_output_file('filament.png')

        """
        """
        for refinement_factor in self.refinement_factors:
                self.config, refinement_factor, self.refinement)
        eval_time = section.getfloat('filament_evaluation_time')
        s_per_day = 86400.0
        zidx = 1
        variable_name = 'tracer2'
        num_tau = 21
        filament_tau = np.linspace(0, 1, num_tau)
        filament_norm = np.zeros((len(resolutions), num_tau))
        fig, ax = plt.subplots()
        for i, refinement_factor in enumerate(self.refinement_factors):
            mesh_name = resolution_to_subdir(resolutions[i])
            ds = xr.open_dataset(f'output_r{refinement_factor:02g}.nc')
            tidx = time_index_from_xtime(ds.xtime.values,
                                         eval_time * s_per_day)
            tracer = ds[variable_name]
            area_cell = ds["areaCell"]
            for j, tau in enumerate(filament_tau):
                cells_above_tau = tracer[tidx, :, zidx] >= tau
                cells_above_tau0 = tracer[0, :, zidx] >= tau
                if np.sum(cells_above_tau0 * area_cell) == 0.:
                    filament_norm[i, j] = np.nan
                else:
                    filament_norm[i, j] = np.divide(
                        np.sum(area_cell * cells_above_tau),
                        np.sum(cells_above_tau0 * area_cell))
            plt.plot(filament_tau, filament_norm[i, :], '.-', label=mesh_name)
        plt.plot([filament_tau[0], filament_tau[-1]], [1., 1.], 'k--')
        ax.set_xlim([filament_tau[0], filament_tau[-1]])
        ax.set_xlabel(r'$\tau$')
        ax.set_ylabel(r'$l_f$')
        plt.title(f'Filament preservation diagnostic for {variable_name}')
        fig.savefig('filament.png', bbox_inches='tight')

        res_array = np.array(resolutions, dtype=float)
        data = np.column_stack((res_array, filament_norm))
        col_headers = ['resolution']
        for tau in filament_tau:
            col_headers.append(f'{tau:g}')
        df = pd.DataFrame(data, columns=col_headers)
        df.to_csv('filament.csv', index=False)

from polaris.ocean.tasks.sphere_transport.resources.flow_types import (
    flow_divergent,
    flow_nondivergent,
    flow_rotation,
)
from polaris.ocean.tasks.sphere_transport.resources.tracer_distributions import (  # noqa: E501
    correlation_fn,
    cosine_bells,
    slotted_cylinders,
    xyztrig,
)

    """
    """
    def __init__(self, component, name, subdir, base_mesh, case_name):
        """

        name : str

        """

        """
        """

        section = config['sphere_transport']

        temperature_array = temperature * xr.ones_like(latCell)

        # tracer1
        tracer1 = xyztrig(lonCell, latCell, sphere_radius)

        # tracer2
        radius = section.getfloat('cosine_bells_radius')
        background_value = section.getfloat('cosine_bells_background')
        amplitude = section.getfloat('cosine_bells_amplitude')
        tracer2 = cosine_bells(lonCell, latCell, radius, background_value,
                               amplitude, sphere_radius)

        # tracer3
            coeff = config.getlist(case_name, 'correlation_coefficients',
                                   dtype=float)
            tracer3 = correlation_fn(tracer2, coeff[0], coeff[1], coeff[2])
        else:
            section = config['sphere_transport']
            radius = section.getfloat('slotted_cylinders_radius')
            background_value = section.getfloat('slotted_cylinders_background')
            amplitude = section.getfloat('slotted_cylinders_amplitude')
            tracer3 = slotted_cylinders(lonCell, latCell, radius,
                                        background_value, amplitude,
                                        sphere_radius)
        _, tracer1_array = np.meshgrid(ds.refZMid.values, tracer1)
        _, tracer2_array = np.meshgrid(ds.refZMid.values, tracer2)
        _, tracer3_array = np.meshgrid(ds.refZMid.values, tracer3)

        ds['tracer1'] = (('nCells', 'nVertLevels',), tracer1_array)
        ds['tracer1'] = ds.tracer1.expand_dims(dim='Time', axis=0)
        ds['tracer2'] = (('nCells', 'nVertLevels',), tracer2_array)
        ds['tracer2'] = ds.tracer2.expand_dims(dim='Time', axis=0)
        ds['tracer3'] = (('nCells', 'nVertLevels',), tracer3_array)
        ds['tracer3'] = ds.tracer3.expand_dims(dim='Time', axis=0)

        if case_name == 'rotation_2d':
            rotation_vector = config.getlist(case_name, 'rotation_vector',
                                             dtype=float)
            vector = np.array(rotation_vector)
            u, v = flow_rotation(lonEdge, latEdge, vector,
                                 vel_pd * s_per_hour, sphere_radius)
        elif case_name == 'divergent_2d':
            section = config[case_name]
            vel_amp = section.getfloat('vel_amp')
            u, v = flow_divergent(0., lonEdge, latEdge,
                                  vel_amp, vel_pd * s_per_hour)
        elif (case_name == 'nondivergent_2d' or
              case_name == 'correlated_tracers_2d'):
            u, v = flow_nondivergent(0., lonEdge, latEdge,
                                     vel_amp, vel_pd * s_per_hour)
        else:
            raise ValueError(f'Unexpected test case name {case_name}')

        normalVelocity = sphere_radius * (u * np.cos(angleEdge) +
                                          v * np.sin(angleEdge))

from matplotlib.lines import Line2D

class MixingAnalysis(Step):
    """

        meshes

    """
        """

            meshes

        """
        super().__init__(component=component, name='mixing_analysis',

        self.add_output_file('triplots.png')

        """
        """
        eval_time = section.getfloat('mixing_evaluation_time')
        zidx = 1
        nrows = int(ceil(len(resolutions) / 2))
        fig, axes = plt.subplots(nrows=nrows, ncols=2, sharex=True,
                                 sharey=True, figsize=(5.5, 7))
            ax = axes[int(i / 2), i % 2]
            _init_triplot_axes(ax)
            ax.set(title=mesh_name)
            if i % 2 == 0:
                ax.set_ylabel("tracer3")
            if int(i / 2) == nrows - 1:
                ax.set_xlabel("tracer2")
            ds = ds.isel(Time=tidx)
            ds = ds.isel(nVertLevels=zidx)
            tracer2 = ds["tracer2"].values
            tracer3 = ds["tracer3"].values
            ax.plot(tracer2, tracer3, '.', markersize=1)
        if i % 2 < 1:
            ax = axes[int(i / 2), 1]
            ax.set_axis_off()
        plt.subplots_adjust(wspace=0.1, hspace=0.1)
        fig.suptitle('Correlated tracers 2-d')
        fig.savefig('triplots.png', bbox_inches='tight')

def _init_triplot_axes(ax):
    lw = 0.4
    topline = Line2D([0.1, 1.0], [0.9, 0.9], color='k',
                     linestyle='-', linewidth=lw)
    midline = Line2D([0.1, 1.0], [0.9, 0.1], color='k',
    rightline = Line2D([1, 1], [0.1, 0.9], color='k',
                       linestyle='-', linewidth=lw)
    leftline = Line2D([0.1, 0.1], [0.1, 0.9], color='k',
                      linestyle='-', linewidth=lw)
    botline = Line2D([0.1, 1.0], [0.1, 0.1], color='k',
    crvx = np.linspace(0.1, 1)
    crvy = -0.8 * np.square(crvx) + 0.9
    ticks = np.array(range(6)) / 5
    ax.plot(crvx, crvy, 'k-', linewidth=1.25 * lw)
    ax.set_xticks(ticks)
    ax.set_yticks(ticks)
    ax.add_artist(topline)
    ax.add_artist(midline)
    ax.add_artist(botline)
    ax.add_artist(rightline)
    ax.add_artist(leftline)
    ax.set_xlim([0, 1.1])
    ax.set_ylim([0, 1.0])
    ax.text(0.98, 0.87, 'Range-preserving\n unmixing', fontsize=8,
            horizontalalignment='right', verticalalignment='top')
    ax.text(0.12, 0.12, 'Range-preserving\n unmixing', fontsize=8,
            horizontalalignment='left', verticalalignment='bottom')
    ax.text(0.5, 0.27, 'Real mixing', rotation=-40., fontsize=8)
    ax.text(0.02, 0.1, 'Overshooting', rotation=90., fontsize=8)
    ax.grid(color='lightgray')

    """

    """
        """

        name : str

        """
        variables_to_plot = dict({'tracer1': 'tracer',
                                  'tracer2': 'tracer',
                                  'tracer3': 'tracer',
                                  'layerThickness': 'h'})
        self.variables_to_plot = variables_to_plot
        for var in variables_to_plot.keys():
            self.add_output_file(f'{var}_init.png')
            self.add_output_file(f'{var}_final.png')
            self.add_output_file(f'{var}_diff.png')

        """
        """
        run_duration = config.getfloat('convergence_forward',
                                       'run_duration')

        variables_to_plot = self.variables_to_plot
        ds_init = ds_init[variables_to_plot.keys()].isel(Time=0, nVertLevels=0)

        # Visualization at halfway around the globe (provided run duration is
        # set to the time needed to circumnavigate the globe)
        tidx = time_index_from_xtime(ds_out.xtime.values,
                                     run_duration * s_per_hour / 2.)
        ds_mid = ds_out[variables_to_plot.keys()].isel(Time=tidx,
                                                       nVertLevels=0)

        # Visualization at all the way around the globe
                                     run_duration * s_per_hour)
        ds_final = ds_out[variables_to_plot.keys()].isel(Time=tidx,
                                                         nVertLevels=0)

        for var, section_name in variables_to_plot.items():
            colormap_section = f'sphere_transport_viz_{section_name}'
                out_filename=f'{var}_init.png',
                plot_land=False,
                central_longitude=180.)
                da=ds_mid[var],
                out_filename=f'{var}_mid.png',
                title=f'{mesh_name} {var} after {run_duration / 48.:g} days',
                da=ds_final[var],
                out_filename=f'{var}_final.png',
                title=f'{mesh_name} {var} after {run_duration / 24.:g} days',
                da=ds_final[var] - ds_init[var],
                out_filename=f'{var}_diff.png',
                colormap_section=f'{colormap_section}_diff',
                title=f'Difference in {mesh_name} {var} from initial '
                      f'condition after {run_duration / 24.:g} days',

from importlib import resources

from netCDF4 import Dataset

from polaris.ocean.vertical.grid_1d.index_tanh_dz import (
    create_index_tanh_dz_grid,
)
from polaris.ocean.vertical.grid_1d.tanh_dz import create_tanh_dz_grid

def generate_1d_grid(config):
    """
    Generate a vertical grid for a test case, using the config options in the
    ``vertical_grid`` section

        Configuration options with parameters used to construct the vertical
        grid

    Returns
    -------
    interfaces : numpy.ndarray
        A 1D array of positive depths for layer interfaces in meters
    """
    section = config['vertical_grid']
    grid_type = section.get('grid_type')
    if grid_type == 'uniform':
        interfaces = _generate_uniform(vert_levels)
    elif grid_type == 'tanh_dz':
        min_layer_thickness = section.getfloat('min_layer_thickness')
        max_layer_thickness = section.getfloat('max_layer_thickness')
        interfaces = create_tanh_dz_grid(vert_levels,
                                         bottom_depth,
                                         min_layer_thickness,
                                         max_layer_thickness)

    elif grid_type == 'index_tanh_dz':
        transition_levels = section.getfloat('transition_levels')
        interfaces = create_index_tanh_dz_grid(
            vert_levels,
            bottom_depth,
            min_layer_thickness,
            max_layer_thickness,
            transition_levels)

    elif grid_type in ['60layerPHC', '80layerE3SMv1', '100layerE3SMv1']:
        interfaces = _read_json(grid_type)
    else:
        raise ValueError(f'Unexpected grid type: {grid_type}')

    if config.has_option('vertical_grid', 'bottom_depth') and \
            grid_type != 'tanh_dz':
        # renormalize to the requested range
        interfaces = (bottom_depth / interfaces[-1]) * interfaces

    return interfaces

def write_1d_grid(interfaces, out_filename):
    """
    write the vertical grid to a file

        MPAS file name for output of vertical grid
    """

    nz = len(interfaces) - 1

    # open a new netCDF file for writing.
    ncfile = Dataset(out_filename, 'w')
    # create the depth_t dimension.
    ncfile.createDimension('nVertLevels', nz)

    refBottomDepth = ncfile.createVariable(
        'refBottomDepth', np.dtype('float64').char, ('nVertLevels',))
    refMidDepth = ncfile.createVariable(
        'refMidDepth', np.dtype('float64').char, ('nVertLevels',))
    refLayerThickness = ncfile.createVariable(
        'refLayerThickness', np.dtype('float64').char, ('nVertLevels',))

    botDepth = interfaces[1:]
    midDepth = 0.5 * (interfaces[0:-1] + interfaces[1:])

    refBottomDepth[:] = botDepth
    refMidDepth[:] = midDepth
    refLayerThickness[:] = interfaces[1:] - interfaces[0:-1]
    ncfile.close()

def add_1d_grid(config, ds):
    """
    Add a 1D vertical grid based on the config options in the ``vertical_grid``
    section to a mesh data set

    The following variables are added to the mesh:
    * ``refTopDepth`` - the positive-down depth of the top of each ref. level
    * ``refZMid`` - the positive-down depth of the middle of each ref. level
    * ``refBottomDepth`` - the positive-down depth of the bottom of each ref.
      level
    * ``refInterfaces`` - the positive-down depth of the interfaces between
      ref. levels (with ``nVertLevels`` + 1 elements).
    There is considerable redundancy between these variables but each is
    sometimes convenient.

        grid

    ds : xarray.Dataset
        A data set to add the grid variables to
    """

    interfaces = generate_1d_grid(config=config)

    ds['refTopDepth'] = ('nVertLevels', interfaces[0:-1])
    ds['refZMid'] = ('nVertLevels', -0.5 * (interfaces[1:] + interfaces[0:-1]))
    ds['refBottomDepth'] = ('nVertLevels', interfaces[1:])
    ds['refInterfaces'] = ('nVertLevelsP1', interfaces)

def _generate_uniform(vert_levels):
    """ Generate uniform layer interfaces between 0 and 1 """
    interfaces = np.linspace(0., 1., vert_levels + 1)

def _read_json(grid_type):
    """ Read the grid interfaces from a json file """

    filename = f'{grid_type}.json'
    with resources.open_text("polaris.ocean.vertical.grid_1d", filename) as \
            data_file:
        data = json.load(data_file)
        interfaces = np.array(data)

from polaris.ocean.vertical.sigma import (
    init_sigma_vertical_coord,
    update_sigma_layer_thickness,
)
from polaris.ocean.vertical.zlevel import (
    init_z_level_vertical_coord,
    update_z_level_layer_thickness,
)
from polaris.ocean.vertical.zstar import (
    init_z_star_vertical_coord,
    update_z_star_layer_thickness,
)

def init_vertical_coord(config, ds):
    """
    Create a vertical coordinate based on the config options in the
    ``vertical_grid`` section and the ``bottomDepth`` and ``ssh`` variables of
    the mesh data set.

    The following new variables will be added to the data set:

      * ``minLevelCell`` - the index of the top valid layer

      * ``maxLevelCell`` - the index of the bottom valid layer

      * ``cellMask`` - a mask of where cells are valid

      * ``layerThickness`` - the thickness of each layer

      * ``restingThickness`` - the thickness of each layer stretched as if
        ``ssh = 0``

      * ``zMid`` - the elevation of the midpoint of each layer

    So far, all supported coordinates make use of a 1D reference vertical grid.
    The following variables associated with that field are also added to the
    mesh:

      * ``refTopDepth`` - the positive-down depth of the top of each ref. level

      * ``refZMid`` - the positive-down depth of the middle of each ref. level

      * ``refBottomDepth`` - the positive-down depth of the bottom of each ref.
        level

      * ``refInterfaces`` - the positive-down depth of the interfaces between
        ref. levels (with ``nVertLevels`` + 1 elements).

      * ``vertCoordMovementWeights`` - the weights (all ones) for coordinate
        movement

        grid

        A data set containing ``bottomDepth`` and ``ssh`` variables used to
        construct the vertical coordinate
    """

    for var in ['bottomDepth', 'ssh']:
        if var not in ds:
            raise ValueError(f'{var} must be added to ds before this call.')

    if 'Time' in ds.ssh.dims:
        # drop it for now, we'll add it back at the end
        ds['ssh'] = ds.ssh.isel(Time=0)

    coord_type = config.get('vertical_grid', 'coord_type')

    if coord_type == 'z-level':
        init_z_level_vertical_coord(config, ds)
    elif coord_type == 'z-star':
        init_z_star_vertical_coord(config, ds)
    elif coord_type == 'sigma':
        init_sigma_vertical_coord(config, ds)
    elif coord_type == 'haney-number':
        raise ValueError('Haney Number coordinate not yet supported.')
    else:
        raise ValueError(f'Unknown coordinate type {coord_type}')

    # recompute the cell mask since min/max indices may have changed
    ds['cellMask'] = _compute_cell_mask(ds.minLevelCell, ds.maxLevelCell,
                                        ds.sizes['nVertLevels'])

    # mask layerThickness and restingThickness
    ds['layerThickness'] = ds.layerThickness.where(ds.cellMask)
    ds['restingThickness'] = ds.restingThickness.where(ds.cellMask)

    # add (back) Time dimension
    ds['ssh'] = ds.ssh.expand_dims(dim='Time', axis=0)
    ds['layerThickness'] = ds.layerThickness.expand_dims(dim='Time', axis=0)
    ds['restingThickness'] = \
        ds.restingThickness.expand_dims(dim='Time', axis=0)

    ds['zMid'] = _compute_zmid_from_layer_thickness(
        ds.layerThickness, ds.ssh, ds.cellMask)

    # fortran 1-based indexing
    ds['minLevelCell'] = ds.minLevelCell + 1
    ds['maxLevelCell'] = ds.maxLevelCell + 1

def update_layer_thickness(config, ds):
    """
    Update the layer thicknesses in ds after the vertical coordinate has
    already been initialized based on the ``bottomDepth`` and ``ssh``
    variables of the mesh data set.

        grid

    """

        update_z_level_layer_thickness(config, ds)
        update_z_star_layer_thickness(config, ds)
        update_sigma_layer_thickness(config, ds)
    else:

def _compute_cell_mask(minLevelCell, maxLevelCell, nVertLevels):
    cellMask = []
    for zIndex in range(nVertLevels):
        mask = np.logical_and(zIndex >= minLevelCell,
                              zIndex <= maxLevelCell)
        cellMask.append(mask)
    cellMaskArray = xr.DataArray(cellMask, dims=['nVertLevels', 'nCells'])
    cellMaskArray = cellMaskArray.transpose('nCells', 'nVertLevels')
    return cellMaskArray

def _compute_zmid_from_layer_thickness(layerThickness, ssh, cellMask):
    """
    Compute zMid from ssh and layerThickness for any vertical coordinate

    layerThickness : xarray.DataArray
        The thickness of each layer

    ssh : xarray.DataArray
        The sea surface height

    cellMask : xarray.DataArray
        A boolean mask of where there are valid cells

    Returns
    -------
    zMid : xarray.DataArray
        The elevation of layer centers
    """

    zTop = ssh.copy()
    nVertLevels = layerThickness.sizes['nVertLevels']
    zMid = []
        mask = cellMask.isel(nVertLevels=zIndex)
        thickness = layerThickness.isel(nVertLevels=zIndex).where(mask, 0.)
        z = (zTop - 0.5 * thickness).where(mask)
        zMid.append(z)
        zTop -= thickness
    zMid = xr.concat(zMid, dim='nVertLevels').transpose('Time', 'nCells',
                                                        'nVertLevels')
    return zMid

import numpy

def alter_bottom_depth(config, bottomDepth, refBottomDepth, maxLevelCell):
    """
    Alter ``bottomDepth`` and ``maxLevelCell`` for full or partial top cells,
    if requested

        grid

    bottomDepth : xarray.DataArray
        The positive-down depth of the seafloor

    refBottomDepth : xarray.DataArray
        A 1D array of positive-down depths of the bottom of each z level

    maxLevelCell : xarray.DataArray
        The zero-based index of the bottom valid level

    Returns
    -------
        The positive-down depth of the seafloor, after alteration

        The zero-based index of the bottom valid level, after alteration
    """
    partial_cell_type = 'none'
    min_pc_fraction = 0.
    if config.has_option('vertical_grid', 'partial_cell_type'):

        partial_cell_type = section.get('partial_cell_type').lower()
        min_pc_fraction = section.getfloat('min_pc_fraction')

    if partial_cell_type == 'full':
        bottomDepth = _compute_full_cells_depth(
            refBottomDepth, maxLevelCell)
    elif partial_cell_type == 'partial':
        bottomDepth, maxLevelCell = _alter_bottom_depth_for_partial_cells(
            bottomDepth, refBottomDepth, maxLevelCell, min_pc_fraction)
    elif partial_cell_type != 'none':
        raise ValueError(f'Unexpected partial cell type {partial_cell_type}')

    return bottomDepth, maxLevelCell

def alter_ssh(config, ssh, refBottomDepth, minLevelCell):
    """
    Alter ``ssh`` and ``maxLevelCell`` for full or partial top cells,

        grid

    minLevelCell : xarray.DataArray
        The zero-based index of the top valid level

    Returns
    -------
        The sea surface height, after alteration

        The zero-based index of the top valid level, after alteration
    """

        ssh = _compute_full_cells_depth(
            refBottomDepth, minLevelCell - 1)
        ssh, minLevelCell = _alter_ssh_for_partial_cells(
            ssh, refBottomDepth, minLevelCell, min_pc_fraction)

    return ssh, minLevelCell

def _compute_full_cells_depth(refBottomDepth, levelIndex):
    """
    Compute the full cell bottom depth given a level index
    """

    depth = refBottomDepth.isel(nVertLevels=levelIndex).where(
        levelIndex >= 0, other=0.)
    return depth

def _alter_bottom_depth_for_partial_cells(bottomDepth, refBottomDepth,
                                          maxLevelCell, min_pc_fraction):
    """
    Alter bottomDepth and maxLevelCell for partial cells
    """

    fullBot = _compute_full_cells_depth(refBottomDepth, maxLevelCell)

    fullTop = _compute_full_cells_depth(refBottomDepth, maxLevelCell - 1)

    fullThickness = fullBot - fullTop

    minBottomDepth = fullBot - (1. - min_pc_fraction) * fullThickness

    minBottomDepthMid = 0.5 * (minBottomDepth + fullTop)

    # where the bottom depth is far too shallow, we're going to fill in the
    # last level
    mask = bottomDepth < minBottomDepthMid
    maxLevelCell = xarray.where(mask, maxLevelCell - 1, maxLevelCell)
    bottomDepth = xarray.where(mask, fullTop, bottomDepth)

    # where the bottom depth only a bit too shallows, we move it deeper
    mask = numpy.logical_and(numpy.logical_not(mask),
                             bottomDepth < minBottomDepth)
    bottomDepth = xarray.where(mask, minBottomDepth, bottomDepth)

def _alter_ssh_for_partial_cells(ssh, refBottomDepth, minLevelCell,
                                 min_pc_fraction):
    """
    Alter ssh and minLevelCell for partial cells
    """

    zBot = -_compute_full_cells_depth(refBottomDepth, minLevelCell)

    zTop = -_compute_full_cells_depth(refBottomDepth, minLevelCell - 1)

    fullThickness = zTop - zBot

    minSSH = zBot + min_pc_fraction * fullThickness

    minSSHMid = 0.5 * (minSSH + zBot)

    # where the SSH is far too deep, we're going to fill the current top level
    mask = ssh < minSSHMid
    minLevelCell = xarray.where(mask, minLevelCell + 1, minLevelCell)
    ssh = xarray.where(mask, zBot, ssh)

    # where the SSH only a bit too deep, we move it shallower
    mask = numpy.logical_and(numpy.logical_not(mask), ssh < minSSH)
    ssh = xarray.where(mask, minSSH, ssh)

from polaris.ocean.vertical.grid_1d import add_1d_grid

def init_sigma_vertical_coord(config, ds):
    """
    Create a sigma (terrain-following) vertical coordinate based on the config
    options in the `vertical_grid`` section and the ``bottomDepth`` and
    ``ssh`` variables of the mesh data set.

        ``ssh = 0``

    The sigma coordinate makes use of a 1D reference vertical grid. The
    following variables associated with that field are also added to the mesh:

        level

        movement

        grid

    """
    add_1d_grid(config, ds)

    ds['vertCoordMovementWeights'] = xarray.ones_like(ds.refBottomDepth)

    valid = ds.bottomDepth > -ds.ssh
    cell_mask, _ = xarray.broadcast(valid, ds.refBottomDepth)
    ds['cellMask'] = cell_mask.transpose('nCells', 'nVertLevels')

    ds['minLevelCell'] = xarray.zeros_like(ds.bottomDepth, dtype=int)
    ds['maxLevelCell'] = (ds.sizes['nVertLevels'] - 1 *
                          xarray.ones_like(ds.bottomDepth, dtype=int))

    resting_ssh = xarray.zeros_like(ds.bottomDepth)

    ds['restingThickness'] = compute_sigma_layer_thickness(
        ds.refInterfaces, resting_ssh, ds.bottomDepth)

    ds['layerThickness'] = compute_sigma_layer_thickness(
        ds.refInterfaces, ds.ssh, ds.bottomDepth)

def update_sigma_layer_thickness(config, ds):
    """
    Update the sigma vertical coordinate layer thicknesses based on the
    ``bottomDepth`` and ``ssh`` variables of the mesh data set.

        grid

    """

def compute_sigma_layer_thickness(ref_interfaces, ssh, bottom_depth,
                                  max_level=None):
    """
    Compute sigma layer thickness by stretching restingThickness based on ssh
    and bottom_depth

    ref_interfaces : xarray.DataArray
        the interfaces of the reference coordinate used to define sigma layer
        spacing.  The interfaces are renormalized to be between 0 and 1

    bottom_depth : xarray.DataArray

    max_level : int, optional
        The maximum number of levels used for the sigma coordinate.  The
        default is ``nVertLevels``.

    Returns
    -------
    layer_thickness : xarray.DataArray
    """
    n_vert_levels = ref_interfaces.sizes['nVertLevelsP1'] - 1
    if max_level is None:
        max_level = n_vert_levels

    # renormalize ref_interfaces to a coordinate between 0 and 1
    stretch = (ref_interfaces - ref_interfaces.min()).values
    stretch = stretch / stretch[max_level]

    column_thickness = ssh + bottom_depth
    valid = bottom_depth > -ssh
    layer_thickness = []

    for z_index in range(max_level):
        thickness = ((stretch[z_index + 1] - stretch[z_index]) *
                     column_thickness)
        thickness = thickness.where(valid, 0.)
        layer_thickness.append(thickness)
    for z_index in range(max_level, n_vert_levels):
        layer_thickness.append(xarray.zeros_like(bottom_depth))
    layer_thickness_array = xarray.DataArray(layer_thickness,
                                             dims=['nVertLevels', 'nCells'])
    layer_thickness_array = layer_thickness_array.transpose(
        'nCells', 'nVertLevels')
    return layer_thickness_array

from polaris.ocean.vertical.partial_cells import alter_bottom_depth, alter_ssh

def init_z_level_vertical_coord(config, ds):
    """
    Create a z-level vertical coordinate based on the config options in the

        ``ssh = 0``

    mesh:

        level

        movement

        grid

    """

    ds['minLevelCell'], ds['maxLevelCell'], ds['cellMask'] = \
        compute_min_max_level_cell(ds.refTopDepth, ds.refBottomDepth, ds.ssh,
                                   ds.bottomDepth)

    ds['bottomDepth'], ds['maxLevelCell'] = alter_bottom_depth(
        config, ds.bottomDepth, ds.refBottomDepth, ds.maxLevelCell)

    ds['ssh'], ds['minLevelCell'] = alter_ssh(
        config, ds.ssh, ds.refBottomDepth, ds.minLevelCell)

    ds['layerThickness'] = compute_z_level_layer_thickness(
        ds.refTopDepth, ds.refBottomDepth, ds.ssh, ds.bottomDepth,
        ds.minLevelCell, ds.maxLevelCell)

    ds['restingThickness'] = compute_z_level_resting_thickness(
        ds.layerThickness, ds.ssh, ds.bottomDepth, ds.minLevelCell,
        ds.maxLevelCell)

def update_z_level_layer_thickness(config, ds):
    """
    Update the z-level vertical coordinate layer thicknesses based on the

        grid

    """

def compute_min_max_level_cell(refTopDepth, refBottomDepth, ssh, bottomDepth,
                               min_vert_levels=None, min_layer_thickness=None):
    """
    Compute ``minLevelCell`` and ``maxLevelCell`` indices as well as a cell
    mask for the given reference grid and top and bottom topography.

    refTopDepth : xarray.DataArray
        A 1D array of positive-down depths of the top of each z level

    Returns
    -------

    """
    if min_layer_thickness is not None:
        valid = bottomDepth + min_layer_thickness * min_vert_levels >= -ssh
    else:
        valid = bottomDepth > -ssh

    aboveTopMask = (refBottomDepth <= -ssh).transpose('nCells', 'nVertLevels')
    aboveBottomMask = (refTopDepth < bottomDepth).transpose(
    aboveBottomMask = numpy.logical_and(aboveBottomMask, valid)

    minLevelCell = (aboveTopMask.sum(dim='nVertLevels')).where(valid, 0)
    maxLevelCell = (aboveBottomMask.sum(dim='nVertLevels') - 1).where(valid, 0)
    if min_vert_levels is not None:
        maxLevelCell = numpy.maximum(maxLevelCell,
                                     minLevelCell + min_vert_levels - 1)
    cellMask = numpy.logical_and(numpy.logical_not(aboveTopMask),
                                 aboveBottomMask)
    cellMask = numpy.logical_and(cellMask, valid)

    return minLevelCell, maxLevelCell, cellMask

def compute_z_level_layer_thickness(refTopDepth, refBottomDepth, ssh,
                                    bottomDepth, minLevelCell, maxLevelCell):
    """
    Compute z-level layer thickness from ssh and bottomDepth

    Returns
    -------
        The thickness of each layer (level)
    """

    nVertLevels = refBottomDepth.sizes['nVertLevels']
    layerThickness = []
        mask = numpy.logical_and(zIndex >= minLevelCell,
                                 zIndex <= maxLevelCell)
        zTop = numpy.minimum(ssh, -refTopDepth[zIndex])
        zBot = numpy.maximum(-bottomDepth, -refBottomDepth[zIndex])
        thickness = (zTop - zBot).where(mask, 0.)
        layerThickness.append(thickness)
    layerThicknessArray = xarray.DataArray(layerThickness,
                                           dims=['nVertLevels', 'nCells'])
    layerThicknessArray = layerThicknessArray.transpose(
    return layerThicknessArray

def compute_z_level_resting_thickness(layerThickness, ssh, bottomDepth,
                                      minLevelCell, maxLevelCell):
    """
    Compute z-level resting thickness by "unstretching" layerThickness
    based on ssh and bottomDepth

    Returns
    -------
    restingThickness : xarray.DataArray
        The thickness of z-star layers when ssh = 0
    """

    restingThickness = []

    layerStretch = bottomDepth / (ssh + bottomDepth)
        thickness = layerStretch * layerThickness.isel(
            nVertLevels=zIndex)
        thickness = thickness.where(mask, 0.)
        restingThickness.append(thickness)
    restingThicknessArray = xarray.DataArray(restingThickness,
    restingThicknessArray = restingThicknessArray.transpose(
    return restingThicknessArray

from polaris.ocean.vertical.partial_cells import alter_bottom_depth
    compute_min_max_level_cell,
    compute_z_level_layer_thickness,
)

def init_z_star_vertical_coord(config, ds):
    """
    Create a z-star vertical coordinate based on the config options in the

        ``ssh = 0``

    mesh:

        level

        movement

        grid

    """

    restingSSH = xarray.zeros_like(ds.bottomDepth)
    min_vert_levels = config.getint('vertical_grid', 'min_vert_levels')
    min_layer_thickness = config.getfloat('vertical_grid',
                                          'min_layer_thickness')
        compute_min_max_level_cell(ds.refTopDepth, ds.refBottomDepth,
                                   restingSSH, ds.bottomDepth,
                                   min_vert_levels=min_vert_levels,
                                   min_layer_thickness=min_layer_thickness)

    ds['restingThickness'] = compute_z_level_layer_thickness(
        ds.refTopDepth, ds.refBottomDepth, restingSSH, ds.bottomDepth,

    ds['layerThickness'] = _compute_z_star_layer_thickness(
        ds.restingThickness, ds.ssh, ds.bottomDepth, ds.minLevelCell,

def update_z_star_layer_thickness(config, ds):
    """
    Update the z-star vertical coordinate layer thicknesses based on the

        grid

    """

def _compute_z_star_layer_thickness(restingThickness, ssh, bottomDepth,
                                    minLevelCell, maxLevelCell):
    """
    Compute z-star layer thickness by stretching restingThickness based on ssh
    and bottomDepth

    Returns
    -------
    """

    nVertLevels = restingThickness.sizes['nVertLevels']

    layerStretch = (ssh + bottomDepth) / \
        restingThickness.sum(dim='nVertLevels')
        thickness = layerStretch * restingThickness.isel(nVertLevels=zIndex)

from polaris.ocean.viz.transect.horiz import (
    find_planar_transect_cells_and_weights,
    find_spherical_transect_cells_and_weights,
    make_triangle_tree,
    mesh_to_triangles,
)

def compute_transect(x, y, ds_horiz_mesh, layer_thickness, bottom_depth,
                     min_level_cell, max_level_cell, spherical=False):
    """
    build a sequence of quads showing the transect intersecting mpas cells.
    This can be used to plot transects of fields with dimensions ``nCells`` and
    ``nVertLevels`` using :py:func:`polaris.ocean.viz.plot_transect()`

    x : xarray.DataArray
        The x or longitude coordinate of the transect

    y : xarray.DataArray
        The y or latitude coordinate of the transect

    ds_horiz_mesh : xarray.Dataset
        The horizontal MPAS mesh to use for plotting

        The layer thickness at a particular instant in time.
        `layerThickness.isel(Time=tidx)` to select a particular time index
        `tidx` if the original data array contains `Time`.

        the (positive down) depth of the seafloor on the MPAS mesh

    min_level_cell : xarray.DataArray
        the vertical zero-based index of the sea surface on the MPAS mesh

    max_level_cell : xarray.DataArray
        the vertical zero-based index of the bathymetry on the MPAS mesh

    spherical : bool, optional
        Whether the x and y coordinates are latitude and longitude in degrees

    Returns
    -------
    ds_transect : xarray.Dataset
        The transect dataset, see
        :py:func:`polaris.ocean.viz.transect.vert.find_transect_levels_and_weights()`
        for details
    """  # noqa: E501

    ds_tris = mesh_to_triangles(ds_horiz_mesh)

    triangle_tree = make_triangle_tree(ds_tris)

    if spherical:
        ds_horiz_transect = find_spherical_transect_cells_and_weights(
            x, y, ds_tris, ds_horiz_mesh, triangle_tree, degrees=True)
    else:
        ds_horiz_transect = find_planar_transect_cells_and_weights(
            x, y, ds_tris, ds_horiz_mesh, triangle_tree)

    # mask horizontal transect to valid cells (max_level_cell >= 0)
    cell_indices = ds_horiz_transect.horizCellIndices
    seg_mask = max_level_cell.isel(nCells=cell_indices).values >= 0
    node_mask = np.zeros(ds_horiz_transect.sizes['nNodes'], dtype=bool)
    node_mask[0:-1] = seg_mask
    node_mask[1:] = np.logical_or(node_mask[1:], seg_mask)

    ds_horiz_transect = ds_horiz_transect.isel(nSegments=seg_mask,
                                               nNodes=node_mask)

    ds_transect = find_transect_levels_and_weights(
        ds_horiz_transect=ds_horiz_transect, layer_thickness=layer_thickness,
        bottom_depth=bottom_depth, min_level_cell=min_level_cell,
        max_level_cell=max_level_cell)

    ds_transect.compute()

    return ds_transect

def find_transect_levels_and_weights(ds_horiz_transect, layer_thickness,
                                     bottom_depth, min_level_cell,
                                     max_level_cell):
    """
    Construct a vertical coordinate for a transect produced by
    :py:func:`polaris.ocean.viz.transect.horiz.find_spherical_transect_cells_and_weights()`
    or :py:func:`polaris.ocean.viz.transect.horiz.find_planar_transect_cells_and_weights()`.
    Also, compute interpolation weights such that observations at points on the
    original transect and with vertical coordinate ``transectZ`` can be
    bilinearly interpolated to the nodes of the transect.

    ds_horiz_transect : xarray.Dataset
        A dataset that defines nodes of the transect

        layer thicknesses on the MPAS mesh

    Returns
    -------
        A dataset that contains nodes and cells that make up a 2D transect.

        There are ``nSegments`` horizontal and ``nHalfLevels`` vertical
        transect cells (quadrilaterals), bounded by ``nHorizNodes`` horizontal
        and ``nVertNodes`` vertical nodes (corners).

        In addition to the variables and coordinates in the input
        ``ds_transect``, the output dataset contains:

            - ``validCells``, ``validNodes``: which transect cells and nodes
              are valid (above the bathymetry and below the sea surface)

            - zTransectNode: the vertical height of each triangle node
            - ssh, zSeaFloor: the sea-surface height and sea-floor height at
              each node of each transect segment

            - ``cellIndices``: the MPAS-Ocean cell of a given transect segment
            - ``levelIndices``: the MPAS-Ocean vertical level of a given
              transect level

            - ``interpCellIndices``, ``interpLevelIndices``: the MPAS-Ocean
              cells and levels from which the value at a given transect cell is
              interpolated.  This can involve up to
              ``nHorizWeights * nVertWeights = 12`` different cells and levels.
            - interpCellWeights: the weight to multiply each field value by
              to perform interpolation to a transect cell.

            - ``dInterfaceSegment``, ``zInterfaceSegment`` - segments that can
              be used to plot the interfaces between MPAS-Ocean layers

            - ``dCellBoundary``, ``zCellBoundary`` - segments that can
              be used to plot the vertical boundaries between MPAS-Ocean cells

        Interpolation of a DataArray from MPAS cells and levels to transect
        cells can be performed with
        :py:func:`polaris.ocean.viz.transect.vert.interp_mpas_to_transect_cells()`.
        Similarly, interpolation to transect nodes can be performed with
        :py:func:`polaris.ocean.viz.transect.vert.interp_mpas_to_transect_nodes()`.
    if 'Time' in layer_thickness.dims:
        raise ValueError('Please select a single time level in layer '
                         'thickness.')

    ds_transect_cells = ds_horiz_transect.rename({'nNodes': 'nHorizNodes'})

    (z_half_interface, ssh, z_seafloor, interp_cell_indices,
     interp_cell_weights, valid_transect_cells,
     level_indices) = _get_vertical_coordinate(
        ds_transect_cells, layer_thickness, bottom_depth, min_level_cell,
        max_level_cell)

    ds_transect_cells['zTransectNode'] = z_half_interface

    ds_transect_cells['ssh'] = ssh
    ds_transect_cells['zSeafloor'] = z_seafloor

    ds_transect_cells['cellIndices'] = ds_transect_cells.horizCellIndices
    ds_transect_cells['levelIndices'] = level_indices
    ds_transect_cells['validCells'] = valid_transect_cells

    d_interface_seg, z_interface_seg = _get_interface_segments(
        z_half_interface, ds_transect_cells.dNode, valid_transect_cells)

    ds_transect_cells['dInterfaceSegment'] = d_interface_seg
    ds_transect_cells['zInterfaceSegment'] = z_interface_seg

    d_cell_boundary, z_cell_boundary = _get_cell_boundary_segments(
        ssh, z_seafloor, ds_transect_cells.dNode,
        ds_transect_cells.horizCellIndices)

    ds_transect_cells['dCellBoundary'] = d_cell_boundary
    ds_transect_cells['zCellBoundary'] = z_cell_boundary

    interp_level_indices, interp_cell_weights, valid_nodes = \
        _get_interp_indices_and_weights(layer_thickness, interp_cell_indices,
                                        interp_cell_weights, level_indices,
                                        valid_transect_cells)

    ds_transect_cells['interpCellIndices'] = interp_cell_indices
    ds_transect_cells['interpLevelIndices'] = interp_level_indices
    ds_transect_cells['interpCellWeights'] = interp_cell_weights
    ds_transect_cells['validNodes'] = valid_nodes

    dims = ['nSegments', 'nHalfLevels', 'nHorizNodes', 'nVertNodes',
            'nInterfaceSegments', 'nCellBoundaries', 'nHorizBounds',
            'nVertBounds', 'nHorizWeights', 'nVertWeights']
    for dim in ds_transect_cells.dims:
        if dim not in dims:
            dims.insert(0, dim)
    ds_transect_cells = ds_transect_cells.transpose(*dims)

    return ds_transect_cells

def interp_mpas_to_transect_cells(ds_transect, da):
    """
    Interpolate an MPAS-Ocean DataArray with dimensions ``nCells`` by
    ``nVertLevels`` to transect cells

        A dataset that defines an MPAS-Ocean transect, the results of calling
        ``find_transect_levels_and_weights()``

        An MPAS-Ocean field with dimensions `nCells`` and ``nVertLevels``
        (possibly among others)

    Returns
    -------
    da_cells : xarray.DataArray
        The data array interpolated to transect cells with dimensions
        ``nSegments`` and ``nHalfLevels`` (in addition to whatever
        dimensions were in ``da`` besides ``nCells`` and ``nVertLevels``)
    """

    cell_indices = ds_transect.cellIndices
    level_indices = ds_transect.levelIndices

    da_cells = da.isel(nCells=cell_indices, nVertLevels=level_indices)
    da_cells = da_cells.where(ds_transect.validCells)

    return da_cells

def interp_mpas_to_transect_nodes(ds_transect, da):
    """
    ``nVertLevels`` to transect nodes, linearly interpolating fields between
    the closest neighboring cells

    Returns
    -------
    da_nodes : xarray.DataArray
        The data array interpolated to transect nodes with dimensions
        ``nHorizNodes`` and ``nVertNodes`` (in addition to whatever
    """

    interp_cell_indices = ds_transect.interpCellIndices
    interp_level_indices = ds_transect.interpLevelIndices
    interp_cell_weights = ds_transect.interpCellWeights

    da = da.isel(nCells=interp_cell_indices, nVertLevels=interp_level_indices)

    da_nodes = (da * interp_cell_weights).sum(
        dim=('nHorizWeights', 'nVertWeights'))

    da_nodes = da_nodes.where(ds_transect.validNodes)

    return da_nodes

def _get_vertical_coordinate(ds_transect, layer_thickness, bottom_depth,
                             min_level_cell, max_level_cell):
    n_horiz_nodes = ds_transect.sizes['nHorizNodes']
    n_segments = ds_transect.sizes['nSegments']
    n_vert_levels = layer_thickness.sizes['nVertLevels']

    # we assume below that there is a segment (whether valid or invalid)
    # connecting each pair of adjacent nodes
    assert n_horiz_nodes == n_segments + 1

    interp_horiz_cell_indices = ds_transect.interpHorizCellIndices
    interp_horiz_cell_weights = ds_transect.interpHorizCellWeights

    bottom_depth_interp = bottom_depth.isel(nCells=interp_horiz_cell_indices)
    layer_thickness_interp = layer_thickness.isel(
        nCells=interp_horiz_cell_indices)

    cell_mask_interp = _get_cell_mask(interp_horiz_cell_indices,
                                      min_level_cell, max_level_cell,
                                      n_vert_levels)
    layer_thickness_interp = layer_thickness_interp.where(cell_mask_interp, 0.)

    ssh_interp = (-bottom_depth_interp +
                  layer_thickness_interp.sum(dim='nVertLevels'))

    interp_mask = np.logical_and(interp_horiz_cell_indices > 0,
                                 cell_mask_interp)

    interp_cell_weights = interp_mask * interp_horiz_cell_weights
    weight_sum = interp_cell_weights.sum(dim='nHorizWeights')

    cell_indices = ds_transect.horizCellIndices

    valid_cells = _get_cell_mask(cell_indices, min_level_cell, max_level_cell,
                                 n_vert_levels)

    valid_cells = valid_cells.transpose('nSegments', 'nVertLevels').values

    valid_nodes = np.zeros((n_horiz_nodes, n_vert_levels), dtype=bool)
    valid_nodes[0:-1, :] = valid_cells
    valid_nodes[1:, :] = np.logical_or(valid_nodes[1:, :], valid_cells)

    valid_nodes = xr.DataArray(dims=('nHorizNodes', 'nVertLevels'),
                               data=valid_nodes)

    valid_weights = valid_nodes.broadcast_like(interp_cell_weights)
    interp_cell_weights = \
        (interp_cell_weights / weight_sum).where(valid_weights)

    layer_thickness_transect = (layer_thickness_interp *
                                interp_cell_weights).sum(dim='nHorizWeights')

    interp_mask = max_level_cell.isel(nCells=interp_horiz_cell_indices) >= 0
    interp_horiz_cell_weights = interp_mask * interp_horiz_cell_weights
    weight_sum = interp_horiz_cell_weights.sum(dim='nHorizWeights')
    interp_horiz_cell_weights = \
        (interp_horiz_cell_weights / weight_sum).where(interp_mask)

    ssh_transect = (ssh_interp *
                    interp_horiz_cell_weights).sum(dim='nHorizWeights')

    z_bot = ssh_transect - layer_thickness_transect.cumsum(dim='nVertLevels')
    z_mid = z_bot + 0.5 * layer_thickness_transect

    z_half_interfaces = [ssh_transect]
    for z_index in range(n_vert_levels):
        z_half_interfaces.extend([z_mid.isel(nVertLevels=z_index),
                                  z_bot.isel(nVertLevels=z_index)])

    z_half_interface = xr.concat(z_half_interfaces, dim='nVertNodes')
    z_half_interface = z_half_interface.transpose('nHorizNodes', 'nVertNodes')

    z_seafloor = ssh_transect - layer_thickness_transect.sum(
        dim='nVertLevels')

    valid_transect_cells = np.zeros((n_segments, 2 * n_vert_levels),
                                    dtype=bool)
    valid_transect_cells[:, 0::2] = valid_cells
    valid_transect_cells[:, 1::2] = valid_cells
    valid_transect_cells = xr.DataArray(dims=('nSegments', 'nHalfLevels'),
                                        data=valid_transect_cells)

    level_indices = np.zeros(2 * n_vert_levels, dtype=int)
    level_indices[0::2] = np.arange(n_vert_levels)
    level_indices[1::2] = np.arange(n_vert_levels)
    level_indices = xr.DataArray(dims=('nHalfLevels',), data=level_indices)

    return (z_half_interface, ssh_transect, z_seafloor,
            interp_horiz_cell_indices, interp_cell_weights,
            valid_transect_cells, level_indices)

def _get_cell_mask(cell_indices, min_level_cell, max_level_cell,
                   n_vert_levels):
    level_indices = xr.DataArray(data=np.arange(n_vert_levels),
                                 dims='nVertLevels')
    min_level_cell = min_level_cell.isel(nCells=cell_indices)
    max_level_cell = max_level_cell.isel(nCells=cell_indices)

    cell_mask = np.logical_and(
        level_indices >= min_level_cell,
        level_indices <= max_level_cell)

    cell_mask = np.logical_and(cell_mask, cell_indices >= 0)

    return cell_mask

def _get_interface_segments(z_half_interface, d_node, valid_transect_cells):

    d = d_node.broadcast_like(z_half_interface)
    z_interface = z_half_interface.values[:, 0::2]
    d = d.values[:, 0::2]

    n_segments = valid_transect_cells.sizes['nSegments']
    n_half_levels = valid_transect_cells.sizes['nHalfLevels']
    n_vert_levels = n_half_levels // 2

    valid_segs = np.zeros((n_segments, n_vert_levels + 1), dtype=bool)
    valid_segs[:, 0:-1] = valid_transect_cells.values[:, 1::2]
    valid_segs[:, 1:] = np.logical_or(valid_segs[:, 1:],
                                      valid_transect_cells.values[:, 0::2])

    n_interface_segs = np.count_nonzero(valid_segs)

    d_seg = np.zeros((n_interface_segs, 2))
    z_seg = np.zeros((n_interface_segs, 2))
    d_seg[:, 0] = d[0:-1, :][valid_segs]
    d_seg[:, 1] = d[1:, :][valid_segs]
    z_seg[:, 0] = z_interface[0:-1, :][valid_segs]
    z_seg[:, 1] = z_interface[1:, :][valid_segs]

    d_seg = xr.DataArray(dims=('nInterfaceSegments', 'nHorizBounds'),
                         data=d_seg)

    z_seg = xr.DataArray(dims=('nInterfaceSegments', 'nHorizBounds'),
                         data=z_seg)

    return d_seg, z_seg

def _get_cell_boundary_segments(ssh, z_seafloor, d_node, cell_indices):

    n_horiz_nodes = d_node.sizes['nHorizNodes']

    cell_boundary = np.ones(n_horiz_nodes, dtype=bool)
    cell_boundary[1:-1] = cell_indices.values[0:-1] != cell_indices.values[1:]

    n_cell_boundaries = np.count_nonzero(cell_boundary)

    d_seg = np.zeros((n_cell_boundaries, 2))
    z_seg = np.zeros((n_cell_boundaries, 2))
    d_seg[:, 0] = d_node.values[cell_boundary]
    d_seg[:, 1] = d_seg[:, 0]
    z_seg[:, 0] = ssh[cell_boundary]
    z_seg[:, 1] = z_seafloor[cell_boundary]

    d_seg = xr.DataArray(dims=('nCellBoundaries', 'nVertBounds'), data=d_seg)

    z_seg = xr.DataArray(dims=('nCellBoundaries', 'nVertBounds'), data=z_seg)

def _get_interp_indices_and_weights(layer_thickness, interp_cell_indices,
                                    interp_cell_weights, level_indices,
                                    valid_transect_cells):
    n_horiz_nodes = interp_cell_indices.sizes['nHorizNodes']
    n_vert_nodes = 2 * n_vert_levels + 1
    n_vert_weights = 2

    interp_level_indices = -1 * np.ones((n_vert_nodes, n_vert_weights),
                                        dtype=int)
    interp_level_indices[1:, 0] = level_indices.values
    interp_level_indices[0:-1, 1] = level_indices.values

    interp_level_indices = xr.DataArray(dims=('nVertNodes', 'nVertWeights'),
                                        data=interp_level_indices)

    half_level_thickness = 0.5 * layer_thickness.isel(
        nCells=interp_cell_indices, nVertLevels=interp_level_indices)
    half_level_thickness = half_level_thickness.where(
        interp_level_indices >= 0, other=0.)

    # vertical weights are proportional to the half-level thickness
    interp_cell_weights = half_level_thickness * interp_cell_weights.isel(
        nVertLevels=interp_level_indices)

    valid_nodes = np.zeros((n_horiz_nodes, n_vert_nodes), dtype=bool)
    valid_nodes[0:-1, 0:-1] = valid_transect_cells
    valid_nodes[1:, 0:-1] = np.logical_or(valid_nodes[1:, 0:-1],
                                          valid_transect_cells)
    valid_nodes[0:-1, 1:] = np.logical_or(valid_nodes[0:-1, 1:],
    valid_nodes[1:, 1:] = np.logical_or(valid_nodes[1:, 1:],

    valid_nodes = xr.DataArray(dims=('nHorizNodes', 'nVertNodes'),

    weight_sum = interp_cell_weights.sum(dim=('nHorizWeights', 'nVertWeights'))
    out_mask = (weight_sum > 0.).broadcast_like(interp_cell_weights)
    interp_cell_weights = (interp_cell_weights / weight_sum).where(out_mask)

    return interp_level_indices, interp_cell_weights, valid_nodes

from mpas_tools.transects import (
    cartesian_to_lon_lat,
    lon_lat_to_cartesian,
    subdivide_great_circle,
    subdivide_planar,
)
from scipy.spatial import cKDTree
from shapely.geometry import LineString, Point

def mesh_to_triangles(ds_mesh):
    """
    Construct a dataset in which each MPAS cell is divided into the triangles
    connecting pairs of adjacent vertices to cell centers.

    ds_mesh : xarray.Dataset
        An MPAS mesh

    Returns
    -------
    ds_tris : xarray.Dataset
        A dataset that defines triangles connecting pairs of adjacent vertices
        to cell centers as well as the cell index that each triangle is in and
        cell indices and weights for interpolating data defined at cell centers
        to triangle nodes.  ``ds_tris`` includes variables ``triCellIndices``,
        the cell that each triangle is part of; ``nodeCellIndices`` and
        ``nodeCellWeights``, the indices and weights used to interpolate from
        MPAS cell centers to triangle nodes; Cartesian coordinates ``xNode``,
        ``yNode``, and ``zNode``; and ``lonNode``` and ``latNode`` in radians.
        ``lonNode`` is guaranteed to be within 180 degrees of the cell center
        corresponding to ``triCellIndices``.  Nodes always have a
        counterclockwise winding.

    """
    n_vertices_on_cell = ds_mesh.nEdgesOnCell.values
    vertices_on_cell = ds_mesh.verticesOnCell.values - 1
    cells_on_vertex = ds_mesh.cellsOnVertex.values - 1

    on_a_sphere = ds_mesh.attrs['on_a_sphere'].strip() == 'YES'
    is_periodic = False
    x_period = None
    y_period = None
    if not on_a_sphere:
        is_periodic = ds_mesh.attrs['is_periodic'].strip() == 'YES'
        if is_periodic:
            x_period = ds_mesh.attrs['x_period']
            y_period = ds_mesh.attrs['y_period']

    kite_areas_on_vertex = ds_mesh.kiteAreasOnVertex.values

    n_triangles = np.sum(n_vertices_on_cell)

    max_edges = ds_mesh.sizes['maxEdges']
    n_cells = ds_mesh.sizes['nCells']
    if ds_mesh.sizes['vertexDegree'] != 3:
        raise ValueError('mesh_to_triangles only supports meshes with '
                         'vertexDegree = 3')

    # find the third vertex for each triangle
    next_vertex = -1 * np.ones(vertices_on_cell.shape, int)
    for i_vertex in range(max_edges):
        valid = i_vertex < n_vertices_on_cell
        invalid = np.logical_not(valid)
        vertices_on_cell[invalid, i_vertex] = -1
        nv = n_vertices_on_cell[valid]
        cell_indices = np.arange(0, n_cells)[valid]
        i_next = np.where(i_vertex < nv - 1, i_vertex + 1, 0)
        next_vertex[:, i_vertex][valid] = (
            vertices_on_cell[cell_indices, i_next])

    valid = vertices_on_cell >= 0
    vertices_on_cell = vertices_on_cell[valid]
    next_vertex = next_vertex[valid]

    # find the cell index for each triangle
    tri_cell_indices, _ = np.meshgrid(np.arange(0, n_cells),
                                      np.arange(0, max_edges),
                                      indexing='ij')
    tri_cell_indices = tri_cell_indices[valid]

    # find list of cells and weights for each triangle node
    node_cell_indices = -1 * np.ones((n_triangles, 3, 3), dtype=int)
    node_cell_weights = np.zeros((n_triangles, 3, 3))

    # the first node is at the cell center, so the value is just the one from
    # that cell
    node_cell_indices[:, 0, 0] = tri_cell_indices
    node_cell_weights[:, 0, 0] = 1.

    # the other 2 nodes are associated with vertices
    node_cell_indices[:, 1, :] = cells_on_vertex[vertices_on_cell, :]
    node_cell_weights[:, 1, :] = kite_areas_on_vertex[vertices_on_cell, :]
    node_cell_indices[:, 2, :] = cells_on_vertex[next_vertex, :]
    node_cell_weights[:, 2, :] = kite_areas_on_vertex[next_vertex, :]

    weight_sum = np.sum(node_cell_weights, axis=2)
    for i_node in range(3):
        node_cell_weights[:, :, i_node] = (
            node_cell_weights[:, :, i_node] / weight_sum)

    ds_tris = xr.Dataset()
    ds_tris['triCellIndices'] = ('nTriangles', tri_cell_indices)
    ds_tris['nodeCellIndices'] = (('nTriangles', 'nNodes', 'nInterp'),
                                  node_cell_indices)
    ds_tris['nodeCellWeights'] = (('nTriangles', 'nNodes', 'nInterp'),
                                  node_cell_weights)

    # get Cartesian and lon/lat coordinates of each node
    for prefix in ['x', 'y', 'z', 'lat', 'lon']:
        out_var = f'{prefix}Node'
        cell_var = f'{prefix}Cell'
        vertex_var = f'{prefix}Vertex'
        coord = np.zeros((n_triangles, 3))
        coord[:, 0] = ds_mesh[cell_var].values[tri_cell_indices]
        coord[:, 1] = ds_mesh[vertex_var].values[vertices_on_cell]
        coord[:, 2] = ds_mesh[vertex_var].values[next_vertex]
        ds_tris[out_var] = (('nTriangles', 'nNodes'), coord)

    # nothing obvious we can do about triangles containing the poles

    if on_a_sphere:
        ds_tris = _fix_periodic_tris(ds_tris, periodic_var='lonNode',
                                     period=2 * np.pi)
    elif is_periodic:
        ds_tris = _fix_periodic_tris(ds_tris, periodic_var='xNode',
                                     period=x_period)
        ds_tris = _fix_periodic_tris(ds_tris, periodic_var='yNode',
                                     period=y_period)

    return ds_tris

def make_triangle_tree(ds_tris):
    """
    Make a KD-Tree for finding triangle edges that are near enough to transect
    segments that they might intersect

        A dataset that defines triangles, the results of calling
        :py:func:`polaris.ocean.viz.transect.horiz.mesh_to_triangles()`

    Returns
    -------
    tree : scipy.spatial.cKDTree
        A tree of edge centers from triangles making up an MPAS mesh
    """

    n_triangles = ds_tris.sizes['nTriangles']
    n_nodes = ds_tris.sizes['nNodes']
    node_coords = np.zeros((n_triangles * n_nodes, 3))
    node_coords[:, 0] = ds_tris.xNode.values.ravel()
    node_coords[:, 1] = ds_tris.yNode.values.ravel()
    node_coords[:, 2] = ds_tris.zNode.values.ravel()

    next_tri, next_node = np.meshgrid(
        np.arange(n_triangles), np.mod(np.arange(n_nodes) + 1, 3),
        indexing='ij')
    nextIndices = n_nodes * next_tri.ravel() + next_node.ravel()

    # edge centers are half way between adjacent nodes (ignoring great-circle
    # distance)
    edgeCoords = 0.5 * (node_coords + node_coords[nextIndices, :])

    tree = cKDTree(data=edgeCoords, copy_data=True)

def find_spherical_transect_cells_and_weights(
        lon_transect, lat_transect, ds_tris, ds_mesh, tree, degrees=True,
        earth_radius=None, subdivision_res=10e3):
    """
    Find "nodes" where the transect intersects the edges of the triangles
    that make up MPAS cells.

    lon_transect : xarray.DataArray
        The longitude of segments making up the transect

    lat_transect : xarray.DataArray
        The latitude of segments making up the transect

        A data set with the full MPAS mesh.

        A tree of edge centers from triangles making up an MPAS mesh, the
        return value from
        :py:func:`polaris.ocean.viz.transect.horiz.make_triangle_tree()`

    degrees : bool, optional
        Whether ``lon_transect`` and ``lat_transect`` are in degrees (as
        opposed to radians).

    subdivision_res : float, optional
        Resolution in m to use to subdivide the transect when looking for
        intersection candidates.  Should be small enough that curvature is
        small.

    earth_radius : float, optional
        The radius of the Earth in meters, taken from the `sphere_radius`
        global attribute if not provided

    Returns
    -------
    ds_out : xarray.Dataset
        A dataset that contains "nodes" where the transect intersects the
        edges of the triangles in ``ds_tris``.  The nodes also includes the two
        end points of the transect, which typically lie within triangles. Each
        internal node (that is, not including the end points) is purposefully
        repeated twice, once for each triangle that node touches.  This allows
        for discontinuous fields between triangles (e.g. if one wishes to plot
        constant values on each MPAS cell).  The Cartesian and lon/lat
        coordinates of these nodes are ``xCartNode``, ``yCartNode``,
        ``zCartNode``, ``lonNode`` and ``latNode``.  The distance along the
        transect of each intersection is ``dNode``. The index of the triangle
        and the first triangle node in ``ds_tris`` associated with each
        intersection node are given by ``horizTriangleIndices`` and
        ``horizTriangleNodeIndices``, respectively. The second node on the
        triangle for the edge associated with the intersection is given by
        ``numpy.mod(horizTriangleNodeIndices + 1, 3)``.

        The MPAS cell that a given node belongs to is given by
        ``horizCellIndices``. Each node also has an associated set of 6
        ``interpHorizCellIndices`` and ``interpHorizCellWeights`` that can be
        used to interpolate from MPAS cell centers to nodes first with
        area-weighted averaging to MPAS vertices and then linear interpolation
        along triangle edges.  Some of the weights may be zero, in which case
        the associated ``interpHorizCellIndices`` will be -1.

        Finally, ``lonTransect`` and ``latTransect`` are included in the
        dataset, along with Cartesian coordinates ``xCartTransect``,
        ``yCartTransect``, `zCartTransect``, and ``dTransect``, the
        great-circle distance along the transect of each original transect
        point.  In order to interpolate values (e.g. observations) from the
        original transect points to the intersection nodes, linear
        interpolation indices ``transectIndicesOnHorizNode`` and weights
        ``transectWeightsOnHorizNode`` are provided.  The values at nodes are
        found by::

          nodeValues = ((transectValues[transectIndicesOnHorizNode] *
                         transectWeightsOnHorizNode)
                        + (transectValues[transectIndicesOnHorizNode+1] *
                           (1.0 - transectWeightsOnHorizNode))
    """
    if earth_radius is None:
        earth_radius = ds_mesh.attrs['sphere_radius']
    buffer = np.maximum(np.amax(ds_mesh.dvEdge.values),
                        np.amax(ds_mesh.dcEdge.values))

    x, y, z = lon_lat_to_cartesian(lon_transect, lat_transect, earth_radius,
                                   degrees)

    node_cell_weights = ds_tris.nodeCellWeights.values
    node_cell_indices = ds_tris.nodeCellIndices.values

    x_node = ds_tris.xNode.values.ravel()
    y_node = ds_tris.yNode.values.ravel()
    z_node = ds_tris.zNode.values.ravel()

    d_transect = np.zeros(lon_transect.shape)

    d_node = None
    x_out = None
    y_out = None
    z_out = None
    tris = None
    nodes = None
    interp_cells = None
    cell_weights = None

    n_horiz_weights = 6

    first = True

    d_start = 0.
    for seg_index in range(len(x) - 1):
        transectv0 = Vector(x[seg_index].values,
                            y[seg_index].values,
                            z[seg_index].values)
        transectv1 = Vector(x[seg_index + 1].values,
                            y[seg_index + 1].values,
                            z[seg_index + 1].values)

        sub_slice = slice(seg_index, seg_index + 2)
        x_sub, y_sub, z_sub, _, _ = subdivide_great_circle(
            x[sub_slice].values, y[sub_slice].values, z[sub_slice].values,
            subdivision_res, earth_radius)

        coords = np.zeros((len(x_sub), 3))
        coords[:, 0] = x_sub
        coords[:, 1] = y_sub
        coords[:, 2] = z_sub
        radius = buffer + subdivision_res

        index_list = tree.query_ball_point(x=coords, r=radius)

        unique_indices = set()
        for indices in index_list:
            unique_indices.update(indices)

        n0_indices_cand = np.array(list(unique_indices))

        if len(n0_indices_cand) == 0:
            continue

        tris_cand = n0_indices_cand // n_nodes
        next_node_index = np.mod(n0_indices_cand + 1, n_nodes)
        n1_indices_cand = n_nodes * tris_cand + next_node_index

        n0_cand = Vector(x_node[n0_indices_cand],
                         y_node[n0_indices_cand],
                         z_node[n0_indices_cand])
        n1_cand = Vector(x_node[n1_indices_cand],
                         y_node[n1_indices_cand],
                         z_node[n1_indices_cand])

        intersect = Vector.intersects(n0_cand, n1_cand, transectv0,
                                      transectv1)

        n0_inter = Vector(n0_cand.x[intersect],
                          n0_cand.y[intersect],
                          n0_cand.z[intersect])
        n1_inter = Vector(n1_cand.x[intersect],
                          n1_cand.y[intersect],
                          n1_cand.z[intersect])

        tris_inter = tris_cand[intersect]
        n0_indices_inter = n0_indices_cand[intersect]
        n1_indices_inter = n1_indices_cand[intersect]

        intersections = Vector.intersection(n0_inter, n1_inter, transectv0,
                                            transectv1)
        intersections = Vector(earth_radius * intersections.x,
                               earth_radius * intersections.y,
                               earth_radius * intersections.z)

        angular_distance = transectv0.angular_distance(intersections)

        d_node_local = d_start + earth_radius * angular_distance

        d_start += earth_radius * transectv0.angular_distance(transectv1)

        node0_inter = np.mod(n0_indices_inter, n_nodes)
        node1_inter = np.mod(n1_indices_inter, n_nodes)

        node_weights = (intersections.angular_distance(n1_inter) /
                        n0_inter.angular_distance(n1_inter))

        weights = np.zeros((len(tris_inter), n_horiz_weights))
        cell_indices = np.zeros((len(tris_inter), n_horiz_weights), int)
        for index in range(3):
            weights[:, index] = (
                node_weights *
                node_cell_weights[tris_inter, node0_inter, index])
            cell_indices[:, index] = (
                node_cell_indices[tris_inter, node0_inter, index])
            weights[:, index + 3] = (
                (1.0 - node_weights) *
                node_cell_weights[tris_inter, node1_inter, index])
            cell_indices[:, index + 3] = (
                node_cell_indices[tris_inter, node1_inter, index])

        if first:
            x_out = intersections.x
            y_out = intersections.y
            z_out = intersections.z
            d_node = d_node_local

            tris = tris_inter
            nodes = node0_inter
            interp_cells = cell_indices
            cell_weights = weights
            first = False
        else:
            x_out = np.append(x_out, intersections.x)
            y_out = np.append(y_out, intersections.y)
            z_out = np.append(z_out, intersections.z)
            d_node = np.append(d_node, d_node_local)

            tris = np.concatenate((tris, tris_inter))
            nodes = np.concatenate((nodes, node0_inter))
            interp_cells = np.concatenate((interp_cells, cell_indices), axis=0)
            cell_weights = np.concatenate((cell_weights, weights), axis=0)

        d_transect[seg_index + 1] = d_start

    epsilon = 1e-6 * subdivision_res
    (d_node, x_out, y_out, z_out, seg_tris, seg_nodes, interp_cells,
     cell_weights, valid_nodes) = _sort_intersections(
        d_node, tris, nodes, x_out, y_out, z_out, interp_cells, cell_weights,
        epsilon)

    lon_out, lat_out = cartesian_to_lon_lat(x_out, y_out, z_out, earth_radius,
                                            degrees)

    valid_segs = seg_tris >= 0
    cell_indices = -1 * np.ones(seg_tris.shape, dtype=int)
    cell_indices[valid_segs] = (
        ds_tris.triCellIndices.values[seg_tris[valid_segs]])

    ds_out = xr.Dataset()
    ds_out['xCartNode'] = (('nNodes',), x_out)
    ds_out['yCartNode'] = (('nNodes',), y_out)
    ds_out['zCartNode'] = (('nNodes',), z_out)
    ds_out['dNode'] = (('nNodes',), d_node)
    ds_out['lonNode'] = (('nNodes',), lon_out)
    ds_out['latNode'] = (('nNodes',), lat_out)

    ds_out['horizTriangleIndices'] = ('nSegments', seg_tris)
    ds_out['horizCellIndices'] = ('nSegments', cell_indices)
    ds_out['horizTriangleNodeIndices'] = (('nSegments', 'nHorizBounds'),
                                          seg_nodes)
    ds_out['interpHorizCellIndices'] = (('nNodes', 'nHorizWeights'),
                                        interp_cells)
    ds_out['interpHorizCellWeights'] = (('nNodes', 'nHorizWeights'),
                                        cell_weights)
    ds_out['validNodes'] = (('nNodes',), valid_nodes)

    transect_indices_on_horiz_node = np.zeros(d_node.shape, dtype=int)
    transect_weights_on_horiz_node = np.zeros(d_node.shape)
    for trans_index in range(len(d_transect) - 1):
        d0 = d_transect[trans_index]
        d1 = d_transect[trans_index + 1]
        mask = np.logical_and(d_node >= d0, d_node < d1)
        transect_indices_on_horiz_node[mask] = trans_index
        transect_weights_on_horiz_node[mask] = (d1 - d_node[mask]) / (d1 - d0)
    # last index will get missed by the mask and needs to be handled as a
    # special case
    transect_indices_on_horiz_node[-1] = len(d_transect) - 2
    transect_weights_on_horiz_node[-1] = 0.0

    ds_out['lonTransect'] = lon_transect
    ds_out['latTransect'] = lat_transect
    ds_out['xCartTransect'] = x
    ds_out['yCartTransect'] = y
    ds_out['zCartTransect'] = z
    ds_out['dTransect'] = (lon_transect.dims, d_transect)
    ds_out['transectIndicesOnHorizNode'] = (('nNodes',),
                                            transect_indices_on_horiz_node)
    ds_out['transectWeightsOnHorizNode'] = (('nNodes',),
                                            transect_weights_on_horiz_node)

    return ds_out

def find_planar_transect_cells_and_weights(
        x_transect, y_transect, ds_tris, ds_mesh, tree, subdivision_res=10e3):
    """

    x_transect : xarray.DataArray
        The x points defining segments making up the transect

    y_transect : xarray.DataArray
        The y points defining segments making up the transect

        small.

    Returns
    -------
        edges of the triangles in ``ds_tris``.  The nodes also include the two
        constant values on each MPAS cell).  The planar coordinates of these
        nodes are ``xNode`` and ``yNode``.  The distance along the transect of
        each intersection is ``dNode``. The index of the triangle and the first
        triangle node in ``ds_tris`` associated with each intersection node are
        given by ``horizTriangleIndices`` and ``horizTriangleNodeIndices``,
        respectively. The second node on the triangle for the edge associated
        with the intersection is given by

        Finally, ``xTransect`` and ``yTransect`` are included in the
        dataset, along with ``dTransect``, the distance along the transect of
        each original transect point.  In order to interpolate values (e.g.
        observations) from the original transect points to the intersection
        nodes, linear interpolation indices ``transectIndicesOnHorizNode`` and
        weights ``transectWeightsOnHorizNode`` are provided.  The values at
        nodes are found by::

    """

    x = x_transect
    y = y_transect

    coordNode = np.zeros((len(x_node), 2))
    coordNode[:, 0] = x_node
    coordNode[:, 1] = y_node

    d_transect = np.zeros(x_transect.shape)

    x_out = np.array([])
    y_out = np.array([])
    tris = None

        x_sub, y_sub, _, _ = subdivide_planar(
            x[sub_slice].values, y[sub_slice].values, subdivision_res)

        start_point = Point(x_transect[seg_index].values,
                            y_transect[seg_index].values)
        end_point = Point(x_transect[seg_index + 1].values,
                          y_transect[seg_index + 1].values)

        segment = LineString([start_point, end_point])

        start_indices = np.array(list(unique_indices))

        if len(start_indices) == 0:
            continue

        tris_cand = start_indices // n_nodes
        next_node_index = np.mod(start_indices + 1, n_nodes)
        end_indices = n_nodes * tris_cand + next_node_index

        intersecting_nodes = list()
        tris_inter_list = list()
        x_intersection_list = list()
        y_intersection_list = list()
        node_weights_list = list()
        node0_inter_list = list()
        node1_inter_list = list()
        distances_list = list()

        for index in range(len(start_indices)):
            start = start_indices[index]
            end = end_indices[index]

            node0 = Point(coordNode[start, 0], coordNode[start, 1])
            node1 = Point(coordNode[end, 0], coordNode[end, 1])

            edge = LineString([node0, node1])
            if segment.intersects(edge):
                point = segment.intersection(edge)
                intersecting_nodes.append((node0, node1, start, end, edge))

                if isinstance(point, LineString):
                    raise ValueError('A triangle edge exactly coincides with '
                                     'a transect segment and I can\'t handle '
                                     'that case.  Try moving the transect a '
                                     'tiny bit.')
                elif not isinstance(point, Point):
                    raise ValueError(f'Unexpected intersection type {point}')

                x_intersection_list.append(point.x)
                y_intersection_list.append(point.y)

                start_to_intersection = LineString([start_point, point])

                weight = (LineString([point, node1]).length /
                          LineString([node0, node1]).length)

                node_weights_list.append(weight)
                node0_inter_list.append(np.mod(start, n_nodes))
                node1_inter_list.append(np.mod(end, n_nodes))
                distances_list.append(start_to_intersection.length)
                tris_inter_list.append(tris_cand[index])

        distances = np.array(distances_list)
        x_intersection = np.array(x_intersection_list)
        y_intersection = np.array(y_intersection_list)
        node_weights = np.array(node_weights_list)
        node0_inter = np.array(node0_inter_list, dtype=int)
        node1_inter = np.array(node1_inter_list, dtype=int)
        tris_inter = np.array(tris_inter_list, dtype=int)

        d_node_local = d_start + distances

        d_start += segment.length

        if first:
            x_out = x_intersection
            y_out = y_intersection

        else:
            x_out = np.append(x_out, x_intersection)
            y_out = np.append(y_out, y_intersection)

    z_out = np.zeros(x_out.shape)

        epsilon)

    ds_out['xNode'] = (('nNodes',), x_out)
    ds_out['yNode'] = (('nNodes',), y_out)

    transect_indices_on_horiz_node = np.zeros(d_node.shape, int)

    ds_out['xTransect'] = x
    ds_out['yTransect'] = y
    ds_out['dTransect'] = (x_transect.dims, d_transect)

def interp_mpas_horiz_to_transect_nodes(ds_transect, da):
    """
    Interpolate a 2D (``nCells``) MPAS DataArray to transect nodes, linearly
    interpolating fields between the closest neighboring cells

    ds_transect : xr.Dataset
        A dataset that defines an MPAS transect, the results of calling
        ``find_spherical_transect_cells_and_weights()`` or
        ``find_planar_transect_cells_and_weights()``

    da : xr.DataArray
        An MPAS 2D field with dimensions `nCells`` (possibly among others)

    Returns
    -------
    da_nodes : xr.DataArray
        ``nNodes`` (in addition to whatever dimensions were in ``da`` besides
        ``nCells``)
    """
    interp_cell_indices = ds_transect.interpHorizCellIndices
    interp_cell_weights = ds_transect.interpHorizCellWeights
    da = da.isel(nCells=interp_cell_indices)
    da_nodes = (da * interp_cell_weights).sum(dim='nHorizWeights')

def _sort_intersections(d_node, tris, nodes, x_out, y_out, z_out, interp_cells,
                        cell_weights, epsilon):
    """ sort nodes by distance and define segment between them """

    sort_indices = np.argsort(d_node)
    d_sorted = d_node[sort_indices]

    # make a list of indices for each unique value of d
    d = d_sorted[0]
    unique_d_indices = [sort_indices[0]]
    unique_d_all_indices = [[sort_indices[0]]]
    for index, next_d, in zip(sort_indices[1:], d_sorted[1:]):
        if next_d - d < epsilon:
            # this d value is effectively the same as the last, so we'll treat
            # it as the same
            unique_d_all_indices[-1].append(index)
        else:
            # this is a new d, so we'll add to a new list
            d = next_d
            unique_d_indices.append(index)
            unique_d_all_indices.append([index])

    # there is a segment between each unique d, though some are invalid (do
    # not correspond to a triangle)
    seg_tris_list = list()
    seg_nodes_list = list()

    index0 = unique_d_indices[0]
    indices0 = unique_d_all_indices[0]
    d0 = d_node[index0]

    indices = [index0]
    ds = [d0]
    for seg_index in range(len(unique_d_all_indices) - 1):
        indices1 = unique_d_all_indices[seg_index + 1]
        index1 = unique_d_indices[seg_index + 1]
        d1 = d_node[index1]

        # are there any triangles in common between this d value and the next?
        tris0 = tris[indices0]
        tris1 = tris[indices1]
        both = set(tris0).intersection(set(tris1))

        if len(both) > 0:
            tri = both.pop()
            seg_tris_list.append(tri)
            indices.append(index1)
            ds.append(d1)

            # the triangle nodes are the 2 corresponding to the same triangle
            # in the original list
            index0 = indices0[np.where(tris0 == tri)[0][0]]
            index1 = indices1[np.where(tris1 == tri)[0][0]]
            seg_nodes_list.append([nodes[index0], nodes[index1]])
        else:
            # this is an invalid segment so we need to insert and extra invalid
            # node to allow for proper masking
            seg_tris_list.extend([-1, -1])
            seg_nodes_list.extend([[-1, -1], [-1, -1]])
            indices.extend([index0, index1])
            ds.extend([0.5 * (d0 + d1), d1])

        index0 = index1
        indices0 = indices1
        d0 = d1

    indices = np.array(indices, dtype=int)
    d_node = np.array(ds, dtype=float)
    seg_tris = np.array(seg_tris_list, dtype=int)
    seg_nodes = np.array(seg_nodes_list, dtype=int)

    valid_nodes = np.ones(len(indices), dtype=bool)
    valid_nodes[1:-1] = np.logical_or(seg_tris[0:-1] >= 0,
                                      seg_tris[1:] > 0)

    x_out = x_out[indices]
    y_out = y_out[indices]
    z_out = z_out[indices]

    interp_cells = interp_cells[indices, :]
    cell_weights = cell_weights[indices, :]

    return (d_node, x_out, y_out, z_out, seg_tris, seg_nodes, interp_cells,
            cell_weights, valid_nodes)

def _fix_periodic_tris(ds_tris, periodic_var, period):
    """
    make sure the given node coordinate on tris is within one period of the
    cell center
    """
    coord_node = ds_tris[periodic_var].values
    coord_cell = coord_node[:, 0]
    copy_pos = np.zeros(coord_cell.shape, dtype=bool)
    copy_neg = np.zeros(coord_cell.shape, dtype=bool)
    for i_node in [1, 2]:
        mask = coord_node[:, i_node] - coord_cell > 0.5 * period
        copy_pos = np.logical_or(copy_pos, mask)
        coord_node[:, i_node][mask] = coord_node[:, i_node][mask] - period
        mask = coord_node[:, i_node] - coord_cell < -0.5 * period
        copy_neg = np.logical_or(copy_neg, mask)
        coord_node[:, i_node][mask] = coord_node[:, i_node][mask] + period

    pos_indices = np.nonzero(copy_pos)[0]
    neg_indices = np.nonzero(copy_neg)[0]
    tri_indices = np.append(np.append(np.arange(0, n_triangles),
                                      pos_indices), neg_indices)

    ds_new = xr.Dataset(ds_tris)
    ds_new[periodic_var] = (('nTriangles', 'nNodes'), coord_node)
    ds_new = ds_new.isel(nTriangles=tri_indices)
    coord_node = ds_new[periodic_var].values

    pos_slice = slice(n_triangles, n_triangles + len(pos_indices))
    coord_node[pos_slice, :] = coord_node[pos_slice, :] + period
    neg_slice = slice(n_triangles + len(pos_indices),
                      n_triangles + len(pos_indices) + len(neg_indices))
    coord_node[neg_slice, :] = coord_node[neg_slice, :] - period
    return ds_new

    """
    A step for comparing outputs between steps in a single-column restart run

    variables : list of str
        The variables to validate
    """
    def __init__(self, component, step_subdirs, indir, variables,
                 restart_filename):
        """

        indir : str
            the directory the step is in, to which the name of the step will
            be appended

        variables : list of str
            The variables to validate

        restart_filename : str
            The relative path to the restart file to compare in the 2 subdirs
        """

                filename=f'{subdir}_restart.nc',
                target=f'../{subdir}/{restart_filename}')

        self.variables = variables

        """
        Compare the variables in the outputs of two previous steps with each
        other
        """
        all_pass = compare_variables(variables=self.variables,
            raise ValueError(f'Validation failed comparing restart between '

    """
    A step for plotting the results of a single column test
    """
        """

        indir : str
        """
            filename='output.2000.nc',
            target='../forward/output/output.2000.nc')

        """
        """
        style_filename = str(
            imp_res.files('polaris.viz') / 'polaris.mplstyle')
        plt.style.use(style_filename)
        ds = xr.open_dataset('output.2000.nc', decode_times=False)
        daysSinceStartOfSim = ds.daysSinceStartOfSim.values
        snowVolumeCell = ds.snowVolumeCell.values
        iceVolumeCell = ds.iceVolumeCell.values
        surfaceTemperatureCell = ds.surfaceTemperatureCell.values

        fig, axis = plt.subplots(figsize=(8, 8))

        axis.plot(daysSinceStartOfSim, surfaceTemperatureCell,
                  color='green', label='surfaceTemperature')
        axis.set_ylabel('Temperature (C)')
        axis.set_xlabel('Days')
        axis.set_xlim(0, daysSinceStartOfSim[-1])
        axis.set_ylim(None, 0)
        axis.set_title('MPAS_Seaice single cell')

        axis2 = axis.twinx()

        axis2.plot(daysSinceStartOfSim, iceVolumeCell,
                   color='red', label='iceVolume')
        axis2.plot(daysSinceStartOfSim, snowVolumeCell,
                   color='blue', label='snowVolume')
        axis2.set_ylabel('Thickness (m)')
        axis2.set_ylim(0, None)

        plt.savefig('single_cell.pdf')
        plt.savefig('single_cell.png', dpi=300)

#!/usr/bin/env python3

import os
import sys

import polaris.run.serial as run_serial
from polaris import cache, list, setup, suite
from polaris.version import __version__

    """
    Entry point for the main script ``polaris``
    """

        description="Perform polaris operations",
        usage='''
polaris <command> [<args>]

The available polaris commands are:
    list    List the available test cases
    setup   Set up a test case
    suite   Manage a regression test suite
    serial  Run a suite, test case or step in task serial

 To get help on an individual command, run:

    polaris <command> --help
    ''')

    parser.add_argument('command', help='command to run')
    parser.add_argument('-v', '--version',
                        action='version',
                        version=f'polaris {__version__}',
                        help="Show version number and exit")
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(0)

    args = parser.parse_args(sys.argv[1:2])

    commands = {'list': list.main,
                'setup': setup.main,
                'suite': suite.main,
                'serial': run_serial.main}

    # only allow the "polaris cache" command if we're on Anvil or Chrysalis
    allow_cache = ('POLARIS_MACHINE' in os.environ and
                   os.environ['POLARIS_MACHINE'] in ['anvil', 'chrysalis'])

    if allow_cache:
        commands['cache'] = cache.main

    if args.command not in commands:
        print(f'Unrecognized command {args.command}')
        exit(1)

    # call the function associated with the requested command
    commands[args.command]()

if __name__ == "__main__":
    main()

import os
import sys
from datetime import datetime

def update_cache(step_paths, date_string=None, dry_run=False):
    """
    Cache one or more polaris output files for use in a cached variant of the
    test case or step

    step_paths : list of str
        The relative path of the original (uncached) steps from the base work
        directory

    date_string : str, optional
        The datestamp (YYMMDD) to use on the files.  Default is today's date.

    dry_run : bool, optional
        Whether this is a dry run (producing the json file but not copying
        files to the LCRC server)
    """
    if 'POLARIS_MACHINE' not in os.environ:
        machine = None
        invalid = True
    else:
        invalid = machine not in ['anvil', 'chrysalis']

    if invalid:
        raise ValueError('You must cache files from either Anvil or Chrysalis')

    if date_string is None:
        date_string = datetime.now().strftime("%y%m%d")

    # make a dictionary with components as keys, and lists of steps as values
    steps: Dict[str, List[Step]] = dict()
    for path in step_paths:
        with open(f'{path}/step.pickle', 'rb') as handle:
            step = pickle.load(handle)

        component = step.component.name

        if component in steps:
            steps[component].append(step)
        else:
            steps[component] = [step]

    # now, iterate over cores and steps
    for component in steps:
        database_root = config.get('paths', 'database_root')
        cache_root = f'{database_root}/{component}/polaris_cache'

        package = f'polaris.{component}'
        try:
            with open(f'{component}_cached_files.json') as data_file:
                cached_files = json.load(data_file)
            # we don't have a local version of the file yet, let's see if
            # there's a remote one for this component
            try:
                pkg_file = imp_res.files(package).joinpath('cached_files.json')
                with pkg_file.open('r') as data_file:
                    cached_files = json.load(data_file)
                # no cached files yet for this core
                cached_files = dict()

        for step in steps[component]:
            # load the step from its pickle file

            step_path = step.path

            for output in step.outputs:
                output = os.path.basename(output)
                out_filename = os.path.join(step_path, output)
                # remove the component from the file path
                target = out_filename[len(component) + 1:]
                path, ext = os.path.splitext(target)
                target = f'{path}.{date_string}{ext}'
                cached_files[out_filename] = target

                print(out_filename)
                print(f'  ==> {target}')
                output_path = f'{cache_root}/{target}'
                print(f'  copy to: {output_path}')
                print()
                if not dry_run:
                    directory = os.path.dirname(output_path)
                    try:
                        os.makedirs(directory)
                    except FileExistsError:
                        pass
                    shutil.copyfile(out_filename, output_path)

        out_filename = f'{component}_cached_files.json'
        with open(out_filename, 'w') as data_file:
            json.dump(cached_files, data_file, indent=4)

        description='Cache the output files from one or more steps for use in '
                    'a cached variant of the step',
        prog='polaris cache')
    parser.add_argument("-i", "--orig_steps", nargs='+', dest="orig_steps",
                        type=str,
                        help="The relative path of the original (uncached) "
                             "steps from the base work directory.",
    parser.add_argument("-d", "--date_string", dest="date_string", type=str,
                        help="The datestamp (YYMMDD) to use on the files.  "
                             "Default is today's date.",
                        metavar="DATE")
    parser.add_argument("-r", "--dry_run", dest="dry_run",
                        help="Whether this is a dry run (producing the json "
                             "file but not copying files to the LCRC server).",
                        action="store_true")

    update_cache(step_paths=args.orig_steps, date_string=args.date_string,
                 dry_run=args.dry_run)

from typing import List

from polaris import Component
# import new components here
from polaris.ocean import Ocean
from polaris.seaice import SeaIce

def get_components():
    """
    Get a list of components, which in turn contain lists tasks

    Returns
    -------
    components : list of polaris.Component
        A list of components containing all available tasks
    """
    # add new components here
    components: List[Component] = [
        Ocean(),
        SeaIce(),
    ]
    return components

import inspect

def log_method_call(method, logger):
    """
    Log the module path and file path of a call to a method, e.g.::

      polaris calling: polaris.landice.tasks.dome.decomposition_test.DecompositionTest.run()  # noqa: E501
        in /turquoise/usr/projects/climate/mhoffman/mpas/polaris/polaris/landice/tasks/dome/decomposition_test/__init__.py  # noqa: E501

    method : method
        The method of a class that will be run immediately following this call

    logger: logging.Logger
        The logger to log the method path and file path to
    """
    if not inspect.ismethod(method):
        raise ValueError('The "method" argument must be a method')

    method_name = method.__name__

    # the class of whatever object the method belongs to might not be the
    # class where the method is implemented.  The "child" class is the class
    # of the object itself (typically a specific test case or step).  The
    # "actual" class where the method is implemented could be that class or
    # one of its parents.

    # get the "child" class and its location (import sequence) from the method
    child_class = method.__self__.__class__
    child_location = f'{child_class.__module__}.{child_class.__name__}'

    # iterate over the classes that the child class descends from to find the
    # first one that actually implements the given method.
    actual_class = None
    # inspect.getmro() returns a list of classes the child class descends from,
    # starting with the child class itself and going "back" to the "object"
    # class that all python classes descend from.
    for cls in inspect.getmro(child_class):
        if method.__name__ in cls.__dict__:
            actual_class = cls
            break

    if actual_class is None:
        raise ValueError(f'Hmm, it seems that none of {child_location} or its '
                         f'parent classes implements\n the {method_name} '
                         f'method, how strange!')
    actual_location = f'{actual_class.__module__}.{actual_class.__name__}'

    # not every class is defined in a file (e.g. python intrinsics) but we
    # expect they always are in polaris.  Still we'll check to make sure.
    try:
        class_file = inspect.getfile(actual_class)
    except TypeError:
        class_file = None

    # log what we found out
    logger.info(f'polaris calling: {child_location}.{method_name}()')
    if child_location != actual_location:
        # okay, so we're inheriting this method from somewhere else.  Better
        # let the user know.
        logger.info(f'  inherited from: {actual_location}.{method_name}()')
    if class_file is not None:
        logger.info(f'  in {class_file}')

def log_function_call(function, logger):
    """
    Log the module path and file path of a call to a function, e.g.::

      polaris calling: polaris.parallel.set_cores_per_node()
        in /home/xylar/code/polaris/polaris/polaris/parallel.py

    function : function
        The function that will be run immediately following this call

        The logger to log the function path and file path to
    """
    if not inspect.isfunction(function):
        raise ValueError('The "function" argument must be a function')

    filename = inspect.getfile(function)

    logger.info(f'polaris calling: '
                f'{function.__module__}.{function.__name__}()')
    logger.info(f'  in {filename}')

def compute_planar_hex_nx_ny(lx, ly, resolution):
    """
    Compute number of grid cells in each direction for the uniform, hexagonal
    planar mesh with the given physical sizes and resolution.  The resulting
    ``nx`` and ``ny`` account for the staggered nature of the hexagonal grid
    in the y direction, and are appropriate for passing to
    :py:func:`mpas_tools.planar_hex.make_planar_hex_mesh()`.

    lx : float
        The size of the domain in km in the x direction

    ly : float
        The size of the domain in km in the y direction

        The resolution of the mesh (distance between cell centers) in km

    Returns
    -------
    nx : int
        The number of grid cells in the x direction

    ny : int
        The number of grid cells in the y direction
    """
    # these could be hard-coded as functions of specific supported
    # resolutions but it is preferable to make them algorithmic like here
    # for greater flexibility
    nx = max(2 * int(0.5 * lx / resolution + 0.5), 4)
    # factor of 2/sqrt(3) because of hexagonal mesh
    ny = max(2 * int(0.5 * ly * (2. / np.sqrt(3)) / resolution + 0.5), 4)
    return nx, ny

def area_for_field(ds_mesh, field):
    """
    Get the appropriate area (on cells, vertices or edges) for the given
    field

        The MPAS mesh

    field : xarray.DataArray
        The field whose dimensions determine which area to return

    Returns
    -------
    area : xarray.DataArray
        The area on cells, vertices or edges
    """
    if 'nCells' in field.dims:
        area = ds_mesh.areaCell
    elif 'nEdges' in field.dims:
        area = 0.25 * ds_mesh.dcEdge * ds_mesh.dvEdge
    elif 'nVertices' in field.dims:
        vertex_degree = ds_mesh.sizes['vertexDegree']
        area = xr.zeros_like(ds_mesh.xVertex)
        for ivert in range(vertex_degree):
            cov = ds_mesh.cellsOnVertex.isel(vertexDegree=ivert) - 1
            mask = cov >= 0
            kite_area = ds_mesh.kiteAreasOnVertex.isel(vertexDegree=ivert)
            area = area + kite_area.where(mask, other=0.)
    else:
        raise ValueError('The field does not have any of the expected '
                         'horizontal dimensions for an MPAS mesh')

    return area

def cell_mask_to_edge_mask(ds_mesh, cell_mask):
    """Convert a cell mask to edge mask using mesh connectivity information

    True corresponds to valid cells and False are invalid cells

    cell_mask : xarray.DataArray
        The cell mask we want to convert to an edge mask

    Returns
    -------
    edge_mask : xarray.DataArray
        The edge mask corresponding to the input cell mask
    """

    # test if any are False
    if ~cell_mask.any():
        return ds_mesh.nEdges > -1

    # zero index the connectivity array
    cells_on_edge = (ds_mesh.cellsOnEdge - 1)

    # using nCells (dim) instead of indexToCellID since it's already 0 indexed
    masked_cells = ds_mesh.nCells.where(~cell_mask, drop=True).astype(int)

    # use inverse so True/False convention matches input cell_mask
    edge_mask = ~cells_on_edge.isin(masked_cells).any("TWO")

    return edge_mask

def time_index_from_xtime(xtime, dt_target, start_xtime=None):
    """
    Determine the time index closest to the target time

    xtime : numpy.ndarray of numpy.char
        Times in the dataset

    dt_target : float
        Time in seconds since the first time in the list of ``xtime`` values

    start_xtime : str, optional
        The start time, the first entry in ``xtime`` by default

    Returns
    -------
    time_index : int
        Index in xtime that is closest to dt_target
    """
    dt = time_since_start(xtime, start_xtime)
    time_index = np.argmin(np.abs(np.subtract(dt, dt_target)))
    return time_index

def time_since_start(xtime, start_xtime=None):
    """
    Determine the time elapsed since the start of the simulation

    Returns
    -------
    dt : numpy.ndarray
        The elapsed time in seconds corresponding to each entry in xtime
    """
    if start_xtime is None:
        start_xtime = xtime[0].decode()

    t0 = datetime.datetime.strptime(start_xtime, '%Y-%m-%d_%H:%M:%S')
    dt = np.zeros((len(xtime),))
    for idx, xt in enumerate(xtime):
        t = datetime.datetime.strptime(xt.decode(),
                                       '%Y-%m-%d_%H:%M:%S')
        dt[idx] = (t - t0).total_seconds()
    return dt

import os
import sys

def write(work_dir, tasks, config=None):
    """
    Write a file with provenance, such as the git version, conda packages,
    command, and tasks, to the work directory

        The path to the work directory where the tasks will be set up

        A dictionary describing all of the tasks and their steps

        Configuration options for this task, a combination of user configs
    """
    polaris_git_version = None
    if os.path.exists('.git'):
        try:
            args = ['git', 'describe', '--tags', '--dirty', '--always']
            polaris_git_version = subprocess.check_output(args).decode('utf-8')
            polaris_git_version = polaris_git_version.strip('\n')
        except subprocess.CalledProcessError:
            pass

    if config is None:
        # this is a call to clean and we don't need to document the component
        # version
        component_git_version = None
    else:
        component_git_version = _get_component_git_version(config)

    try:
        args = ['conda', 'list']
        conda_list = subprocess.check_output(args).decode('utf-8')
    except subprocess.CalledProcessError:
        conda_list = None

    calling_command = ' '.join(sys.argv)

    try:
        os.makedirs(work_dir)
        pass

    provenance_path = f'{work_dir}/provenance'
    if os.path.exists(provenance_path):
        provenance_file = open(provenance_path, 'a')
        provenance_file.write('\n')
    else:
        provenance_file = open(provenance_path, 'w')

    provenance_file.write('**************************************************'
                          '*********************\n')
    if polaris_git_version is not None:
        provenance_file.write(
            f'polaris git version: {polaris_git_version}\n\n')
    if component_git_version is not None:
            f'component git version: {component_git_version}\n\n')
    provenance_file.write(f'command: {calling_command}\n\n')
    provenance_file.write('tasks:\n')

        prefix = '  '
        lines = list()
        to_print = {'path': task.path,
                    'name': task.name,
                    'component': task.component.name,
                    'subdir': task.subdir}
        for key in to_print:
            key_string = f'{key}: '.ljust(15)
            lines.append(f'{prefix}{key_string}{to_print[key]}')
        lines.append(f'{prefix}steps:')
            if step.name == step.subdir:
                lines.append(f'{prefix} - {step.name}')
            else:
                lines.append(f'{prefix} - {step.name}: {step.subdir}')
        lines.append('')
        print_string = '\n'.join(lines)

        provenance_file.write(f'{print_string}\n')

    if conda_list is not None:
        provenance_file.write('conda list:\n')
        provenance_file.write(f'{conda_list}\n')

    provenance_file.close()

def _get_component_git_version(config):

    component_path = config.get('paths', 'component_path')

    if not os.path.exists(component_path):

    os.chdir(component_path)

    try:
        args = ['git', 'describe', '--tags', '--dirty', '--always']
        component_git_version = subprocess.check_output(args).decode('utf-8')
        component_git_version = component_git_version.strip('\n')
    os.chdir(cwd)

    return component_git_version

import sys

from polaris.setup import setup_tasks

def setup_suite(component, suite_name, work_dir, config_file=None,
                copy_executable=False, clean=False, model=None):
    """
    Set up a suite of tasks

    component : str
        The component ('ocean', 'landice', etc.) of the suite

        The name of the suite.  A file ``<suite_name>.txt`` must exist
        within the core's ``suites`` package that lists the paths of the tasks
        in the suite

    config_file : str, optional
        Configuration file with custom options for setting up and running
        tasks

    work_dir : str, optional

        Whether to copy the MPAS executable to the work directory

        setting up the suite

    """

    text = imp_res.files(f'polaris.{component}.suites').joinpath(
        f'{suite_name}.txt').read_text()

    tasks, cached = _parse_suite(text)

    setup_tasks(work_dir, tasks, config_file=config_file, machine=machine,
                baseline_dir=baseline_dir, component_path=component_path,
                suite_name=suite_name, cached=cached,
                copy_executable=copy_executable, clean=clean, model=model)

        description='Set up a regression suite', prog='polaris suite')
    parser.add_argument("-c", "--component", dest="component",
                        help="The component for the suite.",
                        metavar="COMPONENT", required=True)
    parser.add_argument("-t", "--task_suite", dest="task_suite",
                        help="Path to file containing a suite to setup.",
                        metavar="SUITE", required=True)
                        help="If set, script will setup the suite in "
                        "work_dir rather in this script's location.",
                             "before setting up the suite.")

    setup_suite(component=args.component, suite_name=args.task_suite,
                work_dir=args.work_dir, config_file=args.config_file,
                machine=args.machine, baseline_dir=args.baseline_dir,
                component_path=args.component_path,
                copy_executable=args.copy_executable, clean=args.clean,
                model=args.model)

def _parse_suite(text):
    """ Parse the text of a file defining a suite """

    tasks: List[str] = list()
    cached: List[List[str]] = list()
    for task in text.split('\n'):
        task = task.strip()
        if len(task) == 0 or task.startswith('#'):
            # a blank line or comment
            continue

        if task == 'cached':
            cached[-1] = ['_all']
        elif task.startswith('cached:'):
            steps = task[len('cached:'):].strip().split(' ')
            cached[-1].extend(steps)
        else:
            cached.append(list())

    return tasks, cached

import os

def compare_variables(variables, filename1, filename2, logger, l1_norm=0.0,
                      l2_norm=0.0, linf_norm=0.0, quiet=True):
    """
    compare variables in the two files

    variables : list
        A list of variable names to compare

    filename1 : str
        The relative path to a file within the ``work_dir``.  If ``filename2``
        is also given, comparison will be performed with ``variables`` in that
        file.  If a baseline directory was provided when setting up the
        test case, the ``variables`` will be compared between this test case
        and the same relative filename in the baseline version of the test
        case.

    filename2 : str
        The relative path to another file within the ``work_dir`` if comparing
        between files within the current test case.  If a baseline directory
        was provided, the ``variables`` from this file will also be compared
        with those in the corresponding baseline file.

        The logger to log validation output to

    l1_norm : float, optional
        The maximum allowed L1 norm difference between the variables in
        ``filename1`` and ``filename2``.  To skip L1 norm check, pass None.

    l2_norm : float, optional
        The maximum allowed L2 norm difference between the variables in
        ``filename1`` and ``filename2``.  To skip L2 norm check, pass None.

    linf_norm : float, optional
        The maximum allowed L-Infinity norm difference between the variables in
        ``filename1`` and ``filename2``.  To skip Linf norm check, pass None.

        Whether to print detailed information.  If quiet is False, the norm
        tolerance values being compared against will be printed when the
        comparison is made.  This is generally desirable when using nonzero
        norm tolerance values.

    Returns
    -------
    all_pass : bool
        Whether all variables passed the validation checks

    """

    for filename in [filename1, filename2]:
        if not os.path.exists(filename):
            logger.error(f'File {filename} does not exist.')
            return False

    ds1 = xr.open_dataset(filename1)
    ds2 = xr.open_dataset(filename2)

    all_pass = True

    for variable in variables:
        all_found = True
        for ds, filename in [(ds1, filename1), (ds2, filename2)]:
            if variable not in ds:
                logger.error(f'Variable {variable} not in {filename}.')
                all_found = False
        if not all_found:
            all_pass = False
            continue

        da1 = ds1[variable]
        da2 = ds2[variable]

        if not np.all(da1.dims == da2.dims):
            logger.error(f"Dimensions for variable {variable} don't match "
                         f"between files {filename1} and {filename2}.")
            continue

        all_match = True
        for dim in da1.sizes:
            if da1.sizes[dim] != da2.sizes[dim]:
                logger.error(f"Field sizes for variable {variable} don't "
                             f"match files {filename1} and {filename2}.")
                all_match = False
        if not all_match:
            continue

            print("    Pass thresholds are:")
            if l1_norm is not None:
                print(f"       L1: {l1_norm:16.14e}")
            if l2_norm is not None:
                print(f"       L2: {l2_norm:16.14e}")
            if linf_norm is not None:
                print(f"       L_Infinity: {linf_norm:16.14e}")
        variable_pass = True
        if 'Time' in da1.dims:
            time_range = range(0, da1.sizes['Time'])
            time_str = ', '.join([f'{j}' for j in time_range])
            print(f'{variable.ljust(20)} Time index: {time_str}')
            for time_index in time_range:
                slice1 = da1.isel(Time=time_index)
                slice2 = da2.isel(Time=time_index)
                result = _compute_norms(slice1, slice2, quiet, l1_norm,
                                        l2_norm, linf_norm,
                                        time_index=time_index)
                variable_pass = variable_pass and result

        else:
            print(f'{variable}')
            result = _compute_norms(da1, da2, quiet, l1_norm, l2_norm,
                                    linf_norm)
            variable_pass = variable_pass and result

        # ANSI fail text: https://stackoverflow.com/a/287944/7728169
        start_fail = '\033[91m'
        start_pass = '\033[92m'
        end = '\033[0m'
        pass_str = f'{start_pass}PASS{end}'
        fail_str = f'{start_fail}FAIL{end}'

        if variable_pass:
            print(f'  {pass_str} {filename1}\n')
        else:
            print(f'  {fail_str} {filename1}\n')
        print(f'       {filename2}\n')
        all_pass = all_pass and variable_pass

    return all_pass

def _compute_norms(da1, da2, quiet, max_l1_norm, max_l2_norm, max_linf_norm,
                   time_index=None):
    """ Compute norms between variables in two DataArrays """

    da1 = _rename_duplicate_dims(da1)
    da2 = _rename_duplicate_dims(da2)

    result = True
    diff = np.abs(da1 - da2).values.ravel()
    # skip entries where one field or both are a fill value
    diff = diff[np.isfinite(diff)]

    l1_norm = np.linalg.norm(diff, ord=1)
    l2_norm = np.linalg.norm(diff, ord=2)
    linf_norm = np.linalg.norm(diff, ord=np.inf)

    if time_index is None:
        diff_str = ''
    else:
        diff_str = f'{time_index:d}: '

    if max_l1_norm is not None:
        if max_l1_norm < l1_norm:
            result = False
    diff_str = f'{diff_str} l1: {l1_norm:16.14e} '

    if max_l2_norm is not None:
        if max_l2_norm < l2_norm:
    diff_str = f'{diff_str} l2: {l2_norm:16.14e} '

    if max_linf_norm is not None:
        if max_linf_norm < linf_norm:
    diff_str = f'{diff_str} linf: {linf_norm:16.14e} '

    if not quiet or not result:
        print(diff_str)

    return result

def _rename_duplicate_dims(da):
    dims = list(da.dims)
    new_dims = list(dims)
    duplicates = False
    for index, dim in enumerate(dims):
        if dim in dims[index + 1:]:
            duplicates = True
            suffix = 2
            for other_index, other in enumerate(dims[index + 1:]):
                if other == dim:
                    new_dims[other_index + index + 1] = f'{dim}_{suffix}'
                    suffix += 1

    if not duplicates:
        return da

    da = xr.DataArray(data=da.values, dims=new_dims)
    return da

import os

import matplotlib
from matplotlib.colors import LogNorm

def plot_horiz_field(ds_mesh, field, out_file_name=None,  # noqa: C901
                     ax=None, title=None, t_index=None, z_index=None,
                     vmin=None, vmax=None, show_patch_edges=False,
                     cmap=None, cmap_set_under=None, cmap_set_over=None,
                     cmap_scale='linear', cmap_title=None, figsize=None,
                     vert_dim='nVertLevels', field_mask=None,
                     descriptor=None, antialiased=False,
                     transect_x=None, transect_y=None, transect_color='black',
                     transect_start='red', transect_end='green',
                     transect_linewidth=2., transect_markersize=12.):
    """
    Plot a horizontal field from a planar domain using x,y coordinates at a
    single time and depth slice.

        A data set containing horizontal mesh variables

        The data array to plot

    out_file_name : str, optional
        The path to which the plot image should be written

    ax : matplotlib.axes.Axes
        Axes to plot to if making a multi-panel figure

        The title of the plot

    vmin : float, optional
        The minimum values for the colorbar

    vmax : float, optional
        The maximum values for the colorbar

    show_patch_edges : boolean, optional
        If true, patches will be plotted with visible edges

    t_index: int, optional
        The indices of ``Time`` axes to select for plotting. The default is 0
        (initial time)

    z_index: int, optional
        The indices of ``nVertLevels`` axes to select for plotting. The default
        is 0 (top level)

    cmap : Colormap or str, optional
        A color map to plot

    cmap_set_under : str or None, optional
        A color for low out-of-range values

    cmap_set_over : str or None, optional
        A color for upper out-of-range values

    cmap_scale : {'log', 'linear'}, optional
        Whether the colormap is logarithmic or linear

    cmap_title : str, optional
        Title for color bar

        The width and height of the figure in inches. Default is determined
        based on the aspect ratio of the domain.

    vert_dim : str, optional
        Name of the vertical dimension

    field_mask : xarray.DataArray, optional
        A ``bool`` mask indicating where the ``field`` is valid.

    descriptor : mosaic.Descriptor, optional
        Descriptor from a previous call to ``plot_horiz_field()``

    transect_x : numpy.ndarray or xarray.DataArray, optional
        The x coordinates of a transect to plot on the

    transect_y : numpy.ndarray or xarray.DataArray, optional
        The y coordinates of a transect

    transect_color : str, optional
        The color of the transect line

    transect_start : str or None, optional
        The color of a dot marking the start of the transect

    transect_end : str or None, optional
        The color of a dot marking the end of the transect

    transect_linewidth : float, optional
        The width of the transect line

    transect_markersize : float, optional
        The size of the transect start and end markers

    Returns
    -------
    descriptor : mosaic.Descriptor
        For reuse with future plots. Patches are cached, so the Descriptor only
        needs to be created once per mesh file.
    """
    if (transect_x is None) != (transect_y is None):
        raise ValueError('You must supply both transect_x and transect_y or '
                         'neither')

    create_fig = True
    if ax is not None:
        create_fig = False

    if create_fig:
        if out_file_name is None:
            out_file_name = f'{field.name}.png'
        try:
            os.makedirs(os.path.dirname(out_file_name))
        except OSError:
            pass

    if title is None:
        title = field.name

    if 'Time' in field.dims and t_index is None:
        t_index = 0
    if t_index is not None:
        field = field.isel(Time=t_index)
    if vert_dim in field.dims and z_index is None:
        z_index = 0
    if z_index is not None:
        field = field.isel({vert_dim: z_index})

    if descriptor is None:

    pcolor_kwargs = dict(cmap=None, edgecolor='face', norm=None, vmin=vmin,
                         vmax=vmax, antialiased=antialiased)

    if cmap is not None:
        if isinstance(cmap, str):
            cmap = matplotlib.colormaps[cmap]
        if cmap_set_under is not None:
            cmap.set_under(cmap_set_under)
        if cmap_set_over is not None:
            cmap.set_over(cmap_set_over)

        pcolor_kwargs['cmap'] = cmap

    if show_patch_edges:
        pcolor_kwargs['edgecolor'] = 'black'
        pcolor_kwargs['linewidth'] = 0.25

    if cmap_scale == 'log':
        pcolor_kwargs['norm'] = LogNorm(
            vmin=max(1e-10, vmin), vmax=vmax, clip=False
        )

    if figsize is None:
        width = ds_mesh.xCell.max() - ds_mesh.xCell.min()
        length = ds_mesh.yCell.max() - ds_mesh.yCell.min()
        aspect_ratio = width.values / length.values
        fig_width = 4
        legend_width = fig_width / 5
        figsize = (fig_width + legend_width, fig_width / aspect_ratio)

        plt.figure(figsize=figsize)
        ax = plt.subplot(111)

    if field_mask is not None:

        if field_mask.shape != field.shape:
            raise ValueError(f"The shape of `field_mask`: {field_mask.shape} "
                             f"does match shape of `field array`: "
                             f"{field.shape} make sure both arrays are defined"
                             f" at the same location")

        if np.any(~field_mask):
            field = field.where(field_mask)

    collection = mosaic.polypcolor(ax, descriptor, field, **pcolor_kwargs)

    ax.set_xlabel('x (km)')
    ax.set_ylabel('y (km)')
    ax.set_aspect('equal')
    ax.autoscale(tight=True)
    # scale ticks to be in kilometers
    ax.xaxis.set_major_formatter(lambda x, pos: f'{x / 1e3:g}')
    ax.yaxis.set_major_formatter(lambda x, pos: f'{x / 1e3:g}')

    cbar = plt.colorbar(collection, extend='both', shrink=0.7, ax=ax)
    if cmap_title is not None:
        cbar.set_label(cmap_title)

    if transect_x is not None:
        ax.plot(transect_x, transect_y, color=transect_color,
                linewidth=transect_linewidth)
        if transect_start is not None:
            ax.plot(transect_x[0], transect_y[0], '.', color=transect_start,
                    markersize=transect_markersize)
        if transect_end is not None:
            ax.plot(transect_x[-1], transect_y[-1], '.', color=transect_end,
        plt.title(title)
        plt.savefig(out_file_name, bbox_inches='tight', pad_inches=0.2)

def use_mplstyle():
    """
    Use the Polaris matplotlib style file
    """
    style_filename = str(
        imp_res.files('polaris.viz') / 'polaris.mplstyle')
    plt.style.use(style_filename)

from polaris.ocean.tasks.baroclinic_channel import add_baroclinic_channel_tasks
from polaris.ocean.tasks.barotropic_gyre import add_barotropic_gyre_tasks
from polaris.ocean.tasks.cosine_bell import add_cosine_bell_tasks
from polaris.ocean.tasks.geostrophic import add_geostrophic_tasks
from polaris.ocean.tasks.ice_shelf_2d import add_ice_shelf_2d_tasks
from polaris.ocean.tasks.inertial_gravity_wave import (
    add_inertial_gravity_wave_tasks,
)
from polaris.ocean.tasks.internal_wave import add_internal_wave_tasks
from polaris.ocean.tasks.isomip_plus import add_isomip_plus_tasks
from polaris.ocean.tasks.manufactured_solution import (
    add_manufactured_solution_tasks,
)
from polaris.ocean.tasks.single_column import add_single_column_tasks
from polaris.ocean.tasks.sphere_transport import add_sphere_transport_tasks

class Ocean(Component):
    """
    The collection of all test case for the MPAS-Ocean core
    """

        """
        Construct the collection of MPAS-Ocean test cases
        """
        super().__init__(name='ocean')

        # planar: please keep these in alphabetical order
        add_baroclinic_channel_tasks(component=self)
        add_barotropic_gyre_tasks(component=self)
        add_ice_shelf_2d_tasks(component=self)
        add_inertial_gravity_wave_tasks(component=self)
        add_internal_wave_tasks(component=self)
        add_isomip_plus_tasks(component=self, mesh_type='planar')
        add_manufactured_solution_tasks(component=self)

        # single column
        add_single_column_tasks(component=self)

        # spherical: please keep these in alphabetical order
        add_cosine_bell_tasks(component=self)
        add_geostrophic_tasks(component=self)
        add_isomip_plus_tasks(component=self, mesh_type='spherical')
        add_sphere_transport_tasks(component=self)

        """

        """
        section = config['ocean']
        model = section.get('model')
        configs = {'mpas-ocean': 'mpas_ocean.cfg',
                   'omega': 'omega.cfg'}
        if model not in configs:
            raise ValueError(f'Unknown ocean model {model} in config options')

        config.add_from_package('polaris.ocean', configs[model])

import os

from polaris.mesh.spherical import (
    IcosahedralMeshStep,
    QuasiUniformSphericalMeshStep,
)

def add_spherical_base_mesh_step(component, resolution, icosahedral):
    """
    Add a shared step for creating spherical base mesh with the given
    resolution to the ocean component (if one has not already been added)

        The ocean component that the step will be added to

        The resolution in km of the mesh

        Whether the mesh is a subdivided icosahedral mesh, as opposed to

    Returns
    -------
    base_mesh : polaris.Step
        The base mesh step

        The name of the mesh (e.g. '240km')
    """
    mesh_name = resolution_to_subdir(resolution)

    if icosahedral:
        prefix = 'icos'
    else:
        prefix = 'qu'

    name = f'{prefix}_base_mesh_{mesh_name}'
    subdir = f'spherical/{prefix}/base_mesh/{mesh_name}'

    if subdir in component.steps:
        base_mesh = component.steps[subdir]

    else:
            base_mesh = IcosahedralMeshStep(
                component=component, name=name, subdir=subdir,
                cell_width=resolution)
        else:
            base_mesh = QuasiUniformSphericalMeshStep(

        # add default config options for spherical meshes
        config_filename = f'{base_mesh.name}.cfg'
        filepath = os.path.join(base_mesh.subdir, config_filename)
        config.add_from_package('polaris.mesh', 'spherical.cfg')
        base_mesh.set_shared_config(config)

        component.add_step(base_mesh)

    return base_mesh, mesh_name

def get_time_interval_string(days=None, seconds=None):
    """
    Convert a time interval in days and/or seconds to a string for use in a
    model config option.  If both are provided, they will be added

    days : float, optional
        A time interval in days

    seconds : float, optional
        A time interval in seconds

    Returns
    -------
    time_str : str
        The time as a string in the format "DDDD_HH:MM:SS.SS"

    """
    sec_per_day = 86400
    total = 0.
    if seconds is not None:
        total += seconds
    if days is not None:
        total += sec_per_day * days

    day_part = int(total / sec_per_day)
    sec_part = total - day_part * sec_per_day
    sec_decimal = sec_part - np.floor(sec_part)
    # https://stackoverflow.com/a/1384565/7728169
    seconds_str = time.strftime('%H:%M:%S', time.gmtime(sec_part))
    time_str = f'{day_part:04d}_{seconds_str}.{int(sec_decimal * 1e3):03d}'
    return time_str

import csv

def compute_rpe(mesh_filename, initial_state_filename, output_filenames):
    """
    Computes the reference (resting) potential energy for the whole domain

        Name of the netCDF file containing the MPAS horizontal mesh variables

    initial_state_filename : str
        Name of the netCDF file containing the initial state

    output_filenames : list
        List of netCDF files containing output of forward steps

    Returns
    -------
    rpe : numpy.ndarray
        the reference potential energy of size ``Time`` x ``len(output_files)``
    """
    num_files = len(output_filenames)
    if num_files == 0:
        raise ValueError('Must provide at least one output filename')

    ds_mesh = xr.open_dataset(mesh_filename)
    ds_init = xr.open_dataset(initial_state_filename)
    nVertLevels = ds_init.sizes['nVertLevels']

    xEdge = ds_mesh.xEdge
    yEdge = ds_mesh.yEdge
    areaCell = ds_mesh.areaCell
    minLevelCell = ds_init.minLevelCell - 1
    maxLevelCell = ds_init.maxLevelCell - 1
    bottomDepth = ds_init.bottomDepth

    areaCellMatrix = np.tile(areaCell, (nVertLevels, 1)).transpose()
    bottomMax = np.max(bottomDepth.values)
    yMin = np.min(yEdge.values)
    yMax = np.max(yEdge.values)
    xMin = np.min(xEdge.values)
    xMax = np.max(xEdge.values)
    areaDomain = (yMax - yMin) * (xMax - xMin)

    vert_index = \
        xr.DataArray.from_dict({'dims': ('nVertLevels',),
                                'data': np.arange(nVertLevels)})

    cell_mask = np.logical_and(vert_index >= minLevelCell,
                               vert_index <= maxLevelCell)
    cell_mask = np.swapaxes(cell_mask, 0, 1)

    with xr.open_dataset(output_filenames[0]) as ds:
        nt = ds.sizes['Time']
        xtime = ds.xtime.values

    rpe = np.ones((num_files, nt))

    for file_index, out_filename in enumerate(output_filenames):

        ds = xr.open_dataset(out_filename)

        hFull = ds.layerThickness
        densityFull = ds.density

        for time_index in range(nt):

            h = hFull[time_index, :, :].values
            vol = np.multiply(h, areaCellMatrix)
            density = densityFull[time_index, :, :].values
            density_1D = density[cell_mask]
            vol_1D = vol[cell_mask]

            # Density sorting in ascending order
            sorted_ind = np.argsort(density_1D)
            density_sorted = density_1D[sorted_ind]
            vol_sorted = vol_1D[sorted_ind]

            thickness = np.divide(vol_sorted.tolist(), areaDomain)

            # RPE computation
            z = np.append([0.], -np.cumsum(thickness))
            zMid = z[0:-1] - 0.5 * thickness + bottomMax
            rpe1 = gravity * np.multiply(
                np.multiply(density_sorted, zMid),
                vol_sorted)

            rpe[file_index, time_index] = np.sum(rpe1) / np.sum(areaCell)

        ds.close()

    with open('rpe.csv', 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        headers = ['time'] + \
            [f'output_{index + 1}' for index in range(num_files)]
        writer.writerow(headers)
            time = xtime[time_index].astype(str)
            row = [time] + [f'{rpe_val:g}' for rpe_val in rpe[:, time_index]]
            writer.writerow(row)

    return rpe

from polaris.ocean.tasks.baroclinic_channel.decomp import Decomp
from polaris.ocean.tasks.baroclinic_channel.default import Default
from polaris.ocean.tasks.baroclinic_channel.init import Init
from polaris.ocean.tasks.baroclinic_channel.restart import Restart
from polaris.ocean.tasks.baroclinic_channel.rpe import Rpe
from polaris.ocean.tasks.baroclinic_channel.threads import Threads

def add_baroclinic_channel_tasks(component):
    """
    Add tasks for different baroclinic channel tests to the ocean component

    """
    for resolution in [10., 4., 1.]:
        resdir = resolution_to_subdir(resolution)
        resdir = f'planar/baroclinic_channel/{resdir}'

        config_filename = 'baroclinic_channel.cfg'
        config = PolarisConfigParser(filepath=f'{resdir}/{config_filename}')
        config.add_from_package('polaris.ocean.tasks.baroclinic_channel',
                                'baroclinic_channel.cfg')

        init = Init(component=component, resolution=resolution, indir=resdir)

        default = Default(component=component, resolution=resolution,
                          indir=resdir, init=init)
        default.set_shared_config(config, link=config_filename)
        component.add_task(default)

        if resolution == 10.:
            decomp = Decomp(component=component, resolution=resolution,
                            indir=resdir, init=init)
            decomp.set_shared_config(config, link=config_filename)
            component.add_task(decomp)

            restart = Restart(component=component, resolution=resolution,
                              indir=resdir, init=init)
            restart.set_shared_config(config, link=config_filename)
            component.add_task(restart)

            threads = Threads(component=component, resolution=resolution,
            threads.set_shared_config(config, link=config_filename)
            component.add_task(threads)

        component.add_task(Rpe(component=component, resolution=resolution,
                               indir=resdir, init=init, config=config))

from polaris.ocean.tasks.baroclinic_channel.validate import Validate

class Decomp(Task):
    """
    A baroclinic channel decomposition task, which makes sure the model
    produces identical results on 1 and 4 cores.
    """

    def __init__(self, component, resolution, indir, init):
        """
        Create the task

        indir : str

        """

        super().__init__(component=component, name='decomp', indir=indir)

        subdirs = list()
        for procs in [4, 8]:
            name = f'{procs}proc'

            self.add_step(Forward(
                ntasks=procs, min_tasks=procs, openmp_threads=1,
                resolution=resolution, run_time_steps=3,
                graph_target=f'{init.path}/culled_graph.info'))
            subdirs.append(name)
        self.add_step(Validate(component=component, step_subdirs=subdirs,
                               indir=self.subdir))

from polaris.ocean.tasks.baroclinic_channel.viz import Viz

class Default(Task):
    """
    The default baroclinic channel test case simply creates the mesh and
    initial condition, then performs a short forward run on 4 cores.
    """

        """

        indir : str

        """
        super().__init__(component=component, name='default', indir=indir)

            Forward(component=component, indir=self.subdir, ntasks=None,
                    min_tasks=None, openmp_threads=1, resolution=resolution,
                    run_time_steps=3,
                    graph_target=f'{init.path}/culled_graph.info'))

            Viz(component=component, indir=self.subdir))

from polaris.ocean.tasks.baroclinic_channel.restart.restart_step import (
    RestartStep,
)

class Restart(Task):
    """
    A baroclinic channel restart test case, which makes sure the model
    produces identical results with one longer run and two shorter runs with a
    restart in between.
    """

        """

        indir : str

        """
        super().__init__(component=component, name='restart', indir=indir)

        full = RestartStep(component=component, resolution=resolution,
                           name='full_run', indir=self.subdir, init=init)
        self.add_step(full)

        restart = RestartStep(component=component, resolution=resolution,
                              name='restart_run', indir=self.subdir,
                              init=init)
        restart.add_dependency(full, full.name)
        self.add_step(restart)

        self.add_step(Validate(component=component,
                               step_subdirs=['full_run', 'restart_run'],

class Threads(Task):
    """
    A baroclinic channel thread test case, which makes sure the model produces
    identical results with 1 and 2 threads.
    """

        """

        indir : str

        """

        super().__init__(component=component, name='threads', indir=indir)

        for openmp_threads in [1, 2]:
            name = f'{openmp_threads}thread'
                component=component, name=name, indir=self.subdir, ntasks=4,
                min_tasks=4, openmp_threads=openmp_threads,

from polaris.ocean.convergence.spherical import SphericalConvergenceForward

class Forward(SphericalConvergenceForward):
    """
    A step for performing forward ocean component runs as part of the cosine
    bell test case
    """

    def __init__(self, component, name, subdir, mesh, init,
                 refinement_factor, refinement='both'):
        """

        name : str

        mesh : polaris.Step

        """
        package = 'polaris.ocean.tasks.cosine_bell'
        validate_vars = ['normalVelocity', 'tracer1']
                         mesh=mesh, init=init, package=package,
                         validate_vars=validate_vars,
                         graph_target=f'{mesh.path}/graph.info',

        """
        """

    """
    """

        """

        name : str

        """
        package = 'polaris.ocean.tasks.geostrophic'
        validate_vars = ['temperature',
                         'salinity',
                         'layerThickness',
                         'normalVelocity']

from polaris.ocean.tasks.ice_shelf_2d.default import Default
from polaris.ocean.tasks.ice_shelf_2d.init import Init

def add_ice_shelf_2d_tasks(component):
    """
    Add tasks for different ice shelf 2-d tests to the ocean component

    """
    for resolution in [5., 2.]:
        for coord_type in ['z-star', 'z-level']:
            basedir = f'planar/ice_shelf_2d/{resdir}/{coord_type}'

            config_filename = 'ice_shelf_2d.cfg'
            config = PolarisConfigParser(
                filepath=f'{basedir}/{config_filename}')
            config.add_from_package('polaris.ocean.ice_shelf',
                                    'ssh_adjustment.cfg')
            config.add_from_package('polaris.ocean.tasks.ice_shelf_2d',
                                    config_filename)
            config.set('vertical_grid', 'coord_type', coord_type)

            init = Init(component=component, resolution=resolution,
                        indir=basedir)
            init.set_shared_config(config, link=config_filename)

            default = Default(component=component, resolution=resolution,
                              indir=basedir, init=init, config=config)
            default.set_shared_config(config, link=config_filename)
            component.add_task(default)

                              indir=basedir, init=init, config=config,
                              include_viz=True)

                              include_restart=True)

                              include_tides=True)

        # The only test case that makes sense with the single_layer coordinate
        # type is the one with barotropic tidal_forcing
        for coord_type in ['single_layer']:

            config.set('vertical_grid', 'coord_type', 'z-level')
            config.set('vertical_grid', 'vert_levels', '1')
            config.set('vertical_grid', 'partial_cell_type', 'None')

from polaris.ocean.ice_shelf import IceShelfTask
from polaris.ocean.tasks.ice_shelf_2d.forward import Forward
from polaris.ocean.tasks.ice_shelf_2d.ssh_forward import SshForward
from polaris.ocean.tasks.ice_shelf_2d.validate import Validate
from polaris.ocean.tasks.ice_shelf_2d.viz import Viz

class Default(IceShelfTask):
    """
    The default ice shelf 2d test case simply creates the mesh and
    initial condition, then performs a short forward run. Optionally,
    tidal forcing can be added or visualization and restart steps.

    """

    def __init__(self, component, resolution, indir, init, config,
                 include_viz=False, include_restart=False,
                 include_tides=False):
        """

        indir : str
            The directory the task is in, to which the test case name will be
            added

        include_viz : bool, optional

        include_restart : bool, optional
            Include restart and validation steps to test restart capabilities

        include_tides : bool, optional
            Include tidal forcing in the forward step
        """
        if include_tides and include_restart:
            raise ValueError('Restart test is not compatible with tidal '
                             'forcing')
        if include_tides:
            base_name = 'default_tidal_forcing'
        else:
            base_name = 'default'

        test_name = base_name
        test_subdir = f'{indir}/{base_name}'
            test_name = f'{base_name}_with_viz'
            test_subdir = f'{indir}/{base_name}/with_viz'
        if include_restart:
            test_name = f'{base_name}_with_restart'
            test_subdir = f'{indir}/{base_name}/with_restart'

        ssh_dir = indir
        forward_dir = f'{indir}/{base_name}'

        # Put the ssh adjustment steps in indir rather than subdir
        super().__init__(component=component, min_resolution=resolution,
                         name=test_name, subdir=test_subdir, sshdir=ssh_dir)

        last_adjust_step = self.setup_ssh_adjustment_steps(
            mesh_filename=f'{init.path}/culled_mesh.nc',
            graph_target=f'{init.path}/culled_graph.info',
            init_filename=f'{init.path}/output.nc',
            config=config, config_filename='ice_shelf_2d.cfg',
            ForwardStep=SshForward,
            package='polaris.ocean.tasks.ice_shelf_2d')

        forward_path = f'{forward_dir}/forward'
        if forward_path in component.steps:
            forward_step = component.steps[forward_path]
            symlink = 'forward'
        else:
            forward_step = Forward(component=component,
                                   indir=forward_dir,
                                   ntasks=None, min_tasks=None,
                                   openmp_threads=1, resolution=resolution,
                                   mesh=init, init=last_adjust_step,
                                   tidal_forcing=include_tides)
            forward_step.set_shared_config(config, link='ice_shelf_2d.cfg')
        self.add_step(forward_step, symlink=symlink)

            restart_path = f'{test_subdir}/restart'
            if restart_path in component.steps:
                restart_step = component.steps[restart_path]
                symlink = 'restart'
            else:
                restart_step = Forward(component=component,
                                       indir=test_subdir,
                                       ntasks=None, min_tasks=None,
                                       openmp_threads=1, resolution=resolution,
                                       mesh=init, init=last_adjust_step,
                                       do_restart=True)
                restart_step.set_shared_config(config, link='ice_shelf_2d.cfg')
                restart_step.add_dependency(forward_step, forward_step.name)
            self.add_step(restart_step, symlink=symlink)
            self.add_step(Validate(component=component,
                                   step_subdirs=['forward', 'restart'],
                                   indir=test_subdir))

            self.add_step(
                Viz(component=component, indir=test_subdir, mesh=init,
                    init=last_adjust_step))

from polaris.ocean.tasks.internal_wave.default import Default
from polaris.ocean.tasks.internal_wave.init import Init
from polaris.ocean.tasks.internal_wave.rpe import Rpe

def add_internal_wave_tasks(component):
    """
    Add tasks for different internal wave tests to the ocean component

    """
    config_filename = 'internal_wave.cfg'
    base_dir = 'planar/internal_wave'
    config = PolarisConfigParser(filepath=f'{base_dir}/{config_filename}')
    config.add_from_package('polaris.ocean.tasks.internal_wave',
                            config_filename)

    init = Init(component=component, indir=base_dir)
    init.set_shared_config(config, link=config_filename)

    for vadv_method in ['standard', 'vlr']:
        default = Default(component=component, indir=base_dir, init=init,
                          vadv_method=vadv_method)

        component.add_task(Rpe(component=component, indir=base_dir, init=init,
                               vadv_method=vadv_method, config=config))

from polaris.ocean.tasks.internal_wave.viz import Viz

    """
    The default test case for the internal wave test
    """

    def __init__(self, component, indir, init, vadv_method='standard'):
        """

        indir : str

        """
        super().__init__(component=component, name=f'{vadv_method}/default',

            Forward(component=component, init=init, indir=self.subdir,
                    ntasks=None, min_tasks=None, openmp_threads=1,
                    run_time_steps=3, vadv_method=vadv_method))

            Viz(component=component, indir=self.subdir), run_by_default=False)

from polaris.ocean.tasks.single_column.cvmix import CVMix
from polaris.ocean.tasks.single_column.ideal_age import IdealAge

def add_single_column_tasks(component):
    """
    Add tasks for various single-column tests

    """
    component.add_task(CVMix(component=component))
    component.add_task(IdealAge(component=component))

import os

from polaris.ocean.tasks.single_column.forward import Forward
from polaris.ocean.tasks.single_column.init import Init
from polaris.ocean.tasks.single_column.viz import Viz

class CVMix(Task):
    """
    The CVMix single-column test case creates the mesh and initial condition,
    then performs a short forward run testing vertical mixing on 1 core.
    """
    def __init__(self, component):
        """
        """
        name = 'cvmix'
        subdir = os.path.join('single_column', name)
            Init(component=component, indir=self.subdir))

        validate_vars = ['temperature', 'salinity', 'layerThickness',
            Forward(component=component, indir=self.subdir, ntasks=1,
                    min_tasks=1, openmp_threads=1,
                    validate_vars=validate_vars))

        self.config.add_from_package(
            'polaris.ocean.tasks.single_column',
            'single_column.cfg')

    """
    channel test cases.

    resources_fixed : bool
        Whether resources were set already and shouldn't be updated
        algorithmically

    dt : float

    """
    def __init__(self, component, name='forward', subdir=None, indir=None,
                 validate_vars=None):
        """

        name : str

            provided)
        """
                         openmp_threads=openmp_threads)

                            target='../init/initial_state.nc')
                            target='../init/culled_graph.info')

        self.add_yaml_file('polaris.ocean.tasks.single_column',

        self.add_output_file(filename='output.nc', validate_vars=validate_vars)

        self.resources_fixed = (ntasks is not None)

import os

class IdealAge(Task):
    """
    The ideal-age single-column test case creates the mesh and initial
    condition, then performs a short forward run evolving the ideal-age tracer
    on 1 core.
    """
    def __init__(self, component, ideal_age=True):
        """
        """
        name = 'ideal_age'
            Init(component=component, indir=self.subdir, ideal_age=ideal_age))

        validate_vars = ['temperature', 'salinity', 'iAge']
        step = Forward(component=component, indir=self.subdir, ntasks=1,
                       min_tasks=1, openmp_threads=1,
                       validate_vars=validate_vars)

        step.add_yaml_file('polaris.ocean.tasks.single_column.ideal_age',

        self.add_step(step)

            Viz(component=component, indir=self.subdir, ideal_age=ideal_age))

    """
    A step for performing forward ocean component runs as part of the sphere
    transport test case
    """

    def __init__(self, component, name, subdir, base_mesh, init,
                 case_name, refinement_factor, refinement):
        """

        name : str

        """
        package = 'polaris.ocean.tasks.sphere_transport'
        flow_id = {'rotation_2d': 1,
                   'nondivergent_2d': 2,
                   'divergent_2d': 3,
                   'correlated_tracers_2d': 4}
        namelist_options = {
            'mpas-ocean': {
                'config_transport_tests_flow_id': flow_id[case_name]
            }
        }
        validate_vars = ['normalVelocity', 'tracer1', 'tracer2', 'tracer3']
                         mesh=base_mesh,
                         init=init, package=package,
                         options=namelist_options,
                         graph_target=f'{base_mesh.path}/graph.info',

from polaris.ocean.viz.transect.vert import (
    interp_mpas_to_transect_cells,
    interp_mpas_to_transect_nodes,
)

def plot_transect(ds_transect, mpas_field=None, out_filename=None, ax=None,
                  title=None, vmin=None, vmax=None, colorbar_label=None,
                  cmap=None, figsize=(12, 6), dpi=200, method='flat',
                  outline_color='black', ssh_color=None, seafloor_color=None,
                  interface_color=None, cell_boundary_color=None,
                  linewidth=1.0, color_start_and_end=False,
                  start_color='red', end_color='green'):
    """
    plot a transect showing the field on the MPAS-Ocean mesh and save to a file

        A transect dataset from
        :py:func:`polaris.ocean.viz.compute_transect()`

    mpas_field : xarray.DataArray
        The MPAS-Ocean 3D field to plot

    out_filename : str, optional
        The png file to write out to

    title : str

        The colorbar label, or ``None`` if no colorbar is to be included.
        Use an empty string to display a colorbar without a label.

    cmap : str, optional
        The name of a colormap to use

        The dots per inch of the image

    method : {'flat', 'bilinear'}, optional
        The type of interpolation to use in plots.  ``flat`` means constant
        values over each MPAS cell.  ``bilinear`` means smooth interpolation
        between horizontally between cell centers and vertical between the
        middle of layers.

    outline_color : str or None, optional
        The color to use to outline the transect or ``None`` for no outline

    ssh_color : str or None, optional
        The color to use to plot the SSH (sea surface height) or ``None`` if
        not plotting the SSH (except perhaps as part of the outline)

    seafloor_color : str or None, optional
        The color to use to plot the seafloor depth or ``None`` if not plotting
        the seafloor depth (except perhaps as part of the outline)

    interface_color : str or None, optional
        The color to use to plot interfaces between layers or ``None`` if
        not plotting the layer interfaces

    cell_boundary_color : str or None, optional
        The color to use to plot vertical boundaries between cells or ``None``
        if not plotting cell boundaries.  Typically, ``cell_boundary_color``
        will be used along with ``interface_color`` to outline cells both
        horizontally and vertically.

    linewidth : float, optional
        The width of outlines, interfaces and cell boundaries

    color_start_and_end : bool, optional
        Whether to color the left and right axes of the transect, which is
        useful if the transect is also being plotted in an inset or on top of
        a horizontal field

    start_color : str, optional
        The color of left axis marking the start of the transect if
        ``plot_start_end == True``

    end_color : str, optional
        The color of right axis marking the end of the transect if
    """

    if ax is None and out_filename is None:
        raise ValueError('One of ax or out_filename must be supplied')

    create_fig = ax is None

    z = ds_transect.zTransectNode
    x = 1e-3 * ds_transect.dNode.broadcast_like(z)

    if mpas_field is not None:
        if method == 'flat':
            transect_field = interp_mpas_to_transect_cells(ds_transect,
                                                           mpas_field)
            shading = 'flat'
        elif method == 'bilinear':
            transect_field = interp_mpas_to_transect_nodes(ds_transect,
            shading = 'gouraud'
        else:
            raise ValueError(f'Unsupported method: {method}')

        pc = ax.pcolormesh(x.values, z.values, transect_field.values,
                           shading=shading, cmap=cmap, vmin=vmin, vmax=vmax,
                           zorder=0)
        ax.autoscale(tight=True)
        if colorbar_label is not None:
            plt.colorbar(pc, extend='both', shrink=0.7, ax=ax,
                         label=colorbar_label)

    _plot_interfaces(ds_transect, ax, interface_color, cell_boundary_color,
                     ssh_color, seafloor_color, color_start_and_end,
                     start_color, end_color, linewidth)

    _plot_outline(x, z, ds_transect.validNodes, ax, outline_color,
                  linewidth)

    ax.set_xlabel('transect distance (km)')
    ax.set_ylabel('z (m)')

        if title is not None:
        plt.savefig(out_filename, dpi=dpi, bbox_inches='tight', pad_inches=0.2)

def _plot_interfaces(ds_transect, ax, interface_color, cell_boundary_color,
                     start_color, end_color, linewidth):
    if cell_boundary_color is not None:
        x_bnd = 1e-3 * ds_transect.dCellBoundary.values.T
        z_bnd = ds_transect.zCellBoundary.values.T
        ax.plot(x_bnd, z_bnd, color=cell_boundary_color, linewidth=linewidth,
                zorder=1)

    if interface_color is not None:
        x_int = 1e-3 * ds_transect.dInterfaceSegment.values.T
        z_int = ds_transect.zInterfaceSegment.values.T
        ax.plot(x_int, z_int, color=interface_color, linewidth=linewidth,
                zorder=2)

    if ssh_color is not None:
        valid = ds_transect.validNodes.any(dim='nVertNodes')
        x_ssh = 1e-3 * ds_transect.dNode.values
        z_ssh = ds_transect.ssh.where(valid).values
        ax.plot(x_ssh, z_ssh, color=ssh_color, linewidth=linewidth, zorder=4)

    if seafloor_color is not None:
        x_floor = 1e-3 * ds_transect.dNode.values
        z_floor = ds_transect.zSeafloor.where(valid).values
        ax.plot(x_floor, z_floor, color=seafloor_color, linewidth=linewidth,
                zorder=5)

    if color_start_and_end:
        ax.spines['left'].set_color(start_color)
        ax.spines['left'].set_linewidth(4 * linewidth)
        ax.spines['right'].set_color(end_color)
        ax.spines['right'].set_linewidth(4 * linewidth)

def _plot_outline(x, z, valid_nodes, ax, outline_color, linewidth,
                  epsilon=1e-6):
    if outline_color is not None:
        # add a buffer of invalid values around the edge of the domain
        valid = np.zeros((x.shape[0] + 2, x.shape[1] + 2), dtype=float)
        z_buf = np.zeros(valid.shape, dtype=float)
        x_buf = np.zeros(valid.shape, dtype=float)

        valid[1:-1, 1:-1] = valid_nodes.astype(float)
        z_buf[1:-1, 1:-1] = z.values
        z_buf[0, 1:-1] = z_buf[1, 1:-1]
        z_buf[-1, 1:-1] = z_buf[-2, 1:-1]
        z_buf[:, 0] = z_buf[:, 1]
        z_buf[:, -1] = z_buf[:, -2]

        x_buf[1:-1, 1:-1] = x.values
        x_buf[0, 1:-1] = x_buf[1, 1:-1]
        x_buf[-1, 1:-1] = x_buf[-2, 1:-1]
        x_buf[:, 0] = x_buf[:, 1]
        x_buf[:, -1] = x_buf[:, -2]

        ax.contour(x_buf, z_buf, valid, levels=[1. - epsilon],
                   colors=outline_color, linewidths=linewidth, zorder=3)

from polaris.seaice.tasks.single_column import add_single_column_tasks

class SeaIce(Component):
    """
    The collection of all test case for the MPAS-Seaice component
    """

        """
        Construct the collection of MPAS-Seaice test cases
        """
        super().__init__(name='seaice')

        # please keep these in alphabetical order

from polaris.seaice.tasks.single_column.exact_restart import ExactRestart
from polaris.seaice.tasks.single_column.standard_physics import StandardPhysics

    """
    Add various tasks that define single-column tests

    component : polaris.seaice.Seaice
        the component that that the tasks will be added to
    """
    component.add_task(
        StandardPhysics(component=component))
        ExactRestart(component=component))

import os

from polaris.seaice.tasks.single_column.exact_restart.validate import Validate
from polaris.seaice.tasks.single_column.forward import Forward

class ExactRestart(Task):
    """
    A single-column restart test case, which makes sure  the model produces
    identical results with one longer run and two shorter runs with a restart
    in between.
    """

        """

        component : polaris.seaice.Seaice
            the component that that the task belongs to
        """
        name = 'exact_restart'

        validate_vars = ['iceAreaCategory',
                         'iceVolumeCategory',
                         'snowVolumeCategory',
                         'surfaceTemperature',
                         'iceEnthalpy',
                         'iceSalinity',
                         'snowEnthalpy',
                         'levelIceArea',
                         'levelIceVolume',
                         'pondArea',
                         'pondDepth',
                         'pondLidThickness',
                         'uVelocity',
                         'vVelocity',
                         'freezeOnset',
                         'snowfallRate',
                         'pondSnowDepthDifference',
                         'pondLidMeltFluxFraction',
                         'solarZenithAngleCosine',
                         'shortwaveScalingFactor',
                         'shortwaveVisibleDirectDown',
                         'shortwaveVisibleDiffuseDown',
                         'shortwaveIRDirectDown',
                         'shortwaveIRDiffuseDown',
                         'oceanStressCellU',
                         'oceanStressCellV',
                         'seaSurfaceTemperature',
                         'freezingMeltingPotential',
                         'airOceanDragCoefficientRatio']

        step = Forward(component=component, name='full_run', indir=self.subdir)
        step.add_output_file(
            filename='restarts/restart.2000-01-01_12.00.00.nc',
            validate_vars=validate_vars)
            filename='restarts/restart.2000-01-02_00.00.00.nc',

        step.add_namelist_file(
            package='polaris.seaice.tasks.single_column.exact_restart',
            namelist='namelist.full')
        step.add_streams_file(
            streams='streams.full')

        step = Forward(component=component, name='restart_run',
                       indir=self.subdir)
        step.add_input_file(
            target='../../full_run/restarts/restart.2000-01-01_12.00.00.nc')

            filename='restarts/restart.2000-01-02_00.00.00.nc')

            namelist='namelist.restart')
            streams='streams.restart')

        subdirs = ['full_run', 'restart_run']
        restart_filename = 'restarts/restart.2000-01-02_00.00.00.nc'
                               indir=self.subdir, variables=validate_vars,
                               restart_filename=restart_filename))

from polaris import ModelStep

class Forward(ModelStep):
    """
    A step for staging a mesh for “single column” test cases
    """
    def __init__(self, component, name='forward', indir=None):
        """

          The name of the step

        """
        super().__init__(component=component, name=name, indir=indir,
                         ntasks=1, min_tasks=1, openmp_threads=1)

        self.add_input_file(filename='grid.nc',
                            target='grid_sc_71.35_-156.5.nc',
                            database='domains/domain_sc_71.35_-156.5/')

        for year in range(1948, 2008):
                filename=f'forcing/atmosphere_forcing_six_hourly.{year}.nc',
                target=f'LYq_six_hourly.{year}.nc',
                database='domains/domain_sc_71.35_-156.5/')

        self.add_input_file(filename='forcing/atmosphere_forcing_monthly.nc',
                            target='LYq_monthly.nc',

        self.add_input_file(filename='forcing/ocean_forcing_monthly.nc',
                            target='oceanmixed_ice_depth_sc_71.35_-156.5.nc',

            filename='forcing/snicar_optics_5bnd_snow_and_aerosols.nc',
            target='snicar_optics_5bnd_snow_and_aerosols.nc',
            database='domains/domain_sc_71.35_-156.5/')

            filename='forcing/standard_optics_mpas_seaice.nc',
            target='standard_optics_mpas_seaice.nc',

        self.add_namelist_file(
            package='polaris.seaice.tasks.single_column',
            namelist='namelist.seaice')

        self.add_streams_file(
            streams='streams.seaice')

import os

from polaris.seaice.tasks.single_column.standard_physics.viz import Viz

class StandardPhysics(Task):
    """
    The single-column standard physics test case creates the mesh and initial
    condition, then performs a short forward run.
    """

        """

        """
        name = 'standard_physics'
        step = Forward(component=component, indir=self.subdir)
            package='polaris.seaice.tasks.single_column.standard_physics',
        variables = ['iceAreaCell', 'iceVolumeCell', 'snowVolumeCell',
                     'surfaceTemperatureCell', 'shortwaveDown', 'longwaveDown']
        step.add_output_file(filename='output/output.2000.nc',
                             validate_vars=variables)
        self.add_step(Viz(component=component, indir=self.subdir))