JustRelax.jl

Appearance

Stokes ​

State containers, PT and DYREL damping coefficients, solve!, and the stress, principal-stress, and stress-rotation kernels shared by the standard and variational Stokes formulations.

JustRelax.DYREL Type
julia
struct DYREL{T, F}

Structure containing parameters and arrays for the DYREL (Dynamic Relaxation) solver.

Fields

  • γ_eff: Effective penalty parameter.

  • Dx, Dy, Dz: Diagonal preconditioners for velocity updates in x, y, (and z) directions.

  • λmaxVx, λmaxVy, λmaxVz: Maximum eigenvalues for stability calculation.

  • dVxdτ, dVydτ, dVzdτ: Pseudo-time step related damping terms.

  • dτVx, dτVy, dτVz: Pseudo-time steps for velocity fields.

  • dVx, dVy, dVz: Velocity increments for the current iteration.

  • βVx, βVy, βVz: Damping coefficients for momentum equation.

  • cVx, cVy, cVz: Damping coefficients related to dynamic relaxation.

  • αVx, αVy, αVz: Scaling factors for damping.

  • ηb: Bulk viscosity field.

  • P_num: Numerical pressure scratch field.

  • Rx0, Ry0, Rz0: Velocity residual history scratch fields.

  • CFL: Courant-Friedrichs-Lewy number.

  • ϵ: General convergence tolerance.

  • ϵ_vel: Velocity convergence tolerance.

  • c_fact: Damping scaling factor.

source
JustRelax.StressParticles Type
julia
StressParticles{backend, nNormal, nShear, T}

Particle-borne deviatoric stress and vorticity: the normal components τ_normal, the shear components τ_shear, and the vorticity components ω, each a tuple of particle cell arrays. Carrying the old stress on the particles instead of on the grid keeps it attached to the material as it advects and rotates.

Build one from the particles it follows with StressParticles(particles), advance it with rotate_stress!, and write it back onto stokes.τ_o with stress2grid!.

source
JustRelax.unwrap Method
julia
unwrap(x::StressParticles)

Flatten x into a single tuple (τ_normal..., τ_shear..., ω...) of its underlying particle cell arrays.

source
JustRelax.JustRelax2D.Displacement Method
julia
Displacement(nx::Integer, ny::Integer, nz::Integer)

Create the displacement arrays for the Stokes solver in 3D.

Fields

  • Ux: Displacement in x direction at their staggered location

  • Uy: Displacement in y direction at their staggered location

  • Uz: Displacement in z direction at their staggered location

source
JustRelax.JustRelax2D.Displacement Method
julia
Displacement(nx::Integer, ny::Integer)

Create the displacement arrays for the Stokes solver in 2D.

Fields

  • Ux: Displacement in x direction at their staggered location

  • Uy: Displacement in y direction at their staggered location

source
JustRelax.JustRelax2D.PrincipalStress Method
julia
PrincipalStress(ni::NTuple{N, Integer}) where {N}

Create the principal stress arrays for the Stokes solver in 2D or 3D with the extents given by ni (nx x ny or `nx x ny x nz``).

Fields

  • σ1: First principal stress

  • σ2: Second principal stress

  • σ3: Third principal stress (only in 3D). In 2D it is a placeholder array of size (2, 1, 1).

source
JustRelax.JustRelax2D.Residual Method
julia
Residual(nx::Integer, ny::Integer, nz::Integer)

Create the residual arrays for the Stokes solver in 3D.

Fields

  • Rx: Residual for the x-momentum equation

  • Ry: Residual for the y-momentum equation

  • Rz: Residual for the z-momentum equation

  • RP: Residual for the continuity equation

source
JustRelax.JustRelax2D.Residual Method
julia
Residual(nx::Integer, ny::Integer)

Create the residual arrays for the Stokes solver in 2D.

Fields

  • Rx: Residual for the x-momentum equation

  • Ry: Residual for the y-momentum equation

  • RP: Residual for the continuity equation

source
JustRelax.JustRelax2D.StokesArrays Method
julia
StokesArrays(ni::NTuple{N,Integer}) where {N}

Create the Stokes arrays object in 2D or 3D.

Fields

  • P: Pressure field

  • P0: Previous pressure field

  • ∇V: Velocity gradient

  • V: Velocity fields

  • Q: Volumetric source/sink term e.g. ΔV/V_tot [m³/m³]

  • U: Displacement fields

  • ω: Vorticity field

  • τ: Stress tensors

  • τ_o: Old stress tensors

  • ε: Strain rate tensors

  • ε_pl: Plastic strain rate tensors

  • EII_pl: Second invariant of the accumulated plastic strain

  • viscosity: Viscosity fields

  • R: Residual fields

  • Δε: Strain increment tensor

  • ∇U: Displacement gradient

  • λ : plastic multiplier @ centers

  • λv : plastic multiplier @ vertices

  • ΔPψ : pressure correction in dilatant case

source
JustRelax.JustRelax2D.SymmetricTensor Method
julia
SymmetricTensor(nx::Integer, ny::Integer, nz::Integer)

Create the symmetric tensor arrays for the Stokes solver in 3D.

Fields

  • xx: xx component of the tensor at cell centers

  • yy: yy component of the tensor at cell centers

  • zz: zz component of the tensor at cell centers

  • xx_v: xx component of the tensor at vertices

  • yy_v: yy component of the tensor at vertices

  • zz_v: zz component of the tensor at vertices

  • xy: xy component of the tensor at vertices

  • yz: yz component of the tensor at vertices

  • xz: xz component of the tensor at vertices

  • yz_c: yz component of the tensor at cell centers

  • xz_c: xz component of the tensor at cell centers

  • xy_c: xy component of the tensor at cell centers

  • II: second invariant of the tensor at cell centers

source
JustRelax.JustRelax2D.SymmetricTensor Method
julia
SymmetricTensor(nx::Integer, ny::Integer)

Create the symmetric tensor arrays for the Stokes solver in 2D.

Fields

  • xx: xx component of the tensor at cell centers

  • yy: yy component of the tensor at cell centers

  • xx_v: xx component of the tensor at vertices

  • yy_v: yy component of the tensor at vertices

  • xy: xy component of the tensor at vertices

  • xy_c: xy component of the tensor at cell centers

  • II: second invariant of the tensor at cell centers

source
JustRelax.JustRelax2D.Velocity Method
julia
Velocity(nx::Integer, ny::Integer, nz::Integer)

Create the velocity arrays for the Stokes solver in 3D.

Fields

  • Vx: Velocity in x direction (nx + 1, ny + 2, nz + 2)

  • Vy: Velocity in y direction (nx + 2, ny + 1, nz + 2)

  • Vz: Velocity in z direction (nx + 2, ny + 2, nz + 1)

source
JustRelax.JustRelax2D.Velocity Method
julia
Velocity(nx::Integer, ny::Integer)

Create the velocity arrays for the Stokes solver in 2D.

Fields

  • Vx: Velocity in x direction (nx + 1, ny + 2)

  • Vy: Velocity in y direction (nx + 2, ny + 1)

source
JustRelax.JustRelax2D.Viscosity Method
julia
Viscosity(ni::NTuple{N, Integer}) where {N}

Create the viscosity arrays for the Stokes solver in 2D or 3D with the extents given by ni (nx x ny or `nx x ny x nz``).

