Create your own Snek5000 solver, packaging your Nek5000 user source code

Tip

You can now use the Copier template <https://github.com/snek5000/template-snek5000-solver/> to render the following packaging structure effortlessly!

This document describes how to package a Nek5000 source code in your own repository via snek5000’s API. What is described is the bare minimum to get started. Assume your files are named canonical.usr, canonical.par etc. and you wish to package it under a name say snek5000_canonical.

Note

The package name does not have to contain snek5000 or canonical. It could be anything. Make sure the short_name variable (in __init__.py) and entry point key (in setup.cfg) matches.

One first has to create a directory structure, (preferably in a git / mercurial repository) like this:

.
├── LICENSE
├── README.md
├── pyproject.toml
└── src
    └── snek5000_canonical
        ├── Snakefile
        ├── __init__.py
        ├── canonical.usr.f
        ├── output.py
        └── solver.py

You can copy the directory docs/examples/snek5000-canonical from the Snek5000 repository, and then modify the directory names and the files as needed.

Note

The Nek5000 .usr files have to contain Fortran 77 code in strict form layout. In snek5000 solvers, one can rename this file with the .usr.f extension so that editors can correctly highlight the Fortran code.

pyproject.toml

[build-system]
requires = ["setuptools"]
build-backend = "setuptools.build_meta"

[project]
name = "snek5000-canonical"
version = "0.0.1"
description = "An empty solver to demonstrate the packaging layout of a typical Snek5000 solver"
readme = "README.md"
requires-python = ">=3.8"
license = {file = "LICENSE"}
authors = [
    {name = "Ashwin Vishnu" }
]
keywords = ["Snek5000", "Fluidsim", "Nek5000", "CFD"]
dependencies = [
    "snek5000 >= 0.9.0"
]

[project.optional-dependencies]
tests = [
    "nox",
    "pytest",
    "pytest-cov"
]

[project.entry-points."snek5000.solvers"]
canonical = "snek5000_canonical.solver"

[tool.setuptools.packages.find]
where = ["src"]

[tool.setuptools.package-data]
snek5000_canonical = ["*.usr.f", "*.par.cfg", "Snakefile", "etc/*.yml"]

src/snek5000_canonical/__init__.py

from .output import Output
from .solver import Simul

short_name = "canonical"

__all__ = ["Simul", "Output", "short_name"]

src/snek5000_canonical/solver.py

from snek5000.info import InfoSolverMake
from snek5000.solvers.base import SimulNek

# To use KTH Framework import SimulKTH instead
# from snek5000.solvers.kth import SimulKTH


class InfoSolverCanonical(InfoSolverMake):
    """Contain the information on a :class:`snek5000_canonical.solver.Simul`
    instance.

    """

    def _init_root(self):
        from . import short_name

        super()._init_root()
        self.module_name = "snek5000_canonical.solver"
        self.class_name = "Simul"
        self.short_name = short_name

        self.classes.Output.module_name = "snek5000_canonical.output"
        self.classes.Output.class_name = "OutputCanonical"

        # To solve for the temperature, one needs to add a [TEMPERATURE]
        # section in the .par file. It can be done like this:
        self.par_sections_disabled.remove("temperature")


class SimulCanonical(SimulNek):
    """A solver which compiles and runs using a Snakefile."""

    InfoSolver = InfoSolverCanonical

    @classmethod
    def create_default_params(cls):
        """Set default values of parameters as given in reference
        implementation.

        """
        params = super().create_default_params()
        # Re-define default values for parameters here, if necessary
        # following ``canonical.par``, ``canonical.box`` and ``SIZE`` files

        # Extend with new default parameters here, for example:
        # params.nek.velocity._set_attrib("advection", True)

        return params


Simul = SimulCanonical

Note

As you might have guessed all additional files except for Fortran code can be generated by snek5000 and need not be packaged. See snek5000.operators.Operators and snek5000.solvers.kth.SimulKTH to see how to set the params instead.

src/snek5000_canonical/output.py

from snek5000.output.base import Output as OutputBase
from snek5000.resources import get_base_templates

box, makefile_usr, size = get_base_templates()


