Level sets and volume fractions using Chmy

.JuliaFormatter.toml

@@ -1,5 +1,5 @@
 style = "yas"
-margin = 140
+margin = 165
 align_assignment = true
 align_conditional = true
 whitespace_ops_in_indices = false

Project.toml

@@ -21,7 +21,7 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Adapt = "4"
 ElasticArrays = "1"
 GeometryBasics = "0.4"
-HDF5 = "0.17"
+HDF5 = "0.16, 0.17"
 KernelAbstractions = "0.9"
 LightXML = "0.9"
 MPI = "0.20"

scripts_future_API/Project.toml

@@ -1,8 +1,6 @@
 [deps]
-CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+Chmy = "33a72cf0-4690-46d7-b987-06506c2248b9"
 FastIce = "e0de9f13-a007-490e-b696-b07d031015ca"
-FastIceTools = "14532042-4700-46e0-8dec-55d517bc1ae0"
-FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"

scripts_future_API/generate_synthetic.jl

@@ -0,0 +1,29 @@
+using FastIce.Geometries
+
+nx, ny     = 128, 128
+lx, ly     = 5.0, 5.0
+zmin, zmax = 0.0, 1.0
+synth_dem  = make_synthetic(nx, ny, lx, ly, zmin, zmax, :turtle)
+
+@views function visme(synth_dem)
+    x = LinRange(synth_dem.domain.xmin, synth_dem.domain.xmax, nx + 1)
+    y = LinRange(synth_dem.domain.ymin, synth_dem.domain.ymax, ny + 1)
+    bed  = synth_dem.z_bed
+    surf = synth_dem.z_surf
+
+    surf2 = copy(surf)
+    surf2[surf.<bed] .= NaN
+
+    fig = Figure(; size=(1000, 400), fontsize=22)
+    ax1 = Axis3(fig[1, 1]; aspect=(2, 2, 1), azimuth=-π / 8, elevation=π / 5)
+    ax2 = Axis(fig[1, 2]; aspect=DataAspect())
+
+    surface!(ax1, x, y, bed; colormap=:turbo)
+    surface!(ax1, x, y, surf2; colormap=:turbo)
+    plot!(ax2, x, bed[:, ceil(Int, length(y) / 2)])
+    plot!(ax2, x, surf2[:, ceil(Int, length(y) / 2)])
+
+    return display(fig)
+end
+
+visme(synth_dem)

scripts_future_API/test_levelsets.jl

@@ -1,78 +1,78 @@
-using CUDA
-using FileIO
-using FastIceTools
-using JLD2
-using KernelAbstractions
-using HDF5
+# using CUDA
+# using FileIO
+# using FastIceTools
+# using JLD2
+# using KernelAbstractions
+# using HDF5
 
-using FastIce.Grids
-using FastIce.Fields
-using FastIce.LevelSets
-using FastIce.Architectures
-using FastIce.Writers
+# using FastIce.Grids
+# using FastIce.Fields
+# using FastIce.LevelSets
+# using FastIce.Architectures
+# using FastIce.Writers
 
-vavilov_path = "../data/Vavilov/vavilov.jld2"
-synthetic_data = "../data/synthetic.jld2"
+# vavilov_path = "../data/Vavilov/vavilov.jld2"
+# synthetic_data = "../data/synthetic.jld2"
 
-# Select backend (CPU(), CUDABackend())
-backend = CUDABackend()
-arch = Architecture(backend)
+# # Select backend (CPU(), CUDABackend())
+# backend = CUDABackend()
+# arch = Architecture(backend)
 
-"""
-    load_dem_on_GPU(path::String, arch::Architecture)
+# """
+#     load_dem_on_GPU(path::String, arch::Architecture)
 
