Merge pull request DJ4Earth#19 from swilliamson7/main

Small edits

Author: swilliamson7

.DS_Store

@@ -6,61 +6,61 @@
 # RHS_terms separate from the states was no longer a good idea. 
 
 
-function advance(u_v_eta::gyre_vector, 
-        grid::Grid, 
-        rhs::RHS_terms, 
-        params::Params, 
-        interp::Interps, 
-        grad::Derivatives, 
-        advec::Advection
-    ) 
+# function advance(u_v_eta::gyre_vector, 
+#         grid::Grid, 
+#         rhs::RHS_terms, 
+#         params::Params, 
+#         interp::Interps, 
+#         grad::Derivatives, 
+#         advec::Advection
+#     ) 
 
-    nx = grid.nx 
-    dt = params.dt
+#     nx = grid.nx 
+#     dt = params.dt
 
-    # we now use RK4 as the timestepper, here I'm storing the coefficients needed for this 
-    rk_a = [1/6, 1/3, 1/3, 1/6]
-    rk_b = [1/2, 1/2, 1.]
+#     # we now use RK4 as the timestepper, here I'm storing the coefficients needed for this 
+#     rk_a = [1/6, 1/3, 1/3, 1/6]
+#     rk_b = [1/2, 1/2, 1.]
 
-    rhs.umid .= u_v_eta.u
-    rhs.vmid .= u_v_eta.v
-    rhs.etamid .= u_v_eta.eta
+#     rhs.umid .= u_v_eta.u
+#     rhs.vmid .= u_v_eta.v
+#     rhs.etamid .= u_v_eta.eta
 
-    rhs.u0 .= u_v_eta.u
-    rhs.v0 .= u_v_eta.v
-    rhs.eta0 .= u_v_eta.eta
+#     rhs.u0 .= u_v_eta.u
+#     rhs.v0 .= u_v_eta.v
+#     rhs.eta0 .= u_v_eta.eta
 
-    rhs.u1 .= u_v_eta.u
-    rhs.v1 .= u_v_eta.v
-    rhs.eta1 .= u_v_eta.eta
+#     rhs.u1 .= u_v_eta.u
+#     rhs.v1 .= u_v_eta.v
+#     rhs.eta1 .= u_v_eta.eta
 
-    for j in 1:4
+#     @inbounds for j in 1:4
 
-        comp_u_v_eta_t(nx, rhs, params, interp, grad, advec)
+#         comp_u_v_eta_t(nx, rhs, params, interp, grad, advec)
 
-        if j < 4
-            rhs.u1 .= rhs.umid .+ rk_b[j] .* dt .* rhs.u_t
-            rhs.v1 .= rhs.vmid .+ rk_b[j] .* dt .* rhs.v_t
-            rhs.eta1 .= rhs.etamid .+ rk_b[j] .* dt .* rhs.eta_t
-        end
+#         if j < 4
+#             rhs.u1 .= rhs.umid .+ rk_b[j] .* dt .* rhs.u_t
+#             rhs.v1 .= rhs.vmid .+ rk_b[j] .* dt .* rhs.v_t
+#             rhs.eta1 .= rhs.etamid .+ rk_b[j] .* dt .* rhs.eta_t
+#         end
 
-        rhs.u0 .= rhs.u0 .+ rk_a[j] .* dt .* rhs.u_t
-        rhs.v0 .= rhs.v0 .+ rk_a[j] .* dt .* rhs.v_t 
-        rhs.eta0 .= rhs.eta0 .+ rk_a[j] .* dt .* rhs.eta_t 
+#         rhs.u0 .= rhs.u0 .+ rk_a[j] .* dt .* rhs.u_t
+#         rhs.v0 .= rhs.v0 .+ rk_a[j] .* dt .* rhs.v_t 
+#         rhs.eta0 .= rhs.eta0 .+ rk_a[j] .* dt .* rhs.eta_t 
 
-    end
+#     end
 
-    @assert all(x -> x < 7.0, rhs.u0)
-    @assert all(x -> x < 7.0, rhs.v0)
-    @assert all(x -> x < 7.0, rhs.eta0)
+#     @assert all(x -> x < 7.0, rhs.u0)
+#     @assert all(x -> x < 7.0, rhs.v0)
+#     @assert all(x -> x < 7.0, rhs.eta0)
 
-    copyto!(u_v_eta.u, rhs.u0)
-    copyto!(u_v_eta.v, rhs.v0)
-    copyto!(u_v_eta.eta, rhs.eta0)
+#     copyto!(u_v_eta.u, rhs.u0)
+#     copyto!(u_v_eta.v, rhs.v0)
+#     copyto!(u_v_eta.eta, rhs.eta0)
 