Fields

  • η: Viscosity at cell centers

  • ηv: Viscosity at vertices

  • η_vep: Viscosity for visco-elastic-plastic rheology

  • ητ: Pseudo-transient viscosity for stress update

source
JustRelax.JustRelax2D.Vorticity Method
julia
Vorticity(nx::Integer, ny::Integer, nz::Integer)

Create the vorticity arrays for the Stokes solver in 3D.

Fields

  • yz: Vorticity component yz at their staggered location

  • xz: Vorticity component xz at their staggered location

  • xy: Vorticity component xy at their staggered location

source
JustRelax.JustRelax2D.Vorticity Method
julia
Vorticity(nx::Integer, ny::Integer)

Create the vorticity arrays for the Stokes solver in 2D.

Fields

  • xy: Vorticity component xy at vertices
source
JustRelax.JustRelax2D.solve! Method
julia
solve!(stokes::StokesArrays, pt_stokes, grid, flow_bcs, ρg, phase_ratios, rheology, args, dt, igg; kwargs...)

Solve the 2D viscoelastoplastic Stokes equations to pseudo-transient convergence, updating stokes (velocity, pressure, stress, viscosity) in place for one physical time step dt. This is the general, multi-phase call form used by most models; dispatch on the type of the sixth argument also accepts, for simpler/benchmark setups:

  • a single-phase rheology::GeoParams.MaterialParams in place of phase_ratios/rheology (drop the phase_ratios argument), or

  • constant K (bulk modulus) or K, G (bulk and shear modulus) fields in place of phase_ratios/rheology/args, for linear (visco)elastic problems with no material rheology.

grid may also be replaced by the grid spacing di alone (a NTuple/NamedTuple).

Arguments

  • stokes: solver state allocated with StokesArrays.

  • pt_stokes: pseudo-transient coefficients, e.g. from PTStokesCoeffs.

  • grid: the model Geometry.

  • flow_bcs: velocity/displacement boundary conditions.

  • ρg: buoyancy forcing (ρgx, ρgy).

  • phase_ratios: per-cell/per-node phase fractions (a JustPIC.PhaseRatios).

  • rheology: one GeoParams.MaterialParams per phase.

  • args: auxiliary fields (e.g. temperature T, pressure P) used by the constitutive updates.

  • dt: physical time step.

  • igg: the distributed-grid context (IGG).

Keyword arguments

  • iterMax = 50e3, iterMin = 100: min/max pseudo-transient iterations.

  • nout = 500: check convergence every nout iterations.

  • viscosity_cutoff = (-Inf, Inf): clamp bounds for the effective viscosity.

  • viscosity_relaxation = 1e-2: relaxation factor for nonlinear viscosity updates.

  • λ_relaxation = 0.2: relaxation factor for the plastic multiplier.

  • strain_increment = false: accumulate strain increments instead of overwriting.

  • free_surface = false: enable free-surface stabilization.

  • b_width = (4, 4, 0): halo width used when overlapping communication and computation.

  • verbose = true: print convergence progress.

Dispatches on the CPU/CUDA/AMDGPU backend selected by stokes.

source
JustRelax.JustRelax2D.solve_VariationalStokes! Method
julia
solve_VariationalStokes!(backend::BackendTrait, stokes::JustRelax.StokesArrays, args...; kwargs)

Stokes solver entry point for variational Stokes solvers. This function dispatches to the appropriate implementation based on the backend provided in the function call.

source
JustRelax.JustRelax2D.solve_VariationalStokes! Method
julia
solve_VariationalStokes!(stokes::JustRelax.StokesArrays, args...; kwargs...)