-Load digital elevation map of surface and bedrock from (jld2) file, set dimensions of simulation, 
-initiate grids, copy data on gpu.
-"""
-function load_dem_on_GPU(path::String, arch::Architecture)
-    data     = load(path)
-    x        = data["DataElevation"].x
-    y        = data["DataElevation"].y
-    R        = data["DataElevation"].rotation
-    z_surf   = permutedims(data["DataElevation"].z_surf)
-    z_bed    = permutedims(data["DataElevation"].z_bed)
-    domain   = data["DataElevation"].domain
-    nx       = length(x) - 1
-    ny       = length(y) - 1
-    nz       = 10
-    z        = LinRange(domain.zmin, domain.zmax, nz)
-    Ψ_grid   = CartesianGrid(; origin=(0.0, 0.0, 0.0), extent=(1.0, 1.0, 1.0), size=(nx, ny, nz))
-    Ψ        = Field(arch, Ψ_grid, Vertex())
-    dem_grid = CartesianGrid(; origin=(0.0, 0.0), extent=(1.0, 1.0), size=(nx, ny))
-    dem_bed  = Field(arch, dem_grid, Vertex())
-    dem_surf = Field(arch, dem_grid, Vertex())
-    set!(dem_bed, z_bed)
-    set!(dem_surf, z_surf)
-    return Ψ, dem_surf, dem_bed, dem_grid, Ψ_grid
-end
+# Load digital elevation map of surface and bedrock from (jld2) file, set dimensions of simulation,
+# initiate grids, copy data on gpu.
+# """
+# function load_dem_on_GPU(path::String, arch::Architecture)
+#     data     = load(path)
+#     x        = data["DataElevation"].x
+#     y        = data["DataElevation"].y
+#     R        = data["DataElevation"].rotation
+#     z_surf   = permutedims(data["DataElevation"].z_surf)
+#     z_bed    = permutedims(data["DataElevation"].z_bed)
+#     domain   = data["DataElevation"].domain
+#     nx       = length(x) - 1
+#     ny       = length(y) - 1
+#     nz       = 10
+#     z        = LinRange(domain.zmin, domain.zmax, nz)
+#     Ψ_grid   = CartesianGrid(; origin=(0.0, 0.0, 0.0), extent=(1.0, 1.0, 1.0), size=(nx, ny, nz))
+#     Ψ        = Field(arch, Ψ_grid, Vertex())
+#     dem_grid = CartesianGrid(; origin=(0.0, 0.0), extent=(1.0, 1.0), size=(nx, ny))
+#     dem_bed  = Field(arch, dem_grid, Vertex())
+#     dem_surf = Field(arch, dem_grid, Vertex())
+#     set!(dem_bed, z_bed)
+#     set!(dem_surf, z_surf)
+#     return Ψ, dem_surf, dem_bed, dem_grid, Ψ_grid
+# end
 
-function load_synth_dem_on_GPU(path::String, arch::Architecture)
-    data   = load(path)
-    z_surf = data["z_surf"]
-    z_bed  = data["z_bed"]
-    nx     = size(z_bed)[1] - 1
-    ny     = size(z_bed)[2] - 1
-    nz     = 100
-    lz = maximum(z_bed) - minimum(z_bed)
-    Ψ_grid   = CartesianGrid(; origin=(0.0, 0.0, minimum(z_bed)), extent=(lz, lz, lz), size=(nx, ny, nz))
-    Ψ        = Field(arch, Ψ_grid, Vertex())
-    dem_grid = CartesianGrid(; origin=(0.0, 0.0), extent=(lz, lz), size=(nx, ny))
-    dem_bed  = Field(arch, dem_grid, Vertex())
-    dem_surf = Field(arch, dem_grid, Vertex())
-    set!(dem_bed, z_bed)
-    set!(dem_surf, z_surf)
-    return Ψ, dem_surf, dem_bed, dem_grid, Ψ_grid
-end
+# function load_synth_dem_on_GPU(path::String, arch::Architecture)
+#     data   = load(path)
+#     z_surf = data["z_surf"]
+#     z_bed  = data["z_bed"]
+#     nx     = size(z_bed)[1] - 1
+#     ny     = size(z_bed)[2] - 1
+#     nz     = 100
+#     lz = maximum(z_bed) - minimum(z_bed)
+#     Ψ_grid   = CartesianGrid(; origin=(0.0, 0.0, minimum(z_bed)), extent=(lz, lz, lz), size=(nx, ny, nz))
+#     Ψ        = Field(arch, Ψ_grid, Vertex())
+#     dem_grid = CartesianGrid(; origin=(0.0, 0.0), extent=(lz, lz), size=(nx, ny))
+#     dem_bed  = Field(arch, dem_grid, Vertex())
+#     dem_surf = Field(arch, dem_grid, Vertex())
+#     set!(dem_bed, z_bed)
+#     set!(dem_surf, z_surf)
+#     return Ψ, dem_surf, dem_bed, dem_grid, Ψ_grid
+# end
 
-# Ψ, dem_surf, dem_bed, dem_grid, Ψ_grid = load_dem_on_GPU(vavilov_path, arch);
-Ψ, dem_surf, dem_bed, dem_grid, Ψ_grid = load_synth_dem_on_GPU(synthetic_data, arch);
+# # Ψ, dem_surf, dem_bed, dem_grid, Ψ_grid = load_dem_on_GPU(vavilov_path, arch);
+# Ψ, dem_surf, dem_bed, dem_grid, Ψ_grid = load_synth_dem_on_GPU(synthetic_data, arch);
 
