Fix Stokes model

scripts_future_API/tm_stokes_mpi_wip.jl

@@ -14,7 +14,7 @@ const SBC = BoundaryCondition{Slip}
 
 using KernelAbstractions
 
-# using GLMakie
+using GLMakie
 
 using FastIce.Distributed
 using MPI
@@ -30,7 +30,7 @@ ebg = 1.0
 
 topo = details(arch)
 
-size_l = (32, 32, 32)
+size_l = (64, 64, 64)
 size_g = global_grid_size(topo, size_l)
 
 if global_rank(topo) == 0
@@ -73,7 +73,7 @@ physics = (rheology=GlensLawRheology(1),)
 other_fields = (A=Field(backend, grid_l, Center()),)
 
 init_incl(x, y, z, x0, y0, z0, r, Ai, Am) = ifelse((x - x0)^2 + (y - y0)^2 + (z - z0)^2 < r^2, Ai, Am)
-set!(other_fields.A, grid_l, init_incl; parameters=(x0=0.0, y0=0.0, z0=0.5, r=0.2, Ai=1e-1, Am=1.0))
+set!(other_fields.A, grid_l, init_incl; parameters=(x0=0.0, y0=0.0, z0=1.0, r=0.2, Ai=1e-1, Am=1.0))
 
 model = IsothermalFullStokesModel(;
                                   arch,
@@ -83,22 +83,37 @@ model = IsothermalFullStokesModel(;
                                   iter_params,
                                   other_fields)
 
-# if global_rank(topo) == 0
-    # fig = Figure(; resolution=(1200, 1000), fontsize=32)
-    # axs = (Pr=Axis(fig[1, 1][1, 1]; aspect=DataAspect(), xlabel="x", ylabel="z", title="Pr"),
-    #        Vx=Axis(fig[1, 2][1, 1]; aspect=DataAspect(), xlabel="x", ylabel="z", title="Vx"),
-    #        Vy=Axis(fig[2, 1][1, 1]; aspect=DataAspect(), xlabel="x", ylabel="z", title="Vy"),
-    #        Vz=Axis(fig[2, 2][1, 1]; aspect=DataAspect(), xlabel="x", ylabel="z", title="Vz"))
-
-    # plt = (Pr=heatmap!(axs.Pr, xcenters(grid_l), zcenters(grid_l), interior(model.fields.Pr)[:, size(grid_l, 2)÷2, :]; colormap=:turbo),
-    #        Vx=heatmap!(axs.Vx, xvertices(grid_l), zcenters(grid_l), interior(model.fields.V.x)[:, size(grid_l, 2)÷2, :]; colormap=:turbo),
-    #        Vy=heatmap!(axs.Vy, xcenters(grid_l), zcenters(grid_l), interior(model.fields.V.y)[:, size(grid_l, 2)÷2, :]; colormap=:turbo),
-    #        Vz=heatmap!(axs.Vz, xcenters(grid_l), zvertices(grid_l), interior(model.fields.V.z)[:, size(grid_l, 2)÷2, :]; colormap=:turbo))
-    # Colorbar(fig[1, 1][1, 2], plt.Pr)
-    # Colorbar(fig[1, 2][1, 2], plt.Vx)
-    # Colorbar(fig[2, 1][1, 2], plt.Vy)
-    # Colorbar(fig[2, 2][1, 2], plt.Vz)
-# end
+if global_rank(topo) == 0
+    fig = Figure(; resolution=(1200, 1000), fontsize=32)
+    axs = (Pr=Axis(fig[1, 1][1, 1]; aspect=DataAspect(), xlabel="x", ylabel="z", title="Pr"),
+           Vx=Axis(fig[1, 2][1, 1]; aspect=DataAspect(), xlabel="x", ylabel="z", title="Vx"),
+           Vy=Axis(fig[2, 1][1, 1]; aspect=DataAspect(), xlabel="x", ylabel="z", title="Vy"),
+           Vz=Axis(fig[2, 2][1, 1]; aspect=DataAspect(), xlabel="x", ylabel="z", title="Vz"))
+
+    Pr_g = zeros(size(grid_g))
+    Vx_g = zeros(size(grid_g))
+    Vy_g = zeros(size(grid_g))
+    Vz_g = zeros(size(grid_g))
+
+    Pr_v = zeros(size(grid_l))
+    Vx_v = zeros(size(grid_l))
+    Vy_v = zeros(size(grid_l))
+    Vz_v = zeros(size(grid_l))
+
+    plt = (Pr=heatmap!(axs.Pr, xcenters(grid_g), zcenters(grid_g), @view Pr_g[:, size(grid_g, 2)÷2, :]; colormap=:turbo),
+           Vx=heatmap!(axs.Vx, xcenters(grid_g), zcenters(grid_g), @view Vx_g[:, size(grid_g, 2)÷2, :]; colormap=:turbo),
+           Vy=heatmap!(axs.Vy, xcenters(grid_g), zcenters(grid_g), @view Vy_g[:, size(grid_g, 2)÷2, :]; colormap=:turbo),
+           Vz=heatmap!(axs.Vz, xcenters(grid_g), zcenters(grid_g), @view Vz_g[:, size(grid_g, 2)÷2, :]; colormap=:turbo))
+    Colorbar(fig[1, 1][1, 2], plt.Pr)
+    Colorbar(fig[1, 2][1, 2], plt.Vx)
+    Colorbar(fig[2, 1][1, 2], plt.Vy)
+    Colorbar(fig[2, 2][1, 2], plt.Vz)
+else
+    Pr_g = nothing
+    Vx_g = nothing
+    Vy_g = nothing
+    Vz_g = nothing
+end
 
 # set!(model.fields.Pr, 0.0)
 # foreach(x -> set!(x, 0.0), model.fields.τ)
@@ -108,43 +123,43 @@ foreach(x -> fill!(parent(x), 0.0), model.fields.τ)
 foreach(x -> fill!(parent(x), 0.0), model.fields.V)
 
 KernelLaunch.apply_all_boundary_conditions!(arch, grid_l, model.boundary_conditions.stress)
-println("$(global_rank(topo)) applied stress BCs")
 
 set!(model.fields.V.x, grid_l, psh_x)
 set!(model.fields.V.y, grid_l, psh_y)
 set!(model.fields.V.z, 0.0)
 
-# println("at rank $(global_rank(topo)) bcs $(typeof(model.boundary_conditions.velocity))")
-println("at rank $(global_rank(topo)) topo $topo")
-MPI.Barrier(cartesian_communicator(topo))
-
 KernelLaunch.apply_all_boundary_conditions!(arch, grid_l, model.boundary_conditions.velocity)
-println("$(global_rank(topo)) applied velocity BCs")
 
 set!(model.fields.η, other_fields.A)
-extrapolate!(model.fields.η)
-
-# if global_rank(topo) == 0
-    # display(fig)
-# end
+KernelLaunch.apply_all_boundary_conditions!(arch, grid_l, model.boundary_conditions.rheology)
 
-println("Hi from $(global_rank(topo))")
-MPI.Barrier(cartesian_communicator(topo))
+if global_rank(topo) == 0
+    display(fig)
+end
 
 for it in 1:nt
     advance_iteration!(model, 0.0, 1.0; async=false)
-    if it % nviz == 0# && global_rank(topo) == 0
+    if it % nviz == 0
+        @views avx(A) = 0.5.*(A[1:end-1,:,:] .+ A[2:end,:,:])
+        @views avy(A) = 0.5.*(A[:,1:end-1,:] .+ A[:,2:end,:])
+        @views avz(A) = 0.5.*(A[:,:,1:end-1] .+ A[:,:,2:end])
+
+        comm = cartesian_communicator(topo)
+
+        gather!(Pr_g, interior(model.fields.Pr), comm)
+        gather!(Vx_g, avx(interior(model.fields.V.x)), comm)
+        gather!(Vy_g, avy(interior(model.fields.V.y)), comm)
+        gather!(Vz_g, avz(interior(model.fields.V.z)), comm)
         if global_rank(topo) == 0
-            println("it = $it/$nt")
+            plt.Pr[3][] = @view Pr_g[:, size(grid_g, 2)÷2, :]
+            plt.Vx[3][] = @view Vx_g[:, size(grid_g, 2)÷2, :]
+            plt.Vy[3][] = @view Vy_g[:, size(grid_g, 2)÷2, :]
+            plt.Vz[3][] = @view Vz_g[:, size(grid_g, 2)÷2, :]
+            yield()
         end
-        # plt.Pr[3][] = interior(model.fields.Pr)[:, size(grid_l, 2)÷2, :]
-        # plt.Vx[3][] = interior(model.fields.V.x)[:, size(grid_l, 2)÷2, :]
-        # plt.Vy[3][] = interior(model.fields.V.y)[:, size(grid_l, 2)÷2, :]
-        # plt.Vz[3][] = interior(model.fields.V.z)[:, size(grid_l, 2)÷2, :]
-        # yield()
     end
 end
 
-sleep(30)
+sleep(120)
 
 MPI.Finalize()

src/BoundaryConditions/BoundaryConditions.jl

@@ -4,6 +4,7 @@ export FieldBoundaryCondition, BoundaryConditionsBatch
 export apply_boundary_conditions!, apply_all_boundary_conditions!
 
 export DirichletBC, HalfCell, FullCell
+export ExtrapolateBC
 export ContinuousBC, DiscreteBC
 export BoundaryFunction, DiscreteBoundaryFunction, ContinuousBoundaryFunction
 
@@ -23,6 +24,7 @@ include("field_boundary_conditions.jl")
 include("utils.jl")
 include("boundary_function.jl")
 include("dirichlet_bc.jl")
+include("extrapolate_bc.jl")
 include("hide_boundaries.jl")
 
 struct BoundaryConditionsBatch{F,BC}
@@ -38,6 +40,5 @@ end
 end
 
 apply_boundary_conditions!(side, val, arch, grid, ::Nothing; kwargs...) = nothing
-apply_boundary_conditions!(side, val, arch, grid, fields, ::Nothing; kwargs...) = nothing
 
 end

src/BoundaryConditions/extrapolate_bc.jl

@@ -0,0 +1,9 @@
+# Boundary conditions that extrapolates the values of the field outside of the domain
+struct ExtrapolateBC <: FieldBoundaryCondition end
+
+@inline function _apply_field_boundary_condition!(side, dim, grid, f, loc, Ibc, ::ExtrapolateBC)
+    I = _bc_index(dim, side, loc, size(f), Ibc)
+    DI = _bc_offset(Val(ndims(grid)), dim, side)
+    @inbounds f[I] = 2 * f[I+DI] - f[I+2DI]
+    return
+end

src/BoundaryConditions/field_boundary_conditions.jl

@@ -1,10 +1,8 @@
-const FieldOrNothing = Union{FieldBoundaryCondition,Nothing}
-
 function apply_boundary_conditions!(::Val{S}, ::Val{D},
                                     arch::Architecture,
                                     grid::CartesianGrid,
                                     fields::NTuple{N,Field},
-                                    conditions::NTuple{N,FieldOrNothing}; async=true) where {S,D,N}
+                                    conditions::NTuple{N,FieldBoundaryCondition}; async=true) where {S,D,N}
     _validate_fields(fields, D, S)
     sizes = ntuple(ifield -> remove_dim(Val(D), size(parent(fields[ifield]))), Val(N))
     halos = ntuple(ifield -> remove_dim(Val(D), halo(fields[ifield])), Val(N))