Solve the 2D volume-fraction variational Stokes problem with matrix-free pseudo-transient iterations.

ϕ carries liquid weights at pressure cells, stress vertices, and staggered velocity faces. A pressure degree of freedom is retained only when its cell and all four surrounding velocity faces are connected to liquid:

julia
             Vy[i, j+1]
                   o
                   |
    Vx[i, j]  o--- p[i,j] ---o  Vx[i+1, j]
                   |
                   o
             Vy[i, j]

   inactive face => pressure row and disconnected velocity row eliminated

Zero-weight rows are written as zero instead of being solved with air material properties. Positive sliver fractions remain active; their velocity diagonal uses the bounded face mass max(ϕ_face, 0.1).

The free_surface keyword enables the density-gradient correction in the vertical momentum row. In this solver it is included implicitly in the local face diagonal, so the physical timestep does not create an explicit feedback instability.

source
JustRelax.JustRelax2D.solve_DYREL! Method
julia
solve_DYREL!(
    stokes, ρg, dyrel, flow_bcs, phase_ratios, rheology, args, grid, dt, igg;
    kwargs...,
)

Solve the Stokes system with the self-tuned dynamic relaxation (DYREL) method.

Arguments (in the following order)

  • stokes: JustRelax.StokesArrays containing the simulation fields.

  • ρg: buoyancy forces arrays.

  • dyrel: DYREL-specific parameters and fields.

  • flow_bcs: AbstractFlowBoundaryConditions defining velocity boundary conditions.

  • phase_ratios: JustPIC.PhaseRatios for material phase tracking.

  • rheology: Material properties and rheological laws.

  • args: Tuple of additional arguments needed to update viscosity, stress, and buoyancy forces.

  • grid: Geometry object carrying grid spacing and staggered-grid coordinates. A legacy 2D spacing tuple or named tuple is also accepted and converted to a uniform Geometry.

  • dt: Time step.

  • igg: IGG object for global grid information (MPI).

Keyword Arguments

  • viscosity_cutoff: Limits for viscosity (min, max). Default: (-Inf, Inf).

  • viscosity_relaxation: Relaxation factor for viscosity updates. Default: 1.0e-2.

  • λ_relaxation_DR: Relaxation factor for dynamic relaxation. Default: 1.

  • λ_relaxation_PH: Relaxation factor for Powell-Hestenes iterations. Default: 1.

  • iterMax: Maximum number of iterations for each dynamic-relaxation solve. Default: 50.0e3.

  • total_iterMax: Maximum number of total dynamic-relaxation iterations. Default: 50.0e3.

  • nout: Output frequency for residuals. Default: 100.

  • rel_drop: Relative residual drop tolerance. Default: 1.0e-2.

  • verbose_PH: Print Powell-Hestenes iteration info. Default: true.

  • verbose_DR: Print Dynamic Relaxation iteration info. Default: true.

  • linear_viscosity: Whether to use linear viscosity. Default: false.

  • free_surface: Include the density-gradient free-surface stabilization term. Default: false.

source
JustRelax.JustRelax2D.RockRatio Method
julia
RockRatio(nx, ny, nz)

Create a RockRatio object for a 3D grid with dimensions nx x ny x nz on a staggered grid.

source
JustRelax.JustRelax2D.RockRatio Method
julia
RockRatio(nx, ny)

Create a RockRatio object for a 2D grid with dimensions nx x ny on a staggered grid.

source
JustRelax.JustRelax2D._update_rock_ratio! Method
julia
_update_rock_ratio!(ϕ, ratio, air_phase)

Inner kernel of update_rock_ratio that clamps the computed rock ratio to the range [0, 1] for the given ratio and air_phase.

source
JustRelax.JustRelax2D.compute_air_ratio Method
julia
compute_air_ratio(phase_ratio, air_phase, inds...)

Compute the air ratio at the given indices based on the phase_ratio and air_phase.

source
JustRelax.JustRelax2D.compute_rock_ratio Method
julia
compute_rock_ratio(phase_ratio, air_phase, inds...)

Compute the rock ratio at the given indices based on the phase_ratio and air_phase.

source
JustRelax.JustRelax2D.isvalid_c Method
julia
isvalid_v(ϕ::JustRelax.RockRatio, inds...)

Check if ϕ.vertex[inds...] is a not a nullspace in 3D.

source
JustRelax.JustRelax2D.isvalid_c Method
julia
isvalid_c(ϕ::JustRelax.RockRatio, inds...)

Check whether the 2D pressure degree of freedom is connected to liquid.

The cell-centred pressure row is retained only when ϕ.center[i,j] and all four adjacent velocity faces are active:

julia
             Vy[i, j+1]
                   o
                   |
    Vx[i, j]  o--- p[i,j] ---o  Vx[i+1, j]
                   |
                   o
             Vy[i, j]

This is the local null-space elimination used by the matrix-free reduced system. It is intentionally stricter than testing the centre fraction alone.

Arguments

  • ϕ::JustRelax.RockRatio: The RockRatio object to check against.

  • inds: Cartesian indices to check.

source
JustRelax.JustRelax2D.isvalid_v Method
julia
isvalid_v(ϕ::JustRelax.RockRatio, inds...)

Check if ϕ.vertex[inds...] is a not a nullspace in 3D.

source
JustRelax.JustRelax2D.isvalid_v Method
julia
isvalid_v(ϕ::JustRelax.RockRatio, inds...)

