Stokes integrators

stokes_integrators

Solver classes targetting systems dealing with momentum conservation.

This module provides a fine-tuned abstract base class, StokesSolverBase, from which efficient numerical solvers can be derived, such as StokesSolver for the Stokes system of conservation equations and ViscoelasticStokesSolver for an incremental displacement formulation of the latter using a Maxwell rheology. The module also exposes a function, create_stokes_nullspace, to automatically generate null spaces compatible with PETSc solvers and a class, BoundaryNormalStressSolver, to solve for the normal stress acting on a domain boundary. Typically, users instantiate the StokesSolver class by providing a relevant set of arguments and then call the solve method to request a solver update.

iterative_outer_stokes_solver_parameters = MappingProxyType({'mat_type': 'matfree', 'ksp_type': 'preonly', 'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'schur', 'pc_fieldsplit_schur_type': 'full', 'fieldsplit_1': MappingProxyType({'ksp_type': 'fgmres', 'ksp_rtol': 0.0001, 'ksp_max_it': 200, 'pc_type': 'python', 'pc_python_type': 'firedrake.MassInvPC', 'Mp_pc_type': 'ksp', 'Mp_ksp_ksp_rtol': 1e-05, 'Mp_ksp_ksp_type': 'cg', 'Mp_ksp_pc_type': 'sor'})}) module-attribute

Default iterative solver parameters for solution of Stokes system.

We configure the Schur complement approach as described in Section of 4.3 of Davies et al. (2022), using PETSc's fieldsplit preconditioner type, which provides a class of preconditioners for mixed problems that allows a user to apply different preconditioners to different blocks of the system.

The fieldsplit_0 entries configure solver options for the velocity block. The linear systems associated with this matrix are solved using a combination of the Conjugate Gradient (cg) method and an algebraic multigrid preconditioner (gamg).

The fieldsplit_1 entries contain solver options for the Schur complement solve itself. For preconditioning, we approximate the Schur complement matrix with a mass matrix scaled by viscosity, with the viscosity sourced from the approximation object provided to Stokes solver. Since this preconditioner step involves an iterative solve, the Krylov method used for the Schur complement needs to be of flexible type, and we use the Flexible Generalized Minimal Residual (fgmres) method by default.

Note

G-ADOPT will use the above default iterative solver parameters if the argument solver_parameters="iterative" is provided or in 3-D if the solver_parameters argument is omitted. To make modifications to these default values, the most convenient approach is to provide the modified values as a dictionary via the solver_parameters_extra argument. This dictionary can also hold new pairs of keys and values to extend the default ones. .

direct_stokes_solver_parameters = MappingProxyType({'mat_type': 'aij', 'ksp_type': 'preonly', 'pc_type': 'lu', 'pc_factor_mat_solver_type': 'mumps'}) module-attribute

Default direct solver parameters for solution of Stokes system.

We configure the direct solver to use the LU (lu) factorisation provided by the MUMPS package.

Note

G-ADOPT will use the above default direct solver parameters if the argument solver_parameters="direct" is provided or in 2-D if the solver_parameters argument is omitted. To make modifications to these default values, the most convenient approach is to provide the modified values as a dictionary via the solver_parameters_extra argument. This dictionary can also hold new pairs of keys and values to extend the default ones.

newton_stokes_solver_parameters = MappingProxyType({'snes_type': 'newtonls', 'snes_linesearch_type': 'l2', 'snes_max_it': 100, 'snes_atol': 1e-10, 'snes_rtol': 1e-05}) module-attribute

Default non-linear solver parameters for solution of Stokes system.

We use a setup based on Newton's method (newtonls) with a secant line search over the L2-norm of the function.

Note

G-ADOPT will use the above default non-linear solver parameters in conjunction with the above iterative or default solver parameters if the viscosity, mu, provided to the approximation depends on the solver's solution. To make modifications to these default values, the most convenient approach is to provide the modified values as a dictionary via the solver_parameters_extra argument. This dictionary can also hold new pairs of keys and values to extend the default ones.

coupled_gia_solver_parameters = MappingProxyType({'mat_type': 'matfree', 'ksp_type': 'gmres', 'ksp_rtol': 0.001, 'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'symmetric_multiplicative', 'fieldsplit_1': MappingProxyType({'ksp_type': 'cg', 'pc_type': 'python', 'pc_python_type': 'firedrake.AssembledPC', 'assembled_pc_type': 'sor', 'ksp_rtol': 1e-05})}) module-attribute

Default iterative solver parameters for CoupledInternalVariableSolver (GIA problems).

Uses a Newton SNES outer loop with a GMRES/fieldsplit preconditioned inner solve. SNES is always active: for a Newtonian rheology (exponent=1) the power-law factor is identically 1 and the system is linear, so SNES converges in a single iteration at negligible extra cost. For power-law rheology (exponent > 1) the full Newton iteration is required. The snes_rtol is set looser than ksp_rtol so that the single-iteration Newtonian case is accepted without additional SNES iterations.

StokesSolverBase(solution, approximation, /, *, dt=None, theta=0.5, additional_forcing_term=None, bcs={}, quad_degree=6, solver_parameters=None, solver_parameters_extra=None, gpu_extra_parameters=None, J=None, constant_jacobian=False, nullspace=None, transpose_nullspace=None, near_nullspace=None)

Bases: SolverConfigurationMixin, ABC

Solver for a system involving mass and momentum conservation.

Parameters:

Name	Type	Description	Default
solution	Function	Firedrake function representing the field over the mixed Stokes space	required
approximation	BaseApproximation	G-ADOPT approximation defining terms in the system of equations	required
dt	float | None	Float specifying the time step if the system involves a coupled time integration	None
theta	float	Float defining the theta scheme parameter used in a coupled time integration	0.5
additional_forcing_term	Form | None	Firedrake form specifying an additional term contributing to the residual	None
bcs	dict[int | str, dict[str, Any]]	Dictionary specifying boundary conditions (identifier, type, and value)	{}
quad_degree	int	Integer denoting the quadrature degree	6
solver_parameters	ConfigType | str | None	Dictionary of PETSc solver options or string matching one of the default sets	None
solver_parameters_extra	ConfigType | None	Dictionary of PETSc solver options used to update the default G-ADOPT options	None
gpu_extra_parameters	ConfigType | None	Dictionary of additional GPU-specific settings that StokesSolver will translate into appropriate PETSc solver options	None
J	Function | None	Firedrake function representing the Jacobian of the mixed Stokes system	None
constant_jacobian	bool	Boolean specifying whether the Jacobian of the system is constant	False
nullspace	MixedVectorSpaceBasis | None	A MixedVectorSpaceBasis for the operator's kernel	None
transpose_nullspace	MixedVectorSpaceBasis | None	A MixedVectorSpaceBasis for the kernel of the operator's transpose	None
near_nullspace	MixedVectorSpaceBasis | None	A MixedVectorSpaceBasis for the operator's smallest eigenmodes (e.g. rigid body modes)	None

Valid keys for boundary conditions

Condition	Type	Description
u	Strong	Solution
ux	Strong	Solution along the first Cartesian axis
uy	Strong	Solution along the second Cartesian axis
uz	Strong	Solution along the third Cartesian axis
un	Weak	Solution along the normal to the boundary
stress	Weak	Traction across the boundary
normal_stress	Weak	Stress component normal to the boundary
free_surface	Weak	Free-surface characteristics of the boundary

Valid keys describing the free surface boundary

Argument	Required	Description
delta_rho_fs	Yes (d)	Density contrast along the free surface
RaFS	Yes (nd)	Rayleigh number (free-surface density contrast)
variable_rho_fs	No	Account for buoyancy effects on interior density

Classic theta values for coupled implicit time integration

Theta	Scheme
0.5	Crank-Nicolson method
1.0	Backward Euler method

Note: Such a coupling can arise in the case of a free-surface implementation.

Source code in g-adopt/gadopt/stokes_integrators.py

def __init__(
    self,
    solution: fd.Function,
    approximation: BaseApproximation,
    /,
    *,
    dt: float | None = None,
    theta: float = 0.5,
    additional_forcing_term: fd.Form | None = None,
    bcs: dict[int | str, dict[str, Any]] = {},
    quad_degree: int = 6,
    solver_parameters: ConfigType | str | None = None,
    solver_parameters_extra: ConfigType | None = None,
    gpu_extra_parameters: ConfigType | None = None,
    J: fd.Function | None = None,
    constant_jacobian: bool = False,
    nullspace: fd.MixedVectorSpaceBasis | None = None,
    transpose_nullspace: fd.MixedVectorSpaceBasis | None = None,
    near_nullspace: fd.MixedVectorSpaceBasis | None = None,
) -> None:
    self.solution = solution
    self.approximation = approximation
    self.dt = dt
    self.theta = theta
    self.additional_forcing_term = additional_forcing_term
    self.bcs = bcs
    self.quad_degree = quad_degree
    self.J = J
    self.constant_jacobian = constant_jacobian
    self.nullspace = nullspace
    self.transpose_nullspace = transpose_nullspace
    self.near_nullspace = near_nullspace

    self.solution_old = self.solution.copy(deepcopy=True)
    self.solution_space = self.solution.function_space()
    self.mesh = self.solution_space.mesh()
    self.k = upward_normal(self.mesh)

    self.is_mixed_space = isinstance(
        self.solution_space.topological, fd.functionspaceimpl.MixedFunctionSpace
    )
    if self.is_mixed_space:
        self.tests = fd.TestFunctions(self.solution_space)
        self.solution_split = fd.split(solution)
        self.solution_old_split = fd.split(self.solution_old)
    else:
        self.test = fd.TestFunction(self.solution_space)
        self.solution_split = (solution,)
        self.solution_old_split = (self.solution_old,)

    self.solution_theta_split = [
        self.theta * sol + (1 - self.theta) * sol_old
        for sol, sol_old in zip(self.solution_split, self.solution_old_split)
    ]
    self.tests = fd.TestFunctions(self.solution_space)

    self.rho_continuity = self.approximation.rho_continuity()
    self.equations = []  # G-ADOPT's Equation instances

    self.set_boundary_conditions()
    self.set_equations()
    self.set_form()
    self.set_solver_options(solver_parameters, solver_parameters_extra, gpu_extra_parameters)
    self.set_solver()

set_boundary_conditions()

Sets strong and weak boundary conditions.

Source code in g-adopt/gadopt/stokes_integrators.py

def set_boundary_conditions(self) -> None:
    """Sets strong and weak boundary conditions."""
    self.strong_bcs = []
    self.weak_bcs = {}

    if self.is_mixed_space:
        bc_map = {"u": self.solution_space.sub(0)}
    else:
        bc_map = {"u": self.solution_space}
    if is_cartesian(self.mesh):
        bc_map["ux"] = bc_map["u"].sub(0)
        bc_map["uy"] = bc_map["u"].sub(1)
        if self.mesh.geometric_dimension == 3:
            bc_map["uz"] = bc_map["u"].sub(2)

    for bc_id, bc in self.bcs.items():
        weak_bc = defaultdict(float)

        for bc_type, val in bc.items():
            match bc_type:
                case "u" | "ux" | "uy" | "uz":
                    self.strong_bcs.append(
                        fd.DirichletBC(bc_map[bc_type], val, bc_id)
                    )
                case "free_surface":
                    if not hasattr(self, "set_free_surface_boundary"):
                        raise ValueError(
                            "This solver does not implement a free surface."
                        )
                    weak_bc["normal_stress"] += self.set_free_surface_boundary(
                        val, bc_id
                    )
                case _:
                    weak_bc[bc_type] += val

        self.weak_bcs[bc_id] = weak_bc

set_free_surface_boundary(params_fs, bc_id)

Sets the given boundary as a free surface.

This method calculates the normal stress at the free surface boundary. In the coupled approach, it also sets a zero-interior strong condition away from that boundary and populates the free_surface_map dictionary used to calculate the free-surface contribution to the momentum weak form.

Parameters:

Name	Type	Description	Default
params_fs	dict[str, int | bool]	Dictionary holding information about the free surface boundary	required
bc_id	int | str	Integer representing the index of the mesh boundary	required

Returns:

Type	Description
Expr	UFL expression for the normal stress at the free surface boundary

Source code in g-adopt/gadopt/stokes_integrators.py

def set_free_surface_boundary(
    self, params_fs: dict[str, int | bool], bc_id: int | str
) -> Expr:
    """Sets the given boundary as a free surface.

    This method calculates the normal stress at the free surface boundary. In the
    coupled approach, it also sets a zero-interior strong condition away from that
    boundary and populates the `free_surface_map` dictionary used to calculate the
    free-surface contribution to the momentum weak form.

    Args:
      params_fs:
        Dictionary holding information about the free surface boundary
      bc_id:
        Integer representing the index of the mesh boundary

    Returns:
      UFL expression for the normal stress at the free surface boundary
    """

set_equations() abstractmethod

Sets Equation instances for each equation in the system.

Equations must be ordered like solutions in the mixed space.

Source code in g-adopt/gadopt/stokes_integrators.py

@abc.abstractmethod
def set_equations(self):
    """Sets Equation instances for each equation in the system.

    Equations must be ordered like solutions in the mixed space.
    """

set_form()

Sets the weak form including linear and bilinear terms.

Source code in g-adopt/gadopt/stokes_integrators.py

def set_form(self) -> None:
    """Sets the weak form including linear and bilinear terms."""
    self.F = sum(
        eq.residual(sol) for eq, sol in zip(self.equations, self.solution_split)
    )

set_solver_options(solver_preset, solver_extras, gpu_extras, iterative_preset=iterative_outer_stokes_solver_parameters, direct_preset=direct_stokes_solver_parameters)

Sets PETSc solver options.

Source code in g-adopt/gadopt/stokes_integrators.py

def set_solver_options(
    self,
    solver_preset: ConfigType | str | None,
    solver_extras: ConfigType | None,
    gpu_extras: ConfigType | None,
    iterative_preset: ConfigType = iterative_outer_stokes_solver_parameters,
    direct_preset: ConfigType = direct_stokes_solver_parameters,
) -> None:
    """Sets PETSc solver options."""
    # Application context for the inverse mass matrix preconditioner
    self.appctx = {"mu": self.approximation.mu / self.rho_continuity}

    if isinstance(solver_preset, Mapping):
        self.add_to_solver_config(solver_preset)
        self.add_to_solver_config(solver_extras)
        self.register_update_callback(self.set_solver)
        return

    iterative_solve = self._use_iterative_solve(solver_preset)
    if iterative_solve:
        self.add_to_solver_config(iterative_preset)
    else:
        self.add_to_solver_config(direct_preset)

    if "snes_type" not in self.solver_parameters:
        if not depends_on(self.approximation.mu, self.solution):
            # If a set of solver parameters already has a "snes_type", do not
            # override it.
            self.add_to_solver_config({"snes_type": "ksponly"})
        else:
            self.add_to_solver_config(newton_stokes_solver_parameters)

    device_type = get_device_type(gpu_extras)
    if iterative_solve:
        self._configure_iterative_solver(device_type, gpu_extras)

    self._set_monitoring_config(device_type, gpu_extras, iterative_solve)

    self.add_to_solver_config(solver_extras)
    self.register_update_callback(self.set_solver)

set_solver()

Sets up the Firedrake variational problem and solver.

Source code in g-adopt/gadopt/stokes_integrators.py

def set_solver(self) -> None:
    """Sets up the Firedrake variational problem and solver."""
    if self.additional_forcing_term is not None:
        self.F += self.additional_forcing_term

    if self.constant_jacobian:
        warn(
            "Constant Jacobian specified for the Stokes system; please ensure that"
            "the viscosity does not vary in time."
        )

        trial = fd.TrialFunction(self.solution_space)
        F = fd.replace(self.F, {self.solution: trial})
        a, L = fd.lhs(F), fd.rhs(F)

        self.problem = fd.LinearVariationalProblem(
            a, L, self.solution, bcs=self.strong_bcs, constant_jacobian=True
        )
        self.solver = fd.LinearVariationalSolver(
            self.problem,
            solver_parameters=self.solver_parameters,
            nullspace=self.nullspace,
            transpose_nullspace=self.transpose_nullspace,
            near_nullspace=self.near_nullspace,
            appctx=self.appctx,
            options_prefix=self.name,
        )
    else:
        self.problem = fd.NonlinearVariationalProblem(
            self.F, self.solution, bcs=self.strong_bcs, J=self.J
        )
        self.solver = fd.NonlinearVariationalSolver(
            self.problem,
            solver_parameters=self.solver_parameters,
            nullspace=self.nullspace,
            transpose_nullspace=self.transpose_nullspace,
            near_nullspace=self.near_nullspace,
            appctx=self.appctx,
            options_prefix=self.name,
        )

solve()

Solves the system.

Source code in g-adopt/gadopt/stokes_integrators.py

def solve(self) -> None:
    """Solves the system."""
    self.solver.solve()
    self.solution_old.assign(self.solution)

StokesSolver(solution, approximation, T=0.0, /, **kwargs)

Bases: StokesSolverBase

Solver for the Stokes system.

Parameters:

Name	Type	Description	Default
solution	Function	Firedrake function representing the field over the mixed Stokes space	required
approximation	BaseApproximation	G-ADOPT approximation defining terms in the system of equations	required
T	Function | float	Firedrake function representing the temperature field	0.0
dt		Float quantifying the time step used in a coupled time integration	required
theta		Float quantifying the implicit contribution in a coupled time integration	required
additional_forcing_term		Firedrake form specifying an additional term contributing to the residual	required
bcs		Dictionary specifying boundary conditions (identifier, type, and value)	required
quad_degree		Integer denoting the quadrature degree	required
solver_parameters		Dictionary of PETSc solver options or string matching one of the default sets	required
solver_parameters_extra		Dictionary of PETSc solver options used to update the default G-ADOPT options	required
gpu_extra_parameters		Dictionary of additional GPU-specific settings that StokesSolver will translate into appropriate PETSc solver options	required
J		Firedrake function representing the Jacobian of the mixed Stokes system	required
constant_jacobian		Boolean specifying whether the Jacobian of the system is constant	required
nullspace		A MixedVectorSpaceBasis for the operator's kernel	required
transpose_nullspace		A MixedVectorSpaceBasis for the kernel of the operator's transpose	required
near_nullspace		A MixedVectorSpaceBasis for the operator's smallest eigenmodes (e.g. rigid body modes)	required

Source code in g-adopt/gadopt/stokes_integrators.py

def __init__(
    self,
    solution: fd.Function,
    approximation: BaseApproximation,
    T: fd.Function | float = 0.0,
    /,
    **kwargs,
) -> None:
    self.T = T

    self.eta_ind = 2
    self.free_surface_map = {}
    self.free_surface_equations = []
    super().__init__(solution, approximation, **kwargs)

force_on_boundary(subdomain_id, **kwargs)

Computes the force acting on a boundary.

Parameters:

Name	Type	Description	Default
subdomain_id	int | str	The subdomain ID of a physical boundary.	required

Returns:

Type	Description
Function	The force acting on the boundary.

Source code in g-adopt/gadopt/stokes_integrators.py

def force_on_boundary(self, subdomain_id: int | str, **kwargs) -> fd.Function:
    """Computes the force acting on a boundary.

    Arguments:
      subdomain_id: The subdomain ID of a physical boundary.

    Returns:
      The force acting on the boundary.

    """
    if not hasattr(self, "BoundaryNormalStressSolvers"):
        self.BoundaryNormalStressSolvers = {}

    if subdomain_id not in self.BoundaryNormalStressSolvers:
        self.BoundaryNormalStressSolvers[subdomain_id] = BoundaryNormalStressSolver(
            self, subdomain_id, **kwargs
        )

    return self.BoundaryNormalStressSolvers[subdomain_id].solve()

ViscoelasticStokesSolver(solution, approximation, stress_old, displacement, /, *, dt, **kwargs)

Bases: StokesSolverBase

Solves the Stokes system assuming an incompressible Maxwell viscoelastic rheology.

Parameters:

Name	Type	Description	Default
solution	Function	Firedrake function representing the field over the mixed Stokes space	required
approximation	BaseGIAApproximation	G-ADOPT approximation defining terms in the system of equations	required
stress_old	Function	Firedrake function representing the deviatoric stress at the previous time step	required
displacement	Function	Firedrake function representing the total displacement	required
dt	float	Float quantifying the time step used in a coupled time integration	required
theta		Float quantifying the implicit contribution in a coupled time integration	required
additional_forcing_term		Firedrake form specifying an additional term contributing to the residual	required
bcs		Dictionary specifying boundary conditions (identifier, type, and value)	required
quad_degree		Integer denoting the quadrature degree	required
solver_parameters		Dictionary of PETSc solver options or string matching one of the default sets	required
solver_parameters_extra		Dictionary of PETSc solver options used to update the default G-ADOPT options	required
J		Firedrake function representing the Jacobian of the mixed Stokes system	required
constant_jacobian		Boolean specifying whether the Jacobian of the system is constant	required
nullspace		A MixedVectorSpaceBasis for the operator's kernel	required
transpose_nullspace		A MixedVectorSpaceBasis for the kernel of the operator's transpose	required
near_nullspace		A MixedVectorSpaceBasis for the operator's smallest eigenmodes (e.g. rigid body modes)	required

Source code in g-adopt/gadopt/stokes_integrators.py

def __init__(
    self,
    solution: fd.Function,
    approximation: BaseGIAApproximation,
    stress_old: fd.Function,
    displacement: fd.Function,
    /,
    *,
    dt: float,
    **kwargs,
) -> None:

    warn(
        '''This solver is being phased out of G-ADOPT. We recommend using
    `InternalVariableSolver` for viscoelastic applications in G-ADOPT.
    ''')

    self.stress_old = stress_old  # Deviatoric stress from previous time step
    self.displacement = displacement  # Total displacement
    # Replace approximation's viscosity with effective viscosity
    approximation.mu = approximation.effective_viscosity(dt)
    super().__init__(solution, approximation, dt=dt, **kwargs)
    # Scaling factor for the previous stress
    self.stress_scale = self.approximation.prefactor_prestress(self.dt)

set_equations()

Sets up UFL forms for the viscoelastic Stokes equations residual.

Source code in g-adopt/gadopt/stokes_integrators.py

def set_equations(self) -> None:
    """Sets up UFL forms for the viscoelastic Stokes equations residual."""
    u, p = self.solution_split
    stress = self.approximation.stress(u, stress_old=self.stress_old, dt=self.dt)
    source = self.approximation.buoyancy(self.displacement) * self.k
    eqs_attrs = [
        {"p": p, "stress": stress, "source": source},
        {"u": u, "rho_continuity": self.rho_continuity},
    ]

    for i in range(len(stokes_terms)):
        self.equations.append(
            Equation(
                self.tests[i],
                self.solution_space[i],
                stokes_terms[i],
                eq_attrs=eqs_attrs[i],
                approximation=self.approximation,
                bcs=self.weak_bcs,
                quad_degree=self.quad_degree,
                # Scaling factor roughly size of mantle Maxwell time to make sure
                # that dimensional solve converges with strong bcs in parallel
                scaling_factor=1e-10,
            )
        )

InternalVariableSolver(solution, approximation, /, *, internal_variables, dt, **kwargs)

Bases: StokesSolverBase

Solver for internal variable viscoelastic formulation.

Parameters:

Name	Type	Description	Default
solution	Function	Firedrake function representing the displacement solution	required
approximation	BaseGIAApproximation	G-ADOPT approximation defining terms in the system of equations	required
internal_variables	Function | list[Function]	Internal variable functions accounting for viscoelastic stress relaxation. The number of internal variables provided must be consistent with the number of viscosity and shear modulus fields provided to the approximation.	required
dt	float	Float quantifying the time step used for time integration	required
additional_forcing_term		Firedrake form specifying an additional term contributing to the residual	required
bcs		Dictionary specifying boundary conditions (identifier, type, and value)	required
quad_degree		Integer denoting the quadrature degree	required
solver_parameters		Dictionary of PETSc solver options or string matching one of the default sets	required
solver_parameters_extra		Dictionary of PETSc solver options used to update the default G-ADOPT options	required
J		Firedrake function representing the Jacobian of the mixed Stokes system	required
constant_jacobian		Boolean specifying whether the Jacobian of the system is constant	required
nullspace		A MixedVectorSpaceBasis for the operator's kernel	required
transpose_nullspace		A MixedVectorSpaceBasis for the kernel of the operator's transpose	required
near_nullspace		A MixedVectorSpaceBasis for the operator's smallest eigenmodes (e.g. rigid body modes)	required

Source code in g-adopt/gadopt/stokes_integrators.py

def __init__(
    self,
    solution: fd.Function,
    approximation: BaseGIAApproximation,
    /,
    *,
    internal_variables: fd.Function | list[fd.Function],
    dt: float,
    **kwargs,
) -> None:

    # Convert `internal_variables` to a list in the case
    # for Maxwell Rheology where there is only one internal variable
    if not isinstance(internal_variables, list):
        internal_variables = [internal_variables]
    self.internal_variables = internal_variables

    # The substituted displacement operator has this effective shear
    # coefficient, which also supplies its penalty and preconditioner scale.
    approximation.mu = approximation.effective_viscosity(dt)

    super().__init__(solution, approximation, dt=dt, **kwargs)

update_m(m, maxwell_time)

Calculates updated internal variable using Backward Euler formula

Parameters:

Name	Type	Description	Default
m	Function	Firedrake function representing the current value of the internal variable	required
maxwell_time	Function | Expr	Firedrake function or UFL expression for Maxwell time associated with m terms in the system of equations	required

Returns:

Type	Description
Expr	UFL expression for the updated internal variable using Backward Euler

Source code in g-adopt/gadopt/stokes_integrators.py

def update_m(
    self, m: fd.Function, maxwell_time: fd.Function | Expr
) -> Expr:
    """Calculates updated internal variable using Backward Euler formula

    Args:
      m:
        Firedrake function representing the current value of the internal variable
      maxwell_time:
        Firedrake function or UFL expression for Maxwell time associated with m
        terms in the system of equations

    Returns:
        UFL expression for the updated internal variable using Backward Euler
    """
    m_new = (m + self.dt / maxwell_time * self.strain) / (1 + self.dt / maxwell_time)
    return m_new

CoupledInternalVariableSolver(solution, approximation, /, *, dt, scaling_factor=1, **kwargs)

Bases: StokesSolverBase

Solver for coupled internal variable viscoelastic formulation.

Solves the momentum equation and internal variable evolution equations simultaneously as a single coupled nonlinear variational problem.

The internal variable evolution equation is:

dm/dt = deviatoric_strain(u)/tau - m/tau

discretised using backward Euler:

(m_new - m_old)/dt + m_new/tau - deviatoric_strain(u_new)/tau = 0

Only backward Euler time-stepping is currently implemented. Both the strain rate and the internal variable m are evaluated at the new time level, which is consistent with a fully implicit (backward Euler) discretisation. A proper theta scheme would require theta-weighting of both the strain rate and the sink term. In the future, this should be unified with Irksome, either by using Irksome directly or by using Irksome's Dt operator together with expand_time_derivatives and UFL replace.

Parameters:

Name	Type	Description	Default
solution	Function	Firedrake function representing the field over the mixed space (displacement and internal variables)	required
approximation	BaseGIAApproximation	G-ADOPT approximation defining terms in the system of equations	required
dt	float	Float quantifying the time step used for time integration	required
scaling_factor	float	A constant factor used to rescale residual terms.	1
additional_forcing_term		Firedrake form specifying an additional term contributing to the residual	required
bcs		Dictionary specifying boundary conditions (identifier, type, and value)	required
quad_degree		Integer denoting the quadrature degree	required
solver_parameters		Dictionary of PETSc solver options or string matching one of the default sets	required
solver_parameters_extra		Dictionary of PETSc solver options used to update the default G-ADOPT options	required
J		Firedrake function representing the Jacobian of the mixed Stokes system	required
constant_jacobian		Boolean specifying whether the Jacobian of the system is constant	required
nullspace		A MixedVectorSpaceBasis for the operator's kernel	required
transpose_nullspace		A MixedVectorSpaceBasis for the kernel of the operator's transpose	required
near_nullspace		A MixedVectorSpaceBasis for the operator's smallest eigenmodes (e.g. rigid body modes)	required

Source code in g-adopt/gadopt/stokes_integrators.py

def __init__(
    self,
    solution: fd.Function,
    approximation: BaseGIAApproximation,
    /,
    *,
    dt: float,
    scaling_factor: float = 1,
    **kwargs,
) -> None:

    # At fixed internal variables, the displacement block is elastic. Use
    # `mu0` for its Nitsche penalty and preconditioner scale.
    approximation.mu = approximation.mu0
    self.scaling_factor = scaling_factor

    super().__init__(solution, approximation, dt=dt, theta=self._theta, **kwargs)

set_solver_options(solver_preset, solver_extras, gpu_extras, iterative_preset=coupled_gia_solver_parameters | newton_stokes_solver_parameters, direct_preset=direct_stokes_solver_parameters | newton_stokes_solver_parameters)

Sets PETSc solver options for the coupled GIA system.

Overrides the base class to ensure SNES Newton is always active. For Newtonian rheology (exponent=1) the power-law factor is identically 1 and the system is linear, so SNES converges in a single iteration. For power-law rheology (exponent > 1) full Newton iteration is performed.

When solver_preset is a Mapping it is honoured verbatim. The string preset "direct" uses direct_stokes_solver_parameters plus Newton SNES. The string preset "iterative" and the default None both use coupled_gia_solver_parameters, which already includes Newton SNES and the GIA-specific fieldsplit preconditioner. iterative_stokes_solver_parameters is intentionally not used here: its Schur-complement structure is designed for the standard Stokes system, not the larger coupled GIA block.

Source code in g-adopt/gadopt/stokes_integrators.py

def set_solver_options(
    self,
    solver_preset: ConfigType | str | None,
    solver_extras: ConfigType | None,
    gpu_extras: ConfigType | None,
    iterative_preset: ConfigType = coupled_gia_solver_parameters | newton_stokes_solver_parameters,
    direct_preset: ConfigType = direct_stokes_solver_parameters | newton_stokes_solver_parameters
) -> None:
    """Sets PETSc solver options for the coupled GIA system.

    Overrides the base class to ensure SNES Newton is always active.
    For Newtonian rheology (exponent=1) the power-law factor is identically
    1 and the system is linear, so SNES converges in a single iteration.
    For power-law rheology (exponent > 1) full Newton iteration is performed.

    When solver_preset is a Mapping it is honoured verbatim. The string
    preset "direct" uses direct_stokes_solver_parameters plus Newton SNES.
    The string preset "iterative" and the default None both use
    coupled_gia_solver_parameters, which already includes Newton SNES and
    the GIA-specific fieldsplit preconditioner. iterative_stokes_solver_parameters
    is intentionally not used here: its Schur-complement structure is designed
    for the standard Stokes system, not the larger coupled GIA block.
    """
    super().set_solver_options(
        solver_preset, solver_extras, gpu_extras, iterative_preset, direct_preset
    )

BoundaryNormalStressSolver(stokes_solver, subdomain_id, **kwargs)

Bases: SolverConfigurationMixin

A class for calculating surface forces acting on a boundary.

This solver computes topography on boundaries using the equation:

\[ h = sigma_(rr) / (g delta rho) \]

where \(sigma_(rr)\) is defined as:

\[ sigma_(rr) = [-p I + 2 mu (nabla u + nabla u^T)] . hat n . hat n \]

Instead of assuming a unit normal vector \(hat n\), this solver uses FacetNormal from Firedrake to accurately determine the normal vectors, which is particularly useful for complex meshes like icosahedron meshes in spherical simulations.

Parameters:

Name	Type	Description	Default
stokes_solver	StokesSolver	The Stokes solver providing the necessary fields for calculating stress.	required
subdomain_id	str | int	The subdomain ID of a physical boundary.	required
**kwargs		Optional keyword arguments. You can provide: - solver_parameters (dict[str, str | float]): Parameters to control the variational solver. If not provided, defaults are chosen based on whether the Stokes solver is direct or iterative. - solver_parameters_extra: Dictionary of PETSc solver options used to update the default solver options. If not provided, defaults to None (i.e. do not modify default settings)	{}

Source code in g-adopt/gadopt/stokes_integrators.py

def __init__(self, stokes_solver: StokesSolver, subdomain_id: int | str, **kwargs):
    # pressure and velocity together with viscosity are needed
    self.u, self.p, *self.eta = stokes_solver.solution.subfunctions

    # geometry
    self.mesh = stokes_solver.mesh
    self.dim = self.mesh.geometric_dimension

    # approximation tells us if we need to consider compressible formulation or not
    self.approximation = stokes_solver.approximation

    # Domain id that we want to use for boundary force
    self.subdomain_id = subdomain_id

    self._kwargs = kwargs

    solver_parameters = self._kwargs.get(
        "solver_parameters",
        BoundaryNormalStressSolver.iterative_solver_parameters
        if stokes_solver.is_iterative_solver()
        else BoundaryNormalStressSolver.direct_solve_parameters,
    )

    self.add_to_solver_config(solver_parameters)
    self.add_to_solver_config(self._kwargs.get("solver_parameters_extra"))
    self.register_update_callback(self.setup_solver)
    self.setup_solver()

solve()

Solves a linear system for the force and applies necessary boundary conditions.

Returns:

Type	Description
	The modified force after solving the linear system and applying boundary conditions.

Source code in g-adopt/gadopt/stokes_integrators.py

def solve(self):
    """
    Solves a linear system for the force and applies necessary boundary conditions.

    Returns:
        The modified force after solving the linear system and applying boundary conditions.
    """
    # Solve for the force
    self.solver.solve()

    # Take the average out
    vave = fd.assemble(self.force * self.ds) / fd.assemble(1 * self.ds)
    self.force.assign(self.force - vave)

    # Re-apply the zero condition everywhere except for the boundary
    self.interior_null_bc.apply(self.force)

    return self.force