diff --git a/.gitignore b/.gitignore index e5d062c..2f3f6d3 100644 --- a/.gitignore +++ b/.gitignore @@ -9,4 +9,5 @@ parameter.txt # It records a fixed state of all packages used by the project. As such, it should not be # committed for packages, but should be committed for applications that require a static # environment. -Manifest.toml \ No newline at end of file +Manifest.toml + diff --git a/src/ShallowWaters.jl b/src/ShallowWaters.jl index 0d01ab3..fda5303 100644 --- a/src/ShallowWaters.jl +++ b/src/ShallowWaters.jl @@ -8,9 +8,9 @@ module ShallowWaters include("grid.jl") include("constants.jl") include("forcing.jl") + include("preallocate.jl") include("model_setup.jl") include("initial_conditions.jl") - include("preallocate.jl") include("time_integration.jl") include("ghost_points.jl") diff --git a/src/default_parameters.jl b/src/default_parameters.jl index 4ef41c3..3f54fe3 100644 --- a/src/default_parameters.jl +++ b/src/default_parameters.jl @@ -70,6 +70,15 @@ α::Real=2. # lateral boundary condition parameter # 0 free-slip, 0<α<2 partial-slip, 2 no-slip + # PARAMETERS FOR ADJOINT METHOD + data_steps::StepRange{Int,Int} = 0:1:0 # Timesteps where data exists + data::Array{Float32, 1} = [0.] # model data + J::Float64 = 0. # Placeholder for cost function evaluation + j::Int = 1 # For keeping track of the entry in data + + # CHECKPOINTING VARIABLES + i::Int = 0 # Placeholder for current timestep, needed for Checkpointing.jl + # MOMENTUM ADVECTION OPTIONS adv_scheme::String="ArakawaHsu" # "Sadourny" or "ArakawaHsu" dynamics::String="nonlinear" # "linear" or "nonlinear" @@ -273,4 +282,4 @@ Creates a Parameter struct with following options and default values run_id::Int=-1 # Output with a specific run id init_interpolation::Bool=true # Interpolate the initial conditions in case grids don't match? """ -Parameter +Parameter \ No newline at end of file diff --git a/src/ghost_points.jl b/src/ghost_points.jl index 162b584..11c38ff 100644 --- a/src/ghost_points.jl +++ b/src/ghost_points.jl @@ -32,6 +32,42 @@ function add_halo( u::Array{T,2}, return u,v,η,sst end +""" Extends the matrices u,v,η,sst with a halo of ghost points for boundary conditions.""" +function add_halo( u::Array{T,2}, + v::Array{T,2}, + η::Array{T,2}, + sst::Array{T,2}, + G::Grid, + P::Parameter, + C::Constants) where {T<:AbstractFloat} + + @unpack nx,ny,nux,nuy,nvx,nvy = G + @unpack halo,haloη,halosstx,halossty = G + + # Add zeros to satisfy kinematic boundary conditions + u = cat(zeros(T,halo,nuy),u,zeros(T,halo,nuy),dims=1) + u = cat(zeros(T,nux+2*halo,halo),u,zeros(T,nux+2*halo,halo),dims=2) + + v = cat(zeros(T,halo,nvy),v,zeros(T,halo,nvy),dims=1) + v = cat(zeros(T,nvx+2*halo,halo),v,zeros(T,nvx+2*halo,halo),dims=2) + + η = cat(zeros(T,haloη,ny),η,zeros(T,haloη,ny),dims=1) + η = cat(zeros(T,nx+2*haloη,haloη),η,zeros(T,nx+2*haloη,haloη),dims=2) + + sst = cat(zeros(T,halosstx,ny),sst,zeros(T,halosstx,ny),dims=1) + sst = cat(zeros(T,nx+2*halosstx,halossty),sst,zeros(T,nx+2*halosstx,halossty),dims=2) + + # SCALING + @unpack scale,scale_sst = C + u *= scale + v *= scale + sst *= scale_sst + + ghost_points!(u,v,η,P,C) + ghost_points_sst!(sst,P,G) + return u,v,η,sst +end + """Cut off the halo from the prognostic variables.""" function remove_halo( u::Array{T,2}, v::Array{T,2}, @@ -51,6 +87,94 @@ function remove_halo( u::Array{T,2}, return ucut,vcut,ηcut,sstcut end +"""Cut off the halo from the prognostic variables.""" +function remove_halo( u::Array{T,2}, + v::Array{T,2}, + η::Array{T,2}, + sst::Array{T,2}, + G::Grid, + C::Constants + ) where {T<:AbstractFloat} + + @unpack halo,haloη,halosstx,halossty = G + @unpack scale_inv,scale_sst = C + + # undo scaling as well + @views ucut = scale_inv*u[halo+1:end-halo,halo+1:end-halo] + @views vcut = scale_inv*v[halo+1:end-halo,halo+1:end-halo] + @views ηcut = η[haloη+1:end-haloη,haloη+1:end-haloη] + @views sstcut = sst[halosstx+1:end-halosstx,halossty+1:end-halossty]/scale_sst + + return ucut,vcut,ηcut,sstcut +end + +"""Decide on boundary condition P.bc which ghost point function to execute.""" +function ghost_points!( u::AbstractMatrix, + v::AbstractMatrix, + η::AbstractMatrix, + P::Parameter, + C::Constants) + + @unpack bc,Tcomm = P + @unpack one_minus_α = C + + if bc == "periodic" + ghost_points_u_periodic!(Tcomm,u,one_minus_α) + ghost_points_v_periodic!(Tcomm,v) + ghost_points_η_periodic!(Tcomm,η) + else + ghost_points_u_nonperiodic!(u,one_minus_α) + ghost_points_v_nonperiodic!(v,one_minus_α) + ghost_points_η_nonperiodic!(η) + end +end + +"""Decide on boundary condition P.bc which ghost point function to execute.""" +function ghost_points_uv!( u::AbstractMatrix, + v::AbstractMatrix, + P::Parameter, + C::Constants) + + @unpack bc,Tcomm = P + @unpack one_minus_α = C + + if bc == "periodic" + ghost_points_u_periodic!(Tcomm,u,one_minus_α) + ghost_points_v_periodic!(Tcomm,v) + else + ghost_points_u_nonperiodic!(u,one_minus_α) + ghost_points_v_nonperiodic!(v,one_minus_α) + end +end + +"""Decide on boundary condition P.bc which ghost point function to execute.""" +function ghost_points_η!( η::AbstractMatrix, + P::Parameter) + + @unpack bc, Tcomm = P + + if bc == "periodic" + ghost_points_η_periodic!(Tcomm,η) + else + ghost_points_η_nonperiodic!(η) + end +end + +"""Decide on boundary condition P.bc which ghost point function to execute.""" +function ghost_points_sst!( sst::AbstractMatrix, + P::Parameter, + G::Grid) + + @unpack bc,Tcomm = P + @unpack halosstx,halossty = G + + if bc == "periodic" + ghost_points_sst_periodic!(Tcomm,sst,halosstx,halossty) + else + ghost_points_sst_nonperiodic!(sst,halosstx,halossty) + end +end + """ Copy ghost points for u from inside to the halo in the nonperiodic case. """ function ghost_points_u_nonperiodic!(u::AbstractMatrix{T},one_minus_α::T) where T n,m = size(u) diff --git a/src/initial_conditions.jl b/src/initial_conditions.jl index 5e25269..ff48952 100644 --- a/src/initial_conditions.jl +++ b/src/initial_conditions.jl @@ -1,15 +1,3 @@ -""" - P = ProgVars{T}(u,v,η,sst) - -Struct containing the prognostic variables u,v,η and sst. -""" -struct PrognosticVars{T<:AbstractFloat} - u::Array{T,2} # u-velocity - v::Array{T,2} # v-velocity - η::Array{T,2} # sea surface height / interface displacement - sst::Array{T,2} # tracer / sea surface temperature -end - """Zero generator function for Grid G as argument.""" function PrognosticVars{T}(G::Grid) where {T<:AbstractFloat} @unpack nux,nuy,nvx,nvy,nx,ny = G @@ -23,12 +11,12 @@ function PrognosticVars{T}(G::Grid) where {T<:AbstractFloat} return PrognosticVars{T}(u,v,η,sst) end -function initial_conditions(::Type{T},S::ModelSetup) where {T<:AbstractFloat} +function initial_conditions(::Type{T},G::Grid,P::Parameter,C::Constants) where {T<:AbstractFloat} ## PROGNOSTIC VARIABLES U,V,η - @unpack nux,nuy,nvx,nvy,nx,ny = S.grid - @unpack initial_cond = S.parameters - @unpack Tini = S.parameters + @unpack nux,nuy,nvx,nvy,nx,ny = G + @unpack initial_cond = P + @unpack Tini = P if initial_cond == "rest" @@ -38,14 +26,14 @@ function initial_conditions(::Type{T},S::ModelSetup) where {T<:AbstractFloat} elseif initial_cond == "ncfile" - @unpack initpath,init_run_id,init_starti = S.parameters - @unpack init_interpolation = S.parameters - @unpack nx,ny = S.grid + @unpack initpath,init_run_id,init_starti = P + @unpack init_interpolation = P + @unpack nx,ny = G - inirunpath = joinpath(initpath,"run"*@sprintf("%04d",init_run_id)) + # inirunpath = joinpath(initpath,"run"*@sprintf("%04d",init_run_id)) # take starti time step from existing netcdf files - ncstring = joinpath(inirunpath,"u.nc") + ncstring = joinpath(initpath,"u.nc") ncu = NetCDF.open(ncstring) if init_starti == -1 # replace -1 with length of time dimension @@ -54,10 +42,10 @@ function initial_conditions(::Type{T},S::ModelSetup) where {T<:AbstractFloat} u = ncu.vars["u"][:,:,init_starti] - ncv = NetCDF.open(joinpath(inirunpath,"v.nc")) + ncv = NetCDF.open(joinpath(initpath,"v.nc")) v = ncv.vars["v"][:,:,init_starti] - ncη = NetCDF.open(joinpath(inirunpath,"eta.nc")) + ncη = NetCDF.open(joinpath(initpath,"eta.nc")) η = ncη.vars["eta"][:,:,init_starti] # remove singleton time dimension @@ -118,10 +106,10 @@ function initial_conditions(::Type{T},S::ModelSetup) where {T<:AbstractFloat} ## SST - @unpack SSTmin, SSTmax, SSTw, SSTϕ = S.parameters - @unpack SSTwaves_nx,SSTwaves_ny,SSTwaves_p = S.parameters - @unpack sst_initial,scale = S.parameters - @unpack x_T,y_T,Lx,Ly = S.grid + @unpack SSTmin, SSTmax, SSTw, SSTϕ = P + @unpack SSTwaves_nx,SSTwaves_ny,SSTwaves_p = P + @unpack sst_initial,scale = P + @unpack x_T,y_T,Lx,Ly = G xx_T,yy_T = meshgrid(x_T,y_T) @@ -136,7 +124,7 @@ function initial_conditions(::Type{T},S::ModelSetup) where {T<:AbstractFloat} elseif sst_initial == "flat" sst = fill(SSTmin,size(xx_T)) elseif sst_initial == "rect" - @unpack sst_rect_coords = S.parameters + @unpack sst_rect_coords = P x0,x1,y0,y1 = sst_rect_coords sst = fill(SSTmin,size(xx_T)) @@ -159,7 +147,7 @@ function initial_conditions(::Type{T},S::ModelSetup) where {T<:AbstractFloat} η = T.(Tini.(η)) #TODO SST INTERPOLATION - u,v,η,sst = add_halo(u,v,η,sst,S) + u,v,η,sst = add_halo(u,v,η,sst,G,P,C) return PrognosticVars{T}(u,v,η,sst) end diff --git a/src/model_setup.jl b/src/model_setup.jl index 7cbfb1b..36ceeaa 100644 --- a/src/model_setup.jl +++ b/src/model_setup.jl @@ -1,6 +1,30 @@ -struct ModelSetup{T<:AbstractFloat,Tprog<:AbstractFloat} +mutable struct PrognosticVars{T<:AbstractFloat} + u::Array{T,2} # u-velocity + v::Array{T,2} # v-velocity + η::Array{T,2} # sea surface height / interface displacement + sst::Array{T,2} # tracer / sea surface temperature +end + +struct DiagnosticVars{T,Tprog} + RungeKutta::RungeKuttaVars{Tprog} + Tendencies::TendencyVars{Tprog} + VolumeFluxes::VolumeFluxVars{T} + Vorticity::VorticityVars{T} + Bernoulli::BernoulliVars{T} + Bottomdrag::BottomdragVars{T} + ArakawaHsu::ArakawaHsuVars{T} + Laplace::LaplaceVars{T} + Smagorinsky::SmagorinskyVars{T} + SemiLagrange::SemiLagrangeVars{T} + PrognosticVarsRHS::PrognosticVars{T} # low precision version +end + +mutable struct ModelSetup{T<:AbstractFloat,Tprog<:AbstractFloat} parameters::Parameter grid::Grid{T,Tprog} constants::Constants{T,Tprog} forcing::Forcing{T} -end + Prog::PrognosticVars{Tprog} + Diag::DiagnosticVars{T, Tprog} + t::Int # SW: I believe this has something to do with Checkpointing, need to verify +end \ No newline at end of file diff --git a/src/output.jl b/src/output.jl index a30589d..af2c299 100644 --- a/src/output.jl +++ b/src/output.jl @@ -192,38 +192,29 @@ function get_run_id_path(S::ModelSetup) @unpack output,outpath,get_id_mode = S.parameters if output - runlist = filter(x->startswith(x,"run"),readdir(outpath)) - existing_runs = [parse(Int,id[4:end]) for id in runlist] + + pattern = r"run_\d\d\d\d" # run_???? in regex + runlist = filter(x->startswith(x,pattern),readdir(outpath)) + runlist = filter(x->endswith( x,pattern),runlist) + existing_runs = [parse(Int,id[5:end]) for id in runlist] + + # get the run id from existing folders if length(existing_runs) == 0 # if no runfolder exists yet - runpath = joinpath(outpath,"run0000") - mkdir(runpath) - return 0,runpath - else # create next folder - if get_id_mode == "fill" # find the smallest gap in runfolders - run_id = gap(existing_runs) - runpath = joinpath(outpath,"run"*@sprintf("%04d",run_id)) - mkdir(runpath) - - elseif get_id_mode == "specific" # specify the run_id as input argument - @unpack run_id = S.parameters - runpath = joinpath(outpath,"run"*@sprintf("%04d",run_id)) - try # create folder if not existent - mkdir(runpath) - catch # else rm folder and create new one - rm(runpath,recursive=true) - mkdir(runpath) - end - - elseif get_id_mode == "continue" # find largest folder and count one up - run_id = maximum(existing_runs)+1 - runpath = joinpath(outpath,"run"*@sprintf("%04d",run_id)) - mkdir(runpath) - else - throw(error("Order '$get_id_mode' is not valid for get_run_id_path(), choose continue or fill.")) - end - return run_id,runpath + run_id = 1 # start with run_0001 + else + run_id = maximum(existing_runs)+1 # next run gets id +1 end + + id = @sprintf("%04d",run_id) + + run_id2 = string("run_",id) + run_path = joinpath(outpath,run_id2) + @assert !(run_id2 in readdir(outpath)) "Run folder $run_path already exists." + mkdir(run_path) # actually create the folder + + return run_id, run_path + else return 0,"no runpath" end -end +end \ No newline at end of file diff --git a/src/preallocate.jl b/src/preallocate.jl index 889bd7a..6e23611 100644 --- a/src/preallocate.jl +++ b/src/preallocate.jl @@ -448,22 +448,6 @@ function SemiLagrangeVars{T}(G::Grid) where {T<:AbstractFloat} halosstx=halosstx,halossty=halossty) end -################################################################### - -struct DiagnosticVars{T,Tprog} - RungeKutta::RungeKuttaVars{Tprog} - Tendencies::TendencyVars{Tprog} - VolumeFluxes::VolumeFluxVars{T} - Vorticity::VorticityVars{T} - Bernoulli::BernoulliVars{T} - Bottomdrag::BottomdragVars{T} - ArakawaHsu::ArakawaHsuVars{T} - Laplace::LaplaceVars{T} - Smagorinsky::SmagorinskyVars{T} - SemiLagrange::SemiLagrangeVars{T} - PrognosticVarsRHS::PrognosticVars{T} # low precision version -end - """Preallocate the diagnostic variables and return them as matrices in structs.""" function preallocate( ::Type{T}, ::Type{Tprog}, diff --git a/src/run_model.jl b/src/run_model.jl index cb02620..396edf7 100644 --- a/src/run_model.jl +++ b/src/run_model.jl @@ -10,7 +10,7 @@ julia> u,v,η,sst = run_model(Float64,nx=200,output=true) ``` """ function run_model(::Type{T}=Float32; # number format - kwargs... # all additional parameters + kwargs... # all additional parameters ) where {T<:AbstractFloat} P = Parameter(T=T;kwargs...) @@ -29,11 +29,14 @@ function run_model(::Type{T},P::Parameter) where {T<:AbstractFloat} G = Grid{T,Tprog}(P) C = Constants{T,Tprog}(P,G) F = Forcing{T}(P,G) - S = ModelSetup{T,Tprog}(P,G,C,F) - Prog = initial_conditions(Tprog,S) + Prog = initial_conditions(Tprog,G,P,C) Diag = preallocate(T,Tprog,G) - Prog = time_integration(Prog,Diag,S) + # one structure with everything already inside + S = ModelSetup{T,Tprog}(P,G,C,F,Prog,Diag,0) + Prog = time_integration(S) + return Prog -end + +end \ No newline at end of file diff --git a/src/time_integration.jl b/src/time_integration.jl index 1bcd508..53979f0 100644 --- a/src/time_integration.jl +++ b/src/time_integration.jl @@ -1,7 +1,8 @@ """Integrates ShallowWaters forward in time.""" -function time_integration( Prog::PrognosticVars{Tprog}, - Diag::DiagnosticVars{T,Tprog}, - S::ModelSetup{T,Tprog}) where {T<:AbstractFloat,Tprog<:AbstractFloat} +function time_integration(S::ModelSetup{T,Tprog}) where {T<:AbstractFloat,Tprog<:AbstractFloat} + + Diag = S.Diag + Prog = S.Prog @unpack u,v,η,sst = Prog @unpack u0,v0,η0 = Diag.RungeKutta @@ -171,7 +172,7 @@ function time_integration( Prog::PrognosticVars{Tprog}, axb!(v1,Δt_Δn,dv) # if compensated - # axb!(du_sum,Δt_Δn,du) + # axb!(du_sum,Δt_Δn,du) # axb!(dv_sum,Δt_Δn,dv) # end end