Check if ϕ.vertex[inds...] is a not a nullspace 2D.

Arguments

  • ϕ::JustRelax.RockRatio: The RockRatio object to check against.

  • inds: Cartesian indices to check.

source
JustRelax.JustRelax2D.isvalid_velocity Method
julia
isvalid_velocity(ϕ::JustRelax.RockRatio, inds...)

Check if the velocity components at the given indices are not nullspaces in 3D.

source
JustRelax.JustRelax2D.isvalid_velocity Method
julia
isvalid_velocity(ϕ::JustRelax.RockRatio, inds...)

Check if the velocity components at the given indices are not nullspaces in 2D.

source
JustRelax.JustRelax2D.isvalid_vx Method
julia
isvalid_vx(ϕ::JustRelax.RockRatio, inds...)

Check if ϕ.Vx[inds...] is a not a nullspace.

Arguments

  • ϕ::JustRelax.RockRatio: The RockRatio object to check against.

  • inds: Cartesian indices to check.

source
JustRelax.JustRelax2D.isvalid_vz Method
julia
isvalid_vz(ϕ::JustRelax.RockRatio, inds...)

Check if ϕ.Vz[inds...] is a not a nullspace.

Arguments

  • ϕ::JustRelax.RockRatio: The RockRatio object to check against.

  • inds: Cartesian indices to check.

source
JustRelax.JustRelax2D.isvalid_xy Method
julia
isvalid_xy(ϕ, inds...)

Check if the xy shear component at the given indices is not a nullspace.

source
JustRelax.JustRelax2D.isvalid_xz Method
julia
isvalid_xz(ϕ, inds...)

Check if the xz shear component at the given indices is not a nullspace.

source
JustRelax.JustRelax2D.isvalid_yz Method
julia
isvalid_yz(ϕ, inds...)

Check if the yz shear component at the given indices is not a nullspace.

source
JustRelax.JustRelax2D.update_rock_ratio! Method
julia
update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase)

Update the rock ratio ϕ based on the provided phase_ratios and air_phase.

Arguments

  • ϕ::JustRelax.RockRatio: The rock ratio object to be updated.

  • phase_ratios: The ratios of different phases present.

  • air_phase: The phase representing air.

source
JustRelax.JustRelax2D.update_rock_ratio! Method
julia
update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase)

Update the rock ratio ϕ for a 3D grid based on the provided phase_ratios and air_phase.

Arguments

  • ϕ::JustRelax.RockRatio: The rock ratio object to be updated.

  • phase_ratios: The ratios of different phases present.

  • air_phase: The phase representing air.

source
JustRelax.JustRelax2D.update_rock_ratio_cv! Method
julia
update_rock_ratio_cv!(ϕ, ratio_center, ratio_vertex, air_phase)

Update the rock ratio for both center and vertex values based on the provided ratio_center, ratio_vertex, and air_phase.

source
JustRelax.JustRelax2D.accumulate_tensor! Method
julia
accumulate_tensor!(II, A::JustRelax.SymmetricTensor, dt)

Accumulate the second invariant of the symmetric tensor A over a time step: II[I] += dt * A_II[I].

The tensor components are gathered onto the cell centers where II lives, so A may hold its shear components on the vertices. Used to integrate the deviatoric plastic strain rate ε_pl into EII_pl; the volumetric counterpart is accumulate_vol!.

source
JustRelax.JustRelax2D.accumulate_vol! Method
julia
accumulate_vol!(EVol_pl, ε_vol_pl, dt)

Accumulate the volumetric plastic strain over a time step: EVol_pl[I] += dt * ε_vol_pl[I].

ε_vol_pl is the volumetric plastic strain rate at cell centers (= λ · (-dQ/dP), set inside the stress kernel), and is the volumetric counterpart of the deviatoric ε_pl. EVol_pl is the running invariant accumulated through time and is distinct from EII_pl (which integrates the second invariant of the deviatoric plastic strain rate via accumulate_tensor!).

source
JustRelax.JustRelax2D.tensor_invariant! Method
julia
tensor_invariant!(A::JustRelax.SymmetricTensor)

Compute the tensor invariant of the given symmetric tensor A.

Arguments

  • A::JustRelax.SymmetricTensor: The input symmetric tensor.
source
JustRelax.JustRelax2D.compute_principal_stresses! Method
julia
compute_principal_stresses!(stokes, σ::PrincipalStress)

In-place version of compute_principal_stresses, writing into a pre-allocated σ.

source
JustRelax.JustRelax2D.compute_principal_stresses Method
julia
compute_principal_stresses(backend, stokes::StokesArrays)

Compute the principal deviatoric stresses (eigenvalues and eigenvectors of the stress tensor) at cell centers from stokes, returning a new PrincipalStress.

source
JustRelax.JustRelax2D.compute_P! Method

compute_P!(P, P0, RP, ∇V, Q, ΔT, η, rheology::NTuple{N,MaterialParams}, phase_ratio::C, dt, r, θ_dτ)

Compute the pressure field P and the residual RP for the compressible case. This function introduces thermal stresses after the implementation of Kiss et al. (2023).

Arguments

  • P: pressure field

  • RP: residual field

  • ∇V: divergence of the velocity field

  • Q: volumetric source/sink term which should have the properties of dV/V_tot [m³/m³] normalized per cell, default is zero.

  • ΔT: temperature difference on the cell center, to account for thermal stresses. The thermal expansivity α is computed from the material parameters.

  • η: viscosity field

  • rheology: material parameters

  • phase_ratio: phase field

  • melt_fraction: melt fraction field, used for the thermal expansion coefficient if provided

  • dt: time step

  • r: relaxation parameter for the pressure update

  • θ_dτ: numerical parameter for the pressure update

