WIP Stokes

scripts_future_API/tm_stokes_wip.jl

@@ -6,81 +6,67 @@ using FastIce.BoundaryConditions
 using FastIce.Models.FullStokes.Isothermal
 using FastIce.Physics
 
+const VBC = BoundaryCondition{Velocity}
+
 using KernelAbstractions
 
 using GLMakie
 
 # physics
 ebg = 1.0
 
-grid = CartesianGrid(
-    origin = (-0.5, -0.5, 0.0),
-    extent = ( 1.0,  1.0, 1.0),
-    size   = (  64,   64, 64 ),
-)
+grid = CartesianGrid(; origin=(-0.5, -0.5, 0.0),
+                     extent=(1.0, 1.0, 1.0),
+                     size=(64, 64, 64))
 
-psh_x(x, _, _) = -x*ebg
-psh_y(_, y, _) =  y*ebg
+psh_x(x, _, _) = -x * ebg
+psh_y(_, y, _) = y * ebg
 
 x_bc = BoundaryFunction(psh_x; reduce_dims=false)
 y_bc = BoundaryFunction(psh_y; reduce_dims=false)
 
-boundary_conditions = (
-    west   = BoundaryCondition{Velocity}(x_bc, y_bc, 0.0),
-    east   = BoundaryCondition{Velocity}(x_bc, y_bc, 0.0),
-    south  = BoundaryCondition{Velocity}(x_bc, y_bc, 0.0),
-    north  = BoundaryCondition{Velocity}(x_bc, y_bc, 0.0),
-    bottom = BoundaryCondition{Velocity}(0.0 , 0.0 , 0.0),
-    top    = BoundaryCondition{Velocity}(0.0 , 0.0 , 0.0),
-)
+boundary_conditions = (x = (VBC(x_bc, y_bc, 0.0), VBC(x_bc, y_bc, 0.0)),
+                       y = (VBC(x_bc, y_bc, 0.0), VBC(x_bc, y_bc, 0.0)),
+                       z = (VBC(0.0, 0.0, 0.0), VBC(0.0, 0.0, 0.0)))
 
 r       = 0.7
 re_mech = 10π
 lτ_re_m = minimum(extent(grid)) / re_mech
-vdτ     = minimum(spacing(grid)) / sqrt(10.1)
+vdτ     = minimum(spacing(grid)) / sqrt(ndims(grid) * 3.1)
 θ_dτ    = lτ_re_m * (r + 4 / 3) / vdτ
 dτ_r    = 1.0 / (θ_dτ + 1.0)
 nudτ    = vdτ * lτ_re_m
 
-iter_params = (
-    η_rel = 1e-1,
-    Δτ = ( Pr = r / θ_dτ, τ = (xx = dτ_r, yy = dτ_r, zz = dτ_r, xy = dτ_r, xz = dτ_r, yz = dτ_r), V = (x = nudτ, y = nudτ, z = nudτ)),
-)
+iter_params = (η_rel=1e-1,
+               Δτ=(Pr=r / θ_dτ, τ=(xx=dτ_r, yy=dτ_r, zz=dτ_r, xy=dτ_r, xz=dτ_r, yz=dτ_r), V=(x=nudτ, y=nudτ, z=nudτ)))
 
 backend = CPU()
 
