Skip to content

Packages

Jump to bottom
Ben Prather edited this page Sep 4, 2024 · 60 revisions

Introduction: Package Components

Packages are Parthenon's way of introducing modularity in the code, separating different physics and features so that they don't interfere with one another or render the code unreadable. Each package has its own C++ namespace for functions, as well as grid variables ("fields"), and a space for parameters/options/globals (a parthenon::Params object). Packages operate by defining functions, which they then assign as "callbacks," to be invoked at particular points during each step (or at various other times, e.g. before each output file is produced).

Initialize()

Every package should have an Initialize function in its namespace, which exists to declare the parameters, fields, and callbacks which a package will use. The function should return a KHARMAPackage object, which is added to a master list of packages accessible from throughout the code.

The KHARMA::ProcessPackages function is responsible for loading packages by calling each Initialize. Remember to add your package to this function in order to load it!

ProcessPackages uses a task list in order to load packages in the correct order. The basic packages which are always loaded are Globals, Driver, GRMHD, Inverter, Flux, Reductions, and Boundaries. A package for magnetic field is loaded if the problem includes magnetic fields, and the Implicit package is loaded if any variables are marked with the "Implicit" flag -- usually for Extended MHD problems (see ImEx driver). Additional packages are loaded as needed, usually because either they are explicitly enabled in the parameters (e.g., Wind), or they are necessary to calculate some variable (e.g., Current, CoordinateOutput).

Params

The Params class acts quite a bit like a Dictionary in Python, and one is allocated for each package. Members of any type can be added with pkg->Add("name", object) -- this includes not only the usual int, Real, etc. but also structs, objects, or even Kokkos Views containing any array-type information which should not be block-attached (though this last use is discouraged in KHARMA if it can be avoided).

Unless marked private, variables can be accessed by nearly and function in KHARMA via the Mesh or MeshBlock pointer, usually called pmb. This is a feature, not a bug -- it allows any function to access parameters from a categorized global repository, and acts as an alternative to dangerous global variables or cumbersome extra function parameters. For example, one should access the fluid adiabatic index gamma via the GRMHD package:

auto gam = pmb->packages.Get("GRMHD")->Param<Real>("gamma");

and this is done all across the code.

Fields and Flags

In additional to global constants that are available through the params object, packages can also include variables that are defined over the entire mesh (also called fields). These can be scalar, vectors or even tensor objects, attached to each cell center, face, or edge of each zone, or free of the grid structure entirely. Each variable carries metadata flags, which contains all the information about how a variable should be allocated and handled (size, datatype, whether the variable should be synchronized at MPI communications, and so on).

Metadata m_real = Metadata({Metadata::Real, Metadata::Cell, Metadata::Derived, Metadata::OneCopy});
pkg->AddField("solve_norm", m_real);

Here we have defined an object of Parthenon's Metadata class called m_real by calling Metadata's constructor. We've provided a MetadataFlag vector that specifies four flags, Metadata::Real tells it that the field defined by this metadata will consist of Reals, Metadata::Cell clarifies that array of Reals will reside at the cell centers, Metadata::Derived explains that the data in this array can be computed on-the-fly from a more fundamental state of variables, and finally Metadata::OneCopy tells Parthenon that any field defined with this metadata will be shared between stages (recall that we have a multistage driver). The Parthenon docs have a page that lists all the metadata flags available. The second line in the code excerpt above defines a field solve_norm in the implicit package with m_real as its metadata.

User flags may be added with Metadata::AddUserFlag("FlagName") and retrieved with Metadata::GetUserFlag("FlagName"). Flags are useful for categorizing variables, to pack only those needed for each kernel.

Callbacks

Each package can additionally register functions to be called at certain points during each step, e.g. to add a source term, update fluxes before they are applied, etc, etc. These callbacks are intended to be the primary way most packages operate, rather than adding code to the different Driver functions explicitly.

