Tests for Distributed architecture with AnelasticDynamics

examples/stationary_parcel_model.jl

@@ -27,8 +27,7 @@ using CairoMakie
 # the parcel moves with the hydrometeors. We use an `ImpenetrableBoundaryCondition()`
 # to ensure moisture is conserved within the parcel.
 
-grid = RectilinearGrid(CPU(); size=(1, 1, 1), x=(0, 1), y=(0, 1), z=(0, 1),
-                       topology=(Periodic, Periodic, Bounded))
+grid = RectilinearGrid(CPU(); size=1, z=(0, 1), topology=(Flat, Flat, Bounded))
 
 constants = ThermodynamicConstants()
 reference_state = ReferenceState(grid, constants; surface_pressure=101325, potential_temperature=300)

src/AnelasticEquations/AnelasticEquations.jl

@@ -22,11 +22,12 @@ using Oceananigans: Oceananigans, CenterField, XFaceField, YFaceField, ZFaceFiel
 using Oceananigans.Architectures: architecture
 using Oceananigans.BoundaryConditions: FieldBoundaryConditions, regularize_field_boundary_conditions, fill_halo_regions!
 using Oceananigans.Fields: set!
-using Oceananigans.Grids: ZDirection, inactive_cell
+using Oceananigans.Grids: topology, ZDirection, inactive_cell, FullyConnected
 using Oceananigans.ImmersedBoundaries: mask_immersed_field!
 using Oceananigans.Models.NonhydrostaticModels: NonhydrostaticModels
 using Oceananigans.Operators: Δzᵃᵃᶜ, Δzᵃᵃᶠ, divᶜᶜᶜ, Δzᶜᶜᶜ, ℑzᵃᵃᶠ, ∂xᶠᶜᶜ, ∂yᶜᶠᶜ, ∂zᶜᶜᶠ
 using Oceananigans.Solvers: Solvers, solve!, FourierTridiagonalPoissonSolver, AbstractHomogeneousNeumannFormulation
+using Oceananigans.DistributedComputations: DistributedFourierTridiagonalPoissonSolver, reconstruct_global_grid
 using Oceananigans.Utils: prettysummary, launch!
 
 using Breeze.Thermodynamics: ReferenceState, MoistureMassFractions, mixture_gas_constant

src/AnelasticEquations/anelastic_pressure_solver.jl

