diff --git a/docs/make.jl b/docs/make.jl
index ad20359f..81bb2d65 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -43,13 +43,14 @@ pages = [
 ]
 
 makedocs(
-   sitename = "ClimaOcean.jl",
-    modules = [ClimaOcean],
-     format = format,
-      pages = pages,
-    doctest = true,
-      clean = true,
-  checkdocs = :exports
+     sitename = "ClimaOcean.jl",
+      modules = [ClimaOcean],
+       format = format,
+        pages = pages,
+      doctest = true,
+        clean = true,
+     warnonly = [:cross_references, :missing_docs],
+    checkdocs = :exports
 )
 
 @info "Clean up temporary .jld2 and .nc output created by doctests or literated examples..."
diff --git a/docs/src/library/internals.md b/docs/src/library/internals.md
index 920f347b..173ade94 100644
--- a/docs/src/library/internals.md
+++ b/docs/src/library/internals.md
@@ -16,3 +16,31 @@ Modules = [ClimaOcean.Diagnostics]
 Public = false
 ```
 
+
+## InitialConditions
+
+```@autodocs
+Modules = [ClimaOcean.InitialConditions]
+Public = false
+```
+
+## DataWrangling
+
+```@autodocs
+Modules = [ClimaOcean.DataWrangling]
+Public = false
+```
+
+## ECCO2
+
+```@autodocs
+Modules = [ClimaOcean.ECCO2]
+Public = false
+```
+
+## Bathymetry
+
+```@autodocs
+Modules = [ClimaOcean.Bathymetry]
+Public = false
+```
diff --git a/docs/src/library/public.md b/docs/src/library/public.md
index 3ba1c84e..65a7fed6 100644
--- a/docs/src/library/public.md
+++ b/docs/src/library/public.md
@@ -17,3 +17,32 @@ Private = false
 Modules = [ClimaOcean.Diagnostics]
 Private = false
 ```