-physics = (rheology = GlensLawRheology(1), )
-other_fields = (
-    A = Field(backend, grid, Center()),
-)
+physics = (rheology=GlensLawRheology(1),)
+other_fields = (A=Field(backend, grid, 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, init_incl; parameters = (x0 = 0.0, y0 = 0.0, z0 = 0.5, r = 0.2, Ai = 1e1, Am = 1.0))
+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, init_incl; parameters=(x0=0.0, y0=0.0, z0=0.5, r=0.2, Ai=1e1, Am=1.0))
 
 model = IsothermalFullStokesModel(;
-    backend,
-    grid,
-    physics,
-    boundary_conditions,
-    iter_params,
-    other_fields
-)
-
-fig = Figure(resolution=(1000,1000), fontsize=32)
-axs = (
-    Pr = Axis(fig[1,1][1,1]; aspect=DataAspect(), xlabel="x", ylabel="y", title="Pr"),
-    Vx = Axis(fig[2,1][1,1]; aspect=DataAspect(), xlabel="x", ylabel="y", title="Vx"),
-    Vy = Axis(fig[2,2][1,1]; aspect=DataAspect(), xlabel="x", ylabel="y", title="Vy"),
-)
-
-plt = (
-    Pr = heatmap!(axs.Pr, xcenters(grid), ycenters(grid), interior(model.fields.Pr)[:, :, size(grid,3)÷2]; colormap=:turbo),
-    Vx = heatmap!(axs.Vx, xvertices(grid), ycenters(grid), interior(model.fields.V.x)[:, :, size(grid,3)÷2]; colormap=:turbo),
-    Vy = heatmap!(axs.Vy, xcenters(grid), yvertices(grid), interior(model.fields.V.y)[:, :, size(grid,3)÷2]; colormap=:turbo),
-)
-Colorbar(fig[1,1][1,2], plt.Pr)
-Colorbar(fig[2,1][1,2], plt.Vx)
-Colorbar(fig[2,2][1,2], plt.Vy)
+                                  backend,
+                                  grid,
+                                  physics,
+                                  boundary_conditions,
+                                  iter_params,
+                                  other_fields)
+
+fig = Figure(; resolution=(1000, 1000), fontsize=32)
+axs = (Pr = Axis(fig[1, 1][1, 1]; aspect=DataAspect(), xlabel="x", ylabel="y", title="Pr"),
+       Vx = Axis(fig[2, 1][1, 1]; aspect=DataAspect(), xlabel="x", ylabel="y", title="Vx"),
+       Vy = Axis(fig[2, 2][1, 1]; aspect=DataAspect(), xlabel="x", ylabel="y", title="Vy"))
+
+plt = (Pr = heatmap!(axs.Pr, xcenters(grid), ycenters(grid), interior(model.fields.Pr)[:, :, size(grid, 3)÷2]; colormap=:turbo),
+       Vx = heatmap!(axs.Vx, xvertices(grid), ycenters(grid), interior(model.fields.V.x)[:, :, size(grid, 3)÷2]; colormap=:turbo),
+       Vy = heatmap!(axs.Vy, xcenters(grid), yvertices(grid), interior(model.fields.V.y)[:, :, size(grid, 3)÷2]; colormap=:turbo))
+Colorbar(fig[1, 1][1, 2], plt.Pr)
+Colorbar(fig[2, 1][1, 2], plt.Vx)
+Colorbar(fig[2, 2][1, 2], plt.Vy)
 
 set!(model.fields.Pr, 0.0)
 foreach(x -> set!(x, 0.0), model.fields.τ)
@@ -95,11 +81,11 @@ set!(model.fields.η, other_fields.A)
 extrapolate!(model.fields.η)
 
 for it in 1:100
-    advance_iteration!(model, 0.0, 1.0; async = false)
+    advance_iteration!(model, 0.0, 1.0; async=false)
     if it % 10 == 0
-        plt.Pr[3][] = interior(model.fields.Pr)[:, :, size(grid,3)÷2]
-        plt.Vx[3][] = interior(model.fields.V.x)[:, :, size(grid,3)÷2]
-        plt.Vy[3][] = interior(model.fields.V.y)[:, :, size(grid,3)÷2]
+        plt.Pr[3][] = interior(model.fields.Pr)[:, :, size(grid, 3)÷2]
+        plt.Vx[3][] = interior(model.fields.V.x)[:, :, size(grid, 3)÷2]
+        plt.Vy[3][] = interior(model.fields.V.y)[:, :, size(grid, 3)÷2]
         yield()
     end
 end

src/Models/full_stokes/isothermal/boundary_conditions.jl

@@ -1,15 +1,18 @@
 struct Traction end
 struct Velocity end
-struct Slip     end
+struct Slip end
 
-struct BoundaryCondition{Kind, N, C}
+"""
+Boundary condition for Stokes problem.
+"""
+struct BoundaryCondition{Kind,N,C}
     components::C
-    BoundaryCondition{Kind}(components::Tuple) where {Kind} = new{Kind, length(components), typeof(components)}(components)
+    BoundaryCondition{Kind}(components::Tuple) where {Kind} = new{Kind,length(components),typeof(components)}(components)
 end
 
 BoundaryCondition{Kind}(components...) where {Kind} = BoundaryCondition{Kind}(components)
 
-const _COORDINATES  = ((:x, 1), (:y, 2), (:z, 3))
+const _COORDINATES = ((:x, 1), (:y, 2), (:z, 3))
 
 for (dim, val) in _COORDINATES
     td = remove_dim(Val(val), _COORDINATES)
@@ -20,7 +23,7 @@ for (dim, val) in _COORDINATES
     N  = Symbol(:τ, dim, dim)
 
     ex1 = Expr(:(=), :Pr, :(DirichletBC{HalfCell}(bc.components[$val])))