The callbacks available are:

        // PHYSICS
        // Recovery of primitive variables from conserved.
        // These can be host-side functions because they are not called from GetFlux()
        // rather, they are called on zone center values once per step only.
        std::function<void(MeshBlockData<Real>*, IndexDomain, bool)> BlockUtoP = nullptr;
        std::function<void(MeshData<Real>*, IndexDomain, bool)> MeshUtoP = nullptr;
        // Allow applying UtoP only/separately for boundary domains after sync/prolong/restrict ops
        // All packages with independent variables should register this for AMR
        std::function<void(MeshBlockData<Real>*, IndexDomain, bool)> BoundaryUtoP = nullptr;
        // On domain boundaries, however, we sometimes need to respect the primitive variables.
        // Currently only the GRMHD primitives (rho, u, uvec) do this
        std::function<void(MeshBlockData<Real>*, IndexDomain, bool)> DomainBoundaryPtoU = nullptr;

        // Going the other way, however, is handled by Flux::{Block,Mesh}PtoU.
        // All PtoU implementations are device-side (called prim_to_flux),
        // so we do not need something like
        //std::function<void(MeshBlockData<Real>*, IndexDomain, bool)> BlockPtoU = nullptr;

        // Source term to add to the conserved variables during each step
        std::function<void(MeshData<Real>*, MeshData<Real>*, IndexDomain)> AddSource = nullptr;

        // Source term to apply to primitive variables, needed for some problems in order
        // to control dissipation (Hubble, turbulence).
        // Must be applied over entire domain!
        std::function<void(MeshBlockData<Real>*)> BlockApplyPrimSource = nullptr;

        // Apply any fixes after the initial fluxes are calculated
        std::function<void(MeshData<Real>*)> FixFlux = nullptr;

        // Apply any floors or limiters specific to the package (that is, on the package's variables)
        std::function<void(MeshBlockData<Real>*, IndexDomain)> BlockApplyFloors = nullptr;
        std::function<void(MeshData<Real>*, IndexDomain)> MeshApplyFloors = nullptr;

        // CONVENIENCE
        // Anything to be done before each step begins -- currently just updating global "in_loop"
        std::function<void(Mesh*, ParameterInput*, const SimTime&)> PreStepWork = nullptr;
        // Anything to be done after every step is fully complete -- usually reductions or preservation of variables
        // Note that most diagnostics should go in "PostStepDiagnosticsMesh" instead
        std::function<void(Mesh*, ParameterInput*, const SimTime&)> PostStepWork = nullptr;

        // Anything to be done just before any outputs (dump files, restarts, history files) are made
        // Usually for filling output-only variables
        // TODO Add MeshUserWorkBeforeOutput to Parthenon
        std::function<void(MeshBlock*, ParameterInput*)> BlockUserWorkBeforeOutput = nullptr;

        // Anything at the very end of simulation. Cleanup, summaries, outputs if you're brave
        std::function<void(Mesh*, ParameterInput*, const SimTime&)> PostExecute = nullptr;

KHARMA Packages

Globals

GRMHD

B field

Floors

Current

Electrons

Implicit

Overview of the solve

The implicit solver is essentially a multidimensional Newton-Raphson scheme. In its current form it finds the root for the residual of the system of balance laws, $\partial_t\boldsymbol{U}(\boldsymbol{P}) + \boldsymbol{\nabla}\boldsymbol{F}(\boldsymbol{P}) = \boldsymbol{S}(\boldsymbol{P})$ where $\boldsymbol{U}$ is the vector of conserved variables that are analytic functions of $\boldsymbol{P}$ - the vector of primitives, $\boldsymbol{F}$ are the corresponding fluxes and $\boldsymbol{S}$ are the source terms. To get a more detailed introduction to hyperbolic PDEs and balance laws, have a look at this page. The residual is then defined as $\boldsymbol{R}(\boldsymbol{P})\equiv\partial_t\boldsymbol{U}(\boldsymbol{P}) + \boldsymbol{\nabla}\boldsymbol{F}(\boldsymbol{P}) - \boldsymbol{S}(\boldsymbol{P})$. This is the cost/loss function that we'd like to minimize over the vector space of the implicitly evolved variables. This solver was necessitated by the introduction of extended MHD (EMHD) in KHARMA which contains implicit source terms.