+
+## InitialConditions
+
+```@autodocs
+Modules = [ClimaOcean.InitialConditions]
+Private = false
+```
+
+## DataWrangling
+
+```@autodocs
+Modules = [ClimaOcean.DataWrangling]
+Private = false
+```
+
+## ECCO2
+
+```@autodocs
+Modules = [ClimaOcean.ECCO2]
+Private = false
+```
+
+## Bathymetry
+
+```@autodocs
+Modules = [ClimaOcean.Bathymetry]
+Private = false
+```
+
diff --git a/examples/analyze_neverworld_simulation.jl b/examples/analyze_neverworld_simulation.jl
deleted file mode 100644
index f76a0035..00000000
--- a/examples/analyze_neverworld_simulation.jl
+++ /dev/null
@@ -1,43 +0,0 @@
-using Oceananigans
-using Oceananigans.Units
-using Printf
-using JLD2
-
-dir = "archive"
-backend = OnDisk()
-bt = FieldTimeSeries("$dir/neverworld_xyz.jld2", "b"; backend)
-t = bt.times
-Nt = length(t)
-grid = bt.grid
-
-Nyears = 20
-t_end = 200 * 360 * day # 200 "years"
-t_start = (200 - Nyears) * 360 * day # 200 "years"
-
-t = bt.times
-n_stop  = searchsortedfirst(t, t_stop)
-n_start = searchsortedfirst(t, t_start) - 1
-
-for season = 1:4
-    n_average_start = n_start + season - 1
-    n_average = n_average_start:4:n_stop
-
-    for name in ("b", "u", "v", "w", "e", "κᶜ")
-        ψt = FieldTimeSeries("$dir/neverworld_xyz.jld2", name; backend)
-        LX, LY, LZ = location(ψt)
-        ψ_avg = Field{LX, LY, LZ}(grid)
-
-        for n in n_average
-            @info name season n
-            ψn = ψt[n]
-            ψ_avg .+= ψn / length(n_average)
-        end
-        
-        filename = string("climatology_", name, "_season_", season, ".jld2")
-        file = jldopen(filename, "a+")
-        file[name] = ψ_avg
-        file["season"] = season
-        close(file)
-    end
-end
-
diff --git a/examples/animate_neverworld_simulation.jl b/examples/animate_neverworld_simulation.jl
deleted file mode 100644
index 822ff075..00000000
--- a/examples/animate_neverworld_simulation.jl
+++ /dev/null
@@ -1,152 +0,0 @@
-using Oceananigans
-using Oceananigans.Units
-using GLMakie
-using Printf
-using Statistics
-
-# Some plotting parameters
-ζlim = 6e-5
-
-# Plot a limited number of years (here, 10)
-Nyears = 3
-t_end = 200 * 360 * day # 200 "years"
-t_start = (200 - Nyears) * 360 * day # 200 "years"
-
-# Load OnDisk time-series without loading massive amount of data
-backend = OnDisk()
-bxyt = FieldTimeSeries("neverworld_xy.jld2", "b"; backend)
-
-# Determine which iterations to load into memory
-t̃ = bxyt.times
-n_end   = searchsortedfirst(t̃, t_end)
-n_start = searchsortedfirst(t̃, t_start) - 1
-
-t = times = t̃[n_start:n_end]
-Nt = length(t)
-
-bxyt = FieldTimeSeries("neverworld_xy.jld2", "b"; times)
-uxyt = FieldTimeSeries("neverworld_xy.jld2", "u"; times)
-vxyt = FieldTimeSeries("neverworld_xy.jld2", "v"; times)
-exyt = FieldTimeSeries("neverworld_xy.jld2", "e"; times)
-κxyt = FieldTimeSeries("neverworld_xy.jld2", "κᶜ"; times)
-
-byzt = FieldTimeSeries("neverworld_zonal_average.jld2", "b"; times)
-eyzt = FieldTimeSeries("neverworld_zonal_average.jld2", "e"; times)
-κyzt = FieldTimeSeries("neverworld_zonal_average.jld2", "κᶜ"; times)
-
-# Hack to set immersed regions to NaN
-for ft in (bxyt, uxyt, exyt, κxyt, byzt, eyzt, κyzt)
-    fp = parent(ft)
-    fp[fp .== 0] .= NaN
-end
-
-fig = Figure(resolution=(2000, 1800))
-
-axbxy = Axis(fig[2, 1], xlabel="Longitude", ylabel="Latitude", title="b(x, y)")
-axζxy = Axis(fig[2, 2], xlabel="Longitude", ylabel="Latitude", title="ζ(x, y)")
-axexy = Axis(fig[2, 3], xlabel="Longitude", ylabel="Latitude", title="e(x, y)")
-
-axNyz = Axis(fig[3, 1], xlabel="Longitude", ylabel="z (m)", title="N²(y, z)")
-axeyz = Axis(fig[3, 2], xlabel="Longitude", ylabel="z (m)", title="e(y, z)")
-axκyz = Axis(fig[3, 3], xlabel="Longitude", ylabel="z (m)", title="κᶜ(y, z)")
-
-slider = Slider(fig[5, 1:3], range=1:Nt, startvalue=Nt)
-n = slider.value
-
-title = @lift begin
-    t_day = t[$n] / day
-    t_year = floor(Int, t_day / 360)
-    yearday = t_day % 360
-    return @sprintf("CATKE Neverworld at % 3d years, % 3d days", t_year, yearday)
-end
-
-Label(fig[0, 1:3], title, fontsize=36)
-
-using Oceananigans.Operators: ζ₃ᶠᶠᶜ
-grid = bxyt.grid
-grid = grid.underlying_grid
-Nx, Ny, Nz = size(grid)
-uxy = Field{Face, Center, Center}(grid, indices=(:, :, 1))
-vxy = Field{Center, Face, Center}(grid, indices=(:, :, 1))
-ζop = KernelFunctionOperation{Face, Face, Center}(ζ₃ᶠᶠᶜ, grid, uxy, vxy)
-ζxy = Field(ζop, indices=(:, :, 1))
-
-ζxyn = @lift begin
-    uxyn = uxyt[$n]
-    vxyn = vxyt[$n]
-    interior(uxy) .= interior(uxyn)
-    interior(vxy) .= interior(vxyn)
-    compute!(ζxy)
-    return interior(ζxy, :, :, 1)
-end
-
-byz = Field{Nothing, Center, Center}(grid)
-N²yz = Field(∂z(byz))
-
-N²yzn = @lift begin
-    byzn = byzt[$n]
-    interior(byz) .= interior(byzn)
-    compute!(N²yz)
-    return interior(N²yz, 1, :, :)
-end
-
-bxyn = @lift interior(bxyt[$n], :, :, 1)
-exyn = @lift interior(exyt[$n], :, :, 1)
-
-byzn = @lift interior(byzt[$n], 1, :, :)
-eyzn = @lift interior(eyzt[$n], 1, :, :)
-κyzn = @lift interior(κyzt[$n], 1, :, :)
-
-κlims = @lift begin
-    κ_mean = mean(filter(!isnan, interior(κyzt[$n])))
-    (κ_mean/2, 3κ_mean)
-end
-
-eyzlims = @lift begin
-    e_mean = mean(filter(!isnan, interior(eyzt[$n])))
-    (e_mean/2, 3e_mean)
-end
-
-exylims = @lift begin
-    e_mean = mean(filter(!isnan, interior(eyzt[$n])))
-    (e_mean/10, e_mean)
-end
-
-x, y, z = nodes(bxyt)
-xf, yf, zf = nodes(grid, Face(), Face(), Face())
-cbkw = (vertical=false, flipaxis=true, tellwidth=false)
-
-hm = heatmap!(axbxy, x, y, bxyn, colorrange=(-0.05, 0.05))
-Colorbar(fig[1, 1], hm; label="Buoyancy", cbkw...)
-
-hm = heatmap!(axζxy, x, y, ζxyn, colormap=:balance, colorrange=(-ζlim, ζlim))
-Colorbar(fig[1, 2], hm; label="Relative vertical vorticity (s⁻¹)", cbkw...)
-
-hm = heatmap!(axexy, x, y, exyn, colormap=:solar, colorrange=exylims)
-Colorbar(fig[1, 3], hm; label="Turbulent kinetic energy (m² s⁻²)", cbkw...)
-
-cbkw = (vertical=false, flipaxis=false, tellwidth=false)
-
-hm = heatmap!(axNyz, y, zf, N²yzn, nan_color=:gray, colorrange=(1e-6, 1e-4))
-Colorbar(fig[4, 1], hm; label="Buoyancy gradient (s⁻²)", cbkw...)
-contour!(axNyz, y, z, byzn, levels=30, color=:white, linewidth=3)
-
-hm = heatmap!(axeyz, y, z, eyzn, nan_color=:gray, colormap=:solar, colorrange=eyzlims)
-Colorbar(fig[4, 2], hm; label="Turbulent kinetic energy (m² s⁻²)", cbkw...)
-contour!(axeyz, y, z, byzn, levels=30, color=:white, linewidth=3)
-
-hm = heatmap!(axκyz, y, zf, κyzn, nan_color=:gray, colormap=:solar, colorrange=κlims)
-Colorbar(fig[4, 3], hm; label="Tracer eddy diffusivity (m² s⁻¹)", cbkw...)
-contour!(axκyz, y, z, byzn, levels=30, color=:white, linewidth=3)
-
-ylims!(axNyz, -4000, 0)
-ylims!(axeyz, -4000, 0)
-ylims!(axκyz, -4000, 0)
-
-display(fig)
-
-record(fig, "catke_neverworld.mp4", 1:Nt, framerate=8) do nn
-     @info "Recording frame $nn of $Nt..."
-     n[] = nn
-end
-
diff --git a/examples/freely_decaying_mediterranean.jl b/examples/freely_decaying_mediterranean.jl
new file mode 100644
index 00000000..83b719c4
--- /dev/null
+++ b/examples/freely_decaying_mediterranean.jl
@@ -0,0 +1,189 @@
+using GLMakie
+using Oceananigans
+using Oceananigans: architecture
+using ClimaOcean
+using ClimaOcean.ECCO2
+using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: CATKEVerticalDiffusivity
+using Oceananigans.Coriolis: ActiveCellEnstrophyConserving
+using Oceananigans.Units
+using Printf
+
+#####
+##### Regional Mediterranean grid 
+#####
+
+# Domain and Grid size
+#
+# We construct a grid that represents the Mediterranean sea, 
+# with a resolution of 1/10th of a degree (roughly 10 km resolution)
+λ₁, λ₂  = ( 0, 42) # domain in longitude
+φ₁, φ₂  = (30, 45) # domain in latitude
+# A stretched vertical grid with a Δz of 1.5 meters in the first 50 meters
+z_faces = stretched_vertical_faces(depth = 5000, 
+                             surface_layer_Δz = 2.5, 
+                             stretching = PowerLawStretching(1.070), 
+                             surface_layer_height = 50)
+
+Nx = 15 * 42 # 1 / 15th of a degree resolution
+Ny = 15 * 15 # 1 / 15th of a degree resolution
+Nz = length(z_faces) - 1
+
+grid = LatitudeLongitudeGrid(CPU();
+                             size = (Nx, Ny, Nz),
+                             latitude  = (φ₁, φ₂),
+                             longitude = (λ₁, λ₂),
+                             z = z_faces,
+                             halo = (7, 7, 7))
+
+# Interpolating the bathymetry onto the grid
+#
+# We regrid the bathymetry onto the grid.
+# we allow a minimum depth of 10 meters (all shallower regions are 
+# considered land) and we use 25 intermediate grids (interpolation_passes = 25)
+# Note that more interpolation passes will smooth the bathymetry
+bottom_height = regrid_bathymetry(grid, 
+                                  height_above_water = 1,
+                                  minimum_depth = 10,
+                                  interpolation_passes = 25)
+
+# Let's use an active cell map to elide computation in inactive cells
+grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map = true)
+
+# Correct oceananigans
+import Oceananigans.Advection: nothing_to_default
+
+nothing_to_default(user_value; default) = isnothing(user_value) ? default : user_value
+
+# Constructing the model
+#
+# We construct a model that evolves two tracers, temperature (:T), salinity (:S)
+# We do not have convection since the simulation is just slumping to equilibrium so we do not need a turbulence closure
+# We select a linear equation of state for buoyancy calculation, and WENO schemes for both tracer and momentum advection.
+# The free surface utilizes a split-explicit formulation with a barotropic CFL of 0.75 based on wave speed.
+model = HydrostaticFreeSurfaceModel(; grid,
+                            momentum_advection = WENOVectorInvariant(),
+                              tracer_advection = WENO(grid; order = 7),
+                                  free_surface = SplitExplicitFreeSurface(; cfl = 0.75, grid),
+                                      buoyancy = SeawaterBuoyancy(),
+                                      tracers  = (:T, :S, :c),
+                                      coriolis = HydrostaticSphericalCoriolis(scheme = ActiveCellEnstrophyConserving()))
+
+# Initializing the model
+#
+# the model can be initialized with custom values or with ecco2 fields.
+# In this case, our ECCO2 dataset has access to a temperature and a salinity
+# field, so we initialize T and S from ECCO2. 
+# We initialize our passive tracer with a surface blob near to the coasts of Libia
+@info "initializing model"
+libia_blob(x, y, z) = z > -20 || (x - 15)^2 + (y - 34)^2 < 1.5 ? 1 : 0
+
+set!(model, T = ECCO2Metadata(:temperature), S = ECCO2Metadata(:salinity), c = libia_blob)
+
+fig = Figure()
+ax  = Axis(fig[1, 1])
+heatmap!(ax, interior(model.tracers.T, :, :, Nz), colorrange = (10, 20), colormap = :thermal)
+ax  = Axis(fig[1, 2])
+heatmap!(ax, interior(model.tracers.S, :, :, Nz), colorrange = (35, 40), colormap = :haline)
+
+simulation = Simulation(model, Δt = 10minutes, stop_time = 10*365days)
+
+function progress(sim) 
+    u, v, w = sim.model.velocities  
+    T, S = sim.model.tracers
+
+    @info @sprintf("Time: %s, Iteration %d, Δt %s, max(vel): (%.2e, %.2e, %.2e), max(T, S): %.2f, %.2f\n",
+                   prettytime(sim.model.clock.time),
+                   sim.model.clock.iteration,
+                   prettytime(sim.Δt),
+                   maximum(abs, u), maximum(abs, v), maximum(abs, w),
+                   maximum(abs, T), maximum(abs, S))
+end
+
+simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))
+
+# Simulation warm up!
+#
+# We have regridded from a coarse solution (1/4er of a degree) to a
+# fine grid (1/15th of a degree). Also, the bathymetry has little mismatches 
+# that might crash our simulation. We warm up the simulation with a little 
+# time step for few iterations to allow the solution to adjust to the new_grid
+# bathymetry
+simulation.Δt = 10
+simulation.stop_iteration = 1000
+run!(simulation)
+
+# Run the real simulation
+#
+# Now that the solution has adjusted to the bathymetry we can ramp up the time
+# step size. We use a `TimeStepWizard` to automatically adapt to a cfl of 0.2
+wizard = TimeStepWizard(; cfl = 0.2, max_Δt = 10minutes, max_change = 1.1)
+
+simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))
+
+# Let's reset the maximum number of iterations
+simulation.stop_iteration = Inf
+
+simulation.output_writers[:surface_fields] = JLD2OutputWriter(model, merge(model.velocities, model.tracers);
+                                                              indices = (:, :, Nz),
+                                                              schedule = TimeInterval(1day),
+                                                              overwrite_existing = true,
+                                                              filename = "med_surface_field")
+
+run!(simulation)
+
+# Record a video
+#
+# Let's read the data and record a video of the Mediterranean Sea's surface
+# (1) Zonal velocity (u)
+# (2) Meridional velocity (v)
+# (3) Temperature (T)
+# (4) Salinity (S)
+u_series = FieldTimeSeries("med_surface_field.jld2", "u")
+v_series = FieldTimeSeries("med_surface_field.jld2", "v")
+T_series = FieldTimeSeries("med_surface_field.jld2", "T")
+S_series = FieldTimeSeries("med_surface_field.jld2", "S")
+c_series = FieldTimeSeries("med_surface_field.jld2", "c")
+iter = Observable(1)
+
+u = @lift begin
+    f = interior(u_series[$iter], :, :, 1)
+    f[f .== 0] .= NaN
+    f
+end
+v = @lift begin
+    f = interior(v_series[$iter], :, :, 1)
+    f[f .== 0] .= NaN
+    f
+end
+T = @lift begin
+    f = interior(T_series[$iter], :, :, 1)
+    f[f .== 0] .= NaN
+    f
+end
+S = @lift begin
+    f = interior(S_series[$iter], :, :, 1)
+    f[f .== 0] .= NaN
+    f
+end
+c = @lift begin
+    f = interior(c_series[$iter], :, :, 1)
+    f[f .== 0] .= NaN
+    f
+end
+
+fig = Figure()
+ax  = Axis(fig[1, 1], title = "surface zonal velocity ms⁻¹")
+heatmap!(u)
+ax  = Axis(fig[1, 2], title = "surface meridional velocity ms⁻¹")
+heatmap!(v)
+ax  = Axis(fig[2, 1], title = "surface temperature ᵒC")
+heatmap!(T)
+ax  = Axis(fig[2, 2], title = "surface salinity psu")
+heatmap!(S)
+ax  = Axis(fig[2, 3], title = "passive tracer -")
+heatmap!(c)
+
+GLMakie.record(fig, "mediterranean_video.mp4", 1:length(u_series.times); framerate = 5) do i
+    @info "recording iteration $i"
+    iter[] = i    
+end
\ No newline at end of file
diff --git a/examples/generate_bathymetry.jl b/examples/generate_bathymetry.jl
index 64ebc579..78d83b01 100644
--- a/examples/generate_bathymetry.jl
+++ b/examples/generate_bathymetry.jl
@@ -35,23 +35,32 @@ display(fig)
 ##### Regional Mediterranean grid 
 #####
 