@@ -16,7 +16,12 @@ function AtmosphereModels.dynamics_pressure_solver(dynamics::AnelasticDynamics,
         # With this method, we are using an approximate solver that
         # will produce a divergent velocity field near terrain.
         FourierTridiagonalPoissonSolver(grid.underlying_grid; tridiagonal_formulation)
-    else # the solver is exact
+    elseif any(T -> T === FullyConnected, topology(grid))
+        # Multi-rank distributed grids have FullyConnected topology in split dimensions.
+        # Use the distributed solver with MPI-aware FFT transposes.
+        global_grid = reconstruct_global_grid(grid)
+        DistributedFourierTridiagonalPoissonSolver(global_grid, grid; tridiagonal_formulation)
+    else # the solver is exact (also works for single-rank Distributed)
         FourierTridiagonalPoissonSolver(grid; tridiagonal_formulation)
     end
 
@@ -95,6 +100,14 @@ function compute_anelastic_source_term!(solver::FourierTridiagonalPoissonSolver,
     return nothing
 end
 
+function compute_anelastic_source_term!(solver::DistributedFourierTridiagonalPoissonSolver, ρŨ, Δt)
+    rhs = solver.storage.zfield
+    grid = solver.local_grid
+    arch = architecture(grid)
+    launch!(arch, grid, :xyz, _compute_anelastic_source_term!, rhs, grid, ρŨ, Δt)
+    return nothing
+end
+
 @kernel function _compute_anelastic_source_term!(rhs, grid, ρŨ, Δt)
     i, j, k = @index(Global, NTuple)
     active = !inactive_cell(i, j, k, grid)

src/AtmosphereModels/AtmosphereModels.jl

@@ -65,6 +65,7 @@ using KernelAbstractions: @kernel, @index
 
 using Oceananigans: Oceananigans, CenterField, fields
 using Oceananigans.BoundaryConditions: FieldBoundaryConditions, regularize_field_boundary_conditions, fill_halo_regions!
+using Oceananigans.Fields: Field
 using Oceananigans.ImmersedBoundaries: mask_immersed_field!
 using Oceananigans.Operators: Δzᶜᶜᶜ, ℑzᵃᵃᶠ
 using Oceananigans.Solvers: Solvers
@@ -93,6 +94,12 @@ include("atmosphere_model.jl")
 include("atmosphere_model_buoyancy.jl")
 include("radiation_interface.jl")
 include("dynamics_kernel_functions.jl")
+
+# Column fields (Field{Nothing, Nothing, <:Any}) have singleton x/y dimensions
+# and cannot participate in MPI halo communication on distributed grids.
+_is_column_field(::Field{Nothing, Nothing}) = true
+_is_column_field(::Any) = false
+
 include("update_atmosphere_model_state.jl")
 include("compute_hydrostatic_pressure.jl")
 

src/AtmosphereModels/Diagnostics/dewpoint_temperature.jl

@@ -50,15 +50,15 @@ the secant iteration convergence.
 ```jldoctest dewpoint
 using Breeze
 
-grid = RectilinearGrid(size=(1, 1, 8), extent=(1, 1, 1e3))
+grid = RectilinearGrid(size=8, z=(0, 1e3), topology=(Flat, Flat, Bounded))
 model = AtmosphereModel(grid; microphysics=SaturationAdjustment())
 set!(model, θ=300, qᵗ=0.01)
 
 T⁺ = DewpointTemperature(model)
 
 # output
 KernelFunctionOperation at (Center, Center, Center)
-├── grid: 1×1×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×3 halo
+├── grid: 1×1×8 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── kernel_function: DewpointTemperatureKernelFunction
 └── arguments: ()
 ```
@@ -70,13 +70,13 @@ T⁺_field = Field(T⁺)
 
 # output
 1×1×8 Field{Center, Center, Center} on RectilinearGrid on CPU
-├── grid: 1×1×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×3 halo
+├── grid: 1×1×8 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── boundary conditions: FieldBoundaryConditions
-│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
+│   └── west: Nothing, east: Nothing, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
 ├── operand: KernelFunctionOperation at (Center, Center, Center)
 ├── status: time=0.0
-└── data: 3×3×14 OffsetArray(::Array{Float64, 3}, 0:2, 0:2, -2:11) with eltype Float64 with indices 0:2×0:2×-2:11
-    └── max=289.062, min=287.475, mean=288.27
+└── data: 1×1×14 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:11) with eltype Float64 with indices 1:1×1:1×-2:11
+    └── max=287.245, min=285.625, mean=286.436
 ```
 """
 function DewpointTemperature(model; tolerance=1e-4, maxiter=10)

src/AtmosphereModels/Diagnostics/potential_temperatures.jl

@@ -112,7 +112,7 @@ computed using the moist air gas constant ``Rᵐ`` and heat capacity ``cᵖᵐ``
 ```jldoctest
 using Breeze
 
-grid = RectilinearGrid(size=(1, 1, 8), extent=(1, 1, 1e3))
+grid = RectilinearGrid(size=8, z=(0, 1e3), topology=(Flat, Flat, Bounded))
 model = AtmosphereModel(grid)
 set!(model, θ=300)
 
@@ -121,12 +121,12 @@ Field(θ)
 
 # output
 1×1×8 Field{Center, Center, Center} on RectilinearGrid on CPU
-├── grid: 1×1×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×3 halo
+├── grid: 1×1×8 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── boundary conditions: FieldBoundaryConditions
-│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
+│   └── west: Nothing, east: Nothing, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
 ├── operand: KernelFunctionOperation at (Center, Center, Center)
 ├── status: time=0.0
-└── data: 3×3×14 OffsetArray(::Array{Float64, 3}, 0:2, 0:2, -2:11) with eltype Float64 with indices 0:2×0:2×-2:11
+└── data: 1×1×14 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:11) with eltype Float64 with indices 1:1×1:1×-2:11
     └── max=300.0, min=300.0, mean=300.0
 ```
 """
@@ -187,7 +187,7 @@ and that ``δᵛ ≈ 0.608``.
 ```jldoctest
 using Breeze
 