source
JustRelax.JustRelax2D.StressParticles Method
julia
StressParticles(particles::Particles)

Allocate the stress and vorticity cell arrays that follow particles, on the same backend and with the same per-cell capacity. Two normal and one shear component in 2-D, three of each in 3-D.

source
JustRelax.JustRelax2D.rotate_stress! Method
julia
rotate_stress!(τ_particles::StressParticles, stokes, particles, dt)

Interpolate the current deviatoric stress stokes.τ and vorticity stokes.ω onto the particles and rotate the particle stress over dt. stokes.ω must hold the vorticity of the current velocity field. Use stress2grid! afterwards to map the rotated stress back onto stokes.τ_o.

source
JustRelax.JustRelax2D.rotate_stress_particles! Method
julia
rotate_stress_particles!(τ::NTuple, ω::NTuple, particles::Particles, dt; method = :matrix)

Rotate the deviatoric stress carried by each active particle over dt with the local vorticity, using GeoParams' elastic stress rotation. τ holds the stress components and ω the vorticity components, as particle cell arrays. method is accepted for call-site compatibility and does not select an algorithm.

source
JustRelax.JustRelax2D.stress2grid! Method
julia
stress2grid!(stokes, τ_particles::StressParticles, particles)

Interpolate the particle stress in τ_particles back onto the old-stress fields stokes.τ_o, normal components onto the cell centers and shear components onto the vertices. Counterpart of rotate_stress!, and the step that hands the rotated stress to the next Stokes solve.

source
JustRelax.JustRelax2D.displacement2velocity! Method
julia
displacement2velocity!(stokes::StokesArrays, dt)
displacement2velocity!(stokes::StokesArrays, dt, flow_bcs::AbstractFlowBoundaryConditions)

Set stokes.V (velocity) to stokes.U / dt (displacement divided by the time step), in place — the inverse of velocity2displacement!. The 3-argument method is a no-op when flow_bcs isa VelocityBoundaryConditions (velocity is already the primary variable).

source
JustRelax.JustRelax2D.velocity2displacement! Method
julia
velocity2displacement!(stokes::StokesArrays, dt)

Set stokes.U (displacement) to stokes.V * dt (velocity times the time step), in place.

source
JustRelax.JustRelax3D.Displacement Method
julia
Displacement(nx::Integer, ny::Integer, nz::Integer)

Create the displacement arrays for the Stokes solver in 3D.

Fields

  • Ux: Displacement in x direction at their staggered location

  • Uy: Displacement in y direction at their staggered location

  • Uz: Displacement in z direction at their staggered location

source
JustRelax.JustRelax3D.Displacement Method
julia
Displacement(nx::Integer, ny::Integer)

Create the displacement arrays for the Stokes solver in 2D.

Fields

  • Ux: Displacement in x direction at their staggered location

  • Uy: Displacement in y direction at their staggered location

source
JustRelax.JustRelax3D.PrincipalStress Method
julia
PrincipalStress(ni::NTuple{N, Integer}) where {N}

Create the principal stress arrays for the Stokes solver in 2D or 3D with the extents given by ni (nx x ny or `nx x ny x nz``).

Fields

  • σ1: First principal stress

  • σ2: Second principal stress

  • σ3: Third principal stress (only in 3D). In 2D it is a placeholder array of size (2, 1, 1).

source
JustRelax.JustRelax3D.Residual Method
julia
Residual(nx::Integer, ny::Integer, nz::Integer)

Create the residual arrays for the Stokes solver in 3D.

Fields

  • Rx: Residual for the x-momentum equation

  • Ry: Residual for the y-momentum equation

  • Rz: Residual for the z-momentum equation

  • RP: Residual for the continuity equation

source
JustRelax.JustRelax3D.Residual Method
julia
Residual(nx::Integer, ny::Integer)

Create the residual arrays for the Stokes solver in 2D.

Fields

  • Rx: Residual for the x-momentum equation

  • Ry: Residual for the y-momentum equation

  • RP: Residual for the continuity equation

source
JustRelax.JustRelax3D.StokesArrays Method
julia
StokesArrays(ni::NTuple{N,Integer}) where {N}

Create the Stokes arrays object in 2D or 3D.

Fields

  • P: Pressure field

  • P0: Previous pressure field

  • ∇V: Velocity gradient

  • V: Velocity fields

  • Q: Volumetric source/sink term e.g. ΔV/V_tot [m³/m³]

  • U: Displacement fields

  • ω: Vorticity field

  • τ: Stress tensors

  • τ_o: Old stress tensors

  • ε: Strain rate tensors

  • ε_pl: Plastic strain rate tensors

  • EII_pl: Second invariant of the accumulated plastic strain

  • viscosity: Viscosity fields

  • R: Residual fields

  • Δε: Strain increment tensor

  • ∇U: Displacement gradient

  • λ : plastic multiplier @ centers

  • λv : plastic multiplier @ vertices

  • ΔPψ : pressure correction in dilatant case

source
JustRelax.JustRelax3D.SymmetricTensor Method
julia
SymmetricTensor(nx::Integer, ny::Integer, nz::Integer)

Create the symmetric tensor arrays for the Stokes solver in 3D.

Fields

  • xx: xx component of the tensor at cell centers

  • yy: yy component of the tensor at cell centers

  • zz: zz component of the tensor at cell centers

  • xx_v: xx component of the tensor at vertices

  • yy_v: yy component of the tensor at vertices

  • zz_v: zz component of the tensor at vertices

  • xy: xy component of the tensor at vertices

  • yz: yz component of the tensor at vertices

  • xz: xz component of the tensor at vertices

  • yz_c: yz component of the tensor at cell centers

  • xz_c: xz component of the tensor at cell centers

  • xy_c: xy component of the tensor at cell centers

  • II: second invariant of the tensor at cell centers

source
JustRelax.JustRelax3D.SymmetricTensor Method
julia
SymmetricTensor(nx::Integer, ny::Integer)

Create the symmetric tensor arrays for the Stokes solver in 2D.

Fields

  • xx: xx component of the tensor at cell centers

  • yy: yy component of the tensor at cell centers

  • xx_v: xx component of the tensor at vertices

  • yy_v: yy component of the tensor at vertices

  • xy: xy component of the tensor at vertices

  • xy_c: xy component of the tensor at cell centers

  • II: second invariant of the tensor at cell centers

source
JustRelax.JustRelax3D.Velocity Method
julia
Velocity(nx::Integer, ny::Integer, nz::Integer)

Create the velocity arrays for the Stokes solver in 3D.

Fields

  • Vx: Velocity in x direction (nx + 1, ny + 2, nz + 2)

  • Vy: Velocity in y direction (nx + 2, ny + 1, nz + 2)

  • Vz: Velocity in z direction (nx + 2, ny + 2, nz + 1)

source
JustRelax.JustRelax3D.Velocity Method
julia
Velocity(nx::Integer, ny::Integer)

Create the velocity arrays for the Stokes solver in 2D.

Fields

  • Vx: Velocity in x direction (nx + 1, ny + 2)

  • Vy: Velocity in y direction (nx + 2, ny + 1)

source
JustRelax.JustRelax3D.Viscosity Method
julia
Viscosity(ni::NTuple{N, Integer}) where {N}

Create the viscosity arrays for the Stokes solver in 2D or 3D with the extents given by ni (nx x ny or `nx x ny x nz``).