+# 1/25th degree resolution
+Nλ = 25 * 55
+Nφ = 25 * 25
+
 grid = LatitudeLongitudeGrid(CPU();
-                             size = (25 * 50, 55 * 50, 1),
+                             size = (Nλ, Nφ, 1),
                              latitude = (25, 50),
                              longitude = (-10, 45),
                              z = (0, 1),
                              halo = (4, 4, 4))
 
-h = regrid_bathymetry(grid, height_above_water=1)
+h_smooth = regrid_bathymetry(grid, height_above_water=1, interpolation_passes = 40)
+h_rough  = regrid_bathymetry(grid, height_above_water=1, interpolation_passes = 1)
 
-λ, φ, z = nodes(h)
+λ, φ, z = nodes(h_smooth)
 
-land = interior(h) .> 0
-interior(h)[land] .= NaN
+land_smooth = interior(h_smooth) .> 0
+interior(h_smooth)[land_smooth] .= NaN
+land_rough = interior(h_rough) .> 0
+interior(h_rough)[land_rough] .= NaN
 
-fig = Figure(resolution=(2400, 1200))
+fig = Figure(resolution=(2400, 800))
 ax = Axis(fig[1, 1])
-heatmap!(ax, λ, φ, interior(h, :, :, 1), nan_color=:white) #, colorrange=(-5000, 0))
+heatmap!(ax, λ, φ, interior(h_smooth, :, :, 1), nan_color=:white) #, colorrange=(-5000, 0))
+ax = Axis(fig[1, 2])
+heatmap!(ax, λ, φ, interior(h_rough, :, :, 1), nan_color=:white) #, colorrange=(-5000, 0))
 
 display(fig)
 