-grid = RectilinearGrid(size=(1, 1, 8), extent=(1, 1, 1e3))
+grid = RectilinearGrid(size=8, z=(0, 1e3), topology=(Flat, Flat, Bounded))
 model = AtmosphereModel(grid)
 set!(model, θ=300, qᵗ=0.01)
 
@@ -196,12 +196,12 @@ Field(θᵛ)
 
 # output
 1×1×8 Field{Center, Center, Center} on RectilinearGrid on CPU
-├── grid: 1×1×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×3 halo
+├── grid: 1×1×8 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── boundary conditions: FieldBoundaryConditions
-│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
+│   └── west: Nothing, east: Nothing, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
 ├── operand: KernelFunctionOperation at (Center, Center, Center)
 ├── status: time=0.0
-└── data: 3×3×14 OffsetArray(::Array{Float64, 3}, 0:2, 0:2, -2:11) with eltype Float64 with indices 0:2×0:2×-2:11
+└── data: 1×1×14 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:11) with eltype Float64 with indices 1:1×1:1×-2:11
     └── max=301.824, min=301.824, mean=301.824
 ```
 
@@ -266,7 +266,7 @@ This is the prognostic thermodynamic variable used in `LiquidIcePotentialTempera
 ```jldoctest
 using Breeze
 
-grid = RectilinearGrid(size=(1, 1, 8), extent=(1, 1, 1e3))
+grid = RectilinearGrid(size=8, z=(0, 1e3), topology=(Flat, Flat, Bounded))
 model = AtmosphereModel(grid)
 set!(model, θ=300)
 
@@ -275,12 +275,12 @@ Field(θˡⁱ)
 
 # output
 1×1×8 Field{Center, Center, Center} on RectilinearGrid on CPU
-├── grid: 1×1×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×3 halo
+├── grid: 1×1×8 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── boundary conditions: FieldBoundaryConditions
-│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
+│   └── west: Nothing, east: Nothing, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
 ├── operand: KernelFunctionOperation at (Center, Center, Center)
 ├── status: time=0.0
-└── data: 3×3×14 OffsetArray(::Array{Float64, 3}, 0:2, 0:2, -2:11) with eltype Float64 with indices 0:2×0:2×-2:11
+└── data: 1×1×14 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:11) with eltype Float64 with indices 1:1×1:1×-2:11
     └── max=300.0, min=300.0, mean=300.0
 ```
 """
@@ -341,7 +341,7 @@ adapted from the derivation in the work by [Durran and Klemp (1982)](@cite Durra
 ```jldoctest
 using Breeze
 
-grid = RectilinearGrid(size=(1, 1, 8), extent=(1, 1, 1e3))
+grid = RectilinearGrid(size=8, z=(0, 1e3), topology=(Flat, Flat, Bounded))
 model = AtmosphereModel(grid)
 set!(model, θ=300, qᵗ=0.01)
 
@@ -350,13 +350,13 @@ Field(θᵉ)
 
 # output
 1×1×8 Field{Center, Center, Center} on RectilinearGrid on CPU
-├── grid: 1×1×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×3 halo
+├── grid: 1×1×8 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── boundary conditions: FieldBoundaryConditions
-│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
+│   └── west: Nothing, east: Nothing, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
 ├── operand: KernelFunctionOperation at (Center, Center, Center)
 ├── status: time=0.0
-└── data: 3×3×14 OffsetArray(::Array{Float64, 3}, 0:2, 0:2, -2:11) with eltype Float64 with indices 0:2×0:2×-2:11
-    └── max=326.162, min=325.849, mean=326.005
+└── data: 1×1×14 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:11) with eltype Float64 with indices 1:1×1:1×-2:11
+    └── max=326.531, min=326.207, mean=326.368
 ```
 
 # References
@@ -423,7 +423,7 @@ calculations in saturated atmospheres.
 ```jldoctest
 using Breeze
 
-grid = RectilinearGrid(size=(1, 1, 8), extent=(1, 1, 1e3))
+grid = RectilinearGrid(size=8, z=(0, 1e3), topology=(Flat, Flat, Bounded))
 model = AtmosphereModel(grid)
 set!(model, θ=300, qᵗ=0.01)
 
@@ -432,13 +432,13 @@ Field(θᵇ)
 
 # output
 1×1×8 Field{Center, Center, Center} on RectilinearGrid on CPU
-├── grid: 1×1×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×3 halo
+├── grid: 1×1×8 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── boundary conditions: FieldBoundaryConditions
-│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
+│   └── west: Nothing, east: Nothing, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
 ├── operand: KernelFunctionOperation at (Center, Center, Center)
 ├── status: time=0.0
-└── data: 3×3×14 OffsetArray(::Array{Float64, 3}, 0:2, 0:2, -2:11) with eltype Float64 with indices 0:2×0:2×-2:11
-    └── max=326.162, min=325.849, mean=326.005
+└── data: 1×1×14 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:11) with eltype Float64 with indices 1:1×1:1×-2:11
+    └── max=326.531, min=326.207, mean=326.368
 ```
 
 # References

src/AtmosphereModels/Diagnostics/static_energy.jl

@@ -51,7 +51,7 @@ This is the prognostic thermodynamic variable used in `StaticEnergyThermodynamic
 ```jldoctest
 using Breeze
 
-grid = RectilinearGrid(size=(1, 1, 8), extent=(1, 1, 1e3))
+grid = RectilinearGrid(size=8, z=(0, 1e3), topology=(Flat, Flat, Bounded))
 model = AtmosphereModel(grid)
 set!(model, θ=300)
 
@@ -60,13 +60,13 @@ Field(e)
 
 # output
 1×1×8 Field{Center, Center, Center} on RectilinearGrid on CPU
-├── grid: 1×1×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×3 halo
+├── grid: 1×1×8 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── boundary conditions: FieldBoundaryConditions
-│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
+│   └── west: Nothing, east: Nothing, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
 ├── operand: KernelFunctionOperation at (Center, Center, Center)
 ├── status: time=0.0
-└── data: 3×3×14 OffsetArray(::Array{Float64, 3}, 0:2, 0:2, -2:11) with eltype Float64 with indices 0:2×0:2×-2:11
-    └── max=3.03055e5, min=3.02663e5, mean=3.02859e5
+└── data: 1×1×14 OffsetArray(::Array{Float64, 3}, 1:1, 1:1, -2:11) with eltype Float64 with indices 1:1×1:1×-2:11
+    └── max=3.02608e5, min=3.02216e5, mean=3.02412e5
 ```
 """
 function StaticEnergy(model, flavor_symbol=:specific)