Fields

  • η: Viscosity at cell centers

  • ηv: Viscosity at vertices

  • η_vep: Viscosity for visco-elastic-plastic rheology

  • ητ: Pseudo-transient viscosity for stress update

source
JustRelax.JustRelax3D.Vorticity Method
julia
Vorticity(nx::Integer, ny::Integer, nz::Integer)

Create the vorticity arrays for the Stokes solver in 3D.

Fields

  • yz: Vorticity component yz at their staggered location

  • xz: Vorticity component xz at their staggered location

  • xy: Vorticity component xy at their staggered location

source
JustRelax.JustRelax3D.Vorticity Method
julia
Vorticity(nx::Integer, ny::Integer)

Create the vorticity arrays for the Stokes solver in 2D.

Fields

  • xy: Vorticity component xy at vertices
source
JustRelax.JustRelax3D.solve! Method
julia
solve!(stokes::StokesArrays, pt_stokes, grid, flow_bcs, ρg, phase_ratios, rheology, args, dt, igg; kwargs...)

Solve the 3D viscoelastoplastic Stokes equations to pseudo-transient convergence, updating stokes (velocity, pressure, stress, viscosity) in place for one physical time step dt. This is the general, multi-phase call form used by most models; dispatch on the type of the sixth argument also accepts, for simpler/benchmark setups:

  • a single-phase rheology::GeoParams.MaterialParams in place of phase_ratios/rheology (drop the phase_ratios argument), or

  • constant K, G (bulk and shear modulus) fields in place of phase_ratios/rheology/args, for linear viscoelastic problems with no material rheology.

grid may also be replaced by the grid spacing di alone (a NTuple/NamedTuple).

Arguments

  • stokes: solver state allocated with StokesArrays.

  • pt_stokes: pseudo-transient coefficients, e.g. from PTStokesCoeffs.

  • grid: the model Geometry.

  • flow_bcs: velocity/displacement boundary conditions.

  • ρg: buoyancy forcing (ρgx, ρgy, ρgz).

  • phase_ratios: per-cell/per-node phase fractions (a JustPIC.PhaseRatios).

  • rheology: one GeoParams.MaterialParams per phase.

  • args: auxiliary fields (e.g. temperature T, pressure P) used by the constitutive updates.

  • dt: physical time step.

  • igg: the distributed-grid context (IGG).

Keyword arguments

  • iterMax = 10e3: maximum pseudo-transient iterations.

  • nout = 500: check convergence every nout iterations.

  • viscosity_relaxation = 1e-2: relaxation factor for nonlinear viscosity updates.

  • viscosity_cutoff = (-Inf, Inf): clamp bounds for the effective viscosity.

  • λ_relaxation = 0.2: relaxation factor for the plastic multiplier.

  • b_width = (4, 4, 4): halo width used when overlapping communication and computation.

  • verbose = true: print convergence progress.

Dispatches on the CPU/CUDA/AMDGPU backend selected by stokes.

source
JustRelax.JustRelax3D.solve_VariationalStokes! Method
julia
solve_VariationalStokes!(stokes::StokesArrays, pt_stokes, grid, flow_bcs, ρg, phase_ratios, ϕ::RockRatio, rheology, args, dt, igg; kwargs...)

