Single structure (#175)

swilliamson7 · milankl · web-flow · commit 033239f1c253 · 2024-04-19T11:53:51.000-05:00
* 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 &lt;milankloewer@gmx.de&gt;
diff --git 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
+Manifest.toml
+
diff --git 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
@@ -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
diff --git 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
@@ -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
@@ -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
diff --git 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
diff --git 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
diff --git a/src/time_integration.jl b/src/time_integration.jl