-    return nothing 
+#     return nothing 
 
-end 
+# end 
 
 function advance(states_rhs::SWM_pde, 
         grid::Grid, 
@@ -70,12 +70,8 @@ function advance(states_rhs::SWM_pde,
         advec::Advection
     ) 
 
-    nx = grid.nx 
-    dt = params.dt
-
-    # we now use RK4 as the timestepper, here I'm storing the coefficients needed for this 
-    rk_a = [1/6, 1/3, 1/3, 1/6]
-    rk_b = [1/2, 1/2, 1.]
+    @unpack nx = grid 
+    @unpack dt, rk_a, rk_b = params 
 
     states_rhs.umid .= states_rhs.u
     states_rhs.vmid .= states_rhs.v
@@ -89,9 +85,9 @@ function advance(states_rhs::SWM_pde,
     states_rhs.v1 .= states_rhs.v
     states_rhs.eta1 .= states_rhs.eta
 
-    for j in 1:4
+    @inbounds for j in 1:4
 
-        comp_u_v_eta_t!(nx, states_rhs, params, interp, grad, advec)
+        comp_u_v_eta_t(nx, states_rhs, params, interp, grad, advec)
 
         if j < 4
             states_rhs.u1 .= states_rhs.umid .+ rk_b[j] .* dt .* states_rhs.u_t
@@ -145,7 +141,7 @@ function integrate(days, nx, ny; Lx = 3840e3, Ly = 3840e3)
     )
 
     @time for t in 1:T
-        advance!(states_rhs, grid, params, interp, grad, advec)
+        advance(states_rhs, grid, params, interp, grad, advec)
         copyto!(states_rhs.u, states_rhs.u0)
         copyto!(states_rhs.v, states_rhs.v0)
         copyto!(states_rhs.eta, states_rhs.eta0)

explicit_solver/advance.jl

@@ -48,24 +48,24 @@ end
 
 # This function will just serve to take the vectors containing state information 
 # and transform them to matrices 
-function vec_to_mat(u, v, eta, grid)
+function vec_to_mat(states, grid)
 
     state_matrices = gyre_matrix(
-        reshape(u, grid.nx-1, grid.ny)',
-        reshape(v, grid.nx, grid.ny-1)',
-        reshape(eta, grid.nx, grid.ny)'
+        reshape(states.u, grid.nx-1, grid.ny)',
+        reshape(states.v, grid.nx, grid.ny-1)',
+        reshape(states.eta, grid.nx, grid.ny)'
     )
 
     return state_matrices
 
 end
 
-function mat_to_vec(um, vm, etam, grid)
+function mat_to_vec(states_mat, grid)
 
     state_vectors = gyre_vector(
-        reshape(um', grid.Nu),
-        reshape(vm', grid.Nv),
-        reshape(etam', grid.NT)
+        reshape(states_mat.u', grid.Nu),
+        reshape(states_mat.v', grid.Nv),
+        reshape(states_mat.eta', grid.NT)
     )
 
     return state_vectors 

explicit_solver/build_grid.jl

@@ -2,57 +2,57 @@
 # computes the advection term (needed for the time derivatives). Two versions of each function
 # are defined, each just depends on what type of structure I'm passing for RHS 
 
-function comp_u_v_eta_t!(
-        nx::Int, 
-        rhs::RHS_terms, 
-        params::Params, 
-        interp::Interps, 
-        grad::Derivatives, 
-        advec::Advection
-    ) 
-
-    rhs.h .= rhs.eta1 .+ params.H 
-
-    rhs.h_u .= interp.ITu * rhs.h
-    rhs.h_v .= interp.ITv * rhs.h
-    rhs.h_q .= interp.ITq * rhs.h
-
-    rhs.U .= rhs.u1 .* rhs.h_u 
-    rhs.V .= rhs.v1 .* rhs.h_v 
-
-    rhs.kinetic .= interp.IuT * (rhs.u1.^2) .+ interp.IvT * (rhs.v1.^2)
-
-    # Kloewer defined new terms q and p corresponding to potential vorticity and 
-    # Bernoulli potential respectively. To avoid errors in my mimic I'm following 
-    # along and doing the same 
-    rhs.q .= (params.coriolis .+ grad.Gvx * rhs.v1 .- grad.Guy * rhs.u1) ./ rhs.h_q 
-    rhs.p .= 0.5 .* rhs.kinetic .+ params.g .* rhs.h
-
-    # bottom friction
-    rhs.kinetic_sq .= sqrt.(rhs.kinetic)
-    rhs.bfric_u .= params.bottom_drag .* ((interp.ITu * rhs.kinetic_sq) .* rhs.u1) ./ rhs.h_u
-    rhs.bfric_v .= params.bottom_drag .* ((interp.ITv * rhs.kinetic_sq) .* rhs.v1) ./ rhs.h_v
-
-    # deal with the advection term 
-    comp_advection(nx, rhs, advec)
-
-    # rhs.Mu .= params.A_h .* (grad.LLu * rhs.u1)
-    # rhs.Mv .= params.A_h .* (grad.LLv * rhs.v1) 
+# function comp_u_v_eta_t(
+#         nx::Int, 
+#         rhs::RHS_terms, 
+#         params::Params, 
+#         interp::Interps, 
+#         grad::Derivatives, 
+#         advec::Advection
+#     ) 
+
+#     rhs.h .= rhs.eta1 .+ params.H 
+
+#     rhs.h_u .= interp.ITu * rhs.h
+#     rhs.h_v .= interp.ITv * rhs.h
+#     rhs.h_q .= interp.ITq * rhs.h
+
+#     rhs.U .= rhs.u1 .* rhs.h_u 
+#     rhs.V .= rhs.v1 .* rhs.h_v 
+
+#     rhs.kinetic .= interp.IuT * (rhs.u1.^2) .+ interp.IvT * (rhs.v1.^2)
+
+#     # Kloewer defined new terms q and p corresponding to potential vorticity and 
+#     # Bernoulli potential respectively. To avoid errors in my mimic I'm following 
+#     # along and doing the same 
+#     rhs.q .= (params.coriolis .+ grad.Gvx * rhs.v1 .- grad.Guy * rhs.u1) ./ rhs.h_q 
+#     rhs.p .= 0.5 .* rhs.kinetic .+ params.g .* rhs.h
+
+#     # bottom friction
+#     rhs.kinetic_sq .= sqrt.(rhs.kinetic)
+#     rhs.bfric_u .= params.bottom_drag .* ((interp.ITu * rhs.kinetic_sq) .* rhs.u1) ./ rhs.h_u
+#     rhs.bfric_v .= params.bottom_drag .* ((interp.ITv * rhs.kinetic_sq) .* rhs.v1) ./ rhs.h_v
+
+#     # deal with the advection term 
+#     comp_advection(nx, rhs, advec)
+
+#     # rhs.Mu .= params.A_h .* (grad.LLu * rhs.u1)
+#     # rhs.Mv .= params.A_h .* (grad.LLv * rhs.v1) 
 
-    rhs.Mu .= (interp.ITu * params.nu) .* (grad.LLu * rhs.u1)
-    rhs.Mv .= (interp.ITv * params.nu) .* (grad.LLv * rhs.v1) 
+#     rhs.Mu .= (interp.ITu * params.nu) .* (grad.LLu * rhs.u1)
+#     rhs.Mv .= (interp.ITv * params.nu) .* (grad.LLv * rhs.v1) 
 
-    rhs.u_t .= rhs.adv_u .- grad.GTx * rhs.p .+ params.wind_stress ./ rhs.h_u .- rhs.Mu .- rhs.bfric_u
+#     rhs.u_t .= rhs.adv_u .- grad.GTx * rhs.p .+ params.wind_stress ./ rhs.h_u .- rhs.Mu .- rhs.bfric_u
 
-    rhs.v_t .= rhs.adv_v .- grad.GTy * rhs.p .- rhs.Mv .- rhs.bfric_v 
+#     rhs.v_t .= rhs.adv_v .- grad.GTy * rhs.p .- rhs.Mv .- rhs.bfric_v 
 
-    rhs.eta_t .= - (grad.Gux * rhs.U .+ grad.Gvy * rhs.V)
+#     rhs.eta_t .= - (grad.Gux * rhs.U .+ grad.Gvy * rhs.V)
 
-    return nothing
+#     return nothing
 
-end 
+# end 
 
-function comp_u_v_eta_t!(nx::Int, 
+function comp_u_v_eta_t(nx::Int, 
         rhs::SWM_pde, 
         params::Params, 
         interp::Interps, 
@@ -83,25 +83,25 @@ function comp_u_v_eta_t!(nx::Int,
     rhs.bfric_v .= params.bottom_drag .* ((interp.ITv * rhs.kinetic_sq) .* rhs.v1) ./ rhs.h_v
 
     # deal with the advection term 
-    comp_advection!(nx, rhs, advec)
+    comp_advection(nx, rhs, advec)
 
     # rhs.Mu .= params.A_h .* (grad.LLu * rhs.u1)
     # rhs.Mv .= params.A_h .* (grad.LLv * rhs.v1) 
 
+    ### Important: Need to apply Laplacian to nu for when nu varies spatially ##################
     rhs.Mu .= (interp.ITu * params.nu) .* (grad.LLu * rhs.u1)
     rhs.Mv .= (interp.ITv * params.nu) .* (grad.LLv * rhs.v1) 
 
     rhs.u_t .= rhs.adv_u .- grad.GTx * rhs.p .+ params.wind_stress ./ rhs.h_u .- rhs.Mu .- rhs.bfric_u
-
     rhs.v_t .= rhs.adv_v .- grad.GTy * rhs.p .- rhs.Mv .- rhs.bfric_v 
-
     rhs.eta_t .= - (grad.Gux * rhs.U .+ grad.Gvy * rhs.V)
 
     return nothing
 
 end 
 
-function comp_advection!(
+# Arakawa and Lamb advection scheme
+function comp_advection(
     nx::Int, 
     rhs::SWM_pde, 
     advec::Advection
@@ -132,29 +132,29 @@ function comp_advection!(
 
 end
 
-function comp_advection(nx::Int, rhs::RHS_terms, advec::Advection)
+# function comp_advection(nx::Int, rhs::RHS_terms, advec::Advection)
 
-    rhs.AL1q .= advec.AL1 * rhs.q 
-    rhs.AL2q .= advec.AL2 * rhs.q 
+#     rhs.AL1q .= advec.AL1 * rhs.q 
+#     rhs.AL2q .= advec.AL2 * rhs.q 
 
-    AL1q_au = @view rhs.AL1q[1:end-nx]
-    AL2q_bu = @view rhs.AL2q[nx+1:end]
-    AL2q_cu = @view rhs.AL2q[1:end-nx]
-    AL1q_du = @view rhs.AL1q[nx+1:end]
+#     AL1q_au = @view rhs.AL1q[1:end-nx]
+#     AL2q_bu = @view rhs.AL2q[nx+1:end]
+#     AL2q_cu = @view rhs.AL2q[1:end-nx]
+#     AL1q_du = @view rhs.AL1q[nx+1:end]
 
-    AL1q_av = @view rhs.AL1q[advec.index_av]
-    AL2q_bv = @view rhs.AL2q[advec.index_bv]
-    AL2q_cv = @view rhs.AL2q[advec.index_cv]
-    AL1q_dv = @view rhs.AL1q[advec.index_dv]
+#     AL1q_av = @view rhs.AL1q[advec.index_av]
+#     AL2q_bv = @view rhs.AL2q[advec.index_bv]
+#     AL2q_cv = @view rhs.AL2q[advec.index_cv]
+#     AL1q_dv = @view rhs.AL1q[advec.index_dv]
 
-    rhs.adv_u .= advec.Seur * (advec.ALeur * rhs.q .* rhs.U) .+ advec.Seul * (advec.ALeul * rhs.q .* rhs.U) .+
-    advec.Sau * (AL1q_au .* rhs.V) .+ advec.Sbu * (AL2q_bu .* rhs.V) .+ 
-    advec.Scu * (AL2q_cu .* rhs.V) .+ advec.Sdu * (AL1q_du .* rhs.V)
+#     rhs.adv_u .= advec.Seur * (advec.ALeur * rhs.q .* rhs.U) .+ advec.Seul * (advec.ALeul * rhs.q .* rhs.U) .+
+#     advec.Sau * (AL1q_au .* rhs.V) .+ advec.Sbu * (AL2q_bu .* rhs.V) .+ 
+#     advec.Scu * (AL2q_cu .* rhs.V) .+ advec.Sdu * (AL1q_du .* rhs.V)
 
-    rhs.adv_v .= advec.Spvu * ((advec.ALpvu * rhs.q) .* rhs.V) .+ advec.Spvd * ((advec.ALpvd * rhs.q) .* rhs.V) .-
-    advec.Sav * (AL1q_av .* rhs.U) .- advec.Sbv * (AL2q_bv .* rhs.U) .- 
-    advec.Scv * (AL2q_cv .* rhs.U) .- advec.Sdv * (AL1q_dv .* rhs.U)
+#     rhs.adv_v .= advec.Spvu * ((advec.ALpvu * rhs.q) .* rhs.V) .+ advec.Spvd * ((advec.ALpvd * rhs.q) .* rhs.V) .-
+#     advec.Sav * (AL1q_av .* rhs.U) .- advec.Sbv * (AL2q_bv .* rhs.U) .- 
+#     advec.Scv * (AL2q_cv .* rhs.U) .- advec.Sdv * (AL1q_dv .* rhs.U)
 
-    return nothing
+#     return nothing
 
-end
+# end