src/Distributed/Distributed.jl

@@ -13,11 +13,11 @@ export CartesianTopology
 export global_rank, shared_rank, node_name, cartesian_communicator, shared_communicator, coordinates
 export dimensions, global_size, node_size
 export global_grid_size, local_grid
-export neighbors, neighbor
+export neighbors, neighbor, has_neighbor
 export override_boundary_conditions
 export gather!
 
-export DistributedBoundaryConditions
+export ExchangeInfo, DistributedBoundaryConditions
 
 struct DistributedMPI end
 

src/Distributed/topology.jl

@@ -46,6 +46,8 @@ neighbors(t::CartesianTopology) = t.neighbors
 
 neighbor(t::CartesianTopology, dim, side) = t.neighbors[dim][side]
 
+has_neighbor(t::CartesianTopology, dim, side) = t.neighbors[dim][side] != MPI.PROC_NULL
+
 global_size(t::CartesianTopology) = MPI.Comm_size(t.cart_comm)
 node_size(t::CartesianTopology) = MPI.Comm_size(t.shared_comm)
 

src/Models/full_stokes/isothermal/boundary_conditions.jl

@@ -43,66 +43,75 @@ for (dim, val) in _COORDINATES
     ex_slip_vel = Expr(:tuple, ex_slip_vel)
 
     @eval begin