src/AtmosphereModels/atmosphere_model.jl

@@ -2,6 +2,7 @@ using ..Thermodynamics: Thermodynamics, ThermodynamicConstants
 
 using Oceananigans: Oceananigans, AbstractModel, Center, CenterField, Clock, Field,
                     Centered, fields, prognostic_fields
+import Oceananigans.Architectures: architecture
 using Oceananigans.Advection: Advection, adapt_advection_order, cell_advection_timescale
 using Oceananigans.AbstractOperations: @at
 using Oceananigans.BoundaryConditions: FieldBoundaryConditions, regularize_field_boundary_conditions
@@ -26,6 +27,35 @@ function validate_tracers(tracers::Tuple)
     return tracers
 end
 
+using Oceananigans.Grids: topology, halo_size, Periodic
+
+# Minimum halo size required for velocity interpolation stencil.
+# The velocity computation uses 2-point interpolation (e.g., ℑxᶠᵃᵃ accesses i-1),
+# and the kernel range is shrunk by 1 on each side, so we need halo >= 2.
+const MINIMUM_VELOCITY_INTERPOLATION_HALO = 2
+
+"""
+Validate that the grid's halo is large enough for the velocity interpolation stencil.
+For periodic dimensions, we need at least 2 halo points because the velocity computation
+uses a 2-point interpolation stencil (accesses i-1) and the kernel range is shrunk by 1.
+"""
+function validate_velocity_interpolation_halo(grid)
+    Hx, Hy, Hz = halo_size(grid)
+    TX, TY, TZ = topology(grid)
+
+    if TX == Periodic && Hx < MINIMUM_VELOCITY_INTERPOLATION_HALO
+        throw(ArgumentError("For periodic x-topology, AtmosphereModel requires halo ≥ $MINIMUM_VELOCITY_INTERPOLATION_HALO in x (got Hx=$Hx)"))
+    end
+
+    if TY == Periodic && Hy < MINIMUM_VELOCITY_INTERPOLATION_HALO
+        throw(ArgumentError("For periodic y-topology, AtmosphereModel requires halo ≥ $MINIMUM_VELOCITY_INTERPOLATION_HALO in y (got Hy=$Hy)"))
+    end
+
+    # Note: z is typically Bounded, so no check needed for Hz
+
+    return nothing
+end
+
 mutable struct AtmosphereModel{Dyn, Frm, Arc, Tst, Grd, Clk, Thm, Mom, Moi, Mfr, Buy,
                                Tmp, Sol, Vel, Trc, Adv, Cor, Frc, Mic, Cnd, Cls, Cfs, Rad} <: AbstractModel{Tst, Arc}
     architecture :: Arc