diff --git a/examples/run_neverworld_simulation.jl b/examples/run_neverworld_simulation.jl
deleted file mode 100644
index 0acddbe3..00000000
--- a/examples/run_neverworld_simulation.jl
+++ /dev/null
@@ -1,135 +0,0 @@
-using Oceananigans
-using Oceananigans.Units
-using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity
-using Oceananigans.ImmersedBoundaries: PartialCellBottom, GridFittedBottom
-using ClimaOcean.IdealizedSimulations: neverworld_simulation
-using ClimaOcean.VerticalGrids: stretched_vertical_faces, PowerLawStretching
-using Printf
-using CUDA
-
-closure = CATKEVerticalDiffusivity(minimum_turbulent_kinetic_energy = 1e-6,
-                                   minimum_convective_buoyancy_flux = 1e-11)
-
-# closure = RiBasedVerticalDiffusivity()
-
-z = stretched_vertical_faces(surface_layer_Δz = 8,
-                             surface_layer_height = 64,
-                             stretching = PowerLawStretching(1.02),
-                             maximum_Δz = 400.0,
-                             minimum_depth = 4000)
-
-simulation = neverworld_simulation(GPU(); z,
-                                   #ImmersedBoundaryType = GridFittedBottom,
-                                   ImmersedBoundaryType = PartialCellBottom,
-                                   horizontal_resolution = 1/8,
-                                   longitude = (0, 60),
-                                   latitude = (-70, 0),
-                                   time_step = 10minutes,
-                                   stop_time = 4 * 360days,
-                                   closure)
-
-model = simulation.model
-grid = model.grid
-
-@show grid
-@show model
-
-start_time = Ref(time_ns())
-previous_model_time = Ref(time(simulation))
-
-function progress(sim) 
-    b = sim.model.tracers.b
-    e = sim.model.tracers.e
-    u, v, w = sim.model.velocities
-
-    msg = @sprintf("Iter: %d, time: %s, extrema(b): (%6.2e, %6.2e)",
-                   iteration(sim), prettytime(sim), minimum(b), maximum(b))
-
-    msg *= @sprintf(", max(e): %6.2e", maximum(e))
-
-    msg *= @sprintf(", max|u|: %6.2e, max|w|: %6.2e",
-                    maximum(maximum(abs, q) for q in (u, v, w)), maximum(abs, w))
-
-    try 
-        κᶜ = sim.model.diffusivity_fields.κᶜ
-        msg *= @sprintf(", max(κᶜ): %6.2e", maximum(κᶜ))
-    catch
-    end
-
-    elapsed = 1e-9 * (time_ns() - start_time[])
-    elapsed_model_time = time(sim) - previous_model_time[]
-    SYPD = (elapsed_model_time/360days) / (elapsed/day)
-
-    msg *= @sprintf(", wall time: %s, SYPD: %.1f", prettytime(elapsed), SYPD)
-    start_time[] = time_ns()
-    previous_model_time[] = time(sim)
-
-    @info msg
-
-    return nothing
-end
-
-simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))
-
-# Set up output
-Nx, Ny, Nz = size(grid)
-Δt = simulation.Δt
-Δt_minutes = round(Int, Δt / minutes) 
-ib_str = grid.immersed_boundary isa PartialCellBottom ? "partial_cells" : "full_cells"
-output_suffix = "$(Nx)_$(Ny)_$(Nz)_dt$(Δt_minutes)_$(ib_str).jld2"
-output_dir = "." #/nobackup1/glwagner/"
-fine_output_frequency = 1day
-i = round(Int, Nx/10) # index for yz-sliced output
-
-z = znodes(grid, Face(), with_halos=true)
-
-K = CUDA.@allowscalar [Nz,
-                       searchsortedfirst(z, -100),
-                       searchsortedfirst(z, -400)]
-
-i = round(Int, Nx/10) # index for yz-sliced output
-
-diffusivity_fields = (; κᶜ = model.diffusivity_fields.κᶜ)
-outputs = merge(model.velocities, model.tracers, diffusivity_fields)
-zonally_averaged_outputs = NamedTuple(n => Average(outputs[n], dims=1) for n in keys(outputs))
-
-simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs;
-                                                  schedule = TimeInterval(fine_output_frequency),
-                                                  filename = joinpath(output_dir, "neverworld_yz_" * output_suffix),
-                                                  indices = (i, :, :),
-                                                  with_halos = true,
-                                                  overwrite_existing = true)
-
-simulation.output_writers[:zonal] = JLD2OutputWriter(model, zonally_averaged_outputs;
-                                                     schedule = TimeInterval(fine_output_frequency),
-                                                     filename = joinpath(output_dir, "neverworld_zonal_average_" * output_suffix),
-                                                     with_halos = true,
-                                                     overwrite_existing = true)
-
-for (n, k) in enumerate(K)
-    name = Symbol(:xy, n)
-    simulation.output_writers[name] = JLD2OutputWriter(model, outputs;
-                                                       schedule = TimeInterval(fine_output_frequency),
-                                                       filename = joinpath(output_dir, "neverworld_xy$(n)_" * output_suffix),
-                                                       indices = (:, :, Nz),
-                                                       with_halos = true,
-                                                       overwrite_existing = true)
-end
-
-simulation.output_writers[:xyz] = JLD2OutputWriter(model, outputs;
-                                                   schedule = TimeInterval(90days),
-                                                   filename = joinpath(output_dir, "neverworld_xyz_" * output_suffix),
-                                                   with_halos = true,
-                                                   overwrite_existing = true)
-
-simulation.output_writers[:checkpointer] = Checkpointer(model,
-                                                        schedule = TimeInterval(360days),
-                                                        dir = output_dir,
-                                                        prefix = "neverworld_$(Nx)_$(Ny)_$(Nz)_checkpoint",
-                                                        cleanup = true)
-
-@info "Running..."
-@show simulation
-
-run!(simulation)
-
diff --git a/experiments/catke_partial_cells/run_neverworld_simulation.jl b/experiments/catke_partial_cells/run_neverworld_simulation.jl
deleted file mode 100644
index bc1fe488..00000000
--- a/experiments/catke_partial_cells/run_neverworld_simulation.jl
+++ /dev/null
@@ -1,141 +0,0 @@
-using Oceananigans
-using Oceananigans.Units
-using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity
-using Oceananigans.ImmersedBoundaries: PartialCellBottom, GridFittedBottom
-using ClimaOcean.IdealizedSimulations: neverworld_simulation
-using ClimaOcean.VerticalGrids: stretched_vertical_faces, PowerLawStretching
-using Printf
-using CUDA
-
-closure = CATKEVerticalDiffusivity(minimum_turbulent_kinetic_energy = 1e-6,
-                                   minimum_convective_buoyancy_flux = 1e-11)
-
-# closure = RiBasedVerticalDiffusivity()
-
-z = stretched_vertical_faces(surface_layer_Δz = 8,
-                             surface_layer_height = 64,
-                             stretching = PowerLawStretching(1.02),
-                             maximum_Δz = 400.0,
-                             minimum_depth = 4000)
-
-simulation = neverworld_simulation(GPU(); z,
-                                   ImmersedBoundaryType = GridFittedBottom,
-                                   #ImmersedBoundaryType = PartialCellBottom,
-                                   horizontal_resolution = 1/4,
-                                   longitude = (0, 60),
-                                   latitude = (-70, 0),
-                                   time_step = 10minutes,
-                                   stop_time = 1 * 360days,
-                                   closure)
-
-model = simulation.model
-grid = model.grid
-
-@show grid
-@show model
-
-start_time = Ref(time_ns())
-previous_model_time = Ref(time(simulation))
-
-_maximum(x) = maximum(parent(x))
-_maximum(f, x) = maximum(f, parent(x))
-_minimum(x) = minimum(parent(x))
-
-function progress(sim) 
-    b = sim.model.tracers.b
-    e = sim.model.tracers.e
-    u, v, w = sim.model.velocities
-
-    msg = @sprintf("Iter: %d, time: %s, extrema(b): (%6.2e, %6.2e)",
-                   iteration(sim), prettytime(sim), _minimum(b), _maximum(b))
-
-    msg *= @sprintf(", max(e): %6.2e", _maximum(e))
-
-    msg *= @sprintf(", max|u|: %6.2e, max|w|: %6.2e",
-                    maximum(_maximum(abs, q) for q in (u, v, w)), _maximum(abs, w))
-
-    try 
-        κᶜ = sim.model.diffusivity_fields.κᶜ
-        msg *= @sprintf(", max(κᶜ): %6.2e", _maximum(κᶜ))
-    catch
-    end
-
-    elapsed = 1e-9 * (time_ns() - start_time[])
-    elapsed_model_time = time(sim) - previous_model_time[]
-    SYPD = (elapsed_model_time/360days) / (elapsed/day)
-
-    msg *= @sprintf(", wall time: %s, SYPD: %.1f", prettytime(elapsed), SYPD)
-    start_time[] = time_ns()
-    previous_model_time[] = time(sim)
-
-    @info msg
-
-    return nothing
-end
-
-simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))
-
-# Set up output
-Nx, Ny, Nz = size(grid)
-Δt = simulation.Δt
-Δt_minutes = round(Int, Δt / minutes) 
-ib_str = grid.immersed_boundary isa PartialCellBottom ? "partial_cells" : "full_cells"
-output_suffix = "$(Nx)_$(Ny)_$(Nz)_dt$(Δt_minutes)_$(ib_str).jld2"
-output_dir = "."
-fine_output_frequency = 1day
-
-z = znodes(grid, Face(), with_halos=true)
-
-K = CUDA.@allowscalar [Nz,
-                       searchsortedfirst(z, -100),
-                       searchsortedfirst(z, -400)]
-
-I = [round(Int, Nx/10), round(Int, Nx/2)] # index for yz-sliced output
-
-diffusivity_fields = (; κᶜ = model.diffusivity_fields.κᶜ)
-outputs = merge(model.velocities, model.tracers, diffusivity_fields)
-zonally_averaged_outputs = NamedTuple(n => Average(outputs[n], dims=1) for n in keys(outputs))
-
-for (n, i) in enumerate(I)
-    name = Symbol(:yz, n)
-    simulation.output_writers[name] = JLD2OutputWriter(model, outputs;
-                                                       schedule = TimeInterval(fine_output_frequency),
-                                                       filename = joinpath(output_dir, "neverworld_yz$(n)_" * output_suffix),
-                                                       indices = (i, :, :),
-                                                       with_halos = true,
-                                                       overwrite_existing = true)
-end
-
-simulation.output_writers[:zonal] = JLD2OutputWriter(model, zonally_averaged_outputs;
-                                                     schedule = TimeInterval(fine_output_frequency),
-                                                     filename = joinpath(output_dir, "neverworld_zonal_average_" * output_suffix),
-                                                     with_halos = true,
-                                                     overwrite_existing = true)
-
-for (n, k) in enumerate(K)
-    name = Symbol(:xy, n)
-    simulation.output_writers[name] = JLD2OutputWriter(model, outputs;
-                                                       schedule = TimeInterval(fine_output_frequency),
-                                                       filename = joinpath(output_dir, "neverworld_xy$(n)_" * output_suffix),
-                                                       indices = (:, :, Nz),
-                                                       with_halos = true,
-                                                       overwrite_existing = true)
-end
-
-simulation.output_writers[:xyz] = JLD2OutputWriter(model, outputs;
-                                                   schedule = TimeInterval(90days),
-                                                   filename = joinpath(output_dir, "neverworld_xyz_" * output_suffix),
-                                                   with_halos = true,
-                                                   overwrite_existing = true)
-
-simulation.output_writers[:checkpointer] = Checkpointer(model,
-                                                        schedule = TimeInterval(360days),
-                                                        dir = output_dir,
-                                                        prefix = "neverworld_$(Nx)_$(Ny)_$(Nz)_checkpoint",
-                                                        cleanup = true)
-
-@info "Running..."
-@show simulation
-
-run!(simulation)
-
diff --git a/experiments/ri_based_partial_cells/run_neverworld_simulation.jl b/experiments/ri_based_partial_cells/run_neverworld_simulation.jl
deleted file mode 100644
index 49441745..00000000
--- a/experiments/ri_based_partial_cells/run_neverworld_simulation.jl
+++ /dev/null
@@ -1,135 +0,0 @@
-using Oceananigans
-using Oceananigans.Units
-using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity
-using Oceananigans.ImmersedBoundaries: PartialCellBottom, GridFittedBottom
-using ClimaOcean.IdealizedSimulations: neverworld_simulation
-using ClimaOcean.VerticalGrids: stretched_vertical_faces, PowerLawStretching
-using Printf
-using CUDA
-
-# closure = CATKEVerticalDiffusivity(minimum_turbulent_kinetic_energy = 1e-6,
-#                                    minimum_convective_buoyancy_flux = 1e-11)
- 
-closure = RiBasedVerticalDiffusivity()
-
-z = stretched_vertical_faces(surface_layer_Δz = 8,
-                             surface_layer_height = 64,
-                             stretching = PowerLawStretching(1.02),
-                             maximum_Δz = 400.0,
-                             minimum_depth = 4000)
-
-simulation = neverworld_simulation(GPU(); z,
-                                   #ImmersedBoundaryType = GridFittedBottom,
-                                   ImmersedBoundaryType = PartialCellBottom,
-                                   horizontal_resolution = 1/8,
-                                   longitude = (0, 60),
-                                   latitude = (-70, 0),
-                                   time_step = 10minutes,
-                                   stop_time = 4 * 360days,
-                                   closure)
-
-model = simulation.model
-grid = model.grid
-
-@show grid
-@show model
-
-start_time = Ref(time_ns())
-previous_model_time = Ref(time(simulation))
-
-function progress(sim) 
-    b = sim.model.tracers.b
-    e = sim.model.tracers.e
-    u, v, w = sim.model.velocities
-
-    msg = @sprintf("Iter: %d, time: %s, extrema(b): (%6.2e, %6.2e)",
-                   iteration(sim), prettytime(sim), minimum(b), maximum(b))
-
-    msg *= @sprintf(", max(e): %6.2e", maximum(e))
-
-    msg *= @sprintf(", max|u|: %6.2e, max|w|: %6.2e",
-                    maximum(maximum(abs, q) for q in (u, v, w)), maximum(abs, w))
-
-    try 
-        κᶜ = sim.model.diffusivity_fields.κᶜ
-        msg *= @sprintf(", max(κᶜ): %6.2e", maximum(κᶜ))
-    catch
-    end
-
-    elapsed = 1e-9 * (time_ns() - start_time[])
-    elapsed_model_time = time(sim) - previous_model_time[]
-    SYPD = (elapsed_model_time/360days) / (elapsed/day)
-
-    msg *= @sprintf(", wall time: %s, SYPD: %.1f", prettytime(elapsed), SYPD)
-    start_time[] = time_ns()
-    previous_model_time[] = time(sim)
-
-    @info msg
-
-    return nothing
-end
-
-simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))
-
-# Set up output
-Nx, Ny, Nz = size(grid)
-Δt = simulation.Δt
-Δt_minutes = round(Int, Δt / minutes) 
-ib_str = grid.immersed_boundary isa PartialCellBottom ? "partial_cells" : "full_cells"
-output_suffix = "$(Nx)_$(Ny)_$(Nz)_dt$(Δt_minutes)_$(ib_str).jld2"
-output_dir = "." #/nobackup1/glwagner/"
-fine_output_frequency = 1day
-i = round(Int, Nx/10) # index for yz-sliced output
-
-z = znodes(grid, Face(), with_halos=true)
-
-K = CUDA.@allowscalar [Nz,
-                       searchsortedfirst(z, -100),
-                       searchsortedfirst(z, -400)]
-
-i = round(Int, Nx/10) # index for yz-sliced output
-
-diffusivity_fields = (; κᶜ = model.diffusivity_fields.κᶜ)
-outputs = merge(model.velocities, model.tracers, diffusivity_fields)
-zonally_averaged_outputs = NamedTuple(n => Average(outputs[n], dims=1) for n in keys(outputs))
-
-simulation.output_writers[:yz] = JLD2OutputWriter(model, outputs;
-                                                  schedule = TimeInterval(fine_output_frequency),
-                                                  filename = joinpath(output_dir, "neverworld_yz_" * output_suffix),
-                                                  indices = (i, :, :),
-                                                  with_halos = true,
-                                                  overwrite_existing = true)
-
-simulation.output_writers[:zonal] = JLD2OutputWriter(model, zonally_averaged_outputs;
-                                                     schedule = TimeInterval(fine_output_frequency),
-                                                     filename = joinpath(output_dir, "neverworld_zonal_average_" * output_suffix),
-                                                     with_halos = true,
-                                                     overwrite_existing = true)
-
-for (n, k) in enumerate(K)
-    name = Symbol(:xy, n)
-    simulation.output_writers[name] = JLD2OutputWriter(model, outputs;
-                                                       schedule = TimeInterval(fine_output_frequency),
-                                                       filename = joinpath(output_dir, "neverworld_xy$(n)_" * output_suffix),
-                                                       indices = (:, :, Nz),
-                                                       with_halos = true,
-                                                       overwrite_existing = true)
-end
-
-simulation.output_writers[:xyz] = JLD2OutputWriter(model, outputs;
-                                                   schedule = TimeInterval(90days),
-                                                   filename = joinpath(output_dir, "neverworld_xyz_" * output_suffix),
-                                                   with_halos = true,
-                                                   overwrite_existing = true)
-
-simulation.output_writers[:checkpointer] = Checkpointer(model,
-                                                        schedule = TimeInterval(360days),
-                                                        dir = output_dir,
-                                                        prefix = "neverworld_$(Nx)_$(Ny)_$(Nz)_checkpoint",
-                                                        cleanup = true)
-
-@info "Running..."
-@show simulation
-
-run!(simulation)
-
diff --git a/src/Bathymetry.jl b/src/Bathymetry.jl
index 37ae85e5..b38470a3 100644
--- a/src/Bathymetry.jl
+++ b/src/Bathymetry.jl
@@ -1,11 +1,18 @@
 module Bathymetry
 
