Merge pull request DJ4Earth#12 from swilliamson7/main

updating code

Author: swilliamson7

.DS_Store

@@ -1,7 +1,6 @@
 [deps]
 Checkpointing = "eb46d486-4f9c-4c3d-b445-a617f2a2f1ca"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
-EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PProf = "e4faabce-9ead-11e9-39d9-4379958e3056"

explicit_solver/.DS_Store

@@ -5,6 +5,7 @@
 # I needed all of the terms that matter to the derivative to appear in a *single* structure, so keeping 
 # RHS_terms separate from the states was no longer a good idea. 
 
+
 function advance(u_v_eta::gyre_vector, 
         grid::Grid, 
         rhs::RHS_terms, 
@@ -110,4 +111,140 @@ function advance(states_rhs::SWM_pde,
 
     return nothing 
 
-end 
+end 
+
+# This function needs to be given 
+#           T - how many days to integrate the model for
+#           nx, ny - how many grid cells in x and y directions, respectively
+#           Lx - E-W length of the domain [meters], has a default but can change if wanted
+#           Ly - N-S length of the domain [meters], also has a default 
+# mostly keeping because I like having a function that just integrates for some amount
+# of time 
+function integrate(T, nx, ny; Lx = 3840e3, Ly = 3840e3)                 
+    
+    grid_params = build_grid(Lx, Ly, nx, ny)
+    gyre_params = def_params(grid_params)
+    
+    # building discrete operators 
+    grad_ops = build_derivs(grid_params)                # discrete gradient operators 
+    interp_ops = build_interp(grid_params, grad_ops)    # discrete interpolation operators (travels between grids)
+    advec_ops = build_advec(grid_params)
+    rhs_terms = RHS_terms(Nu = grid_params.Nu, Nv = grid_params.Nv, NT = grid_params.NT, Nq = grid_params.Nq)
+    
+    # starting from rest ---> all initial conditions are zero 
+    
+    Trun = days_to_seconds(T, gyre_params.dt)
+    
+    uout = zeros(grid_params.Nu) 
+    vout = zeros(grid_params.Nv) 
+    etaout = zeros(grid_params.NT)
+    
+    # u = [zeros(grid_params.Nu)]
+    # v = [zeros(grid_params.Nv)]
+    # eta = [zeros(grid_params.NT)]
+
+    u_v_eta = gyre_vector(uout, vout, etaout)
+    
+    @time for t in 1:Trun
+        advance(u_v_eta, grid_params, rhs_terms, gyre_params, interp_ops, grad_ops, advec_ops)
+    end
+    
+    # u_v_eta_mat = vec_to_mat(u_v_eta.u, u_v_eta.v, u_v_eta.eta, grid_params)
+    
+    return u_v_eta
+    
+end
+
+# ****IMPORTANT**** not yet sure if I'm moving between high and low res grids, need to check with Patrick
+# and come back here if there are issues with how I did it
+
+# This function needs to be given 
+#           days - how many days to integrate the model for
+#           nx_lowres, ny_lowres - grid resolution (number of cells in the x and y directions
+#                    respectively) of the courser grid
+#           Lx, Ly - size of the domain, have a default value but can set 
+#                    manually if needed 
+#           data_steps - which timesteps we want to store data at
+# Theoretically, we should only ever run this function *once*, from there 
+# the data points will be stored as a JLD2 data file 
+function create_data(days, nx_lowres, ny_lowres, data_steps; scaling = 4, Lx = 3840e3, Ly = 3840e3)
+
+    nx_highres = nx_lowres * scaling
+    ny_highres = ny_lowres * scaling 
+
+    grid_highres = build_grid(Lx, Ly, nx_highres, ny_highres)
+    params = def_params(grid_highres)
+
+    # building discrete operators
+    grad = build_derivs(grid_highres)            # discrete gradient operators
+    interp = build_interp(grid_highres, grad)    # discrete interpolation operators (travels between grids)
+    advec = build_advec(grid_highres)
+
+    Trun = days_to_seconds(days, params.dt)
+
+    Nu = grid_highres.Nu
+    Nv = grid_highres.Nv
+    NT = grid_highres.NT
+    Nq = grid_highres.Nq 
+
+    u_v_eta_rhs = SWM_pde(Nu = Nu, 
+        Nv = Nv,
+        NT = NT, 
+        Nq = Nq
+    )
+
+    # In order to compare high res data to low res velocities I'm going to (1) interpolate 
+    # the velocities to the T-grid (cell centers) (2) average the high
+    # res data points down to the low res grid (3) interpolate the low 
+    # res results to the cell centers. Then I'll be comparing apples to apples (hopefully)
+    # will check with Patrick that this is a valid method 
+    
+    # Building the averaging operator needed for step (2) above
+    diag1 = (1 / scaling^2) .* ones(NT)
+    M = spdiagm(NT, NT, 0.0 .* diag1)
+    for k = 1:scaling
+        for j = 1:scaling
+            M += spdiagm(NT, NT, j+(k-1)*nx_highres - 1 => diag1[1:end-j-(k-1)*nx_highres + 1])
+        end
+    end
+    M = M[1:scaling:end, :]
+    index1 = 1:nx_lowres*ny_lowres*scaling
+    for k in 1:ny_lowres
+        index1 = filter(x -> x ∉ [j for j in ((k-1)*nx_lowres*scaling+nx_lowres+1):((k-1)*nx_lowres*scaling+nx_lowres*scaling)], index1)
+    end
+    M = M[index1, :]
+
+    # initializing where to store the data 
+    data = zeros(grid_highres.Nu + grid_highres.Nv + grid_highres.NT, length(data_steps))
+
+    # the steps where we want data in the high res model correspond to (roughly) scaling * t for t 
+    # in the low res model. for simplicity I'm going to keep the times in the low res where I want to have 
+    # data and then just scale them in the for loop to find corresponding high res data points  
+
+    # for an initial effort I'm just going to run pretty course resolution models for both the high and low res 
+
+    if 1 in scaling .* data_steps 
+        data[:, 1] .= [u_v_eta_rhs.u; u_v_eta_rhs.v; u_v_eta_rhs.eta]
+        j = 2
+    else
+        j = 1
+    end
+
+    for t in 2:Trun
+
+        advance(u_v_eta_rhs, grid_highres, params, interp, grad, advec) 
+
+        if t in scaling .* data_steps 
+            data[:, j] .= [u_v_eta_rhs.u; u_v_eta_rhs.v; u_v_eta_rhs.eta]
+            j += 1
+        end
+
+        copyto!(u_v_eta_rhs.u, u_v_eta_rhs.u0)
+        copyto!(u_v_eta_rhs.v, u_v_eta_rhs.v0)
+        copyto!(u_v_eta_rhs.eta, u_v_eta_rhs.eta0)
+
+    end
+    
+    return data, M
+    
+end

explicit_solver/Project.toml

@@ -1,8 +1,8 @@
 # This function builds all the parameters relating to the grid. It takes 
 #       Lx - the length of the domain from E-W [meters]
 #       Ly - the length of the domain from N-S [meters]
-#       Nx - the number of cells in the E-W direction 
-#       Ny - the number of cells in the N-S direction 
+#       nx - the number of cells in the E-W direction 
+#       ny - the number of cells in the N-S direction 
 # and returns a structure with a bunch of different parameters all built
 # from the above inputs. 
 function build_grid(Lx, Ly, nx, ny)
@@ -46,16 +46,6 @@ function build_grid(Lx, Ly, nx, ny)
 
 end
 
-# This function allows me to specify the number of days that we want to run the model, 
-# and convert that into the total steps to take 
-function days_to_seconds(T, dt)
-
-    Trun_seconds = Int(ceil((T * 24 * 3600) / dt))
-
-    return Trun_seconds
-
-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)
@@ -69,3 +59,15 @@ function vec_to_mat(u, v, eta, grid)
     return state_matrices
 
 end
+
+function mat_to_vec(um, vm, etam, grid)
+
+    state_vectors = gyre_vector(
+        reshape(um', grid.Nu),
+        reshape(vm', grid.Nv),
+        reshape(etam', grid.NT)
+    )
+
+    return state_vectors 
+
+end