Solve the 3D viscoelastoplastic Stokes equations to pseudo-transient convergence using the variational (ghost-node-free) formulation, updating stokes in place for one physical time step dt. Cells are weighted by the rock ratio ϕ (see RockRatio), so partially- or fully-air/sticky-air cells contribute less (or not at all) to the momentum balance. grid may be replaced by the grid spacing di alone.

rheology is one GeoParams.MaterialParams per phase; args carries auxiliary fields (e.g. temperature T, pressure P). Keyword tolerances, iteration limits, and relaxation factors default as in _solve_VS!. Dispatches on the CPU/CUDA/AMDGPU backend selected by stokes.

source
JustRelax.JustRelax3D.RockRatio Method
julia
RockRatio(nx, ny, nz)

Create a RockRatio object for a 3D grid with dimensions nx x ny x nz on a staggered grid.

source
JustRelax.JustRelax3D.RockRatio Method
julia
RockRatio(nx, ny)

Create a RockRatio object for a 2D grid with dimensions nx x ny on a staggered grid.

source
JustRelax.JustRelax3D._update_rock_ratio! Method
julia
_update_rock_ratio!(ϕ, ratio, air_phase)

Inner kernel of update_rock_ratio that clamps the computed rock ratio to the range [0, 1] for the given ratio and air_phase.

source
JustRelax.JustRelax3D.compute_air_ratio Method
julia
compute_air_ratio(phase_ratio, air_phase, inds...)

Compute the air ratio at the given indices based on the phase_ratio and air_phase.

source
JustRelax.JustRelax3D.compute_rock_ratio Method
julia
compute_rock_ratio(phase_ratio, air_phase, inds...)

Compute the rock ratio at the given indices based on the phase_ratio and air_phase.

source
JustRelax.JustRelax3D.isvalid_c Method
julia
isvalid_v(ϕ::JustRelax.RockRatio, inds...)

Check if ϕ.vertex[inds...] is a not a nullspace in 3D.

source
JustRelax.JustRelax3D.isvalid_c Method
julia
isvalid_c(ϕ::JustRelax.RockRatio, inds...)

Check whether the 2D pressure degree of freedom is connected to liquid.

The cell-centred pressure row is retained only when ϕ.center[i,j] and all four adjacent velocity faces are active:

julia
             Vy[i, j+1]
                   o
                   |
    Vx[i, j]  o--- p[i,j] ---o  Vx[i+1, j]
                   |
                   o
             Vy[i, j]

This is the local null-space elimination used by the matrix-free reduced system. It is intentionally stricter than testing the centre fraction alone.

Arguments

  • ϕ::JustRelax.RockRatio: The RockRatio object to check against.

  • inds: Cartesian indices to check.

source
JustRelax.JustRelax3D.isvalid_v Method
julia
isvalid_v(ϕ::JustRelax.RockRatio, inds...)

Check if ϕ.vertex[inds...] is a not a nullspace in 3D.

source
JustRelax.JustRelax3D.isvalid_v Method
julia
isvalid_v(ϕ::JustRelax.RockRatio, inds...)

Check if ϕ.vertex[inds...] is a not a nullspace 2D.

Arguments

  • ϕ::JustRelax.RockRatio: The RockRatio object to check against.

  • inds: Cartesian indices to check.

source
JustRelax.JustRelax3D.isvalid_velocity Method
julia
isvalid_velocity(ϕ::JustRelax.RockRatio, inds...)

Check if the velocity components at the given indices are not nullspaces in 3D.

source
JustRelax.JustRelax3D.isvalid_velocity Method
julia
isvalid_velocity(ϕ::JustRelax.RockRatio, inds...)

Check if the velocity components at the given indices are not nullspaces in 2D.

source
JustRelax.JustRelax3D.isvalid_vx Method
julia
isvalid_vx(ϕ::JustRelax.RockRatio, inds...)

Check if ϕ.Vx[inds...] is a not a nullspace.

Arguments

  • ϕ::JustRelax.RockRatio: The RockRatio object to check against.

  • inds: Cartesian indices to check.

source
JustRelax.JustRelax3D.isvalid_vz Method
julia
isvalid_vz(ϕ::JustRelax.RockRatio, inds...)

Check if ϕ.Vz[inds...] is a not a nullspace.

Arguments

  • ϕ::JustRelax.RockRatio: The RockRatio object to check against.

  • inds: Cartesian indices to check.

source
JustRelax.JustRelax3D.isvalid_xy Method
julia
isvalid_xy(ϕ, inds...)

Check if the xy shear component at the given indices is not a nullspace.

source
JustRelax.JustRelax3D.isvalid_xz Method
julia
isvalid_xz(ϕ, inds...)

Check if the xz shear component at the given indices is not a nullspace.

source
JustRelax.JustRelax3D.isvalid_yz Method
julia
isvalid_yz(ϕ, inds...)

Check if the yz shear component at the given indices is not a nullspace.

source
JustRelax.JustRelax3D.update_rock_ratio! Method
julia
update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase)

Update the rock ratio ϕ based on the provided phase_ratios and air_phase.

Arguments

  • ϕ::JustRelax.RockRatio: The rock ratio object to be updated.

  • phase_ratios: The ratios of different phases present.

  • air_phase: The phase representing air.

source
JustRelax.JustRelax3D.update_rock_ratio! Method
julia
update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase)

Update the rock ratio ϕ for a 3D grid based on the provided phase_ratios and air_phase.