-    ex2 = Expr(:(=), N,   :(DirichletBC{HalfCell}(convert(eltype(bc.components[$val]), 0))))
+    ex2 = Expr(:(=), N, :(DirichletBC{HalfCell}(convert(eltype(bc.components[$val]), 0))))
     ex3 = ntuple(Val(length(TN))) do I
         Expr(:(=), TN[I], :(DirichletBC{FullCell}(bc.components[$(td[I][2])])))
     end
@@ -54,59 +57,47 @@ for (dim, val) in _COORDINATES
     end
 end
 
-no_bcs(names) = NamedTuple(f => nothing for f in names)
-
-unique_names(a, b) = Tuple(unique(tuple(a..., b...)))
+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))
+    return BoundaryConditionsBatch(batch_fields, values(bcs))
+end
 
-function expand_boundary_conditions(left, right)
-    names = unique_names(keys(left), keys(right))
+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
 
-    if isempty(names)
-        return nothing
+function make_batches(grid, bcs, fields)
+    ntuple(Val(ndims(grid))) do D
+        ntuple(Val(2)) do S
+            make_batch(grid, bcs[D][S], fields)
+        end
     end
-
-    default = no_bcs(names)
-    left    = merge(default, left)
-    right   = merge(default, right)
-
-    return NamedTuple{names}(zip(left, right))
 end
 
-function make_field_boundary_conditions(bcs)
-    ordering = (
-        (:west  , :east),
-        (:south , :north),
-        (:bottom, :top),
-    )
+function make_field_boundary_conditions(grid::CartesianGrid{N}, fields, logical_boundary_conditions) where {N}
+    ordering = (:x, :y, :z)
 
-    field_bcs = ntuple(Val(length(ordering))) do dim
-        left, right = bcs[ordering[dim]]
+    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)
 
-        stress, velocity = Tuple(zip(left, right))
-
-        stress   = expand_boundary_conditions(stress...)
-        velocity = expand_boundary_conditions(velocity...)
-
-        stress, velocity
+        Tuple(zip(left, right))
     end
 
-    return NamedTuple{(:stress, :velocity)}(zip(field_bcs...))
-end
+    stress, velocity = zip(field_bcs...)
 
-function _apply_bcs!(backend, grid, fields, bcs)
-    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)
-
-    ntuple(Val(length(bcs))) do D
-        if !isnothing(bcs[D])
-            fs = Tuple( field_map[f] for f in eachindex(bcs[D]) )
-            apply_bcs!(Val(D), backend, grid, fs, values(bcs[D]))
-        end
-    end
-    return
+    stress   = make_batches(grid, stress, fields)
+    velocity = make_batches(grid, velocity, fields)
+
+    return (; stress, velocity)
 end

src/Models/full_stokes/isothermal/isothermal.jl

@@ -8,21 +8,20 @@ using FastIce.Grids
 using FastIce.Fields
 using FastIce.BoundaryConditions
 using FastIce.Utils
+using FastIce.KernelLaunch
 
 include("kernels.jl")
 
 include("boundary_conditions.jl")
 
-function default_physics(::Type{T}) where T
-    return (
-        equation_of_state=default(IncompressibleIceEOS{T}),
-        thermal_properties=default(IceThermalProperties{T}),
-        rheology=default(GlensLawRheology{Int64})
-    )
+function default_physics(::Type{T}) where {T}
+    return (equation_of_state=default(IncompressibleIceEOS{T}),
+            thermal_properties=default(IceThermalProperties{T}),
+            rheology=default(GlensLawRheology{Int64}))
 end
 
-struct IsothermalFullStokesModel{Backend,Grid,BC,Physics,IterParams,Fields}
-    backend::Backend
+struct IsothermalFullStokesModel{Arch,Grid,BC,Physics,IterParams,Fields}
+    arch::Arch
     grid::Grid
     boundary_conditions::BC
     physics::Physics
@@ -31,30 +30,22 @@ struct IsothermalFullStokesModel{Backend,Grid,BC,Physics,IterParams,Fields}
 end
 
 function make_fields_mechanics(backend, grid)