+export regrid_bathymetry
+
 using ..DataWrangling: download_progress
 
 using Oceananigans
 using Oceananigans.Architectures: architecture
-using Oceananigans.Grids: halo_size, λnodes
-using Oceananigans.Utils: pretty_filesize
+using Oceananigans.Grids: halo_size, λnodes, φnodes
+using Oceananigans.Grids: x_domain, y_domain
+using Oceananigans.Grids: topology
+using Oceananigans.Utils: pretty_filesize, launch!
+using Oceananigans.Fields: interpolate!
+using Oceananigans.BoundaryConditions
+using KernelAbstractions: @kernel, @index
 
 using NCDatasets
 using Downloads
@@ -22,14 +29,46 @@ using Printf
 Regrid bathymetry associated with the NetCDF file at `path = joinpath(dir, filename)` to `target_grid`.
 If `path` does not exist, then a download is attempted from `joinpath(url, filename)`.
 
-TODO: describe keyword arguments.
+Arguments:
+==========
+
+- target_grid: grid to interpolate onto
+
+Keyword Arguments:
+==================
+
+- height_above_water: limits the maximum height of above-water topography (where h > 0). If
+                      `nothing` the original topography is retained
+
+- minimum_depth: minimum depth for the shallow regions. `h > minimum_depth` will be considered land
+
+- dir: directory of the bathymetry-containing file
+
+- filename: file containing bathymetric data. Must be netcdf with fields:
+            (1) `lat` vector of latitude nodes
+            (2) `lon` vector of longitude nodes
+            (3) `z` matrix of depth values
+
+- interpolation_passes: regridding/interpolation passes. The bathymetry is interpolated in
+                        `interpolation_passes - 1` intermediate steps. With more steps the 
+                        final bathymetry will be smoother.
+                        Example: interpolating from a 400x200 grid to a 100x100 grid in 4 passes will involve
+                        - 400x200 -> 325x175
+                        - 325x175 -> 250x150
+                        - 250x150 -> 175x125
+                        - 175x125 -> 100x100
+                        If _coarsening_ the original grid, linear interpolation in passes is equivalent to 
+                        applying a smoothing filter, with more passes increasing the strength of the filter.
+                        If _refining_ the original grid, additional passes will not help and no intermediate
+                        steps will be performed.
 """
 function regrid_bathymetry(target_grid;
                            height_above_water = nothing,
                            minimum_depth = 0,
                            dir = joinpath(@__DIR__, "..", "data"),
                            url = "https://www.ngdc.noaa.gov/thredds/fileServer/global/ETOPO2022/60s/60s_surface_elev_netcdf", 
-                           filename = "ETOPO_2022_v1_60s_N90W180_surface.nc")
+                           filename = "ETOPO_2022_v1_60s_N90W180_surface.nc",
+                           interpolation_passes = 1)
 
     filepath = joinpath(dir, filename)
 
@@ -48,22 +87,18 @@ function regrid_bathymetry(target_grid;
 
     dataset = Dataset(filepath)
 
+    FT = eltype(target_grid)
+
     φ_data = dataset["lat"][:]
     λ_data = dataset["lon"][:]
-    h_data = dataset["z"][:, :]
+    h_data = convert.(FT, dataset["z"][:, :])
 
     close(dataset)
 
     # Diagnose target grid information
-    Nxt, Nyt, Nzt = size(target_grid)
     arch = architecture(target_grid)
-    λt, φt, zt = nodes(target_grid, Face(), Face(), Face())
-
-    λ₁ = λt[1]
-    λ₂ = λt[Nxt]
-
-    φ₁ = φt[1]
-    φ₂ = φt[Nyt]
+    φ₁, φ₂ = y_domain(target_grid)
+    λ₁, λ₂ = x_domain(target_grid)
 
     # Calculate limiting indices on the bathymetry grid
     i₁ = searchsortedfirst(λ_data, λ₁)
@@ -106,8 +141,8 @@ function regrid_bathymetry(target_grid;
     end
 
     if minimum_depth > 0
-        ocean = h_data .<= 0
-        h_data[ocean] .= min.(-minimum_depth, h_data[ocean])
+        shallow_ocean = h_data .> minimum_depth
+        h_data[shallow_ocean] .= height_above_water
     end
 
     # Build the "native" grid of the bathymetry and make a bathymetry field.
@@ -120,33 +155,61 @@ function regrid_bathymetry(target_grid;
                                         latitude = (φ₁, φ₂),
                                         longitude = (λ₁, λ₂),
                                         z = (0, 1),
-                                        halo = halo_size(target_grid))
+                                        halo = (10, 10, 1))
 
     native_h = Field{Center, Center, Nothing}(native_grid)
     set!(native_h, h_data)
 
-    Nxi = Nxt
-    Nyi = Nyn
-    Nzi = Nzn
+    target_h = interpolate_bathymetry_in_passes(native_h, target_grid; passes = interpolation_passes)
 
-    if parent(parent(λt)) isa StepRangeLen # longitude is equispaced
-        longitude = (λ₁, λ₂)
-    else
-        longitude = λnodes(target_grid, Face(), Center(), Center())
+    return target_h
+end
+
+# Here we can either use `regrid!` (three dimensional version) or `interpolate`
+function interpolate_bathymetry_in_passes(native_h, target_grid; passes = 10)
+    Nλt, Nφt = Nt = size(target_grid)
+    Nλn, Nφn = Nn = size(native_h)
+    
+    if any(Nt[1:2] .> Nn[1:2]) # We are refining the grid (at least in one direction), more passes will not help!
+        target_h = Field{Center, Center, Nothing}(target_grid)
+        interpolate!(target_h, native_h)
+        return target_h
     end
+ 
+    latitude  = y_domain(target_grid)
+    longitude = x_domain(target_grid)
+
+    ΔNλ = floor((Nλn - Nλt) / passes)
+    ΔNφ = floor((Nφn - Nφt) / passes)
+
+    Nλ = [Nλn - ΔNλ * pass for pass in 1:passes-1]
+    Nφ = [Nφn - ΔNφ * pass for pass in 1:passes-1]
 
-    intermediate_grid = LatitudeLongitudeGrid(arch;
-                                              size = (Nxi, Nyi, Nzi),
-                                              latitude = (φ₁, φ₂),
-                                              longitude,
-                                              z = (0, 1),
-                                              halo = halo_size(target_grid))
+    Nλ = Int[Nλ..., Nλt]
+    Nφ = Int[Nφ..., Nφt]
 
-    intermediate_h =  Field{Center, Center, Nothing}(intermediate_grid)
-    regrid!(intermediate_h, native_h)
+    old_h     = native_h
+    TX, TY, _ = topology(target_grid)
+
+    for pass = 1:passes - 1
+        new_size = (Nλ[pass], Nφ[pass], 1)
+
+        @debug "pass number $pass with size $new_size"
+        new_grid = LatitudeLongitudeGrid(architecture(target_grid),
+                                         size = new_size, 
+                                     latitude = (latitude[1],  latitude[2]), 
+                                    longitude = (longitude[1], longitude[2]), 
+                                            z = (0, 1),
+                                     topology = (TX, TY, Bounded))
+
+        new_h = Field{Center, Center, Nothing}(new_grid)
+
+        interpolate!(new_h, old_h)
+        old_h = new_h
+    end
 
     target_h = Field{Center, Center, Nothing}(target_grid)
-    regrid!(target_h, intermediate_h)
+    interpolate!(target_h, old_h)
 
     return target_h
 end
diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl
index d256bf13..61b51ce7 100644
--- a/src/ClimaOcean.jl
+++ b/src/ClimaOcean.jl
@@ -1,9 +1,15 @@
 module ClimaOcean
 
+export regrid_bathymetry
+export stretched_vertical_faces
+export PowerLawStretching, LinearStretching
+export jra55_field_time_series
+export ecco2_field, ECCO2Metadata
+export initialize!
+
 using Oceananigans
 using Oceananigans.Operators: ℑxyᶠᶜᵃ, ℑxyᶜᶠᵃ
 using DataDeps
-using CubicSplines
 
 function __init__(; remove_existing_data=false)
 
@@ -59,13 +65,16 @@ end
 @inline v_immersed_bottom_drag(i, j, k, grid, c, Φ, μ) = @inbounds - μ * Φ.v[i, j, k] * spᶜᶠᶜ(i, j, k, grid, Φ)
 
 include("VerticalGrids.jl")
+include("InitialConditions/InitialConditions.jl")
 include("DataWrangling/DataWrangling.jl")
 include("Bathymetry.jl")
-include("InitialConditions.jl")
 include("Diagnostics.jl")
 include("NearGlobalSimulations/NearGlobalSimulations.jl")
 
+using .VerticalGrids
+using .Bathymetry
 using .DataWrangling: JRA55
 using .DataWrangling: ECCO2
+using .InitialConditions
 
 end # module
diff --git a/src/DataWrangling/DataWrangling.jl b/src/DataWrangling/DataWrangling.jl
index 9ccbc324..fb2a5b4a 100644
--- a/src/DataWrangling/DataWrangling.jl
+++ b/src/DataWrangling/DataWrangling.jl
@@ -8,6 +8,7 @@ using Oceananigans.Architectures: architecture
 using Oceananigans.Grids: node
 using Oceananigans.BoundaryConditions: fill_halo_regions!
 using Oceananigans.Fields: interpolate
+using Oceananigans: pretty_filesize
 using Oceananigans.Utils: launch!
 using KernelAbstractions: @kernel, @index
 
@@ -107,6 +108,7 @@ function save_field_time_series!(fts; path, name, overwrite_existing=false)
     return nothing
 end
 
+include("inpaint_mask.jl")
 include("JRA55.jl")
 include("ECCO2.jl")
 
diff --git a/src/DataWrangling/ECCO2.jl b/src/DataWrangling/ECCO2.jl
index 7145f99d..28bfd255 100644
--- a/src/DataWrangling/ECCO2.jl
+++ b/src/DataWrangling/ECCO2.jl
@@ -1,13 +1,32 @@
 module ECCO2
 
+export ECCO2Metadata, ecco2_field, ecco2_center_mask, adjusted_ecco_tracers, initialize!
+
+using ClimaOcean.DataWrangling: inpaint_mask!
+using ClimaOcean.InitialConditions: three_dimensional_regrid!
+
 using Oceananigans
+using Oceananigans: architecture
 using Oceananigans.BoundaryConditions
+using Oceananigans.Utils
+using KernelAbstractions: @kernel, @index
 using NCDatasets
 using Downloads: download
 
-temperature_filename = "THETA.1440x720x50.19920102.nc"
-salinity_filename = "SALT.1440x720x50.19920102.nc"
-effective_ice_thickness_filename = "SIheff.1440x720.19920102.nc"
+import Oceananigans.Fields: set!
+
+# Ecco field used to set model's initial conditions
+struct ECCO2Metadata
+    name  :: Symbol
+    year  :: Int
+    month :: Int
+    day   :: Int
+end
+
+# We only have 1992 at the moment
+ECCO2Metadata(name::Symbol) = ECCO2Metadata(name, 1992, 1, 2)
+
+filename(data::ECCO2Metadata) = "ecco2_" * string(data.name) * "_$(data.year)$(data.month)$(data.day).nc"
 
 ecco2_short_names = Dict(
     :temperature   => "THETA",
@@ -15,6 +34,12 @@ ecco2_short_names = Dict(
     :effective_ice_thickness => "SIheff"
 )
 
+ecco2_location = Dict(
+    :temperature   => (Center, Center, Center),
+    :salinity      => (Center, Center, Center),
+    :effective_ice_thickness => (Center, Center, Nothing)
+)
+
 ecco2_depth_names = Dict(
     :temperature   => "DEPTH_T",
     :salinity      => "DEPTH_T",
@@ -60,79 +85,199 @@ function construct_vertical_interfaces(ds, depth_name)
     return zf
 end
 
-function ecco2_field(variable_name;
-                     architecture = CPU(),
-                     horizontal_halo = (1, 1),
-                     url = ecco2_urls[variable_name],
-                     filename = ecco2_file_names[variable_name],
-                     short_name = ecco2_short_names[variable_name])
+function empty_ecco2_field(data::ECCO2Metadata;
+                           architecture = CPU(), 
+                           horizontal_halo = (1, 1))
 
-    isfile(filename) || download(url, filename)
+    variable_name = data.name
 
-    ds = Dataset(filename)
+    location = ecco2_location[variable_name]
 
     longitude = (0, 360)
     latitude = (-90, 90)
     TX, TY = (Periodic, Bounded)
 
+    filename = ecco2_file_names[variable_name]
+    
+    ds = Dataset(filename)
+
     if variable_is_three_dimensional[variable_name] 
-        data = ds[short_name][:, :, :, 1]
         depth_name = ecco2_depth_names[variable_name]
-        
-        # The surface layer in three-dimensional ECCO fields is at `k = 1`
-        data = reverse(data, dims = 3)
-        
         z    = construct_vertical_interfaces(ds, depth_name)
-        N    = size(data)
-
         # add vertical halo for 3D fields
         halo = (horizontal_halo..., 1)
-
         LZ   = Center
         TZ   = Bounded
+        N    = (1440, 720, 50)
     else
-        data = ds[short_name][:, :, 1]
-        N    = size(data)
         z    = nothing
         halo = horizontal_halo
         LZ   = Nothing
         TZ   = Flat
+        N    = (1440, 720)
     end
 
-    close(ds)
-
     # Flat in z if the variable is two-dimensional
     grid = LatitudeLongitudeGrid(architecture; halo, size = N, topology = (TX, TY, TZ),
                                  longitude, latitude, z)
 
-    field = Field{Center, Center, LZ}(grid)
-    
+    return Field{location...}(grid)
+end
+
+"""
+    ecco2_field(variable_name;
+                architecture = CPU(),
+                horizontal_halo = (1, 1),
+                user_data = nothing,
+                url = ecco2_urls[variable_name],
+                filename = ecco2_file_names[variable_name],
+                short_name = ecco2_short_names[variable_name])
+
+Retrieve the ecco2 field corresponding to `variable_name`. 
+The data is either:
+(1) retrieved from `filename`,
+(2) dowloaded from `url` if `filename` does not exists,
+(3) filled from `user_data` if `user_data` is provided.
+"""
+function ecco2_field(variable_name;
+                     architecture = CPU(),
+                     horizontal_halo = (1, 1),
+                     user_data = nothing,
+                     year = 1992,
+                     month = 1,
+                     day  = 2, 
+                     url  = ecco2_urls[variable_name],
+                     filename = ecco2_file_names[variable_name],
+                     short_name = ecco2_short_names[variable_name])
+
+    ecco2_data = ECCO2Metadata(variable_name, year, month, day)
+
+    isfile(filename) || download(url, filename)
+
+    if user_data isa Nothing
+        ds = Dataset(filename)
+        
+        if variable_is_three_dimensional[variable_name] 
+            data = ds[short_name][:, :, :, 1]
+            # The surface layer in three-dimensional ECCO fields is at `k = 1`
+            data = reverse(data, dims = 3)
+        else
+            data = ds[short_name][:, :, 1]
+        end        
+    else
+        data = user_data
+    end
+
+    field = empty_ecco2_field(ecco2_data; architecture, horizontal_halo)
+    FT    = eltype(field)
+    data  = convert.(FT, data)
+
     set!(field, data)
     fill_halo_regions!(field)
 
     return field
 end
 
-function ecco2_bottom_height_from_temperature()
-    Tᵢ = ecco2_field(:temperature)
+@kernel function _set_ecco2_mask!(mask, Tᵢ, minimum_value)
+    i, j, k = @index(Global, NTuple)
+    @inbounds mask[i, j, k] = Tᵢ[i, j, k] < minimum_value
+end
+
+"""
+    ecco2_center_mask(architecture = CPU(); minimum_value = Float32(-1e5))
+
+A boolean field where `false` represents a missing value in the ECCO2 :temperature dataset.
+"""
+function ecco2_center_mask(architecture = CPU(); minimum_value = Float32(-1e5))
+    Tᵢ   = ecco2_field(:temperature; architecture)
+    mask = CenterField(Tᵢ.grid, Bool)
+
+    # Set the mask with ones where T is defined
+    launch!(architecture, Tᵢ.grid, :xyz, _set_ecco2_mask!, mask, Tᵢ, minimum_value)
+
+    return mask
+end
+
+"""
+    inpainted_ecco2_field(variable_name; 
+                          architecture = CPU(),
+                          filename = "./inpainted_ecco2_fields.nc",
+                          mask = ecco2_center_mask(architecture))
+    
+Retrieve the ECCO2 field corresponding to `variable_name` inpainted to fill all the
+missing values in the original dataset.
+
+Arguments:
+==========
+
+- `variable_name`: the variable name corresponding to the Dataset.
 