As is the case for the multidimensional Newton-Raphson scheme, the solver computes a Jacobian $\boldsymbol{J}(\boldsymbol{P})$ and solves for the equation $\boldsymbol{J}(\boldsymbol{P})\cdot\boldsymbol{\delta P} = -\boldsymbol{R}(\boldsymbol{P})$ where $\boldsymbol{\delta P}$ is the update to the primitives in one iteration of the solve. The step length $\boldsymbol{\delta P}$ is the direction in which the update must take place; however, it is possible that when one is close to the solution, updating $\boldsymbol{P}$ with $\boldsymbol{\delta P}$ can overstep the solution. To avoid this, the solve includes a backtracking line search where we select a suitable step length $\lambda$ and update the primitives like $\boldsymbol{P} = \boldsymbol{P} + \lambda\boldsymbol{\delta P}$.

NOTE: This is a local solve ie., it minimizes the residual and computes the update to the primitives on a zone-by-zone basis.

The details of the scheme are present in section 3 of this paper.

Package variables and constants

The Implicit package defines a bunch of constants that are relevant to the implicit solve. These are,

  • min_nonlinear_iter (default: 1) - The minimum number of iterations that must be performed by the solver.
  • max_nonlinear_iter (default: 3) - The maximum number of iterations that must be performed by the solver.
  • rootfind_tol (default: 1.e-12) - The tolerance that must be met during the solve to consider it a success.
  • jacobian_delta (default: 4.e-8) - The small change considered for each primitive - $\Delta P$ - to compute the Jacobian.
  • use_qr (default: true) - A boolean that dictates whether the QR decomposition should be considered while solving the system of linear equations. If false, the solver uses LU decomposition. While QR has certain stability benefits, we've found that it can lead to potentially inaccurate results when performing a >= 8x8 solve on GPUs.
  • linesearch (default: true) - A boolean that decides if backtracking linesearch must be performed.

The following arguments are relevant only if linesearch has been enabled.

  • max_linesearch_iter (default: 3) - The maximum number of iterations that must be performed by the linesearch.
  • linesearch_eps (default: 1.e-4) - A parameter that determines if the rate of decrease of $\boldsymbol{R}$ is small enough to exit linesearch.
  • linesearch_lambda (default: 1.0) - The initial value of the step length.

Additionally, the Implicit package defines to two scalar fields,

  • solve_norm - The L2 norm of the residual.
  • solve_fail - It is possible that the solver may not converge to a physically viable solution in zones where the loss landscape is difficult to navigate. In these scenarios, we manually backtrack by setting the step length to 0.1, recording the failure by setting solve_fail to 1, and performing the linesearch. In the event the manual backtrack is not sufficient, we set solve_fail to 2 and exit the solver (for that particular zone). Later, we consider the average of the primitives of the neighboring zones that did converge.

Package tasks and functions

The package contains a single task Step that performs the solve over the implicit variables. This task is called by the ImexDriver after the explicit variables are updated. Since we are only concerned with evolving the implicit variables, we need to reorder the variables in relevant fluid states. We achieve this with the package's get_ordered_names function. It does this by leveraging metadata flags. The function creates a vector of MetadataFlag by picking variables marked with the ImplicitFlag first and then those marked with ExplicitFlag. We then pack the variables in this order in all fluid states via Parthenon's PackVariables function, and finally generate a VarMap to refer to the right variable index. VarMaps are explained in detail here.

The syntax for the rest of the task is much like KHARMA's GetFlux. We first estimate the scratch memory required and launch a Kokkos kernel (aka. parallel dispatch). Inside the outer for loop (across blocks and along X3 and X2) we allocate temporaries using Parthenon's ScratchPad and initialize them using the corresponding MeshBlockPack. In the inner for loop (along X1) we define Kokkos subviews - slices over the scratchpads - that are lightweight and faster. The remainder of the computational body carries out the necessary operations needed to update the primitives,

  1. Compute and save implicit source terms. Currently, we only include source terms for the EMHD variables, but one can just as easily include function calls for other implicit sources.
  2. Compute the Jacobian.
  3. Solve the system of equations (refer this) for the update $\boldsymbol{\delta P}$.
  4. Perform linesearch to estimate the right step length.
  5. Update the primitives, calculate the residual and its L2 norm.

Additionally, the package has two device side functions calc_residual and calc_jacobian that are defined in the package's header file and as the name implies compute the residual and jacobian respectively.

EMHD

DISCLAIMER: This package is in a developmental phase.

