From a19f160d6887d4ae7410bf93ec8a93ce513ae4c9 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 2 Dec 2024 11:44:31 -0700 Subject: [PATCH 01/22] Cleaner syntax for ECCOFieldTimeSeries --- src/DataWrangling/ECCO/ECCO_restoring.jl | 26 ++++++++++++------------ 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/src/DataWrangling/ECCO/ECCO_restoring.jl b/src/DataWrangling/ECCO/ECCO_restoring.jl index 7400d0f7..d76f951b 100644 --- a/src/DataWrangling/ECCO/ECCO_restoring.jl +++ b/src/DataWrangling/ECCO/ECCO_restoring.jl @@ -51,7 +51,7 @@ end Base.length(backend::ECCONetCDFBackend) = backend.length Base.summary(backend::ECCONetCDFBackend) = string("ECCONetCDFBackend(", backend.start, ", ", backend.length, ")") -const ECCONetCDFFTS{N} = FlavorOfFTS{<:Any, <:Any, <:Any, <:Any, <:ECCONetCDFBackend{N}} where N +const ECCOFieldTimeSeries{N} = FlavorOfFTS{<:Any, <:Any, <:Any, <:Any, <:ECCONetCDFBackend{N}} where N new_backend(b::ECCONetCDFBackend{native, cache_data}, start, length) where {native, cache_data} = ECCONetCDFBackend{native, cache_data}(start, length, b.inpainting, b.metadata) @@ -59,7 +59,7 @@ new_backend(b::ECCONetCDFBackend{native, cache_data}, start, length) where {nati on_native_grid(::ECCONetCDFBackend{native}) where native = native cache_inpainted_data(::ECCONetCDFBackend{native, cache_data}) where {native, cache_data} = cache_data -function set!(fts::ECCONetCDFFTS) +function set!(fts::ECCOFieldTimeSeries) backend = fts.backend start = backend.start inpainting = backend.inpainting @@ -105,7 +105,7 @@ function ECCO_times(metadata; start_time = first(metadata).dates) end """ - ECCO_field_time_series(metadata::ECCOMetadata; + ECCOFieldTimeSeries(metadata::ECCOMetadata; grid = nothing, architecture = isnothing(grid) ? CPU() : architecture(grid), time_indices_in_memory = 2, @@ -138,7 +138,7 @@ Keyword Arguments Default: `true`. """ -function ECCO_field_time_series(metadata::ECCOMetadata; +function ECCOFieldTimeSeries(metadata::ECCOMetadata; architecture = CPU(), time_indices_in_memory = 2, time_indexing = Cyclical(), @@ -167,8 +167,8 @@ function ECCO_field_time_series(metadata::ECCOMetadata; return fts end -ECCO_field_time_series(variable_name::Symbol, version=ECCO4Monthly(); kw...) = - ECCO_field_time_series(ECCOMetadata(variable_name, all_ECCO_dates(version), version); kw...) +ECCOFieldTimeSeries(variable_name::Symbol, version=ECCO4Monthly(); kw...) = + ECCOFieldTimeSeries(ECCOMetadata(variable_name, all_ECCO_dates(version), version); kw...) # Variable names for restoreable data struct Temperature end @@ -330,13 +330,13 @@ function ECCORestoring(arch::AbstractArchitecture, throw(ArgumentError("The architecture of ECCORestoring must match the architecture of the grid.")) end - fts = ECCO_field_time_series(metadata; - grid, - architecture = arch, - time_indices_in_memory, - time_indexing, - inpainting, - cache_inpainted_data) + fts = ECCOFieldTimeSeries(metadata; + grid, + architecture = arch, + time_indices_in_memory, + time_indexing, + inpainting, + cache_inpainted_data) # Grab the correct Oceananigans field to restore variable_name = metadata.name From 98d1a8b13dd96f1391441058ee257fc382b35ecc Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 2 Dec 2024 11:49:01 -0700 Subject: [PATCH 02/22] Fix indenting --- src/DataWrangling/ECCO/ECCO_restoring.jl | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/DataWrangling/ECCO/ECCO_restoring.jl b/src/DataWrangling/ECCO/ECCO_restoring.jl index d76f951b..5b89cdc1 100644 --- a/src/DataWrangling/ECCO/ECCO_restoring.jl +++ b/src/DataWrangling/ECCO/ECCO_restoring.jl @@ -106,11 +106,11 @@ end """ ECCOFieldTimeSeries(metadata::ECCOMetadata; - grid = nothing, - architecture = isnothing(grid) ? CPU() : architecture(grid), - time_indices_in_memory = 2, - time_indexing = Cyclical(), - inpainting_iterations = prod(size(metadata)), + grid = nothing, + architecture = isnothing(grid) ? CPU() : architecture(grid), + time_indices_in_memory = 2, + time_indexing = Cyclical(), + inpainting_iterations = prod(size(metadata)), Create a field time series object for ECCO data. @@ -139,12 +139,12 @@ Keyword Arguments """ function ECCOFieldTimeSeries(metadata::ECCOMetadata; - architecture = CPU(), - time_indices_in_memory = 2, - time_indexing = Cyclical(), - inpainting = NearestNeighborInpainting(prod(size(metadata))), - cache_inpainted_data = true, - grid = nothing) + architecture = CPU(), + time_indices_in_memory = 2, + time_indexing = Cyclical(), + inpainting = NearestNeighborInpainting(prod(size(metadata))), + cache_inpainted_data = true, + grid = nothing) # Make sure all the required individual files are downloaded download_dataset(metadata) From 632c94710f6f00c1b418bc7c5ce0c28b7005a4de Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 2 Dec 2024 11:59:40 -0700 Subject: [PATCH 03/22] Make grid positional argument --- src/DataWrangling/ECCO/ECCO_restoring.jl | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/src/DataWrangling/ECCO/ECCO_restoring.jl b/src/DataWrangling/ECCO/ECCO_restoring.jl index 5b89cdc1..395abbdb 100644 --- a/src/DataWrangling/ECCO/ECCO_restoring.jl +++ b/src/DataWrangling/ECCO/ECCO_restoring.jl @@ -138,13 +138,12 @@ Keyword Arguments Default: `true`. """ -function ECCOFieldTimeSeries(metadata::ECCOMetadata; - architecture = CPU(), +function ECCOFieldTimeSeries(metadata::ECCOMetadata, grid=nothing; + architecture = isnothing(grid) ? CPU() : architecture(grid), time_indices_in_memory = 2, time_indexing = Cyclical(), inpainting = NearestNeighborInpainting(prod(size(metadata))), - cache_inpainted_data = true, - grid = nothing) + cache_inpainted_data = true) # Make sure all the required individual files are downloaded download_dataset(metadata) @@ -176,11 +175,10 @@ struct Salinity end struct UVelocity end struct VVelocity end -oceananigans_fieldname = Dict( - :temperature => Temperature(), - :salinity => Salinity(), - :u_velocity => UVelocity(), - :v_velocity => VVelocity()) +const oceananigans_fieldnames = Dict(:temperature => Temperature(), + :salinity => Salinity(), + :u_velocity => UVelocity(), + :v_velocity => VVelocity()) @inline Base.getindex(fields, i, j, k, ::Temperature) = @inbounds fields.T[i, j, k] @inline Base.getindex(fields, i, j, k, ::Salinity) = @inbounds fields.S[i, j, k] @@ -330,8 +328,7 @@ function ECCORestoring(arch::AbstractArchitecture, throw(ArgumentError("The architecture of ECCORestoring must match the architecture of the grid.")) end - fts = ECCOFieldTimeSeries(metadata; - grid, + fts = ECCOFieldTimeSeries(metadata, grid; architecture = arch, time_indices_in_memory, time_indexing, @@ -340,7 +337,7 @@ function ECCORestoring(arch::AbstractArchitecture, # Grab the correct Oceananigans field to restore variable_name = metadata.name - field_name = oceananigans_fieldname[variable_name] + field_name = oceananigans_fieldnames[variable_name] # If we pass the grid we do not need to interpolate # so we can save parameter space by setting the native grid to nothing From c18ba56e0b70717dbe2056a27de8f499da00df90 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 2 Dec 2024 12:35:20 -0700 Subject: [PATCH 04/22] Eliminate architecture kwarg --- src/DataWrangling/ECCO/ECCO_restoring.jl | 27 ++++++++++++------------ 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/src/DataWrangling/ECCO/ECCO_restoring.jl b/src/DataWrangling/ECCO/ECCO_restoring.jl index 395abbdb..9fbeea66 100644 --- a/src/DataWrangling/ECCO/ECCO_restoring.jl +++ b/src/DataWrangling/ECCO/ECCO_restoring.jl @@ -105,12 +105,11 @@ function ECCO_times(metadata; start_time = first(metadata).dates) end """ - ECCOFieldTimeSeries(metadata::ECCOMetadata; - grid = nothing, + ECCOFieldTimeSeries(metadata::ECCOMetadata, grid=nothing; architecture = isnothing(grid) ? CPU() : architecture(grid), time_indices_in_memory = 2, time_indexing = Cyclical(), - inpainting_iterations = prod(size(metadata)), + inpainting = nothing) Create a field time series object for ECCO data. @@ -138,28 +137,28 @@ Keyword Arguments Default: `true`. """ -function ECCOFieldTimeSeries(metadata::ECCOMetadata, grid=nothing; - architecture = isnothing(grid) ? CPU() : architecture(grid), +function ECCOFieldTimeSeries(metadata::ECCOMetadata, arch::AbstractArchitecture=CPU(); kw...) + download_dataset(metadata) + ftmp = empty_ECCO_field(first(metadata); architecture) + grid = ftmp.grid + return ECCOFieldTimeSeries(metadata, grid; kw...) +end + +function ECCOFieldTimeSeries(metadata::ECCOMetadata, grid::AbstractGrid; time_indices_in_memory = 2, time_indexing = Cyclical(), - inpainting = NearestNeighborInpainting(prod(size(metadata))), + inpainting = nothing, cache_inpainted_data = true) # Make sure all the required individual files are downloaded download_dataset(metadata) inpainting isa Int && (inpainting = NearestNeighborInpainting(inpainting)) - - ftmp = empty_ECCO_field(first(metadata); architecture) - on_native_grid = isnothing(grid) - on_native_grid && (grid = ftmp.grid) + backend = ECCONetCDFBackend(time_indices_in_memory, metadata; on_native_grid, inpainting, cache_inpainted_data) times = ECCO_times(metadata) loc = LX, LY, LZ = location(metadata) boundary_conditions = FieldBoundaryConditions(grid, loc) - - backend = ECCONetCDFBackend(time_indices_in_memory, metadata; on_native_grid, inpainting, cache_inpainted_data) - fts = FieldTimeSeries{LX, LY, LZ}(grid, times; backend, time_indexing, boundary_conditions) set!(fts) @@ -248,7 +247,7 @@ end mask = 1, rate = 1, grid = nothing, - inpainting = NearestNeighborInpainting(prod(size(metadata))), + inpainting = NearestNeighborInpainting(Inf), cache_inpainted_data = true) Create a forcing term that restores to values stored in an ECCO field time series. From 4b87a568ec4a0550ae19ccf3281e738215ae6845 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 3 Dec 2024 13:30:14 -0700 Subject: [PATCH 05/22] Clean up ECCORestoring --- src/DataWrangling/ECCO/ECCO_restoring.jl | 118 ++++++++++------------- 1 file changed, 52 insertions(+), 66 deletions(-) diff --git a/src/DataWrangling/ECCO/ECCO_restoring.jl b/src/DataWrangling/ECCO/ECCO_restoring.jl index 9fbeea66..97eff888 100644 --- a/src/DataWrangling/ECCO/ECCO_restoring.jl +++ b/src/DataWrangling/ECCO/ECCO_restoring.jl @@ -1,5 +1,5 @@ using Oceananigans: location -using Oceananigans.Grids: node, on_architecture +using Oceananigans.Grids: AbstractGrid, node, on_architecture using Oceananigans.Fields: interpolate!, interpolate, location, instantiated_location using Oceananigans.OutputReaders: Cyclical, TotallyInMemory, AbstractInMemoryBackend, FlavorOfFTS, time_indices using Oceananigans.Utils: Time @@ -35,7 +35,7 @@ end Adapt.adapt_structure(to, b::ECCONetCDFBackend{N, C}) where {N, C} = ECCONetCDFBackend{N, C}(b.start, b.length, nothing, nothing) """ - ECCONetCDFBackend(length; on_native_grid = false, inpainting = NearestNeighborInpainting(Inf)) + ECCONetCDFBackend(length; on_native_grid=false, inpainting=NearestNeighborInpainting(Inf)) Represent an ECCO FieldTimeSeries backed by ECCO native netCDF files. Each time instance is stored in an individual file. @@ -105,8 +105,7 @@ function ECCO_times(metadata; start_time = first(metadata).dates) end """ - ECCOFieldTimeSeries(metadata::ECCOMetadata, grid=nothing; - architecture = isnothing(grid) ? CPU() : architecture(grid), + ECCOFieldTimeSeries(metadata::ECCOMetadata [, arch_or_grid=CPU() ]; time_indices_in_memory = 2, time_indexing = Cyclical(), inpainting = nothing) @@ -118,13 +117,12 @@ Arguments - `metadata`: `ECCOMetadata` containing information about the ECCO dataset. +- `arch_or_grid`: Either a grid to interpolate ECCO data to, or an `arch`itecture + to use for the native ECCO grid. Default: CPU(). + Keyword Arguments ================= -- `grid`: where ECCO data is interpolated. If `nothing`, the native `ECCO` grid is used. - -- `architecture`: where data is stored. Should only be set if `isnothing(grid)`. - - `time_indices_in_memory`: The number of time indices to keep in memory. Default: 2. - `time_indexing`: The time indexing scheme to use. Default: `Cyclical()`. @@ -191,14 +189,14 @@ Base.summary(::VVelocity) = "v_velocity" struct ECCORestoring{FTS, G, M, V, N} field_time_series :: FTS - grid :: G + native_grid :: G mask :: M variable_name :: V rate :: N end Adapt.adapt_structure(to, p::ECCORestoring) = ECCORestoring(Adapt.adapt(to, p.field_time_series), - Adapt.adapt(to, p.grid), + Adapt.adapt(to, p.native_grid), Adapt.adapt(to, p.mask), Adapt.adapt(to, p.variable_name), Adapt.adapt(to, p.rate)) @@ -212,8 +210,7 @@ Adapt.adapt_structure(to, p::ECCORestoring) = ECCORestoring(Adapt.adapt(to, p.fi # Possibly interpolate ECCO data from the ECCO grid to simulation grid. # Otherwise, simply extract the pre-interpolated data from p.field_time_series. backend = p.field_time_series.backend - interpolating = on_native_grid(backend) - ψ_ecco = maybe_interpolate(Val(interpolating), p.field_time_series, i, j, k, p.grid, grid, time) + ψ_ecco = maybe_interpolate(p.field_time_series, i, j, k, p.native_grid, grid, time) ψ = @inbounds fields[i, j, k, p.variable_name] μ = stateindex(p.mask, i, j, k, grid, clock.time, loc) @@ -222,9 +219,9 @@ Adapt.adapt_structure(to, p::ECCORestoring) = ECCORestoring(Adapt.adapt(to, p.fi return ω * μ * (ψ_ecco - ψ) end -@inline maybe_interpolate(::Val{false}, fts, i, j, k, native_grid, grid, time) = @inbounds fts[i, j, k, time] +@inline maybe_interpolate(fts, i, j, k, ::Nothing, grid, time) = @inbounds fts[i, j, k, time] -@inline function maybe_interpolate(::Val{true}, fts, i, j, k, native_grid, grid, time) +@inline function maybe_interpolate(fts, i, j, k, native_grid, grid, time) times = fts.times data = fts.data time_indexing = fts.time_indexing @@ -237,40 +234,47 @@ end end """ - ECCORestoring([arch=CPU(),] - variable_name::Symbol; - version=ECCO4Monthly(), - dates=all_ECCO_dates(version), - dates = all_ECCO_dates(version), + ECCORestoring(variable_name::Symbol, [ arch_or_grid = CPU(), ]; + version = ECCO4Monthly(), + dates = all_ECCO_dates(version), time_indices_in_memory = 2, time_indexing = Cyclical(), mask = 1, rate = 1, - grid = nothing, inpainting = NearestNeighborInpainting(Inf), cache_inpainted_data = true) -Create a forcing term that restores to values stored in an ECCO field time series. +Build a forcing term that restores to values stored in an ECCO field time series. The restoring is applied as a forcing on the right hand side of the evolution equations calculated as -```julia -F = mask ⋅ rate ⋅ (ECCO_variable - simulation_variable[i, j, k]) +```math +Fψ = μ r (ψ_ECCO - ψ) ``` -where `ECCO_variable` is linearly interpolated in space and time from the ECCO dataset of choice to the -simulation grid and time. + +where ``μ`` is the mask, ``r`` is the restoring rate, ``ψ`` is the simulation variable, +and the ECCO variable ``ψ_ECCO`` is linearly interpolated in space and time from the +ECCO dataset of choice to the simulation grid and time. Arguments ========= -- `arch`: The architecture. Typically `CPU()` or `GPU()`. Default: `CPU()`. - - `variable_name`: The name of the variable to restore. Choices include: - * `:temperature`, - * `:salinity`, - * `:u_velocity`, - * `:v_velocity`, - * `:sea_ice_thickness`, - * `:sea_ice_area_fraction`. + * `:temperature`, + * `:salinity`, + * `:u_velocity`, + * `:v_velocity`, + * `:sea_ice_thickness`, + * `:sea_ice_area_fraction`. + + Note that `ECCOMetadata` may be provided as the first argument instead + of `variable_name`. In this case the `version` and `dates` kwargs (described below) + cannot be provided. + +- `arch_or_grid`: Either the architecture of the simulation, or a grid on which the ECCO data + is pre-interpolated when loaded. If an `arch`itecture is provided, such as + `arch_or_grid = CPU()` or `arch_or_grid = GPU()`, ECCO data + will be interpolated on-the-fly when the forcing tendency is computed. + Default: CPU(). Keyword Arguments ================= @@ -279,56 +283,41 @@ Keyword Arguments - `dates`: The dates to use for the ECCO dataset. Default: `all_ECCO_dates(version)`. -- `time_indices_in_memory`: The number of time indices to keep in memory; trade-off between performance - and memory footprint. +- `time_indices_in_memory`: The number of time indices to keep in memory. The number is chosen based on + a trade-off between increased performance (more indices in memory) and reduced + memory footprint (fewer indices in memory). Default: 2. -- `time_indexing`: The time indexing scheme for the field time series≥ +- `time_indexing`: The time indexing scheme for the field time series. - `mask`: The mask value. Can be a function of `(x, y, z, time)`, an array, or a number. - `rate`: The restoring rate, i.e., the inverse of the restoring timescale (in s⁻¹). -- `time_indices_in_memory:` how many time instances are loaded in memory; the remaining are loaded lazily. - - `inpainting`: inpainting algorithm, see [`inpaint_mask!`](@ref). Default: `NearestNeighborInpainting(Inf)`. -- `grid`: If `isnothing(grid)`, ECCO data is interpolated on-the-fly to the simulation grid. - If `!isnothing(grid)`, ECCO data is pre-interpolated to `grid`. - Default: nothing. - - `cache_inpainted_data`: If `true`, the data is cached to disk after inpainting for later retrieving. Default: `true`. - -It is possible to also pass an `ECCOMetadata` type as the first argument without the need for the -`variable_name` argument and the `version` and `dates` keyword arguments. """ -function ECCORestoring(arch::AbstractArchitecture, - variable_name::Symbol; - version=ECCO4Monthly(), - dates=all_ECCO_dates(version), +function ECCORestoring(variable_name::Symbol, + arch_or_grid = CPU(); + version = ECCO4Monthly(), + dates = all_ECCO_dates(version), kw...) metadata = ECCOMetadata(variable_name, dates, version) - return ECCORestoring(arch, metadata; kw...) + return ECCORestoring(metadata, arch_or_grid; kw...) end -function ECCORestoring(arch::AbstractArchitecture, - metadata::ECCOMetadata; +function ECCORestoring(metadata::ECCOMetadata, + arch_or_grid = CPU(); rate, mask = 1, - grid = nothing, time_indices_in_memory = 2, # Not more than this if we want to use GPU! time_indexing = Cyclical(), inpainting = NearestNeighborInpainting(Inf), cache_inpainted_data = true) - # Validate architecture - if !isnothing(grid) && architecture(grid) != arch - throw(ArgumentError("The architecture of ECCORestoring must match the architecture of the grid.")) - end - - fts = ECCOFieldTimeSeries(metadata, grid; - architecture = arch, + fts = ECCOFieldTimeSeries(metadata, arch_or_grid; time_indices_in_memory, time_indexing, inpainting, @@ -340,15 +329,12 @@ function ECCORestoring(arch::AbstractArchitecture, # If we pass the grid we do not need to interpolate # so we can save parameter space by setting the native grid to nothing - native_grid = isnothing(grid) ? fts.grid : nothing + on_native_grid = arch_or_grid isa AbstractArchitecture + maybe_native_grid = on_native_grid ? fts.grid : nothing - return ECCORestoring(fts, native_grid, mask, field_name, rate) + return ECCORestoring(fts, maybe_native_grid, mask, field_name, rate) end -# Make sure we can call ECCORestoring with architecture as the first positional argument -ECCORestoring(variable_name::Symbol; kw...) = ECCORestoring(CPU(), variable_name; kw...) -ECCORestoring(metadata::ECCOMetadata; kw...) = ECCORestoring(CPU(), metadata; kw...) - function Base.show(io::IO, p::ECCORestoring) print(io, "ECCORestoring:", '\n', "├── restored variable: ", summary(p.variable_name), '\n', From 313740153a692059b2642d94d83c8398e504c9e5 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Tue, 3 Dec 2024 13:37:57 -0700 Subject: [PATCH 06/22] More simplification --- src/DataWrangling/ECCO/ECCO_restoring.jl | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/DataWrangling/ECCO/ECCO_restoring.jl b/src/DataWrangling/ECCO/ECCO_restoring.jl index 97eff888..ecbe3d7b 100644 --- a/src/DataWrangling/ECCO/ECCO_restoring.jl +++ b/src/DataWrangling/ECCO/ECCO_restoring.jl @@ -209,19 +209,20 @@ Adapt.adapt_structure(to, p::ECCORestoring) = ECCORestoring(Adapt.adapt(to, p.fi # Possibly interpolate ECCO data from the ECCO grid to simulation grid. # Otherwise, simply extract the pre-interpolated data from p.field_time_series. - backend = p.field_time_series.backend - ψ_ecco = maybe_interpolate(p.field_time_series, i, j, k, p.native_grid, grid, time) + if p.native_grid isa Nothing + ψ_ecco = @inbounds p.field_time_series[i, j, k] + else + ψ_ecco = interpolate_to_grid(p.field_time_series, i, j, k, p.native_grid, grid, time) + end ψ = @inbounds fields[i, j, k, p.variable_name] μ = stateindex(p.mask, i, j, k, grid, clock.time, loc) - ω = p.rate + r = p.rate - return ω * μ * (ψ_ecco - ψ) + return r * μ * (ψ_ecco - ψ) end -@inline maybe_interpolate(fts, i, j, k, ::Nothing, grid, time) = @inbounds fts[i, j, k, time] - -@inline function maybe_interpolate(fts, i, j, k, native_grid, grid, time) +@inline function interpolate_to_grid(fts, i, j, k, native_grid, grid, time) times = fts.times data = fts.data time_indexing = fts.time_indexing @@ -248,7 +249,7 @@ Build a forcing term that restores to values stored in an ECCO field time series The restoring is applied as a forcing on the right hand side of the evolution equations calculated as ```math -Fψ = μ r (ψ_ECCO - ψ) +Fψ = r μ (ψ_ECCO - ψ) ``` where ``μ`` is the mask, ``r`` is the restoring rate, ``ψ`` is the simulation variable, From 5fc441d228e9c38511605b877a0bc8abd071e9af Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Wed, 4 Dec 2024 08:43:01 +1100 Subject: [PATCH 07/22] fix typo --- src/DataWrangling/ECCO/ECCO_restoring.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/DataWrangling/ECCO/ECCO_restoring.jl b/src/DataWrangling/ECCO/ECCO_restoring.jl index ecbe3d7b..92cb6cb2 100644 --- a/src/DataWrangling/ECCO/ECCO_restoring.jl +++ b/src/DataWrangling/ECCO/ECCO_restoring.jl @@ -137,7 +137,7 @@ Keyword Arguments """ function ECCOFieldTimeSeries(metadata::ECCOMetadata, arch::AbstractArchitecture=CPU(); kw...) download_dataset(metadata) - ftmp = empty_ECCO_field(first(metadata); architecture) + ftmp = empty_ECCO_field(first(metadata); arch) grid = ftmp.grid return ECCOFieldTimeSeries(metadata, grid; kw...) end From 28632e82dc2d6308afe1d8612cd94e3d99df4a91 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Wed, 4 Dec 2024 11:02:05 +1100 Subject: [PATCH 08/22] fix typo --- src/DataWrangling/ECCO/ECCO_restoring.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/DataWrangling/ECCO/ECCO_restoring.jl b/src/DataWrangling/ECCO/ECCO_restoring.jl index 92cb6cb2..749caff5 100644 --- a/src/DataWrangling/ECCO/ECCO_restoring.jl +++ b/src/DataWrangling/ECCO/ECCO_restoring.jl @@ -135,9 +135,9 @@ Keyword Arguments Default: `true`. """ -function ECCOFieldTimeSeries(metadata::ECCOMetadata, arch::AbstractArchitecture=CPU(); kw...) +function ECCOFieldTimeSeries(metadata::ECCOMetadata, architecture::AbstractArchitecture=CPU(); kw...) download_dataset(metadata) - ftmp = empty_ECCO_field(first(metadata); arch) + ftmp = empty_ECCO_field(first(metadata); architecture) grid = ftmp.grid return ECCOFieldTimeSeries(metadata, grid; kw...) end From 5b6c96a052f20206a21807b66017029d8bc09cef Mon Sep 17 00:00:00 2001 From: glwagner Date: Tue, 3 Dec 2024 23:55:35 -0500 Subject: [PATCH 09/22] Bugfix --- src/DataWrangling/ECCO/ECCO_restoring.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/DataWrangling/ECCO/ECCO_restoring.jl b/src/DataWrangling/ECCO/ECCO_restoring.jl index 749caff5..2c370e79 100644 --- a/src/DataWrangling/ECCO/ECCO_restoring.jl +++ b/src/DataWrangling/ECCO/ECCO_restoring.jl @@ -210,7 +210,7 @@ Adapt.adapt_structure(to, p::ECCORestoring) = ECCORestoring(Adapt.adapt(to, p.fi # Possibly interpolate ECCO data from the ECCO grid to simulation grid. # Otherwise, simply extract the pre-interpolated data from p.field_time_series. if p.native_grid isa Nothing - ψ_ecco = @inbounds p.field_time_series[i, j, k] + ψ_ecco = @inbounds p.field_time_series[i, j, k, time] else ψ_ecco = interpolate_to_grid(p.field_time_series, i, j, k, p.native_grid, grid, time) end @@ -249,7 +249,7 @@ Build a forcing term that restores to values stored in an ECCO field time series The restoring is applied as a forcing on the right hand side of the evolution equations calculated as ```math -Fψ = r μ (ψ_ECCO - ψ) +Fψ = r μ (ψ_{ECCO} - ψ) ``` where ``μ`` is the mask, ``r`` is the restoring rate, ``ψ`` is the simulation variable, From 4be5e85438e95e32f65fa973e58f51669f9be986 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 5 Dec 2024 07:23:09 -0700 Subject: [PATCH 10/22] Fix test syntax --- test/test_ecco.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/test_ecco.jl b/test/test_ecco.jl index 9a3bb9c3..b3250828 100644 --- a/test/test_ecco.jl +++ b/test/test_ecco.jl @@ -24,7 +24,7 @@ inpainting = NearestNeighborInpainting(2) for name in (:temperature, :salinity) @info "Testing ECCO_field on $A..." metadata = ECCOMetadata(name, dates, ECCO4Monthly()) - restoring = ECCORestoring(metadata ; rate = 1 / 1000.0, inpainting) + restoring = ECCORestoring(metadata; rate = 1 / 1000.0, inpainting) for datum in metadata @test isfile(metadata_path(datum)) @@ -98,7 +98,7 @@ end southern = (φ₁, φ₂), z = (z₁, 0)) - t_restoring = ECCORestoring(arch, :temperature; + t_restoring = ECCORestoring(:temperature, arch; dates, mask, rate = 1 / 1000.0, @@ -157,7 +157,7 @@ end true end - forcing_T = ECCORestoring(arch, :temperature; + forcing_T = ECCORestoring(:temperature, arch; dates, rate = 1 / 1000.0, inpainting) From 3f7f58f454ab6c0c9125bf435b7a77763fc141ad Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 5 Dec 2024 07:23:53 -0700 Subject: [PATCH 11/22] Bump project --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 32571d19..2a1f7456 100644 --- a/Project.toml +++ b/Project.toml @@ -2,7 +2,7 @@ name = "ClimaOcean" uuid = "0376089a-ecfe-4b0e-a64f-9c555d74d754" license = "MIT" authors = ["Climate Modeling Alliance and contributors"] -version = "0.2.4" +version = "0.3.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" From 64d68a53a2ca5015ce501f3eff72442c2432d1a1 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 5 Dec 2024 15:26:16 -0700 Subject: [PATCH 12/22] Implement mixed layer depth diagnostic --- Project.toml | 1 - examples/ecco_inspect_temperature_salinity.jl | 87 +++++++++++++++ examples/ecco_mixed_layer_depth.jl | 90 ++++++++++++++++ examples/inspect_ecco_data.jl | 97 ----------------- examples/mixed_layer_depth.jl | 101 ++++++++++++++++++ mwe.jl | 11 ++ 6 files changed, 289 insertions(+), 98 deletions(-) create mode 100644 examples/ecco_inspect_temperature_salinity.jl create mode 100644 examples/ecco_mixed_layer_depth.jl delete mode 100644 examples/inspect_ecco_data.jl create mode 100644 examples/mixed_layer_depth.jl create mode 100644 mwe.jl diff --git a/Project.toml b/Project.toml index 2a1f7456..ccfa2505 100644 --- a/Project.toml +++ b/Project.toml @@ -44,7 +44,6 @@ MPI = "0.20" NCDatasets = "0.12, 0.13, 0.14" Oceananigans = "0.94.3 - 0.99" OffsetArrays = "1.14" -OrthogonalSphericalShellGrids = "0.1.8" Scratch = "1" SeawaterPolynomials = "0.3.4" StaticArrays = "1" diff --git a/examples/ecco_inspect_temperature_salinity.jl b/examples/ecco_inspect_temperature_salinity.jl new file mode 100644 index 00000000..4fcaf1ed --- /dev/null +++ b/examples/ecco_inspect_temperature_salinity.jl @@ -0,0 +1,87 @@ +using Oceananigans +using Oceananigans.ImmersedBoundaries: mask_immersed_field! + +using GLMakie +using Printf +using ClimaOcean +using ClimaOcean.DataWrangling.ECCO: ECCO_field, ECCOFieldTimeSeries +using CFTime +using Dates + +arch = CPU() +Nx = 360 ÷ 4 +Ny = 160 ÷ 4 + +z = ClimaOcean.DataWrangling.ECCO.ECCO_z +z = z[20:end] +Nz = length(z) - 1 + +grid = LatitudeLongitudeGrid(arch; z, + size = (Nx, Ny, Nz), + latitude = (-80, 80), + longitude = (0, 360)) + +bottom_height = regrid_bathymetry(grid; + minimum_depth = 10, + interpolation_passes = 5, + major_basins = 1) + +grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height)) + +T = CenterField(grid) +S = CenterField(grid) + +using SeawaterPolynomials: TEOS10EquationOfState +using Oceananigans.BuoyancyModels: buoyancy + +equation_of_state = TEOS10EquationOfState() +sb = SeawaterBuoyancy(; equation_of_state) +tracers = (T=T, S=S) +b = Field(buoyancy(sb, grid, tracers)) + +start = DateTimeProlepticGregorian(1993, 1, 1) +stop = DateTimeProlepticGregorian(1999, 1, 1) +dates = range(start; stop, step=Month(1)) + +Tmeta = ECCOMetadata(:temperature; dates) +Smeta = ECCOMetadata(:salinity; dates) + +Tt = ECCOFieldTimeSeries(Tmeta, grid; time_indices_in_memory=length(dates)) +St = ECCOFieldTimeSeries(Smeta, grid; time_indices_in_memory=length(dates)) + +fig = Figure(size=(900, 1050)) + +axT = Axis(fig[1, 1]) +axS = Axis(fig[2, 1]) +axb = Axis(fig[3, 1]) + +Nt = length(dates) +grid = T.grid +Nz = size(grid, 3) +kslider = Slider(fig[1:3, 0], range=1:Nz, startvalue=Nz, horizontal=false) +nslider = Slider(fig[4, 1:2], range=1:Nt, startvalue=1) +k = kslider.value +n = nslider.value + +Tk = @lift view(Tt[$n], :, :, $k) +Sk = @lift view(St[$n], :, :, $k) + +Δb = @lift begin + parent(T) .= parent(Tt[$n]) + parent(S) .= parent(St[$n]) + compute!(b) + mask_immersed_field!(b, NaN) + Δb = interior(b, :, :, Nz) .- interior(b, :, :, $k) + Δb +end + +hmT = heatmap!(axT, Tk, nan_color=:lightgray, colorrange=(-2, 30), colormap=:thermal) +hmS = heatmap!(axS, Sk, nan_color=:lightgray, colorrange=(31, 37), colormap=:haline) +hmb = heatmap!(axb, Δb, nan_color=:lightgray, colorrange=(0, 1e-3), colormap=:magma) + +Colorbar(fig[1, 2], hmT, label="Temperature (ᵒC)") +Colorbar(fig[2, 2], hmS, label="Salinity (g kg⁻¹)") +Colorbar(fig[3, 2], hmb, label="Buoyancy difference (m s⁻²)") + +display(fig) + diff --git a/examples/ecco_mixed_layer_depth.jl b/examples/ecco_mixed_layer_depth.jl new file mode 100644 index 00000000..90c6e98a --- /dev/null +++ b/examples/ecco_mixed_layer_depth.jl @@ -0,0 +1,90 @@ +#= +using Oceananigans +using Oceananigans.ImmersedBoundaries: mask_immersed_field! + +#include("mixed_layer_depth.jl") + +using GLMakie +using Printf +using ClimaOcean +using ClimaOcean.DataWrangling.ECCO: ECCO_field, ECCOFieldTimeSeries +using CFTime +using Dates + +using SeawaterPolynomials: TEOS10EquationOfState +using Oceananigans.BuoyancyModels: buoyancy + +arch = CPU() +Nx = 360 ÷ 1 +Ny = 160 ÷ 1 + +z = ClimaOcean.DataWrangling.ECCO.ECCO_z +z = z[20:end] +Nz = length(z) - 1 + +grid = LatitudeLongitudeGrid(arch; z, + size = (Nx, Ny, Nz), + latitude = (-80, 80), + longitude = (0, 360)) + +bottom_height = regrid_bathymetry(grid; + minimum_depth = 10, + interpolation_passes = 5, + major_basins = 1) + +grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height)) + +#dates = ClimaOcean.DataWrangling.ECCO.all_ECCO_dates(ClimaOcean.DataWrangling.ECCO.ECCO4Monthly()) +start = DateTimeProlepticGregorian(1993, 1, 1) +stop = DateTimeProlepticGregorian(2017, 1, 1) +dates = range(start; stop, step=Month(1)) + +Tmeta = ECCOMetadata(:temperature; dates) +Smeta = ECCOMetadata(:salinity; dates) + +Tt = ECCOFieldTimeSeries(Tmeta, grid; time_indices_in_memory=2) +St = ECCOFieldTimeSeries(Smeta, grid; time_indices_in_memory=2) +ht = FieldTimeSeries{Center, Center, Nothing}(grid, Tt.times) + +equation_of_state = TEOS10EquationOfState() +sb = SeawaterBuoyancy(; equation_of_state) +tracers = (T=Tt[1], S=St[1]) +h = MixedLayerDepthField(sb, grid, tracers) + +Nt = length(ht) +for n = 1:Nt-1 + local tracers + tracers = (T=Tt[n], S=St[n]) + h.operand.buoyancy_perturbation = buoyancy(sb, grid, tracers) + @show n + @time compute!(h) + parent(ht[n]) .= parent(h) +end +=# + +function titlestr(n) + d = dates[n] + yr = year(d) + mn = monthname(d) + return string("ECCO mixed layer depth on ", mn, " ", yr) +end + +fig = Figure(size=(1500, 800)) +axh = Axis(fig[2, 1], xlabel="Longitude", ylabel="Latitude") + +#slider = Slider(fig[3, 1], range=1:Nt, startvalue=1) +#n = slider.value +n = Observable(1) + +str = @lift titlestr($n) +Label(fig[1, 1], str, tellwidth=false) + +hn = @lift ht[$n] +hm = heatmap!(axh, hn, colorrange=(0, 500), colormap=:magma, nan_color=:lightgray) +Colorbar(fig[2, 2], hm, label="Mixed layer depth (m)") +display(fig) + +record(fig, "ecco_mld.mp4", 1:Nt-1, framerate=4) do nn + @info "Drawing frame $nn of $Nt..." + n[] = nn +end diff --git a/examples/inspect_ecco_data.jl b/examples/inspect_ecco_data.jl deleted file mode 100644 index 3d260436..00000000 --- a/examples/inspect_ecco_data.jl +++ /dev/null @@ -1,97 +0,0 @@ -# # A quick look at ECCO data -# -# ClimaOcean can download and utilize data from the "ECCO" state estimate, -# which stands for "Estimating the Circulation and Climate of the Ocean" --- two! -# -# This example shows how to download three-dimensional temperature and salinity fields -# from ECCO, and makes a short animation to showcase the fields' content. -# -# For this example we need Oceananigans for Field utilities, CairoMakie for plotting, -# Printf for nice labeling, and of course ClimaOcean to actually download and construct -# the ECCO fields. - -using Oceananigans -using CairoMakie -using Printf - -using ClimaOcean.DataWrangling.ECCO: ECCO_field - -# The function `ECCO_field` provided by `ClimaOcean.DataWrangling.ECCO` automatically -# downloads ECCO data, if the data doesn't already exist at the default location. - -T = ECCO_field(:temperature) -S = ECCO_field(:salinity) - -# Next, we massage the ECCO data by inserting NaNs in "land cells", which -# are diagnosed by having an unphysically low temperature. - -Tp = parent(T) -Sp = parent(S) -Sp[Tp .< -10] .= NaN -Tp[Tp .< -10] .= NaN - -# # Plotting ECCO data -# -# We're ready to plot. We'll make an animation -# that depicts how the ECCO data changes with depth. - -fig = Figure(size=(900, 1050)) - -axT = Axis(fig[1, 1]) -axS = Axis(fig[2, 1]) - -# To make an animation that scrolls through the 3D temperature -# and salinity fields, we make an Observable for the vertical index, -# and then construct slices of T, S using the Observable, `k`. - -grid = T.grid -Nz = size(grid, 3) -k = Observable(Nz) - -Tk = @lift view(T, :, :, $k) -Sk = @lift view(S, :, :, $k) - -# Finally, we make a nice plot with a label that displays depth, colorbars, -# and light gray within land cells. - -hmT = heatmap!(axT, Tk, nan_color=:lightgray, colorrange=(-2, 30), colormap=:thermal) -hmS = heatmap!(axS, Sk, nan_color=:lightgray, colorrange=(31, 37), colormap=:haline) - -Colorbar(fig[1, 2], hmT, label="Temperature (ᵒC)") -Colorbar(fig[2, 2], hmS, label="Salinity (psu)") - -z = znodes(grid, Center()) -depth_str = @lift @sprintf("z = %d meters", z[$k]) -text!(axT, 50, 50, text=depth_str, color=:lemonchiffon, justification=:center, fontsize=20) -text!(axS, 50, 50, text=depth_str, color=:lemonchiffon, justification=:center, fontsize=20) - -# # Making the animation -# -# This animation is a little fancy. We start by displaying the surface -# field, then we scroll through depth to the bottom and pause there. -# Next, we scroll back to the surface and pause. - -stillframes = 10 -movingframes = Nz - -record(fig, "ECCO_temperature_salinity.mp4", framerate=4) do io - - [recordframe!(io) for _ = 1:stillframes] - - for kk in Nz:-2:1 - k[] = kk - recordframe!(io) - end - - [recordframe!(io) for _ = 1:stillframes] - - for kk in 1:2:Nz - k[] = kk - recordframe!(io) - end - - [recordframe!(io) for _ = 1:stillframes] -end -nothing #hide - -# ![](ECCO_temperature_salinity.mp4) diff --git a/examples/mixed_layer_depth.jl b/examples/mixed_layer_depth.jl new file mode 100644 index 00000000..e2128878 --- /dev/null +++ b/examples/mixed_layer_depth.jl @@ -0,0 +1,101 @@ +using Oceananigans +using Oceananigans.Architectures: architecture +using Oceananigans.BuoyancyModels: buoyancy +using Oceananigans.Grids: new_data, inactive_cell, znode +using Oceananigans.BoundaryConditions: FieldBoundaryConditions, fill_halo_regions! +using Oceananigans.Fields: FieldStatus +using Oceananigans.Utils: launch! + +import Oceananigans.Fields: compute! + +using KernelAbstractions: @index, @kernel + +mutable struct MixedLayerDepthOperand{FT, B} + buoyancy_perturbation :: B + difference_criterion :: FT +end + +Base.summary(mldo::MixedLayerDepthOperand) = "MixedLayerDepthOperand" + +function MixedLayerDepthOperand(bm, grid, tracers; difference_criterion=1e-4) + buoyancy_perturbation = buoyancy(bm, grid, tracers) + difference_criterion = convert(eltype(grid), difference_criterion) + return MixedLayerDepthOperand(buoyancy_perturbation, difference_criterion) +end + +const MixedLayerDepthField = Field{<:Any, <:Any, <:Any, <:MixedLayerDepthOperand} + +""" + MixedLayerDepthField(bm, grid, tracers; difference_criterion=1e-4) + +""" +function MixedLayerDepthField(bm, grid, tracers; difference_criterion=1e-4) + operand = MixedLayerDepthOperand(bm, grid, tracers; difference_criterion) + loc = (Center, Center, Nothing) + indices = (:, :, :) + bcs = FieldBoundaryConditions(grid, loc) + data = new_data(grid, loc, indices) + recompute_safely = false + status = FieldStatus() + return Field(loc, grid, data, bcs, indices, operand, status) +end + +function compute!(mld::MixedLayerDepthField, time=nothing) + compute_mixed_layer_depth!(mld) + #@apply_regionally compute_mixed_layer_depth!(mld) + fill_halo_regions!(mld) + return mld +end + +function compute_mixed_layer_depth!(mld) + grid = mld.grid + arch = architecture(grid) + + launch!(arch, mld.grid, :xy, + _compute_mixed_layer_depth!, + mld, + grid, + mld.operand.buoyancy_perturbation, + mld.operand.difference_criterion) + + return mld +end + +const c = Center() +const f = Face() + +@kernel function _compute_mixed_layer_depth!(mld, grid, b, Δb★) + i, j = @index(Global, NTuple) + Nz = size(grid, 3) + + Δb = zero(grid) + bN = @inbounds b[i, j, Nz] + mixed = true + k = Nz - 1 + inactive = inactive_cell(i, j, k, grid) + + while !inactive & mixed & (k > 0) + Δb = @inbounds bN - b[i, j, k] + mixed = Δb < Δb★ + k -= 1 + inactive = inactive_cell(i, j, k, grid) + end + + # Linearly interpolate + # z★ = zk + Δz/Δb * (Δb★ - Δb) + zN = znode(i, j, Nz, grid, c, c, c) + zk = znode(i, j, k, grid, c, c, c) + Δz = zN - zk + z★ = zk - Δz/Δb * (Δb★ - Δb) + + # Special case when domain is one grid cell deep + z★ = ifelse(Δb == 0, zN, z★) + + # Apply various criterion + h = -z★ + h = max(h, zero(grid)) + h = min(h, grid.Lz) + + @inbounds mld[i, j, 1] = h +end + diff --git a/mwe.jl b/mwe.jl new file mode 100644 index 00000000..f77545eb --- /dev/null +++ b/mwe.jl @@ -0,0 +1,11 @@ +using ClimaOcean +using NCDatasets + +cachepath = ClimaOcean.DataWrangling.JRA55.download_jra55_cache +filename = "RYF.tas.1990_1991.nc" +filepath = joinpath(cachepath, filename) + +ds = Dataset(filepath) +Nx, Ny, Nt = size(ds["tas"]) +ds["tas"][1, 1, [Nt, 1]] +close(ds) From 6dcf7274a85e6c1976a93df2137c40eaf05bf9c4 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 12 Dec 2024 21:10:00 -0700 Subject: [PATCH 13/22] Shuffle files --- examples/ecco_mixed_layer_depth.jl | 19 +-- src/ClimaOcean.jl | 2 +- src/Diagnostics.jl | 124 ------------------ src/Diagnostics/Diagnostics.jl | 19 +++ .../Diagnostics}/mixed_layer_depth.jl | 12 -- 5 files changed, 25 insertions(+), 151 deletions(-) delete mode 100644 src/Diagnostics.jl create mode 100644 src/Diagnostics/Diagnostics.jl rename {examples => src/Diagnostics}/mixed_layer_depth.jl (85%) diff --git a/examples/ecco_mixed_layer_depth.jl b/examples/ecco_mixed_layer_depth.jl index 90c6e98a..ff3165e5 100644 --- a/examples/ecco_mixed_layer_depth.jl +++ b/examples/ecco_mixed_layer_depth.jl @@ -1,18 +1,14 @@ -#= +using ClimaOcean +using ClimaOcean.Diagnostics: MixedLayerDepthField +using ClimaOcean.DataWrangling.ECCO: ECCO_field, ECCOFieldTimeSeries using Oceananigans -using Oceananigans.ImmersedBoundaries: mask_immersed_field! - -#include("mixed_layer_depth.jl") - using GLMakie using Printf -using ClimaOcean -using ClimaOcean.DataWrangling.ECCO: ECCO_field, ECCOFieldTimeSeries using CFTime using Dates using SeawaterPolynomials: TEOS10EquationOfState -using Oceananigans.BuoyancyModels: buoyancy +using Oceananigans.BuoyancyFormulations: buoyancy arch = CPU() Nx = 360 ÷ 1 @@ -34,9 +30,8 @@ bottom_height = regrid_bathymetry(grid; grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height)) -#dates = ClimaOcean.DataWrangling.ECCO.all_ECCO_dates(ClimaOcean.DataWrangling.ECCO.ECCO4Monthly()) start = DateTimeProlepticGregorian(1993, 1, 1) -stop = DateTimeProlepticGregorian(2017, 1, 1) +stop = DateTimeProlepticGregorian(2003, 1, 1) dates = range(start; stop, step=Month(1)) Tmeta = ECCOMetadata(:temperature; dates) @@ -60,7 +55,6 @@ for n = 1:Nt-1 @time compute!(h) parent(ht[n]) .= parent(h) end -=# function titlestr(n) d = dates[n] @@ -71,9 +65,6 @@ end fig = Figure(size=(1500, 800)) axh = Axis(fig[2, 1], xlabel="Longitude", ylabel="Latitude") - -#slider = Slider(fig[3, 1], range=1:Nt, startvalue=1) -#n = slider.value n = Observable(1) str = @lift titlestr($n) diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index ae1bb084..d1f6b071 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -75,7 +75,7 @@ include("VerticalGrids.jl") include("InitialConditions/InitialConditions.jl") include("DataWrangling/DataWrangling.jl") include("Bathymetry.jl") -include("Diagnostics.jl") +include("Diagnostics/Diagnostics.jl") include("OceanSimulations/OceanSimulations.jl") using .VerticalGrids diff --git a/src/Diagnostics.jl b/src/Diagnostics.jl deleted file mode 100644 index 74dda306..00000000 --- a/src/Diagnostics.jl +++ /dev/null @@ -1,124 +0,0 @@ -module Diagnostics - -export MixedLayerDepthField - -using Oceananigans -using Oceananigans.BuoyancyModels: buoyancy -using Oceananigans.BoundaryConditions: fill_halo_regions! -using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid -using Oceananigans.Architectures: device, architecture -using Oceananigans.Utils: launch! -using Oceananigans.Grids: Center, Face, inactive_node, znode -using Oceananigans.Operators: Δzᶜᶜᶠ - -using KernelAbstractions: @index, @kernel -using KernelAbstractions.Extras.LoopInfo: @unroll - -import Oceananigans.Fields: compute! - -const c = Center() -const f = Face() - -@inline z_bottom(i, j, grid) = znode(c, c, f, i, j, 1, grid) -@inline z_bottom(i, j, grid::ImmersedBoundaryGrid) = @inbounds grid.immersed_boundary.bottom_height[i, j] - -@inline bottom(i, j, grid) = znode(c, c, f, i, j, 1, grid) -##### -##### MixedLayerDepthField -##### - -@kernel function compute_mld!(h, grid, b, Δb) - i, j = @index(Global, NTuple) - - Nz = grid.Nz - z_surface = znode(c, c, f, i, j, Nz+1, grid) - b_surface = @inbounds b[i, j, Nz] # buoyancy at surface - - # Start downwards look - z_ij = z_surface - found_mixed_layer_depth = false - - @unroll for k in Nz-1 : -1 : 1 # scroll from point just below surface - b⁺ = @inbounds b[i, j, k+1] - bᵏ = @inbounds b[i, j, k] - - # If buoyancy decreases downwards, both are > 0 - Δb⁺ = b_surface - b⁺ - Δbᵏ = b_surface - bᵏ - - zᵏ = znode(c, c, c, i, j, k, grid) - Δz⁺ = Δzᶜᶜᶠ(i, j, k+1, grid) - - # Assuming buoyancy decreases downwards - just_below_mixed_layer = (Δb⁺ < Δb) & (Δbᵏ >= Δb) - just_below_mixed_layer *= !found_mixed_layer_depth - - # Linearly interpolate to find mixed layer height - new_z_ij = zᵏ + (Δb - Δbᵏ) / (Δb⁺ - Δbᵏ) * Δz⁺ - - # Replace z_ij if we found a new mixed layer depth - z_ij = ifelse(!found_mixed_layer_depth | just_below_mixed_layer, new_z_ij, z_ij) - end - - # Note "-" since `h` is supposed to be "depth" rather than "height" - h_ij = z_surface - z_ij - @inbounds h[i, j, 1] = ifelse(inactive_node(i, j, Nz, grid, c, c, c), zero(grid), h_ij) -end - -struct MixedLayerDepthOperand{B, FT} - buoyancy_operation :: B - mixed_layer_buoyancy_differential :: FT -end - -Base.summary(op::MixedLayerDepthOperand) = "MixedLayerDepthOperand" - -const MixedLayerDepthField = Field{Center, Center, Nothing, <:MixedLayerDepthOperand} - -""" - MixedLayerDepthField(grid, tracers, buoyancy_model; Δb = 3e-4, field_kw...) - -Return a reduced `Field{Center, Center, Nothing}` that represents -mixed layer depth for `model`, based on a buoyancy differential criterion. -The mixed layer depth is defined as the depth ``h`` for which - -```math -b(z=0) - b(z=-h) = Δb -``` - -This criterion is solved by integrating downwards and linearly interpolating to find `h`, -assuming that ``b`` decreases with depth. - -Keyword arguments -================= - -* `Δb`: Buoyancy differential used to calculate mixed layer depth -* `field_kw`: Keyword arguments passed to `Field`. - -Example -======= - -```julia -h = MixedLayerDepth(model) -compute!(h) # compute mixed layer depth -``` -""" -function MixedLayerDepthField(grid, tracers, buoyancy_model; Δb = 3e-4, field_kw...) - b_op = buoyancy(buoyancy_model, grid, tracers) - operand = MixedLayerDepthOperand(b_op, Δb) - return Field{Center, Center, Nothing}(grid; operand, field_kw...) -end - -MixedLayerDepthField(model; kw...) = MixedLayerDepthField(model.grid, - model.tracers, - model.buoyancy; kw...) - -function compute!(h::MixedLayerDepthField, time=nothing) - arch = architecture(h) - b = h.operand.buoyancy_operation - Δb = h.operand.mixed_layer_buoyancy_differential - launch!(arch, h.grid, :xy, compute_mld!, h, h.grid, b, Δb) - fill_halo_regions!(h) - return h -end - -end # module diff --git a/src/Diagnostics/Diagnostics.jl b/src/Diagnostics/Diagnostics.jl new file mode 100644 index 00000000..539572b0 --- /dev/null +++ b/src/Diagnostics/Diagnostics.jl @@ -0,0 +1,19 @@ +module Diagnostics + +export MixedLayerDepthField, MixedLayerDepthOperand + +using Oceananigans +using Oceananigans.Architectures: architecture +using Oceananigans.BuoyancyModels: buoyancy +using Oceananigans.Grids: new_data, inactive_cell, znode +using Oceananigans.BoundaryConditions: FieldBoundaryConditions, fill_halo_regions! +using Oceananigans.Fields: FieldStatus +using Oceananigans.Utils: launch! + +import Oceananigans.Fields: compute! + +using KernelAbstractions: @index, @kernel + +include("mixed_layer_depth.jl") + +end # module diff --git a/examples/mixed_layer_depth.jl b/src/Diagnostics/mixed_layer_depth.jl similarity index 85% rename from examples/mixed_layer_depth.jl rename to src/Diagnostics/mixed_layer_depth.jl index e2128878..afc509be 100644 --- a/examples/mixed_layer_depth.jl +++ b/src/Diagnostics/mixed_layer_depth.jl @@ -1,15 +1,3 @@ -using Oceananigans -using Oceananigans.Architectures: architecture -using Oceananigans.BuoyancyModels: buoyancy -using Oceananigans.Grids: new_data, inactive_cell, znode -using Oceananigans.BoundaryConditions: FieldBoundaryConditions, fill_halo_regions! -using Oceananigans.Fields: FieldStatus -using Oceananigans.Utils: launch! - -import Oceananigans.Fields: compute! - -using KernelAbstractions: @index, @kernel - mutable struct MixedLayerDepthOperand{FT, B} buoyancy_perturbation :: B difference_criterion :: FT From 653ae7803364b8a6500d94cfef59c3cca0eb1787 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Thu, 12 Dec 2024 21:11:11 -0700 Subject: [PATCH 14/22] Change difference criterion --- src/Diagnostics/mixed_layer_depth.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Diagnostics/mixed_layer_depth.jl b/src/Diagnostics/mixed_layer_depth.jl index afc509be..70b62bb1 100644 --- a/src/Diagnostics/mixed_layer_depth.jl +++ b/src/Diagnostics/mixed_layer_depth.jl @@ -17,7 +17,7 @@ const MixedLayerDepthField = Field{<:Any, <:Any, <:Any, <:MixedLayerDepthOperand MixedLayerDepthField(bm, grid, tracers; difference_criterion=1e-4) """ -function MixedLayerDepthField(bm, grid, tracers; difference_criterion=1e-4) +function MixedLayerDepthField(bm, grid, tracers; difference_criterion=3e-5) operand = MixedLayerDepthOperand(bm, grid, tracers; difference_criterion) loc = (Center, Center, Nothing) indices = (:, :, :) From 382330134ef71f5eef14863f777fd307b421bfc2 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sat, 14 Dec 2024 12:26:08 -0700 Subject: [PATCH 15/22] Deleted Diagnostics.jl --- src/Diagnostics.jl | 124 --------------------------------------------- 1 file changed, 124 deletions(-) delete mode 100644 src/Diagnostics.jl diff --git a/src/Diagnostics.jl b/src/Diagnostics.jl deleted file mode 100644 index 0fad0b8e..00000000 --- a/src/Diagnostics.jl +++ /dev/null @@ -1,124 +0,0 @@ -module Diagnostics - -export MixedLayerDepthField - -using Oceananigans -using Oceananigans.BuoyancyFormulations: buoyancy -using Oceananigans.BoundaryConditions: fill_halo_regions! -using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid -using Oceananigans.Architectures: device, architecture -using Oceananigans.Utils: launch! -using Oceananigans.Grids: Center, Face, inactive_node, znode -using Oceananigans.Operators: Δzᶜᶜᶠ - -using KernelAbstractions: @index, @kernel -using KernelAbstractions.Extras.LoopInfo: @unroll - -import Oceananigans.Fields: compute! - -const c = Center() -const f = Face() - -@inline z_bottom(i, j, grid) = znode(c, c, f, i, j, 1, grid) -@inline z_bottom(i, j, grid::ImmersedBoundaryGrid) = @inbounds grid.immersed_boundary.bottom_height[i, j] - -@inline bottom(i, j, grid) = znode(c, c, f, i, j, 1, grid) -##### -##### MixedLayerDepthField -##### - -@kernel function compute_mld!(h, grid, b, Δb) - i, j = @index(Global, NTuple) - - Nz = grid.Nz - z_surface = znode(c, c, f, i, j, Nz+1, grid) - b_surface = @inbounds b[i, j, Nz] # buoyancy at surface - - # Start downwards look - z_ij = z_surface - found_mixed_layer_depth = false - - @unroll for k in Nz-1 : -1 : 1 # scroll from point just below surface - b⁺ = @inbounds b[i, j, k+1] - bᵏ = @inbounds b[i, j, k] - - # If buoyancy decreases downwards, both are > 0 - Δb⁺ = b_surface - b⁺ - Δbᵏ = b_surface - bᵏ - - zᵏ = znode(c, c, c, i, j, k, grid) - Δz⁺ = Δzᶜᶜᶠ(i, j, k+1, grid) - - # Assuming buoyancy decreases downwards - just_below_mixed_layer = (Δb⁺ < Δb) & (Δbᵏ >= Δb) - just_below_mixed_layer *= !found_mixed_layer_depth - - # Linearly interpolate to find mixed layer height - new_z_ij = zᵏ + (Δb - Δbᵏ) / (Δb⁺ - Δbᵏ) * Δz⁺ - - # Replace z_ij if we found a new mixed layer depth - z_ij = ifelse(!found_mixed_layer_depth | just_below_mixed_layer, new_z_ij, z_ij) - end - - # Note "-" since `h` is supposed to be "depth" rather than "height" - h_ij = z_surface - z_ij - @inbounds h[i, j, 1] = ifelse(inactive_node(i, j, Nz, grid, c, c, c), zero(grid), h_ij) -end - -struct MixedLayerDepthOperand{B, FT} - buoyancy_operation :: B - mixed_layer_buoyancy_differential :: FT -end - -Base.summary(op::MixedLayerDepthOperand) = "MixedLayerDepthOperand" - -const MixedLayerDepthField = Field{Center, Center, Nothing, <:MixedLayerDepthOperand} - -""" - MixedLayerDepthField(grid, tracers, buoyancy_model; Δb = 3e-4, field_kw...) - -Return a reduced `Field{Center, Center, Nothing}` that represents -mixed layer depth for `model`, based on a buoyancy differential criterion. -The mixed layer depth is defined as the depth ``h`` for which - -```math -b(z=0) - b(z=-h) = Δb -``` - -This criterion is solved by integrating downwards and linearly interpolating to find `h`, -assuming that ``b`` decreases with depth. - -Keyword arguments -================= - -* `Δb`: Buoyancy differential used to calculate mixed layer depth -* `field_kw`: Keyword arguments passed to `Field`. - -Example -======= - -```julia -h = MixedLayerDepth(model) -compute!(h) # compute mixed layer depth -``` -""" -function MixedLayerDepthField(grid, tracers, buoyancy_model; Δb = 3e-4, field_kw...) - b_op = buoyancy(buoyancy_model, grid, tracers) - operand = MixedLayerDepthOperand(b_op, Δb) - return Field{Center, Center, Nothing}(grid; operand, field_kw...) -end - -MixedLayerDepthField(model; kw...) = MixedLayerDepthField(model.grid, - model.tracers, - model.buoyancy; kw...) - -function compute!(h::MixedLayerDepthField, time=nothing) - arch = architecture(h) - b = h.operand.buoyancy_operation - Δb = h.operand.mixed_layer_buoyancy_differential - launch!(arch, h.grid, :xy, compute_mld!, h, h.grid, b, Δb) - fill_halo_regions!(h) - return h -end - -end # module From 7dbd8c7fbf5a47c0be1c5aeca68c7e9f20761b18 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sat, 14 Dec 2024 12:34:56 -0700 Subject: [PATCH 16/22] Update to latest Oceananigans --- src/Diagnostics/Diagnostics.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Diagnostics/Diagnostics.jl b/src/Diagnostics/Diagnostics.jl index 539572b0..7f3931f1 100644 --- a/src/Diagnostics/Diagnostics.jl +++ b/src/Diagnostics/Diagnostics.jl @@ -4,7 +4,7 @@ export MixedLayerDepthField, MixedLayerDepthOperand using Oceananigans using Oceananigans.Architectures: architecture -using Oceananigans.BuoyancyModels: buoyancy +using Oceananigans.BuoyancyFormulations: buoyancy using Oceananigans.Grids: new_data, inactive_cell, znode using Oceananigans.BoundaryConditions: FieldBoundaryConditions, fill_halo_regions! using Oceananigans.Fields: FieldStatus From 25f732efbcf0825dc790ed2904d6f1fed07e7511 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sat, 14 Dec 2024 12:40:42 -0700 Subject: [PATCH 17/22] Add a test for MixedLayerDepthField --- test/test_diagnostics.jl | 52 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) create mode 100644 test/test_diagnostics.jl diff --git a/test/test_diagnostics.jl b/test/test_diagnostics.jl new file mode 100644 index 00000000..670e46ff --- /dev/null +++ b/test/test_diagnostics.jl @@ -0,0 +1,52 @@ +include("runtests_setup.jl") + +using SeawaterPolynomials: TEOS10EquationOfState +using Oceananigans.BuoyancyFormulations: buoyancy +using Oceananigans: location +using ClimaOcean.DataWrangling.ECCO: ECCO_field, ECCOFieldTimeSeries +using ClimaOcean.Diagnostics: MixedLayerDepthField, MixedLayerDepthOperand + +@testset "MixedLayerDepthField" begin + for arch in test_architectures + grid = LatitudeLongitudeGrid(arch; + size = (3, 3, 100), + latitude = (0, 30), + longitude = (150, 180), + z = (-1000, 0)) + + bottom_height = regrid_bathymetry(grid; + minimum_depth = 10, + interpolation_passes = 5, + major_basins = 1) + + grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height)) + + start = DateTimeProlepticGregorian(1993, 1, 1) + stop = DateTimeProlepticGregorian(1993, 2, 1) + dates = range(start; stop, step=Month(1)) + + Tmeta = ECCOMetadata(:temperature; dates) + Smeta = ECCOMetadata(:salinity; dates) + + Tt = ECCOFieldTimeSeries(Tmeta, grid; time_indices_in_memory=2) + St = ECCOFieldTimeSeries(Smeta, grid; time_indices_in_memory=2) + + equation_of_state = TEOS10EquationOfState() + sb = SeawaterBuoyancy(; equation_of_state) + tracers = (T=Tt[1], S=St[1]) + h = MixedLayerDepthField(sb, grid, tracers) + + @test h isa Field + @test location(h) == (Center, Center, Nothing) + @test h.operand isa MixedLayerDepthOperand + @test h.operand.buoyancy_perturbation isa KernelFunctionOperation + + compute!(h) + @test h[1, 1, 1] ≈ 16.255836 # m + + tracers = (T=Tt[2], S=St[2]) + h.operand.buoyancy_perturbation = buoyancy(sb, grid, tracers) + compute!(h) + @test h[1, 1, 1] ≈ 9.295890287 # m + end +end From cb3c74d70074b91bf2e79d616b382aa3b5e3493c Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sat, 14 Dec 2024 12:40:51 -0700 Subject: [PATCH 18/22] Add test for MixeDLayerDepthField --- test/runtests.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/test/runtests.jl b/test/runtests.jl index 932f48ef..aca50ba5 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -58,6 +58,7 @@ end if test_group == :simulations || test_group == :all CUDA.set_runtime_version!(v"12.2", local_toolkit = true) # Seems to help in finding the correct CUDA version include("test_simulations.jl") + include("test_diagnostics.jl") end if test_group == :distributed || test_group == :all From 7242be5de84c5a23bc91490f7e9985754f962086 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sat, 14 Dec 2024 13:56:24 -0700 Subject: [PATCH 19/22] Add allowscalar --- test/runtests_setup.jl | 2 ++ test/test_diagnostics.jl | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/test/runtests_setup.jl b/test/runtests_setup.jl index f8d83ebd..fd829728 100644 --- a/test/runtests_setup.jl +++ b/test/runtests_setup.jl @@ -16,6 +16,8 @@ using ClimaOcean.Bathymetry: download_bathymetry_cache using CFTime using Dates +using CUDA: @allowscalar + gpu_test = parse(Bool, get(ENV, "GPU_TEST", "false")) test_architectures = gpu_test ? [GPU()] : [CPU()] diff --git a/test/test_diagnostics.jl b/test/test_diagnostics.jl index 670e46ff..fc64742d 100644 --- a/test/test_diagnostics.jl +++ b/test/test_diagnostics.jl @@ -42,11 +42,11 @@ using ClimaOcean.Diagnostics: MixedLayerDepthField, MixedLayerDepthOperand @test h.operand.buoyancy_perturbation isa KernelFunctionOperation compute!(h) - @test h[1, 1, 1] ≈ 16.255836 # m + @test @allowscalar h[1, 1, 1] ≈ 16.255836 # m tracers = (T=Tt[2], S=St[2]) h.operand.buoyancy_perturbation = buoyancy(sb, grid, tracers) compute!(h) - @test h[1, 1, 1] ≈ 9.295890287 # m + @test @allowscalar h[1, 1, 1] ≈ 9.295890287 # m end end From 6d4ced651634d0f984f413ef4863fd32775c4676 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Sun, 15 Dec 2024 21:18:18 -0700 Subject: [PATCH 20/22] Hmm --- test/test_jra55.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_jra55.jl b/test/test_jra55.jl index 8f83aa68..fd938839 100644 --- a/test/test_jra55.jl +++ b/test/test_jra55.jl @@ -37,7 +37,7 @@ using ClimaOcean.OceanSeaIceModels: PrescribedAtmosphere @info "Testing loading preprocessed JRA55 data on $A..." in_memory_jra55_fts = JRA55_field_time_series(test_name; - time_indices = 1:3, + time_indices = 1:4, architecture = arch, backend = InMemory(2)) From 12ef52175d2ab3d12894405fa6569b4d80b5a9dc Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Mon, 16 Dec 2024 08:09:42 -0700 Subject: [PATCH 21/22] Update src/Diagnostics/mixed_layer_depth.jl Co-authored-by: Simone Silvestri --- src/Diagnostics/mixed_layer_depth.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/Diagnostics/mixed_layer_depth.jl b/src/Diagnostics/mixed_layer_depth.jl index 70b62bb1..48266d53 100644 --- a/src/Diagnostics/mixed_layer_depth.jl +++ b/src/Diagnostics/mixed_layer_depth.jl @@ -82,7 +82,8 @@ const f = Face() # Apply various criterion h = -z★ h = max(h, zero(grid)) - h = min(h, grid.Lz) + H = static_column_depthᶜᶜᵃ(i, j, grid) + h = min(h, H) @inbounds mld[i, j, 1] = h end From 211cdfd9e3e5131de252a60612646002723fc574 Mon Sep 17 00:00:00 2001 From: Gregory Wagner Date: Mon, 16 Dec 2024 11:24:18 -0700 Subject: [PATCH 22/22] Import static_column_depth --- src/Diagnostics/mixed_layer_depth.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/Diagnostics/mixed_layer_depth.jl b/src/Diagnostics/mixed_layer_depth.jl index 48266d53..8da401ea 100644 --- a/src/Diagnostics/mixed_layer_depth.jl +++ b/src/Diagnostics/mixed_layer_depth.jl @@ -1,3 +1,5 @@ +using Oceananigans.Grids: static_column_depthᶜᶜᵃ + mutable struct MixedLayerDepthOperand{FT, B} buoyancy_perturbation :: B difference_criterion :: FT