-    missing_value = Float32(-9.9e22)
+Keyword Arguments:
+==================
 
-    # Construct bottom_height depth by analyzing T
-    Nx, Ny, Nz = size(Tᵢ)
-    bottom_height = ones(Nx, Ny) .* (zf[1] - Δz)
-    zf = znodes(Tᵢ.grid, Face())
+- `architecture`: either `CPU()` or `GPU()`.
+
+- `filename`: the path where to retrieve the data from. If the file does not exist,
+              the data will be retrived from the ECCO2 dataset, inpainted, and
+              saved to `filename`.
+
+- `mask`: the mask used to inpaint the field (see `inpaint_mask!`).
+"""
+function inpainted_ecco2_field(variable_name; 
+                               architecture = CPU(),
+                               filename = "./inpainted_ecco2_fields.nc",
+                               mask = ecco2_center_mask(architecture),
+                               kw...)
     
-    for i = 1:Nx, j = 1:Ny
-        @inbounds for k = Nz:-1:1
-            if Tᵢ[i, j, k] < -10
-                bottom_height[i, j] = zf[k+1]
-                break
-            end
+    if !isfile(filename)
+        f = ecco2_field(variable_name; architecture)
+        
+        # Make sure all values are extended properly
+        @info "In-painting ecco field $variable_name and saving it in $filename"
+        inpaint_mask!(f, mask; kw...)
+
+        ds = Dataset(filename, "c")
+        defVar(ds, string(variable_name), Array(interior(f)), ("lat", "lon", "z"))
+
+        close(ds)
+    else
+        ds = Dataset(filename, "a")
+
+        if haskey(ds, string(variable_name))
+            data = ds[variable_name][:, :, :]
+            f = ecco2_field(variable_name; architecture, user_data = data)
+        else
+            f = ecco2_field(variable_name; architecture)
+            # Make sure all values are inpainted properly
+            @info "In-painting ecco field $variable_name and saving it in $filename"
+            inpaint_mask!(f, mask; kw...)
+
+            defVar(ds, string(variable_name), Array(interior(f)), ("lat", "lon", "z"))
         end
+
+        close(ds)
     end
 
-    return bottom_height
+    return f
+end
+
+function set!(field::Field, ecco2_metadata::ECCO2Metadata; filename="./inpainted_ecco2_fields.nc", kw...)
+
+    # Fields initialized from ECCO2
+    grid = field.grid
+    name = ecco2_metadata.name
+
+    mask = ecco2_center_mask(architecture(grid))
+    
+    f = inpainted_ecco2_field(name; filename, mask,
+                              architecture = architecture(grid),
+                              kw...)
+
+    f_grid = Field(ecco2_location[name], grid)   
+    three_dimensional_regrid!(f_grid, f)
+    set!(field, f_grid)
+
+    return field
 end
 
 end # module
diff --git a/src/DataWrangling/JRA55.jl b/src/DataWrangling/JRA55.jl
index 5cf4d760..785a5766 100644
--- a/src/DataWrangling/JRA55.jl
+++ b/src/DataWrangling/JRA55.jl
@@ -8,17 +8,17 @@ using Downloads: download
 
 # A list of all variables provided in the JRA55 dataset:
 jra55_short_names = (:freshwater_river_flux,
-                        :freshwater_rain_flux,
-                        :freshwater_snow_flux,
-                        :freshwater_iceberg_flux,
-                        :specific_humidity,
-                        :sea_level_pressure,
-                        :relative_humidity,
-                        :downwelling_longwave_radiation,
-                        :downwelling_shortwave_radiation,
-                        :atmospheric_temperature,
-                        :atmospheric_eastward_velocity,
-                        :atmospheric_northward_velocity)
+                     :freshwater_rain_flux,
+                     :freshwater_snow_flux,
+                     :freshwater_iceberg_flux,
+                     :specific_humidity,
+                     :sea_level_pressure,
+                     :relative_humidity,
+                     :downwelling_longwave_radiation,
+                     :downwelling_shortwave_radiation,
+                     :atmospheric_temperature,
+                     :atmospheric_eastward_velocity,
+                     :atmospheric_northward_velocity)
 
 file_names = Dict(
     :freshwater_river_flux           => "RYF.friver.1990_1991.nc",   # Freshwater fluxes from rivers
diff --git a/src/DataWrangling/inpaint_mask.jl b/src/DataWrangling/inpaint_mask.jl
new file mode 100644
index 00000000..c6ea8c45
--- /dev/null
+++ b/src/DataWrangling/inpaint_mask.jl
@@ -0,0 +1,126 @@
+using Oceananigans
+using Oceananigans.BoundaryConditions
+using Oceananigans.Fields: OneField
+using Oceananigans.Grids: peripheral_node
+using Oceananigans.Utils: launch!
+using Oceananigans.Fields: instantiated_location, interior, CenterField
+using Oceananigans.Architectures: architecture, device, GPU
+
+using KernelAbstractions: @kernel, @index
+using KernelAbstractions.Extras.LoopInfo: @unroll
+
+# Maybe we can remove this propagate field in lieu of a diffusion, 
+# Still we'll need to do this a couple of steps on the original grid
+@kernel function _propagate_field!(field, tmp_field)
+    i, j, k = @index(Global, NTuple)
+
+    @inbounds begin
+        nw = field[i - 1, j, k]
+        ns = field[i, j - 1, k]
+        ne = field[i + 1, j, k]
+        nn = field[i, j + 1, k]
+        nb = (nw, ne, nn, ns)
+    end
+
+    counter = 0
+    cumsum  = zero(eltype(field))
+
+    @unroll for n in nb
+        counter += ifelse(isnan(n), 0, 1)
+        cumsum  += ifelse(isnan(n), 0, n)
+    end
+
+    @inbounds tmp_field[i, j, k] = ifelse(cumsum == 0, NaN, cumsum / counter)
+end
+
+@kernel function _substitute_values!(field, tmp_field)
+    i, j, k = @index(Global, NTuple)
+    @inbounds needs_inpainting = isnan(field[i, j, k])
+    @inbounds field[i, j, k] = ifelse(needs_inpainting, tmp_field[i, j, k], field[i, j, k])
+end
+
+@kernel function _nan_mask!(field, mask)
+    i, j, k = @index(Global, NTuple)
+    @inbounds field[i, j, k] = ifelse(mask[i, j, k], NaN, field[i, j, k])
+end
+
+propagate_horizontally!(field, ::Nothing, tmp_field=deepcopy(field); kw...) = field
+
+function propagating(field, mask, iter, max_iter)
+    mask_sum = sum(field; condition=interior(mask))
+    return isnan(mask_sum) && iter < max_iter
+end
+
+""" 
+    propagate_horizontally!(field, mask [, tmp_field=deepcopy(field)]; max_iter = Inf)
+
+Horizontally propagate the values of `field` into the `mask`.
+In other words, cells where `mask[i, j, k] == false` are preserved,
+and cells where `mask[i, j, k] == true` are painted over.
+"""
+function propagate_horizontally!(field, mask, tmp_field=deepcopy(field); max_iter = Inf) 
+    iter  = 0
+    grid  = field.grid
+    arch  = architecture(grid)
+    
+    launch!(arch, grid, :xyz, _nan_mask!, field, mask)
+    fill_halo_regions!(field)
+
+    # Need temporary field to avoid a race condition
+    parent(tmp_field) .= parent(field)
+
+    while propagating(field, mask, iter, max_iter)
+        launch!(arch, grid, :xyz, _propagate_field!,   field, tmp_field)
+        launch!(arch, grid, :xyz, _substitute_values!, field, tmp_field)
+        iter += 1
+        @debug "Propagate pass $iter with sum $(sum(parent(field)))"
+    end
+
+    return field
+end
+
+continue_downwards!(field, ::Nothing) = field
+
+""" 
+    continue_downwards!(field, mask)
+
+Continue downwards a field with missing values within `mask`.
+Cells where `mask[i, k, k] == false` will be preserved.
+"""
+function continue_downwards!(field, mask)
+    arch = architecture(field)
+    grid = field.grid
+    launch!(arch, grid, :xy, _continue_downwards!, field, grid, mask)
+    return field
+end
+
+@kernel function _continue_downwards!(field, grid, mask)
+    i, j = @index(Global, NTuple)
+    Nz = size(grid, 3)
+
+    @unroll for k = Nz-1 : -1 : 1
+        @inbounds field[i, j, k] = ifelse(mask[i, j, k], field[i, j, k+1], field[i, j, k])
+    end
+end
+
+"""
+    inpaint_mask!(field, mask; max_iter = Inf)
+
+Inpaint field within `mask`, using values outside `mask`.
+In other words, regions where `mask[i, j, k] == 1` will be inpainted
+and regions where `mask[i, j, k] == 0` will be preserved.
+
+Arguments
+=========
+    - `field`: `Field` to be inpainted.
+    - `mask`: Boolean-valued `Field`, values where
+              `mask[i, j, k] == true` are inpainted.
+    - `max_iter`: Maximum iterations for inpainting. Non-Inf values mean that
+                  NaN's can occur within the mask.
+"""
+function inpaint_mask!(field, mask; max_iter = Inf)
+    continue_downwards!(field, mask)
+    propagate_horizontally!(field, mask; max_iter)
+    return field
+end
+
diff --git a/src/InitialConditions.jl b/src/InitialConditions.jl
deleted file mode 100644
index 42272758..00000000
--- a/src/InitialConditions.jl
+++ /dev/null
@@ -1,83 +0,0 @@
-module InitialConditions
-
-export continue_downwards!
-
-using Oceananigans
-using Oceananigans.Grids: peripheral_node
-using Oceananigans.Utils: launch!
-using Oceananigans.Fields: instantiated_location, interior, CenterField
-using Oceananigans.Architectures: architecture, device, GPU
-
-using KernelAbstractions: @kernel, @index
-using KernelAbstractions.Extras.LoopInfo: @unroll
-
-function continue_downards!(field)
-    arch = architecture(field)
-    grid = field.grid
-    loc = instantiated_location(field)
-    launch!(arch, grid, :xy, _continue_downwards!, field, loc, grid)
-    return nothing
-end
-
-@kernel function _continue_downwards!(field, (LX, LY, LZ), grid)
-    i, j = @index(Global, NTuple)
-
-    Nz = grid.Nz
-    active_surface = !peripheral_node(i, j, Nz, grid, LX, LY, LZ)
-
-    @unroll for k = Nz-1 : -1 : 1
-        fill_from_above = active_surface & peripheral_node(i, j, k, grid, LX, LY, LZ)
-        @inbounds field[i, j, k] = ifelse(fill_from_above, field[i, j, k+1], field[i, j, k])
-    end
-end
-
-scale_to_diffusivity(vertical_scale) = vertical_scale^2
-scale_to_diffusivity(vertical_scale::Function) = (x, y, z, t) -> vertical_scale(x, y, z, t)^2
-
-function diffuse_tracers!(grid;
-                          tracers,
-                          horizontal_scale = 0,
-                          vertical_scale = 0,
-                          fractional_time_step = 2e-2)
-
-    # Remake tracers without boundary conditions
-    tracers = NamedTuple(name => CenterField(grid; data=tracers[name].data) for name in keys(tracers))
-
-    # Horizontal diffusivities that mix up to t ∼ ℓ² / κ ∼ 1
-    κh = horizontal_scale^2
-    κz = scale_to_diffusivity(vertical_scale)
-
-    # Determine stable time-step
-    Nx, Ny, Nz = size(grid)
-    ϵ = fractional_time_step
-    Az = minimum(grid.Azᶜᶜᵃ[1:Ny])
-
-    if κh == 0
-        Δt = fractional_time_step
-    else
-        Δt = ϵ * Az / κh
-    end
-    @show Nt = ceil(Int, 1 / Δt)
-
-    vitd = VerticallyImplicitTimeDiscretization()
-    vertical_smoothing = VerticalScalarDiffusivity(vitd, κ=κz)
-    horizontal_smoothing = HorizontalScalarDiffusivity(κ=κh)
-
-    smoothing_model = HydrostaticFreeSurfaceModel(; grid, tracers,
-                                                  velocities = PrescribedVelocityFields(),
-                                                  momentum_advection = nothing,
-                                                  tracer_advection = nothing,
-                                                  buoyancy = nothing,
-                                                  closure = (horizontal_smoothing, vertical_smoothing))
-
-    Nt = ceil(Int, 1 / Δt)
-    @info string("Smoothing tracers ", keys(tracers), " with ", Nt, " time steps")
-    smoothing_simulation = Simulation(smoothing_model; Δt, stop_time=1)
-    pop!(smoothing_simulation.callbacks, :nan_checker) 
-    run!(smoothing_simulation)
-
-    return nothing
-end
-
-end # module
-
diff --git a/src/InitialConditions/InitialConditions.jl b/src/InitialConditions/InitialConditions.jl
new file mode 100644
index 00000000..9edc1e7b
--- /dev/null
+++ b/src/InitialConditions/InitialConditions.jl
@@ -0,0 +1,73 @@
+module InitialConditions
+
+export initialize!
+
+using Oceananigans
+using Oceananigans.BoundaryConditions
+using Oceananigans.Fields: OneField
+using Oceananigans.Grids: peripheral_node
+using Oceananigans.Utils: launch!
+using Oceananigans.Fields: instantiated_location, interior, CenterField
+using Oceananigans.Architectures: architecture, device, GPU
+
+using KernelAbstractions: @kernel, @index
+using KernelAbstractions.Extras.LoopInfo: @unroll
+using JLD2
+
+# Implementation of 3-dimensional regridding
+# TODO: move all the following to Oceananigans! 
+
+using Oceananigans.Fields: regrid!
+using Oceananigans.Grids: cpu_face_constructor_x, 
+                          cpu_face_constructor_y, 
+                          cpu_face_constructor_z,
+                          topology
+
+# Should we move this to grids??
+construct_grid(::Type{<:RectilinearGrid}, arch, size, extent, topology) = 
+    RectilinearGrid(arch; size, x = extent[1], y = extent[2], z = extent[2], topology)
+
+construct_grid(::Type{<:LatitudeLongitudeGrid}, arch, size, extent, topology) = 
+    LatitudeLongitudeGrid(arch; size, longitude = extent[1], latitude = extent[2], z = extent[3], topology)
+
+# Regrid a field in three dimensions
+function three_dimensional_regrid!(a, b)
+    target_grid = a.grid isa ImmersedBoundaryGrid ? a.grid.underlying_grid : a.grid
+    source_grid = b.grid isa ImmersedBoundaryGrid ? b.grid.underlying_grid : b.grid 
+
+    topo = topology(target_grid)
+    arch = architecture(target_grid)
+    
+    target_y = yt = cpu_face_constructor_y(target_grid)
+    target_z = zt = cpu_face_constructor_z(target_grid)
+
+    target_size = Nt = size(target_grid)
+
+    source_x = xs = cpu_face_constructor_x(source_grid)
+    source_y = ys = cpu_face_constructor_y(source_grid)
+
+    source_size = Ns = size(source_grid)
+
+    # Start by regridding in z
+    @debug "Regridding in z"
+    zgrid   = construct_grid(typeof(target_grid), arch, (Ns[1], Ns[2], Nt[3]), (xs, ys, zt), topo)
+    field_z = Field(location(b), zgrid)
+    regrid!(field_z, zgrid, source_grid, b)
+
+    # regrid in y 
+    @debug "Regridding in y"
+    ygrid   = construct_grid(typeof(target_grid), arch, (Ns[1], Nt[2], Nt[3]), (xs, yt, zt), topo)
+    field_y = Field(location(b), ygrid);
+    regrid!(field_y, ygrid, zgrid, field_z);
+
+    # Finally regrid in x
+    @debug "Regridding in x"
+    regrid!(a, target_grid, ygrid, field_y)
+
+    return a
+end
+
+include("diffuse_tracers.jl")
+
+end # module
+
diff --git a/src/InitialConditions/diffuse_tracers.jl b/src/InitialConditions/diffuse_tracers.jl
new file mode 100644
index 00000000..b906ad11
--- /dev/null
+++ b/src/InitialConditions/diffuse_tracers.jl
@@ -0,0 +1,80 @@
+
+scale_to_diffusivity(vertical_scale) = vertical_scale^2
+scale_to_diffusivity(vertical_scale::Function) = (x, y, z, t) -> vertical_scale(x, y, z, t)^2
+
+@kernel function _apply_tracer_mask!(tracer, initial_tracer, mask::AbstractArray)
+    i, j, k = @index(Global, NTuple)
+    @inbounds tracer[i, j, k] = ifelse(mask[i, j, k] == 1, initial_tracer[i, j, k], tracer[i, j, k])
+end
+
+@kernel function _apply_tracer_mask!(tracer, initial_tracer, mask::Function)
+    i, j, k = @index(Global, NTuple)
+    @inbounds tracer[i, j, k] = ifelse(mask(i, j, k, grid) == 1, initial_tracer[i, j, k], tracer[i, j, k])
+end
+
+@inline not_an_immersed_cell(i, j, k, grid) = !immersed_cell(i, j, k, grid)
+
+function diffuse_tracers(initial_tracers;
+                         horizontal_scale = 0,
+                         vertical_scale = 0,
+                         fractional_time_step = 2e-2,
+                         mask = not_an_immersed_cell)
+
+    # Remake the grid
+    grid = initial_tracers[1].grid
+
+    # Remake tracers without boundary conditions
+    tracers = NamedTuple(name => CenterField(grid) for name in keys(initial_tracers))
+
+    # Horizontal diffusivities that mix up to t ∼ ℓ² / κ ∼ 1
+    κh = horizontal_scale^2
+    κz = scale_to_diffusivity(vertical_scale)
+
+    # Determine stable time-step
+    Nx, Ny, Nz = size(grid)
+    ϵ = fractional_time_step
+    Az = minimum(grid.Azᶜᶜᵃ[1:Ny])
+
+    if κh == 0
+        Δt = fractional_time_step
+    else
+        Δt = ϵ * Az / κh
+    end
+    @show Nt = ceil(Int, 1 / Δt)
+
+    vitd = VerticallyImplicitTimeDiscretization()
+    vertical_smoothing = VerticalScalarDiffusivity(vitd, κ=κz)
+    horizontal_smoothing = HorizontalScalarDiffusivity(κ=κh)
+
+    smoothing_model = HydrostaticFreeSurfaceModel(; grid, tracers,
+                                                    velocities = PrescribedVelocityFields(),
+                                                    momentum_advection = nothing,
+                                                    tracer_advection = nothing,
+                                                    buoyancy = nothing,
+                                                    closure = (horizontal_smoothing, vertical_smoothing))
+
+    set!(smoothing_model, (tracer[name] = initial_tracer[name] for name in keys(initial_tracers))...)
+
+    Nt = ceil(Int, 1 / Δt)
+    @info string("Smoothing tracers ", keys(tracers), " with ", Nt, " time steps")
+    smoothing_simulation = Simulation(smoothing_model; Δt, stop_time=1)
+
+    # Remove NaN checker
+    pop!(smoothing_simulation.callbacks, :nan_checker) 
+
+    # Restore values to default in the masked region
+    function restore_values(sim)
+        grid    = sim.model.grid
+        tracers = sim.model.tracers
+        for (tracer, initial_tracer) in zip(tracers, initial_tracers)
+            launch!(architecture(grid), grid, :xyz, _apply_tracer_mask!, tracer, initial_tracer, mask)
+        end
+    end
+
+    # Add restoring to initial values
+    simulation.callbacks[:restoring] = Callback(restore_values, IterationInterval(1))
+
+    run!(smoothing_simulation)
+
+    return smoothing_model.tracers
+end
diff --git a/src/VerticalGrids.jl b/src/VerticalGrids.jl
index e3c7d162..220d791c 100644
--- a/src/VerticalGrids.jl
+++ b/src/VerticalGrids.jl
@@ -1,5 +1,7 @@
 module VerticalGrids
 
+export stretched_vertical_faces, PowerLawStretching, LinearStretching
+
 struct PowerLawStretching{T}
     power :: T
 end
@@ -25,12 +27,12 @@ end
                                maximum_Δz = Inf,
                                stretching = PowerLawStretching(1.02),
                                rounding_digits = 1,
-                               minimum_depth = 5000)
+                               depth = 5000)
 
 Return an array of cell interfaces with `surface_layer_Δz` spacing in
 a surface layer of height `surface_layer_height`, and stretched according to
-the function `stretching(Δz_above, z_above)` down to `minimum_depth`.
-The interfaces extends from `Lz = -z[1]` to `0 = z[end]`, where `Lz ≥ minimum_depth`.
+the function `stretching(Δz_above, z_above)` down to `depth`.
+The interfaces extends from `Lz = -z[1]` to `0 = z[end]`, where `Lz ≥ depth`.
 
 The grid spacing `Δz` is limited to be less than `maximum_Δz`.
 The grid is also uniformly-spaced below `constant_bottom_spacing_depth`.
@@ -43,7 +45,7 @@ function stretched_vertical_faces(; surface_layer_Δz = 5.0,
                                     maximum_Δz = Inf,
                                     stretching = PowerLawStretching(1.02),
                                     rounding_digits = 1,
-                                    minimum_depth = 5000)
+                                    depth = 5000)
 
     Δz₀ = surface_layer_Δz
     h₀ = surface_layer_height
@@ -52,7 +54,7 @@ function stretched_vertical_faces(; surface_layer_Δz = 5.0,
     z = [-Δz₀ * (k-1) for k = 1:ceil(h₀ / Δz₀)]
 
     # Generate stretched interior grid
-    Lz₀ = minimum_depth
+    Lz₀ = depth
 
     while z[end] > - Lz₀
         Δz_above = z[end-1] - z[end]
diff --git a/test/test_ecco2.jl b/test/test_ecco2.jl
index 3e482675..278dc2bf 100644
--- a/test/test_ecco2.jl
+++ b/test/test_ecco2.jl
@@ -1,6 +1,7 @@
 include("runtests_setup.jl")
 
-using ClimaOcean: ECCO2
+using ClimaOcean
+using ClimaOcean.ECCO2
 using Oceananigans.Grids: topology
 
 @testset "ECCO2 fields utilities" begin
@@ -38,4 +39,13 @@ using Oceananigans.Grids: topology
         @test Ny == 720
         @test Nz == 1
     end
-end
\ No newline at end of file
+end
+
+@testset "setting a field with ECCO2" begin
+    for arch in test_architectures
+        grid  = LatitudeLongitudeGrid(size = (10, 10, 10), latitude = (-60, -40), longitude = (-5, 5), z = (-200, 0))
+        field = CenterField(grid)
+        set!(field, ECCO2Metadata(:temperature)) 
+        set!(field, ECCO2Metadata(:salinity))
+    end 
+end