Arguments

  • ϕ::JustRelax.RockRatio: The rock ratio object to be updated.

  • phase_ratios: The ratios of different phases present.

  • air_phase: The phase representing air.

source
JustRelax.JustRelax3D.update_rock_ratio_cv! Method
julia
update_rock_ratio_cv!(ϕ, ratio_center, ratio_vertex, air_phase)

Update the rock ratio for both center and vertex values based on the provided ratio_center, ratio_vertex, and air_phase.

source
JustRelax.JustRelax3D.accumulate_tensor! Method
julia
accumulate_tensor!(II, A::JustRelax.SymmetricTensor, dt)

Accumulate the second invariant of the symmetric tensor A over a time step: II[I] += dt * A_II[I].

The tensor components are gathered onto the cell centers where II lives, so A may hold its shear components on the vertices. Used to integrate the deviatoric plastic strain rate ε_pl into EII_pl; the volumetric counterpart is accumulate_vol!.

source
JustRelax.JustRelax3D.accumulate_vol! Method
julia
accumulate_vol!(EVol_pl, ε_vol_pl, dt)

Accumulate the volumetric plastic strain over a time step: EVol_pl[I] += dt * ε_vol_pl[I].

ε_vol_pl is the volumetric plastic strain rate at cell centers (= λ · (-dQ/dP), set inside the stress kernel), and is the volumetric counterpart of the deviatoric ε_pl. EVol_pl is the running invariant accumulated through time and is distinct from EII_pl (which integrates the second invariant of the deviatoric plastic strain rate via accumulate_tensor!).

source
JustRelax.JustRelax3D.tensor_invariant! Method
julia
tensor_invariant!(A::JustRelax.SymmetricTensor)

Compute the tensor invariant of the given symmetric tensor A.

Arguments

  • A::JustRelax.SymmetricTensor: The input symmetric tensor.
source
JustRelax.JustRelax3D.compute_principal_stresses! Method
julia
compute_principal_stresses!(stokes, σ::PrincipalStress)

In-place version of compute_principal_stresses, writing into a pre-allocated σ.

source
JustRelax.JustRelax3D.compute_principal_stresses Method
julia
compute_principal_stresses(backend, stokes::StokesArrays)

Compute the principal deviatoric stresses (eigenvalues and eigenvectors of the stress tensor) at cell centers from stokes, returning a new PrincipalStress.

source
JustRelax.JustRelax3D.compute_P! Method

compute_P!(P, P0, RP, ∇V, Q, ΔT, η, rheology::NTuple{N,MaterialParams}, phase_ratio::C, dt, r, θ_dτ)

Compute the pressure field P and the residual RP for the compressible case. This function introduces thermal stresses after the implementation of Kiss et al. (2023).

Arguments

  • P: pressure field

  • RP: residual field

  • ∇V: divergence of the velocity field

  • Q: volumetric source/sink term which should have the properties of dV/V_tot [m³/m³] normalized per cell, default is zero.

  • ΔT: temperature difference on the cell center, to account for thermal stresses. The thermal expansivity α is computed from the material parameters.

  • η: viscosity field

  • rheology: material parameters

  • phase_ratio: phase field

  • melt_fraction: melt fraction field, used for the thermal expansion coefficient if provided

  • dt: time step

  • r: relaxation parameter for the pressure update

  • θ_dτ: numerical parameter for the pressure update

source
JustRelax.JustRelax3D.StressParticles Method
julia
StressParticles(particles::Particles)

Allocate the stress and vorticity cell arrays that follow particles, on the same backend and with the same per-cell capacity. Two normal and one shear component in 2-D, three of each in 3-D.

source
JustRelax.JustRelax3D.rotate_stress! Method
julia
rotate_stress!(τ_particles::StressParticles, stokes, particles, dt)

Interpolate the current deviatoric stress stokes.τ and vorticity stokes.ω onto the particles and rotate the particle stress over dt. stokes.ω must hold the vorticity of the current velocity field. Use stress2grid! afterwards to map the rotated stress back onto stokes.τ_o.

source
JustRelax.JustRelax3D.rotate_stress_particles! Method
julia
rotate_stress_particles!(τ::NTuple, ω::NTuple, particles::Particles, dt; method = :matrix)

Rotate the deviatoric stress carried by each active particle over dt with the local vorticity, using GeoParams' elastic stress rotation. τ holds the stress components and ω the vorticity components, as particle cell arrays. method is accepted for call-site compatibility and does not select an algorithm.

source
JustRelax.JustRelax3D.stress2grid! Method
julia
stress2grid!(stokes, τ_particles::StressParticles, particles)

Interpolate the particle stress in τ_particles back onto the old-stress fields stokes.τ_o, normal components onto the cell centers and shear components onto the vertices. Counterpart of rotate_stress!, and the step that hands the rotated stress to the next Stokes solve.

source
JustRelax.JustRelax3D.displacement2velocity! Method
julia
displacement2velocity!(stokes::StokesArrays, dt)
displacement2velocity!(stokes::StokesArrays, dt, flow_bcs::AbstractFlowBoundaryConditions)

Set stokes.V (velocity) to stokes.U / dt (displacement divided by the time step), in place — the inverse of velocity2displacement!. The 3-argument method is a no-op when flow_bcs isa VelocityBoundaryConditions (velocity is already the primary variable).

source
JustRelax.JustRelax3D.velocity2displacement! Method
julia
velocity2displacement!(stokes::StokesArrays, dt)

Set stokes.U (displacement) to stokes.V * dt (velocity times the time step), in place.

source