-        function extract_boundary_conditions(::Val{$val}, bc::BoundaryCondition{Traction})
-            return $ex_tr, NamedTuple()
-        end
-
-        function extract_boundary_conditions(::Val{$val}, bc::BoundaryCondition{Velocity})
-            return NamedTuple(), $ex_vel
-        end
+        extract_stress_bc(::Val{$val}, bc::BoundaryCondition{Traction}) = $ex_tr
+        extract_stress_bc(::Val{$val}, bc::BoundaryCondition{Velocity}) = ()
+        extract_stress_bc(::Val{$val}, bc::BoundaryCondition{Slip})     = $ex_slip_tr
 
-        function extract_boundary_conditions(::Val{$val}, bc::BoundaryCondition{Slip})
-            return $ex_slip_tr, $ex_slip_vel
-        end
+        extract_velocity_bc(::Val{$val}, bc::BoundaryCondition{Traction}) = ()
+        extract_velocity_bc(::Val{$val}, bc::BoundaryCondition{Velocity}) = $ex_vel
+        extract_velocity_bc(::Val{$val}, bc::BoundaryCondition{Slip})     = $ex_slip_vel
     end
 end
 
-function make_batch(::CartesianGrid{2}, bcs::NamedTuple{bnames, Tuple{}}, fields::NamedTuple) where {bnames}
-    return nothing
-end
-
-function make_batch(::CartesianGrid{3}, bcs::NamedTuple{bnames, Tuple{}}, fields::NamedTuple) where {bnames}
-    return nothing
-end
-
-function make_batch(::CartesianGrid{2}, bcs::NamedTuple, fields::NamedTuple)
-    field_map = (Pr  = fields.Pr,
-                 τxx = fields.τ.xx, τyy = fields.τ.yy, τxy = fields.τ.xy,
-                 Vx  = fields.V.x, Vy  = fields.V.y)
-    batch_fields = Tuple(field_map[name] for name in eachindex(bcs))
+function make_batch(bcs::NamedTuple, fields::NamedTuple)
+    batch_fields = Tuple(fields[name] for name in eachindex(bcs))
     return BoundaryConditionsBatch(batch_fields, values(bcs))
 end
 
