From 29f56d22c22066f022dc30e7cd7ff53611ea0a71 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Fri, 31 Jan 2025 07:49:15 -0500 Subject: [PATCH] Add some validatieon cases and use Float64 clock --- src/Advection/weno_reconstruction.jl | 2 + .../nonhydrostatic_model.jl | 2 +- src/Simulations/time_step_wizard.jl | 1 + .../hydrostatic_two_dimensional_turbulence.jl | 117 ++++++++++++++++++ .../two_dimensional_turbulence.jl | 51 ++++++++ 5 files changed, 172 insertions(+), 1 deletion(-) create mode 100644 validation/float_precision_tests/hydrostatic_two_dimensional_turbulence.jl create mode 100644 validation/float_precision_tests/two_dimensional_turbulence.jl diff --git a/src/Advection/weno_reconstruction.jl b/src/Advection/weno_reconstruction.jl index 382535729b..0d09f9409e 100644 --- a/src/Advection/weno_reconstruction.jl +++ b/src/Advection/weno_reconstruction.jl @@ -1,3 +1,5 @@ +import Oceananigans + ##### ##### Weighted Essentially Non-Oscillatory (WENO) advection scheme ##### diff --git a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl index ae986f0ec2..b545fbe3b4 100644 --- a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl +++ b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl @@ -112,7 +112,7 @@ Keyword arguments - `auxiliary_fields`: `NamedTuple` of auxiliary fields. Default: `nothing` """ function NonhydrostaticModel(; grid, - clock = Clock{eltype(grid)}(time = 0), + clock = Clock{Float64}(time = 0), advection = Centered(), buoyancy = nothing, coriolis = nothing, diff --git a/src/Simulations/time_step_wizard.jl b/src/Simulations/time_step_wizard.jl index 4e46608fdd..8ee22550d7 100644 --- a/src/Simulations/time_step_wizard.jl +++ b/src/Simulations/time_step_wizard.jl @@ -1,5 +1,6 @@ using Oceananigans: TurbulenceClosures using Oceananigans.Grids: prettysummary, architecture +import Oceananigans mutable struct TimeStepWizard{FT, C, D} cfl :: FT diff --git a/validation/float_precision_tests/hydrostatic_two_dimensional_turbulence.jl b/validation/float_precision_tests/hydrostatic_two_dimensional_turbulence.jl new file mode 100644 index 0000000000..6e88bdd049 --- /dev/null +++ b/validation/float_precision_tests/hydrostatic_two_dimensional_turbulence.jl @@ -0,0 +1,117 @@ +using Oceananigans +using MPI +using Oceananigans.Models.HydrostaticFreeSurfaceModels: VerticalVorticityField +using Printf +using Statistics +using Oceananigans.BoundaryConditions +using Oceananigans.DistributedComputations +using Random +using JLD2 + +# Run with +# +# ```julia +# mpiexec -n 4 julia --project distributed_hydrostatic_turbulence.jl +# ``` + +function run_simulation(nx, ny, arch; topology = (Periodic, Periodic, Bounded)) + grid = RectilinearGrid(arch; topology, size = (Nx, Ny, 10), extent=(4π, 4π, 0.5), halo=(8, 8, 8)) + + bottom(x, y) = (x > π && x < 3π/2 && y > π/2 && y < 3π/2) ? 1.0 : - grid.Lz - 1.0 + grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom); active_cells_map = true) + + model = HydrostaticFreeSurfaceModel(; grid, + momentum_advection = VectorInvariant(vorticity_scheme=WENO(order=9)), + free_surface = SplitExplicitFreeSurface(grid, substeps=10), + tracer_advection = WENO(), + buoyancy = nothing, + coriolis = FPlane(f = 1), + tracers = :c) + + # Scale seed with rank to avoid symmetry + local_rank = MPI.Comm_rank(arch.communicator) + Random.seed!(1234 * (local_rank + 1)) + + set!(model, u = (x, y, z) -> 1-2rand(), v = (x, y, z) -> 1-2rand()) + + mask(x, y, z) = x > 3π/2 && x < 5π/2 && y > 3π/2 && y < 5π/2 + c = model.tracers.c + + set!(c, mask) + + u, v, _ = model.velocities + # ζ = VerticalVorticityField(model) + η = model.free_surface.η + outputs = merge(model.velocities, model.tracers) + + progress(sim) = @info "Iteration: $(sim.model.clock.iteration), time: $(sim.model.clock.time), Δt: $(sim.Δt)" + simulation = Simulation(model, Δt=0.02, stop_time=100.0) + + wizard = TimeStepWizard(cfl = 0.2, max_change = 1.1) + + simulation.callbacks[:progress] = Callback(progress, IterationInterval(100)) + simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10)) + + filepath = "mpi_hydrostatic_turbulence_rank$(local_rank)" + simulation.output_writers[:fields] = + JLD2OutputWriter(model, outputs, filename=filepath, schedule=TimeInterval(0.1), + overwrite_existing=true) + + run!(simulation) + + MPI.Barrier(arch.communicator) +end + +Nx = 32 +Ny = 32 + +arch = Distributed(CPU(), partition = Partition(2, 2)) + +# Run the simulation +run_simulation(Nx, Ny, arch) + +# Visualize the plane +# Produce a video for variable `var` +try + using GLMakie + + function visualize_simulation(var) + iter = Observable(1) + + v = Vector(undef, 4) + V = Vector(undef, 4) + x = Vector(undef, 4) + y = Vector(undef, 4) + + for r in 1:4 + v[r] = FieldTimeSeries("mpi_hydrostatic_turbulence_rank$(r-1).jld2", var) + nx, ny, _ = size(v[r]) + V[r] = @lift(interior(v[r][$iter], 1:nx, 1:ny, 1)) + + x[r] = xnodes(v[r]) + y[r] = ynodes(v[r]) + end + + fig = Figure() + ax = Axis(fig[1, 1]) + for r in 1:4 + heatmap!(ax, x[r], y[r], V[r], colorrange = (-1.0, 1.0)) + end + + GLMakie.record(fig, "hydrostatic_test_" * var * ".mp4", 1:length(v[1].times), framerate = 11) do i + @info "step $i"; + iter[] = i; + end + end + + if MPI.Comm_rank(MPI.COMM_WORLD) == 0 + visualize_simulation("u") + visualize_simulation("v") + visualize_simulation("c") + end +catch err + @info err +end + +MPI.Barrier(arch.communicator) + diff --git a/validation/float_precision_tests/two_dimensional_turbulence.jl b/validation/float_precision_tests/two_dimensional_turbulence.jl new file mode 100644 index 0000000000..5f8107b311 --- /dev/null +++ b/validation/float_precision_tests/two_dimensional_turbulence.jl @@ -0,0 +1,51 @@ +using Oceananigans +using Statistics +using Printf + +Oceananigans.defaults.FloatType = Float64 +Nx = Ny = 512 + +grid = RectilinearGrid(size = (Nx, Ny), + halo = (7, 7), + extent = (Nx, Ny), + topology = (Periodic, Periodic, Flat)) + +clock = Clock{Float64}(time=0) +model = NonhydrostaticModel(; grid, clock, advection=WENO(order=9)) + +ui(x, y) = randn() +set!(model, u=ui, v=ui) + +Δx = minimum_xspacing(grid) +simulation = Simulation(model; Δt=1e-2, stop_iteration=200) +#conjure_time_step_wizard!(simulation, cfl=0.5) + +u, v, w = model.velocities +eop = @at (Center, Center, Center) (u^2 + v^2) / 2 +e = Field(eop) + +wall_clock = Ref(time_ns()) +function progress(sim) + compute!(e) + avg_e = mean(e) + elapsed = 1e-9 * (time_ns() - wall_clock[]) + msg = @sprintf("Iter: %s, time: %.2f, wall clock: %s, ⟨e⟩: %.2e", + iteration(sim), time(sim), prettytime(elapsed), avg_e) + wall_clock[] = time_ns() + @info msg + return nothing +end + +add_callback!(simulation, progress, IterationInterval(10)) + +FT = Oceananigans.defaults.FloatType +outputs = merge(model.velocities, (; e)) +ow = JLD2OutputWriter(model, outputs; + filename = "float_point_test_$(Nx)_$FT.jld2", + schedule = TimeInterval(0.1), + overwrite_existing = true) + +simulation.output_writers[:jld2] = ow + +run!(simulation) +