-    return (
-        Pr = Field(backend, grid, Center(); halo=1),
-        τ  = (
-            xx = Field(backend, grid, Center(); halo=1),
-            yy = Field(backend, grid, Center(); halo=1),
-            zz = Field(backend, grid, Center(); halo=1),
-            xy = Field(backend, grid, (Vertex(), Vertex(), Center()); halo=0),
-            xz = Field(backend, grid, (Vertex(), Center(), Vertex()); halo=0),
-            yz = Field(backend, grid, (Center(), Vertex(), Vertex()); halo=0),
-        ),
-        V = (
-            x = Field(backend, grid, (Vertex(), Center(), Center()); halo=1),
-            y = Field(backend, grid, (Center(), Vertex(), Center()); halo=1),
-            z = Field(backend, grid, (Center(), Center(), Vertex()); halo=1),
-        )
-    )
+    return (Pr = Field(backend, grid, Center(); halo=1),
+            τ  = (xx=Field(backend, grid, Center(); halo=1),
+            yy=Field(backend, grid, Center(); halo=1),
+            zz=Field(backend, grid, Center(); halo=1),
+            xy=Field(backend, grid, (Vertex(), Vertex(), Center()); halo=0),
+            xz=Field(backend, grid, (Vertex(), Center(), Vertex()); halo=0),
+            yz=Field(backend, grid, (Center(), Vertex(), Vertex()); halo=0)),
+            V  = (x=Field(backend, grid, (Vertex(), Center(), Center()); halo=1),
+            y=Field(backend, grid, (Center(), Vertex(), Center()); halo=1),
+            z=Field(backend, grid, (Center(), Center(), Vertex()); halo=1)))
 end
 
-
-
-function IsothermalFullStokesModel(; backend, grid, boundary_conditions, physics=nothing, iter_params, fields=nothing, other_fields=nothing)
+function IsothermalFullStokesModel(; arch, grid, boundary_conditions, physics=nothing, iter_params, fields=nothing, other_fields=nothing)
     if isnothing(fields)
-        mechanic_fields = make_fields_mechanics(backend, grid)
-        rheology_fields = (η = Field(backend, grid, Center(); halo=1),)
+        mechanic_fields = make_fields_mechanics(backend(arch), grid)
+        rheology_fields = (η=Field(backend(arch), grid, Center(); halo=1),)
         fields = merge(mechanic_fields, rheology_fields)
     end
 
@@ -66,47 +57,38 @@ function IsothermalFullStokesModel(; backend, grid, boundary_conditions, physics
         physics = default_physics(eltype(grid))
     end
 
-    boundary_conditions = make_field_boundary_conditions(boundary_conditions)
+    boundary_conditions = make_field_boundary_conditions(grid, fields, boundary_conditions)
 
-    return IsothermalFullStokesModel(backend, grid, boundary_conditions, physics, iter_params, fields)
+    return IsothermalFullStokesModel(arch, grid, boundary_conditions, physics, iter_params, fields)
 end
 
 fields(model::IsothermalFullStokesModel) = model.fields
 grid(model::IsothermalFullStokesModel) = model.grid
 
-function advance_iteration!(model::IsothermalFullStokesModel, t, Δt; async = true)
+function advance_iteration!(model::IsothermalFullStokesModel, t, Δt; async=true)
     (; Pr, τ, V, η) = model.fields
     (; η_rel, Δτ) = model.iter_params
     η_rh = model.physics.rheology
     Δ = NamedTuple{(:x, :y, :z)}(spacing(model.grid))
-    nx, ny, nz = size(model.grid)
-    backend = model.backend
-
-    set_bcs!(bcs) = _apply_bcs!(model.backend, model.grid, model.fields, bcs)
 
-    # stress
+    launch!(model.arch, model.grid, update_σ! => (Pr, τ, V, η, Δτ, Δ);
+            location=Vertex(), boundary_conditions=model.boundary_conditions.stress)
 
-    # launch!(arch, grid, update_σ!, Pr, τ, V, η, Δτ, Δ)
+    launch!(model.arch, model.grid, update_V! => (V, Pr, τ, η, Δτ, Δ);
+            location=Vertex(), boundary_conditions=model.boundary_conditions.velocity)
 
-    update_σ!(backend, 256, (nx+1, ny+1, nz+1))(Pr, τ, V, η, Δτ, Δ)
-    set_bcs!(model.boundary_conditions.stress)
-    # velocity
-
-    # launch!(arch, grid, (update_res_V! => (rV, V, Pr, τ, η, Δτ, Δ), update_V! => (V, rV, dt)); exchangers = exchangers.velocity, boundary_conditions = boundary_conditions.velocity)
-    update_V!(backend, 256, (nx+1, ny+1, nz+1))(V, Pr, τ, η, Δτ, Δ)
-    set_bcs!(model.boundary_conditions.velocity)
     # rheology
-    update_η!(backend, 256, (nx, ny, nz))(η, η_rh, η_rel, model.grid, model.fields)
+    launch!(model.arch, grid, update_η! => (η, η_rh, η_rel, model.grid, model.fields); location=Center())
     extrapolate!(η)
 
-    async || synchronize(backend)
+    async || synchronize(backend(model.arch))
     return
 end
 
 function advance_timestep!(model::IsothermalFullStokesModel, t, Δt)
     # TODO
-    
+
     return
 end
 
-end
+end