-function make_batch(::CartesianGrid{3}, bcs::NamedTuple, fields::NamedTuple)
-    field_map = (Pr  = fields.Pr,
-                 τxx = fields.τ.xx, τyy = fields.τ.yy, τzz = fields.τ.zz,
-                 τxy = fields.τ.xy, τxz = fields.τ.xz, τyz = fields.τ.yz,
-                 Vx  = fields.V.x, Vy  = fields.V.y, Vz  = fields.V.z)
-    batch_fields = Tuple(field_map[name] for name in eachindex(bcs))
-    return BoundaryConditionsBatch(batch_fields, values(bcs))
-end
-
-function make_batches(grid, bcs, fields)
-    ntuple(Val(ndims(grid))) do D
+function make_stress_bc(arch::Architecture{Kind}, ::CartesianGrid{N}, fields, bc) where {Kind,N}
+    ordering = (:x, :y, :z)
+    ntuple(Val(N)) do D
         ntuple(Val(2)) do S
-            make_batch(grid, bcs[D][S], fields)
+            if (Kind == Distributed.DistributedMPI) && has_neighbor(details(arch), D, S)
+                nothing
+            else
+                new_bc = extract_stress_bc(Val(D), bc[ordering[D]][S])
+                if isempty(new_bc)
+                    nothing
+                else
+                    make_batch(new_bc, fields)
+                end
+            end
         end
     end
 end
 
-function make_field_boundary_conditions(grid::CartesianGrid{N}, fields, logical_boundary_conditions) where {N}
+function make_velocity_bc(arch::Architecture{Kind}, ::CartesianGrid{N}, fields::NamedTuple{names}, bc) where {Kind,N,names}
     ordering = (:x, :y, :z)
+    ntuple(Val(N)) do D
+        ntuple(Val(2)) do S
+            if (Kind == Distributed.DistributedMPI) && has_neighbor(details(arch), D, S)
+                new_bc = NamedTuple{names}(ExchangeInfo(Val(S), Val(D), V) for V in fields)
+                make_batch(new_bc, fields)
+            else
+                new_bc = extract_velocity_bc(Val(D), bc[ordering[D]][S])
+                if isempty(new_bc)
+                    nothing
+                else
+                    make_batch(new_bc, fields)
+                end
+            end
+        end
+    end
+end
 
