Skip to content

Commit

Permalink
Single structure (#175)
Browse files Browse the repository at this point in the history
* Update .gitignore

* Adding new runid

Changed the id and path of the output to match the way SpeedyWeather does it

Matching the upstream repo

* add Manifest.toml to .gitignore

* Fixing the output function  on main

Just correcting the error with path versus outpath, restoring the name
initpath in Parameters, all else the same

* adding what's needed to move everything to a single structure

Modified everything so that only a single structure is needed

* removing redundant functions

* removing /my_scripts from gitignore

* removing the checkpointed time integration for now

* re-added a function definition

* Update ghost_points.jl

* removing manifest file

* minor changes based on comments

* Found one place where I used `T` and should have used `Tprog`

* Reverting the changes to default parameters, leaving the addition of some adjoint/checkpointing stuff. I can maybe delete this later if needed

---------

Co-authored-by: Milan Klöwer <milankloewer@gmx.de>
  • Loading branch information
swilliamson7 and milankl authored Apr 19, 2024
1 parent cd0cb6a commit 033239f
Show file tree
Hide file tree
Showing 10 changed files with 214 additions and 89 deletions.
3 changes: 2 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -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
Manifest.toml

2 changes: 1 addition & 1 deletion src/ShallowWaters.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
11 changes: 10 additions & 1 deletion src/default_parameters.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
124 changes: 124 additions & 0 deletions src/ghost_points.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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)
Expand Down
46 changes: 17 additions & 29 deletions src/initial_conditions.jl
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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"

Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)

Expand All @@ -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))
Expand All @@ -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
Expand Down
28 changes: 26 additions & 2 deletions src/model_setup.jl
Original file line number Diff line number Diff line change
@@ -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
51 changes: 21 additions & 30 deletions src/output.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
16 changes: 0 additions & 16 deletions src/preallocate.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down
Loading

0 comments on commit 033239f

Please sign in to comment.