The EMHD package includes anisotropic heat conduction and momentum transport. The theory is detailed in this work and in the theory subsection we present just the governing equations with variable definitions for brevity.

Theory

The evolution equations for the extended magnetohydrodynamics model are,

$$ \nabla_{\mu}(n u^{\mu}) = 0,$$

$$ \nabla_{\mu}T^{\mu}_{\nu} = 0,$$

$$ \nabla_{\nu}^{*}F^{\mu\nu} = 0,$$

$$ \nabla_{\mu}(\tilde{q}u^{\mu}) = -\frac{\tilde{q} - \tilde{q}_{0}}{\tau_R} + \frac{\tilde{q}}{2}\nabla_{\mu}u^{\mu},$$

$$ \nabla_{\mu}(\Delta\tilde{P}u^{\mu}) = -\frac{\Delta\tilde{P} - \Delta\tilde{P}_{0}}{\tau_R} + \frac{\Delta\tilde{P}}{2}\nabla_{\mu}u^{\mu}, $$

where the first three equations are similar to those for ideal magnetohydrodynamics (see here), and the latter two dictate the evolution of the new non-ideal variables - conductive heat flux scalar $q$ and pressure anisotropy $\Delta P$. These equations are derived from the second law of thermodynamics (the derivation is present in section 4 of the paper). Note that $\tilde{q}$ and $\Delta\tilde{P}$ are rescaled variables related to $q$ and $\Delta P$ as,

$$\tilde{q} = q\Big(\frac{\tau_R}{\rho\chi\Theta^2}\Big)^{1/2}; \Delta\tilde{P} = \Delta P\Big(\frac{\tau_R}{\rho\nu\Theta}\Big)^{1/2}.$$

This is done to increase the stability of the numerical implementation of these equations in regions of low density. $\tau_R$ is the relaxation timescale over which we expect the variables to reach their saturation values - $\tilde{q}_0$ and $\Delta\tilde{P}_0$. $\chi$ is the coefficient of thermal diffusivity, $\nu$ is the kinematic viscosity and $\Theta\equiv P/\rho$ is the dimensionless temperature. The saturation values for these variables, which constitute their source terms, are,

$$ q_0\equiv -\rho\chi\hat{b}^{\mu}(\nabla_{\mu}\Theta + \Theta a_{\mu}),$$

$$\Delta P_{0}\equiv 3\rho\nu(\hat{b}^{\mu}\hat{b}^{\nu}\nabla_{\mu}u_{\nu} - \frac{1}{3}\nabla_{\mu}u^{\mu}).$$

Note that the stress-energy tensor $T^{\mu\nu}$ is,

$$ T^{\mu\nu} \equiv (\rho + u + p + b^2)u^{\mu}u^{\nu} + (p + \frac{b^2}{2})g^{\mu\nu} - b^{\mu}b^{\nu} + q(\hat{b}^{\mu}u^{\nu} + u^{\mu}\hat{b}^{\nu}) - \Delta P (\hat{b}^{\mu}\hat{b}^{\nu} - \frac{1}{3}(g^{\mu\nu} + u^{\mu}u^{\nu})). $$

Numerical implementation

The magnetic field variables are updated explicitly prior to the implicit solver. Therefore, a full non-ideal run contains 7 variables to be updated implicitly and it utilizes the newly updated B-field.

We now elucidate how the various source terms are computed ie., explicitly or implicitly. The source terms for the internal energy and velocity field are evaluated explicitly in Flux::AddSource prior to the solve. Since the evolution equations for the rest-mass density and magnetic fields contain no source terms we are done calculating all the source terms for the ideal MHD sector. For the two remaining equations corresponding to the non-ideal equations we selectively evaluate some explicitly and some implicitly. Let's see the source terms for $\tilde{q}$ as an example (the same will apply for $\Delta\tilde{P}$),

$$ S = -\frac{\tilde{q} - \tilde{q}_{0}}{\tau_R} + \frac{\tilde{q}}{2}\nabla_{\mu}u^{\mu}.$$

What is worth noting is that $q_0$ and the final term on the RHS contain the pesky covariant derivative. This means that there are temporal derivatives present in the source term which necessitates handling them implicitly. Up to muliplicative factors of negative sign and scalar fields, the source terms can be expressed like,