@@ -137,6 +167,7 @@ function AtmosphereModel(grid;
 
     # Check halos and throw an error if the grid's halo is too small
     validate_model_halo(grid, momentum_advection, scalar_advection, closure)
+    validate_velocity_interpolation_halo(grid)
 
     # Reduce the advection order in directions that do not have enough grid points
     momentum_advection = validate_momentum_advection(momentum_advection, grid)
@@ -252,6 +283,8 @@ function AtmosphereModel(grid;
     return model
 end
 
+architecture(model::AtmosphereModel) = model.grid.architecture
+
 function Base.summary(model::AtmosphereModel)
     A = nameof(typeof(model.grid.architecture))
     G = nameof(typeof(model.grid))

src/AtmosphereModels/set_to_mean.jl

@@ -16,7 +16,7 @@ function set_to_mean!(ref::ReferenceState, model)
     constants = model.thermodynamic_constants
 
     mean!(ref.temperature, model.temperature)
-    fill_halo_regions!(ref.temperature)
+    fill_halo_regions!(ref.temperature; only_local_halos=true)
 
     mean_mass_fraction!(ref.vapor_mass_fraction, vapor_mass_fraction(model))
     mean_mass_fraction!(ref.liquid_mass_fraction, liquid_mass_fraction(model))
@@ -29,13 +29,13 @@ end
 
 function mean_mass_fraction!(ref_field, field)
     mean!(ref_field, field)
-    fill_halo_regions!(ref_field)
+    fill_halo_regions!(ref_field; only_local_halos=true)
     return nothing
 end
 
 function mean_mass_fraction!(ref_field, ::Nothing)
     interior(ref_field) .= 0
-    fill_halo_regions!(ref_field)
+    fill_halo_regions!(ref_field; only_local_halos=true)
     return nothing
 end
 

src/AtmosphereModels/update_atmosphere_model_state.jl

@@ -60,8 +60,13 @@ function tracer_specific_to_density!(tracers, density)
     return nothing
 end
 
+# Kernel launch indices for computing diagnostic fields (e.g., velocities from momentum).
+# For periodic dimensions, we shrink by 1 on each side because the velocity computation
+# uses interpolation operators that access neighboring points (e.g., ℑxᶠᵃᵃ accesses i-1).
+# The outermost halo points are then filled by fill_halo_regions! via periodic wrapping.
+# This requires halo >= 2 for periodic dimensions to ensure the computed region is non-empty.
 diagnostic_indices(::Bounded, N, H) = 1:N+1
-diagnostic_indices(::Periodic, N, H) = -H+1:N+H
+diagnostic_indices(::Periodic, N, H) = -H+2:N+H-1
 diagnostic_indices(::Flat, N, H) = 1:N
 
 #####
@@ -94,7 +99,9 @@ function compute_velocities!(model::AtmosphereModel)
     # Ensure halos are filled before velocity computation
     # (prognostic field halo fill in update_state! is async)
     density = dynamics_density(model.dynamics)
-    fill_halo_regions!(density)
+    # Column fields (e.g., anelastic reference density) only need local halo fills;
+    # skip MPI communication that would fail on their singleton x/y dimensions.
+    fill_halo_regions!(density; only_local_halos=_is_column_field(density))
     fill_halo_regions!(model.momentum)
 
     launch!(arch, grid, :xyz,