class OutputCanonical(OutputBase):
    template_box = box
    template_makefile_usr = makefile_usr
    template_size = size

    @classmethod
    def _set_info_solver_classes(cls, classes):
        """Set the classes for info_solver.classes.Output

        NOTE: This method is optional, and only required if custom classes are
        to be set.

        """
        super()._set_info_solver_classes(classes)

        # Replace sim.output.phys_fields with a custom class
        classes.PhysFields.module_name = "snek5000_canonical.phys_fields"
        classes.PhysFields.class_name = "PhysFieldsCanonical"

    @property
    def makefile_usr_sources(self):
        """
        Sources for inclusion to makefile_usr.inc
        Dict[directory]  -> list of source files
        """
        return {}

        # For example this is the list of extra files required for the KTH
        # Framework:

        # return {
        #     "toolbox": [
        #         ("frame.f", "FRAMELP"),
        #         ("mntrlog_block.f", "MNTRLOGD"),
        #         ("mntrlog.f", "MNTRLOGD"),
        #         ("mntrtmr_block.f", "MNTRLOGD", "MNTRTMRD"),
        #         ("mntrtmr.f", "MNTRLOGD", "MNTRTMRD", "FRAMELP"),
        #         ("rprm_block.f", "RPRMD"),
        #         ("rprm.f", "RPRMD", "FRAMELP"),
        #         ("io_tools_block.f", "IOTOOLD"),
        #         ("io_tools.f", "IOTOOLD"),
        #         ("chkpoint.f", "CHKPOINTD"),
        #         ("chkpt_mstp.f", "CHKPTMSTPD", "CHKPOINTD"),
        #         ("map2D.f", "MAP2D", "FRAMELP"),
        #         ("stat.f", "STATD", "MAP2D", "FRAMELP"),
        #         ("stat_IO.f", "STATD", "MAP2D", "FRAMELP"),
        #         ("math_tools.f",),
        #     ],
        # }


Output = OutputCanonical

Note

One can also define your own Jinja templates or re-purpose them from snek5000.resources. See our How to page on templates.

src/snek5000_canonical/phys_fields.py

Optionally, add classes for reader and phys_fields objects which would allow customized load, get_var and plot_* methods suited for your case.

import matplotlib.pyplot as plt

from snek5000.output.phys_fields import PhysFields
from snek5000.output.readers.pymech_ import ReaderPymech


class ReaderPymechAvg(ReaderPymech):
    """Horizontally and temporally averaged profiles."""

    tag = "pymech_avg"

    def load(self, prefix="", index="all", t_stat=None, **kwargs):
        ds = super().load(prefix, index, **kwargs)

        avg_dims = ("x", "z", "time")
        # Use xarray's sel method to slice along time and mean method to do the
        # averaging
        self.data = ds.sel(time=slice(t_stat, None)).mean(avg_dims)
        return self.data


class PhysFieldsCanonical(PhysFields):
    @staticmethod
    def _complete_info_solver(info_solver, classes=None):
        PhysFields._complete_info_solver(info_solver, classes=(ReaderPymechAvg,))

    def plot_mean_vel(self, marker="x"):
        """Plot mean velocity profiles against y-coordinate."""
        fix, ax = plt.subplots()
        data = self.data
        ax.plot("ux", "y", "", marker=marker, label="$U$", data=data)
        ax.plot("uy", "y", "", marker=marker, label="$V$", data=data)
        ax.plot("uz", "y", "", marker=marker, label="$W$", data=data)
        ax.set(ylabel="$y$")
        ax.legend()

src/snek5000_canonical/Snakefile

from snek5000 import ensure_env, get_snek_resource
from snek5000_canonical import short_name, Output


configfile: "config_simul.yml"


ensure_env()

# Necessary to pass configuration to other Snakemake modules
Output.update_snakemake_config(config, short_name)


# default rule
rule all:
    input:
        "nek5000",


# shorthand for mesh
rule mesh:
    input:
        f"{short_name}.re2",
        f"{short_name}.ma2",


# compiler and run rules
# ======================
module compiler:
    snakefile:
        get_snek_resource("compiler.smk")
    config:
        config


use rule * from compiler


# I/O rules
# =========
module io:
    snakefile:
        get_snek_resource("io.smk")
    config:
        config


use rule * from io


# internal rules
# ==============
module internal:
    snakefile:
        get_snek_resource("internal.smk")
    config:
        config


use rule * from internal as internal_*