-compute_level_set_from_dem!(arch, Ψ, dem_bed, dem_grid, Ψ_grid)
+# compute_level_set_from_dem!(arch, Ψ, dem_bed, dem_grid, Ψ_grid)
 
-jldopen(synthetic_data, "a+") do file
-    file["level_sets"] = Array(interior(Ψ))
-end
+# jldopen(synthetic_data, "a+") do file
+#     file["level_sets"] = Array(interior(Ψ))
+# end
 
-jldopen(vavilov_path, "a+") do file
-    file["level_sets"] = Array(interior(Ψ))
-end
+# jldopen(vavilov_path, "a+") do file
+#     file["level_sets"] = Array(interior(Ψ))
+# end

scripts_future_API/test_levelsets_volumefractions_mpi.jl

@@ -0,0 +1,125 @@
+using Chmy.Architectures
+using Chmy.Grids
+using Chmy.Fields
+using Chmy.KernelLaunch
+
+using FastIce.LevelSets
+
+using KernelAbstractions
+using CairoMakie
+# using CUDA
+# using AMDGPU
+# AMDGPU.allowscalar(false)
+
+# Select backend
+backend = CPU()
+# backend = CUDABackend()
+# backend = ROCBackend()
+
+using Chmy.Distributed
+using MPI
+MPI.Init()
+
+@views function make_dem(backend=CPU(); nx, ny, lx, ly, zmin, zmax, amp, ω, tanβ, el, gl)
+    arch_2d = Arch(backend)
+    grid_2d = UniformGrid(arch_2d; origin=(-lx/2, -ly/2), extent=(lx, ly), dims=(nx, ny))
+
+    lz = zmax - zmin
+
+    # type = :turtle
+    generate_surf(x, y, gl, lx, ly) = gl * (1.0 - ((1.7 * x / lx + 0.22)^2 + (0.5 * y / ly)^2))
+    generate_bed(x, y, amp, ω, tanβ, el, lx, ly, lz) = lz * (amp * sin(ω * x / lx) * sin(ω * y /ly) + tanβ * x / lx + el + (1.5 * y / ly)^2)
+
+    surf = FunctionField(generate_surf, grid_2d, Vertex(); parameters=(; gl, lx, ly))
+    bed  = FunctionField(generate_bed, grid_2d, Vertex(); parameters=(; amp, ω, tanβ, el, lx, ly, lz))
+
+    return (; bed, surf, grid_2d, lz)
+end
+
+@views function main(backend=CPU())
+    lx, ly = 5.0, 5.0
+    zmin, zmax = 0.0, 1.0
+    nx = ny = 126
+    # synthetic topo
+    amp  = 0.1
+    ω    = 10π
+    tanβ = tan(-π/6)
+    el   = 0.35
+    gl   = 0.9
+
+    dem_data = make_dem(backend; nx, ny, lx, ly, zmin, zmax, amp, ω, tanβ, el, gl)
+
+    run_simulation(backend; nxyz_l=(nx, ny), dem_data)
+    return
+end
+
+@views function run_simulation(backend=CPU(); nxyz_l=(126, 126), dem_data)
+    arch    = Arch(backend, MPI.COMM_WORLD, (0, 0, 0))
+    topo    = topology(arch)
+    me      = global_rank(topo)
+    # geometry
+    lx, ly, lz = extent(dem_data.grid_2d, Vertex())..., dem_data.lz
+    nx, ny     = nxyz_l[1:2]
+    nz         = (length(nxyz_l) < 3) ? ceil(Int, lz / lx * nxyz_l[1]) : nxyz_l[3]
+    @info "nz = $nz (lz = $lz)"
+    dims_l     = (nx, ny, nz)
+    dims_g     = dims_l .* dims(topo)
+    grid       = UniformGrid(arch; origin=(-lx/2, -ly/2, 0.0), extent=(lx, ly, lz), dims=dims_g)
+    launch     = Launcher(arch, grid; outer_width=(16, 8, 4))
+    # init fields
+    Ψ = (na=Field(backend, grid, Vertex()),
+         ns=Field(backend, grid, Vertex()))
+    wt = (na=volfrac_field(backend, grid),
+          ns=volfrac_field(backend, grid))
+    # comput level set
+    compute_levelset_from_dem!(arch, launch, Ψ.na, dem_data.surf, dem_data.grid_2d, grid)
+    compute_levelset_from_dem!(arch, launch, Ψ.ns, dem_data.bed, dem_data.grid_2d, grid)
+    invert_levelset!(arch, launch, Ψ.ns, grid) # invert level set to set what's below the DEM surface as inside
+    # volume fractions
+    for phase in eachindex(Ψ)
+        compute_volfrac_from_levelset!(arch, launch, wt[phase], Ψ[phase], grid)
+    end
+
+    # compute physics or else
+
+    # postprocessing
+    wt_na_c = (me == 0) ? KernelAbstractions.zeros(CPU(), Float64, size(interior(wt.na.c)) .* dims(topo)) : nothing
+    wt_ns_c = (me == 0) ? KernelAbstractions.zeros(CPU(), Float64, size(interior(wt.ns.c)) .* dims(topo)) : nothing
+    gather!(arch, wt_na_c, wt.na.c)
+    gather!(arch, wt_ns_c, wt.ns.c)
+    # visualise
+    if me == 0
+        dem_bed  = Array(interior(dem_data.bed))
+        dem_surf = Array(interior(dem_data.surf))
+        dem_surf[dem_surf .< dem_bed] .= NaN
+        slx = ceil(Int, size(wt_na_c, 1) / 2) # for visu
+        sly = ceil(Int, size(wt_na_c, 2) / 2) # for visu
+        x_g = LinRange(-lx / 2, lx / 2, size(dem_bed, 1))
+        y_g = LinRange(-ly / 2, ly / 2, size(dem_bed, 2))
+
+        fig = Figure(; size=(1000, 800), fontsize=22)
+        axs = (ax1 = Axis3(fig[1, 1][1, 1]; aspect=(2, 2, 1), azimuth=-π / 8, elevation=π / 5),
+               ax2 = Axis(fig[1, 2]; aspect=DataAspect()),
+               ax3 = Axis(fig[2, 1]; aspect=DataAspect()),
+               ax4 = Axis(fig[2, 2]; aspect=DataAspect()),
+               ax5 = Axis(fig[3, 1]; aspect=DataAspect()),
+               ax6 = Axis(fig[3, 2]; aspect=DataAspect()))
+        plt = (p1  = surface!(axs.ax1, x_g, y_g, dem_bed; colormap=:turbo),
+               p1_ = surface!(axs.ax1, x_g, y_g, dem_surf; colormap=:turbo),
+               p2  = plot!(axs.ax2, x_g, dem_bed[:, sly] |> Array),
+               p2_ = plot!(axs.ax2, x_g, dem_surf[:, sly] |> Array),
+               p3  = heatmap!(axs.ax3, wt_na_c[:, sly, :]; colormap=:turbo),
+               p4  = heatmap!(axs.ax4, wt_ns_c[:, sly, :]; colormap=:turbo),
+               p5  = heatmap!(axs.ax5, wt_na_c[slx, :, :]; colormap=:turbo),
+               p6  = heatmap!(axs.ax6, wt_ns_c[slx, :, :]; colormap=:turbo))
+        Colorbar(fig[1, 1][1, 2], plt.p1)
+        Colorbar(fig[2, 2][1, 2], plt.p4)
+        display(fig)
+        # save("levset_$nx.png", fig)
+    end
+    return
+end
+
+main(backend)
+
+MPI.Finalize()

src/FastIce.jl

@@ -25,11 +25,12 @@ https://github.com/PTsolvers/FastIce.jl
 greet(; kwargs...) = printstyled(GREETING; kwargs...)
 greet_fast(; kwargs...) = printstyled(GREETING_FAST; kwargs...)
 
-# core modules
+# core modules (included in alphabetical order)
+include("Geometries/Geometries.jl")
+include("LevelSets/LevelSets.jl")
 include("Logging.jl")
 include("Physics.jl")
 include("Writers.jl")
-include("LevelSets/LevelSets.jl")
 
 # ice flow models
 # include("Models/Models.jl")

src/Geometries/Geometries.jl

@@ -0,0 +1,13 @@
+module Geometries
+
+# export AbstractElevation, AABB, DataElevation, SyntheticElevation
+# export domain, rotated_domain, rotation, generate_elevation, extents, center, dilate
+export SyntheticElevation
+export extents
+export make_synthetic
+
+include("elevation_data.jl")
+include("geometry_helpers.jl")
+include("synthetic_geometries.jl")
+
+end # module