From 1ce659922fbeeb9c64f863c6d4e420eb3caf1acd Mon Sep 17 00:00:00 2001 From: Ivan Utkin Date: Mon, 13 Nov 2023 01:51:02 +0100 Subject: [PATCH] WIP Stokes --- scripts_future_API/tm_stokes_wip.jl | 92 ++++++++----------- .../isothermal/boundary_conditions.jl | 85 ++++++++--------- .../full_stokes/isothermal/isothermal.jl | 80 +++++++--------- 3 files changed, 108 insertions(+), 149 deletions(-) diff --git a/scripts_future_API/tm_stokes_wip.jl b/scripts_future_API/tm_stokes_wip.jl index d532e2ec..db0e37be 100644 --- a/scripts_future_API/tm_stokes_wip.jl +++ b/scripts_future_API/tm_stokes_wip.jl @@ -6,6 +6,8 @@ using FastIce.BoundaryConditions using FastIce.Models.FullStokes.Isothermal using FastIce.Physics +const VBC = BoundaryCondition{Velocity} + using KernelAbstractions using GLMakie @@ -13,74 +15,58 @@ 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 diff --git a/src/Models/full_stokes/isothermal/boundary_conditions.jl b/src/Models/full_stokes/isothermal/boundary_conditions.jl index 926807fc..becd0cda 100644 --- a/src/Models/full_stokes/isothermal/boundary_conditions.jl +++ b/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 diff --git a/src/Models/full_stokes/isothermal/isothermal.jl b/src/Models/full_stokes/isothermal/isothermal.jl index 20ead3e2..28e98499 100644 --- a/src/Models/full_stokes/isothermal/isothermal.jl +++ b/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 \ No newline at end of file +end