$$ S \sim \textcolor{blue}{\frac{\tilde{q}}{\tau_R}} + \bigg(\textcolor{blue}{\frac{\partial_{t}\Theta}{\tau_R}} + \frac{\partial_{i}\Theta}{\tau_R}\bigg) + \bigg(\textcolor{blue}{u^{t}\partial_{t} u_{\mu}} + u^{i}\partial_{i}u_{\mu} + \Gamma_{\mu\nu\kappa}u^{\nu}u^{\kappa}\bigg) + \tilde{q}\bigg(\textcolor{blue}{\partial_{t}u^{t}} + \partial_{i}u^{i} + g^{\mu\nu}\Gamma^{\kappa}_{\mu\nu}u_{\kappa}\bigg).$$

The terms in blue are treated implicitly and the rest are computed explicitly prior to the solve. The first term on the RHS has a factor of $1/{\tau_R}$ and can be a potentially small term. We therefore treat it implicitly (this is done in EMHD::implicit_sources). The remainder of the implicit terms contain a time derivative and NEED to be calculated within the solver as the fluid state is iteratively updated (EMHD::time_derivative_sources). The terms in black are computationally expensive and do not warrant an implicit update; they are therefore evaluated prior to the solve via EMHD::AddSource and passed as arguments to Implicit::Step.

Package variables and constants

Many of the constants in the EMHD package have to do with the closure scheme used to compute the transfer coefficients. We express this as,

$$ \chi = \phi c_{s}^2 \tau_{R}, $$

$$ \nu = \psi c_{s}^2 \tau_{R}, $$

where $\phi$ and $\psi$ are constant of the order of unity. They must be chosen carefully to ensure stability and causality. In the code we refer to them as conduction_alpha and viscosity_alpha respectively.

  • conduction_alpha (default: 1.) - The closure parameter for thermal diffusivity.
  • viscosity_alpha (default: 1.) - The closure parameter for kinematic viscosity.
  • tau (default: 1.) - Relaxation timescale. Must be provided for test problems. It is computed for the torus problems based on instability conditions.
  • closure_type (default: torus) - The closure scheme to be used. This is a misnomer of sorts since ultimately we use the scheme mentioned above but a few additional operations may need to be performed either to compute $\tau_{R}$ or if $\kappa$ or $\eta$ are provided instead (for eg. in the conducting atmosphere and viscous Bondi problem).
  • kappa (default: 1.) - Thermal conductivity.
  • eta (default: 1.) - Dynamic viscosity.
  • higher_order_terms (default: false) - Should second order terms be considered?
  • feedback (default: true) - Should the EMHD sector feedback on to the ideal MHD variables? We generally want this enabled to get the physics right but there may be cases where one might want to disable it (for eg., the viscous Bondi problem or for debugging purposes).

With the introduction of two new variables, the EMHD package has four new scalar fields,

  • prims.q - $\tilde{q}$
  • prims.dP - $\Delta\tilde{P}$
  • cons.q - $\sqrt{-g}\tilde{q}u^{t}$
  • cons.dP - $\sqrt{-g}\Delta\tilde{P}u^{t}$

Package tasks and functions

The EMHD::AddSource task is called by the ImEx driver prior to an implicit update, it computes all the explicit source terms associated with the evolution equations of the non-ideal variables. These source terms contain temperature and velocity gradients (cf. the decomposition of source terms as implicit/explicit above). This means we need to reconstruct these variables at the zone faces to obtain a second order slope value to zone centers. To facilitate this, the EMHD package has a bunch of auxiliary functions present in emhd_utils.hpp.

KHARMA does not cache scalar fields such as $\chi$, $\nu$, $\tau_{R}$, $q$ and $\Delta P$ unlike iharm3d, and instead computes these on the fly. The set_parameters function in the header file computes the transport coefficients based on the closure scheme, while the convert_prims_to_q_dP function calculates $q$ and $\Delta P$ from $\tilde{q}$ and $\Delta\tilde{P}$ respectively. Note that these differ when higher order terms are enabled. Finally, the EMHD package has it's own version of calc_tensor that is called when source terms or fluxes need to be estimated.

Clone this wiki locally