-    field_bcs = ntuple(Val(N)) do dim
-        left, right = logical_boundary_conditions[ordering[dim]]
-        left  = extract_boundary_conditions(Val(dim), left)
-        right = extract_boundary_conditions(Val(dim), right)
-        Tuple(zip(left, right))
+function make_rheology_bc(arch::Architecture{Kind}, ::CartesianGrid{N}, η) where {Kind,N}
+    ntuple(Val(N)) do D
+        ntuple(Val(2)) do S
+            if (Kind == Distributed.DistributedMPI) && has_neighbor(details(arch), D, S)
+                BoundaryConditionsBatch((η,), (ExchangeInfo(Val(S), Val(D), η),))
+            else
+                BoundaryConditionsBatch((η,), (ExtrapolateBC(),))
+            end
+        end
     end
+end
 
-    stress, velocity = zip(field_bcs...)
-    stress   = make_batches(grid, stress, fields)
-    velocity = make_batches(grid, velocity, fields)
+function make_field_boundary_conditions(arch::Architecture, grid::CartesianGrid, fields, bc)
+    stress_fields   = (; Pr=fields.Pr, NamedTuple{Symbol.(:τ, keys(fields.τ))}(values(fields.τ))...)
+    velocity_fields = NamedTuple{Symbol.(:V, keys(fields.V))}(values(fields.V))
 
-    return (; stress, velocity)
+    return (stress   = make_stress_bc(arch, grid, stress_fields, bc),
+            velocity = make_velocity_bc(arch, grid, velocity_fields, bc),
+            rheology = make_rheology_bc(arch, grid, fields.η))
 end

src/Models/full_stokes/isothermal/isothermal.jl

@@ -31,7 +31,18 @@ struct IsothermalFullStokesModel{Arch,Grid,BC,Physics,IterParams,Fields}
     fields::Fields
 end
 
-function make_fields_mechanics(backend, grid)
+function make_fields_mechanics(backend, grid::CartesianGrid{2})
+    return (Pr=Field(backend, grid, Center(); halo=1),
+            # deviatoric stress
+            τ=(xx=Field(backend, grid, Center(); halo=1),
+               yy=Field(backend, grid, Center(); halo=1),
+               xy=Field(backend, grid, (Vertex(), Vertex(), Center()))),
+            # velocity
+            V=(x=Field(backend, grid, (Vertex(), Center(), Center()); halo=1),
+               y=Field(backend, grid, (Center(), Vertex(), Center()); halo=1)))
+end
+
+function make_fields_mechanics(backend, grid::CartesianGrid{3})
     return (Pr=Field(backend, grid, Center(); halo=1),
             # deviatoric stress
             τ=(xx=Field(backend, grid, Center(); halo=1),
@@ -61,13 +72,7 @@ function IsothermalFullStokesModel(; arch, grid, boundary_conditions, physics=no
         physics = default_physics(eltype(grid))
     end
 
-    boundary_conditions = make_field_boundary_conditions(grid, fields, boundary_conditions)
-
-    if arch isa Architecture{Distributed.DistributedMPI}
-        stress = override_boundary_conditions(arch, boundary_conditions.stress)
-        velocity = override_boundary_conditions(arch, boundary_conditions.velocity; exchange=true)
-        boundary_conditions = (; stress, velocity)
-    end
+    boundary_conditions = make_field_boundary_conditions(arch, grid, fields, boundary_conditions)
 
     return IsothermalFullStokesModel(arch, grid, boundary_conditions, physics, iter_params, fields)
 end
@@ -88,8 +93,8 @@ function advance_iteration!(model::IsothermalFullStokesModel, t, Δt; async=true
             location=Vertex(), boundary_conditions=model.boundary_conditions.velocity)
 
     # rheology
-    launch!(model.arch, model.grid, update_η! => (η, η_rh, η_rel, model.grid, model.fields); location=Center())
-    extrapolate!(η)
+    launch!(model.arch, model.grid, update_η! => (η, η_rh, η_rel, model.grid, model.fields);
+            location=Center(), boundary_conditions=model.boundary_conditions.rheology)
 
     async || synchronize(backend(model.arch))
     return

src/Physics.jl

@@ -1,14 +1,11 @@
 module Physics
 
-export AbstractPhysics
 export IncompressibleIceEOS, IceThermalProperties
-export IceRheology, GlensLawRheology, τ
+export IceRheology, GlensLawRheology
 export default
 
 using FastIce.GridOperators
 
-abstract type AbstractPhysics end
-
 struct IncompressibleIceEOS{T}
     density::T
     heat_capacity::T