Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
58 commits
Select commit Hold shift + click to select a range
7241c49
Unify JRA55 onto generic DatasetBackend; add async PrefetchingBackend
simone-silvestri Apr 17, 2026
32ef1f1
formatting changes
simone-silvestri Apr 17, 2026
9886fad
fix code fragilities
simone-silvestri Apr 17, 2026
8c2f5d5
some changes
simone-silvestri Apr 17, 2026
3984c18
Merge branch 'main' into ss/dataset-backend-prefetch
simone-silvestri Apr 20, 2026
120ab36
Merge branch 'ss/dataset-backend-prefetch' of github.com:NumericalEar…
simone-silvestri Apr 20, 2026
fde7baf
fix all tests and examples
simone-silvestri Apr 20, 2026
bf25439
more test fixes
simone-silvestri Apr 20, 2026
6b2830d
just pass the native grid
simone-silvestri Apr 20, 2026
009327d
make sure we use 2 sizes and 2 halos
simone-silvestri Apr 20, 2026
cd97aa9
inner backend
simone-silvestri Apr 20, 2026
d9ec6e0
pop flat elements
simone-silvestri Apr 20, 2026
b8786d7
put a guard on the time_indices_in_memory
simone-silvestri Apr 20, 2026
a06a656
these variables have a different size
simone-silvestri Apr 20, 2026
e7a6070
discontinue this example
simone-silvestri Apr 20, 2026
58c73ad
fix surface flux tests
simone-silvestri Apr 20, 2026
2efbaad
fix tests
simone-silvestri Apr 20, 2026
0425328
Update src/DataWrangling/restoring.jl
simone-silvestri Apr 20, 2026
3b42167
Update src/DataWrangling/prefetching_backend.jl
simone-silvestri Apr 20, 2026
7199622
Merge branch 'main' into ss/dataset-backend-prefetch
simone-silvestri Apr 22, 2026
b7d3744
Merge branch 'main' into ss/dataset-backend-prefetch
simone-silvestri Apr 22, 2026
00e4273
Fix prefetch predictor off-by-one (slide by Nm - 1, not Nm)
simone-silvestri Apr 22, 2026
92d98e7
Recover from failed prefetch instead of killing the simulation
simone-silvestri Apr 23, 2026
ea52120
remove dataset prefetch
simone-silvestri Apr 23, 2026
ab8aa12
remove JRA55NetCDFBackend
simone-silvestri Apr 23, 2026
67d7ba0
Merge branch 'main' into ss/dataset-backend-prefetch
simone-silvestri Apr 23, 2026
8dea43e
make sure we download corrupted files again
simone-silvestri Apr 23, 2026
f472c3d
try like this
simone-silvestri Apr 23, 2026
9b0c77e
Merge branch 'main' into ss/dataset-backend-prefetch
simone-silvestri Apr 23, 2026
af301f4
fix downloader
simone-silvestri Apr 23, 2026
64bb057
Merge branch 'ss/dataset-backend-prefetch' of github.com:NumericalEar…
simone-silvestri Apr 23, 2026
f37a3c0
Merge branch 'main' into ss/dataset-backend-prefetch
simone-silvestri Apr 24, 2026
fe90124
Merge branch 'main' into ss/dataset-backend-prefetch
simone-silvestri Apr 27, 2026
c03a97e
bounding_box -> region
simone-silvestri Apr 27, 2026
953d543
Merge branch 'main' into ss/dataset-backend-prefetch
simone-silvestri Apr 27, 2026
c49cb34
reenable the single column example
simone-silvestri Apr 27, 2026
a2289da
typo
simone-silvestri Apr 27, 2026
6abba9a
huge refactor
simone-silvestri Apr 28, 2026
7b76d7c
restore CI
simone-silvestri Apr 28, 2026
a05d4dc
more fixes
simone-silvestri Apr 28, 2026
1de664a
Merge branch 'main' into ss/dataset-backend-prefetch
simone-silvestri Apr 28, 2026
e60e249
small bugfix
simone-silvestri Apr 28, 2026
2af644d
Merge branch 'ss/dataset-backend-prefetch' of github.com:NumericalEar…
simone-silvestri Apr 28, 2026
52ef0f4
correct WOA climatology
simone-silvestri Apr 28, 2026
7058a00
fix uniform cells `restrict`
simone-silvestri Apr 28, 2026
d3454dd
fix longitude and latitude names
simone-silvestri Apr 28, 2026
fd577ca
some changes
simone-silvestri Apr 28, 2026
4613e68
use Oceananigans' method
simone-silvestri Apr 28, 2026
85681ae
bump version
simone-silvestri Apr 29, 2026
93b7440
Merge remote-tracking branch 'origin/main' into ss/dataset-backend-pr…
simone-silvestri Apr 29, 2026
3fa6f53
remove extra comment
simone-silvestri Apr 29, 2026
958cffa
fix test surface fluxes
simone-silvestri Apr 29, 2026
f69483b
fix the tests
simone-silvestri Apr 29, 2026
bd6188d
Merge branch 'main' into ss/dataset-backend-prefetch
simone-silvestri May 1, 2026
521eb0a
fix tests
simone-silvestri May 1, 2026
8c6d701
fix restrict to snap out
simone-silvestri May 1, 2026
568a9f4
fix the BBox longitudes into native longitudes
simone-silvestri May 1, 2026
5d1e95d
fix nothing longitudes
simone-silvestri May 1, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "NumericalEarth"
uuid = "904d977b-046a-4731-8b86-9235c0d1ef02"
license = "MIT"
version = "0.3.0"
version = "0.3.1"
authors = ["NumericalEarth contributors"]

[deps]
Expand Down
4 changes: 2 additions & 2 deletions examples/generate_surface_fluxes.jl
Original file line number Diff line number Diff line change
Expand Up @@ -45,9 +45,9 @@ save("ECCO_continents.png", fig)
# - downwelling longwave radiation
#
# We load in memory only the first two time indices, corresponding to January 1st
# (at 00:00 AM and 03:00 AM), by using `JRA55NetCDFBackend(2)`.
# (at 00:00 AM and 03:00 AM), by using `time_indices_in_memory = 2`.

atmosphere = JRA55PrescribedAtmosphere(; backend = JRA55NetCDFBackend(2))
atmosphere = JRA55PrescribedAtmosphere(; time_indices_in_memory = 2)
ocean = ocean_simulation(grid, closure=nothing)

# Now that we have an atmosphere and ocean, we `set!` the ocean temperature and salinity
Expand Down
7 changes: 5 additions & 2 deletions examples/inspect_JRA55_data.jl
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,11 @@ using Oceananigans
using Oceananigans.Units
using Printf

Qswt = NumericalEarth.JRA55.JRA55FieldTimeSeries(:downwelling_shortwave_radiation)
rht = NumericalEarth.JRA55.JRA55FieldTimeSeries(:relative_humidity)
using NumericalEarth.DataWrangling: Metadata
using NumericalEarth.JRA55: RepeatYearJRA55

Qswt = FieldTimeSeries(Metadata(:downwelling_shortwave_radiation; dataset=RepeatYearJRA55()); time_indices_in_memory=8)
rht = FieldTimeSeries(Metadata(:specific_humidity; dataset=RepeatYearJRA55()); time_indices_in_memory=8)
Comment thread
simone-silvestri marked this conversation as resolved.

function lonlat2xyz(lons::AbstractVector, lats::AbstractVector)
x = [cosd(lat) * cosd(lon) for lon in lons, lat in lats]
Expand Down
6 changes: 3 additions & 3 deletions examples/meridional_heat_transport_ecco.jl
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,9 @@ set!(ocean.model, T=ecco_temperature, S=ecco_salinity)
set!(sea_ice.model, h=ecco_sea_ice_thickness, ℵ=ecco_sea_ice_concentration)

radiation = Radiation(arch)
atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(80),
include_rivers_and_icebergs = false)
esm = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation)
atmosphere = JRA55PrescribedAtmosphere(arch; time_indices_in_memory = 80)
land = JRA55PrescribedLand(arch; time_indices_in_memory = 80)
esm = OceanSeaIceModel(ocean, sea_ice; atmosphere, land, radiation)

simulation = Simulation(esm; Δt=20minutes, stop_time=5*365days)

Expand Down
7 changes: 3 additions & 4 deletions examples/near_global_ocean_simulation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -102,11 +102,10 @@ radiation = Radiation(arch)
# The JRA55 dataset provides atmospheric data such as temperature, humidity, and winds
# to calculate turbulent fluxes using bulk formulae, see [`InterfaceComputations`](@ref NumericalEarth.EarthSystemModels.InterfaceComputations).
# The number of snapshots that are loaded into memory is determined by
# the `backend`. Here, we load 41 snapshots at a time into memory.
# `time_indices_in_memory`. Here, we load 41 snapshots at a time into memory.

jra55_backend = JRA55NetCDFBackend(41)
atmosphere = JRA55PrescribedAtmosphere(arch; backend=jra55_backend)
land = JRA55PrescribedLand(arch; backend=jra55_backend)
atmosphere = JRA55PrescribedAtmosphere(arch; time_indices_in_memory = 41)
land = JRA55PrescribedLand(arch; time_indices_in_memory = 41)

# ## The coupled simulation

Expand Down
5 changes: 2 additions & 3 deletions examples/one_degree_simulation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -96,9 +96,8 @@ set!(sea_ice.model, h=ecco_sea_ice_thickness, ℵ=ecco_sea_ice_concentration)

# We force the simulation with a JRA55-do atmospheric reanalysis.
radiation = Radiation(arch)
jra55_backend = JRA55NetCDFBackend(80)
atmosphere = JRA55PrescribedAtmosphere(arch; backend=jra55_backend)
land = JRA55PrescribedLand(arch; backend=jra55_backend)
atmosphere = JRA55PrescribedAtmosphere(arch; time_indices_in_memory = 80)
land = JRA55PrescribedLand(arch; time_indices_in_memory = 80)

# ### Coupled simulation

Expand Down
5 changes: 2 additions & 3 deletions examples/single_column_os_papa_simulation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -63,10 +63,9 @@ set!(ocean.model, T=Metadatum(:temperature, dataset=GLORYSMonthly(), region=col)
# We build a `JRA55PrescribedAtmosphere` at the same location as the single-colunm grid
# which is based on the JRA55 reanalysis.

atmosphere = JRA55PrescribedAtmosphere(longitude = λ★,
latitude = φ★,
atmosphere = JRA55PrescribedAtmosphere(region = Column(λ★, φ★),
end_date = DateTime(1990, 1, 31), # Last day of the simulation
backend = InMemory())
time_indices_in_memory = 1000)

# This builds a representation of the atmosphere on the small grid

Expand Down
2 changes: 1 addition & 1 deletion examples/veros_ocean_forced_simulation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ ocean.set_forcing = set_forcing_tke_only
# This includes 2-meter wind velocity, temperature, humidity, downwelling longwave and shortwave
# radiation, as well as freshwater fluxes.

atmos = JRA55PrescribedAtmosphere(; backend = JRA55NetCDFBackend(10))
atmos = JRA55PrescribedAtmosphere(; time_indices_in_memory = 10)

# The coupled ocean--atmosphere model.
# We use the default radiation model and we do not couple an ice model for simplicity.
Expand Down
2 changes: 1 addition & 1 deletion experiments/arctic_simulation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ set!(sea_ice.model, h=Metadatum(:sea_ice_thickness; dataset),
##### A Prescribed Atmosphere model
#####

atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(40))
atmosphere = JRA55PrescribedAtmosphere(arch; time_indices_in_memory=40)
radiation = Radiation()

#####
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ salinity = ECCOMetadata(:salinity; dir="./")
ice_thickness = ECCOMetadata(:sea_ice_thickness; dir="./")
ice_concentration = ECCOMetadata(:sea_ice_concentration; dir="./")

atmosphere = JRA55PrescribedAtmosphere(arch, backend=JRA55NetCDFBackend(20))
atmosphere = JRA55PrescribedAtmosphere(arch, time_indices_in_memory=20)
radiation = Radiation(ocean_albedo = LatitudeDependentAlbedo(), sea_ice_albedo=0.6)

set!(ocean.model, T=temperature, S=salinity)
Expand Down
2 changes: 1 addition & 1 deletion experiments/flux_climatology/flux_climatology.jl
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,7 @@ heat_capacity(ocean::Simulation{<:PrescribedOcean}) = 3995.6
##### A prescribed atmosphere...
#####

atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(1000))
atmosphere = JRA55PrescribedAtmosphere(arch; time_indices_in_memory=1000)

#####
##### A prescribed earth...
Expand Down
3 changes: 1 addition & 2 deletions experiments/one_degree_simulation/debug_tides.jl
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,8 @@ using Statistics

import SPICE

backend = JRA55NetCDFBackend(41)
ρᵒᶜ = 1020
pa = JRA55_field_time_series(:sea_level_pressure; backend)
pa = JRA55_field_time_series(:sea_level_pressure; time_indices_in_memory=41)
grid = pa.grid

# dt = 30 * 60 # minutes
Expand Down
3 changes: 1 addition & 2 deletions experiments/one_degree_simulation/generate_tidal_forcing.jl
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,7 @@ using Dates
using GLMakie
import SPICE

backend = JRA55NetCDFBackend(41)
pa = JRA55_field_time_series(:sea_level_pressure; backend)
pa = JRA55_field_time_series(:sea_level_pressure; time_indices_in_memory=41)

pa_with_tides = FieldTimeSeries{Center, Center, Nothing}(pa.grid, pa.times,
path = "sea_level_pressure_plus_tides.jld2",
Expand Down
2 changes: 1 addition & 1 deletion experiments/one_degree_simulation/one_degree_simulation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ set!(ocean.model, T=ECCOMetadata(:temperature; dates=first(dates)),
S=ECCOMetadata(:salinity; dates=first(dates)))

radiation = Radiation(arch)
atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(41))
atmosphere = JRA55PrescribedAtmosphere(arch; time_indices_in_memory=41)
coupled_model = OceanOnlyModel(ocean; atmosphere, radiation)

simulation = Simulation(coupled_model; Δt=10minutes, stop_iteration=100)
Expand Down
2 changes: 2 additions & 0 deletions src/DataWrangling/DataWrangling.jl
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,9 @@ default_mask_value(dataset) = NaN

# Fundamentals
include("metadata.jl")
include("set_region_data.jl")
include("metadata_field.jl")
include("dataset_backend.jl")
include("metadata_field_time_series.jl")
include("inpainting.jl")
include("restoring.jl")
Expand Down
45 changes: 10 additions & 35 deletions src/DataWrangling/ECCO/ECCO.jl
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,7 @@ using NumericalEarth.DataWrangling:
location,
compute_mask,
inpaint_mask!,
set_metadata_field!,
extract_column!
set_metadata_field!

using KernelAbstractions: @kernel, @index

Expand All @@ -51,6 +50,8 @@ import NumericalEarth.DataWrangling:
metaprefix,
longitude_interfaces,
latitude_interfaces,
longitude_name,
latitude_name,
z_interfaces,
is_three_dimensional,
inpainted_metadata_path,
Expand Down Expand Up @@ -113,6 +114,13 @@ longitude_interfaces(::ECCODataset) = (0, 360)
longitude_interfaces(::ECCO4Monthly) = (-180, 180)
latitude_interfaces(::ECCODataset) = (-90, 90)

longitude_name(::Metadata{<:ECCODataset}) = "LONGITUDE_T"
latitude_name(::Metadata{<:ECCODataset}) = "LATITUDE_T"
longitude_name(::Metadata{<:ECCO4Monthly}) = "longitude"
latitude_name(::Metadata{<:ECCO4Monthly}) = "latitude"
longitude_name(::Metadata{<:ECCO4DarwinMonthly}) = "longitude"
latitude_name(::Metadata{<:ECCO4DarwinMonthly}) = "latitude"

z_interfaces(::ECCODataset) = [
-6128.75,
-5683.75,
Expand Down Expand Up @@ -371,37 +379,4 @@ inpainted_metadata_path(metadata::ECCOMetadatum) = joinpath(metadata.dir, inpain

include("ECCO_atmosphere.jl")

#####
##### Column Field for ECCO datasets (which always download globally)
#####

using Oceananigans.BoundaryConditions: fill_halo_regions!

const ECCOColumnMetadatum = Metadatum{<:ECCODataset, <:Any, <:Column}

function Oceananigans.Fields.Field(metadata::ECCOColumnMetadatum, arch=CPU();
inpainting = default_inpainting(metadata),
mask = nothing,
halo = (3, 3, 3),
cache_inpainted_data = true)

download_dataset(metadata)
column_grid = native_grid(metadata, arch; halo)

# Build a full-grid Field without a region to load the global data
global_metadatum = Metadatum(metadata.name;
dataset = metadata.dataset,
date = metadata.dates)

intermediate_field = Field(global_metadatum, arch; inpainting, mask, halo, cache_inpainted_data)
fill_halo_regions!(intermediate_field)

# Extract the column
_, _, LZ = location(metadata)
column_field = Field{Nothing, Nothing, LZ}(column_grid)
extract_column!(column_field, intermediate_field, metadata.region)

return column_field
end

end # Module
29 changes: 11 additions & 18 deletions src/DataWrangling/ECCO/ECCO_atmosphere.jl
Original file line number Diff line number Diff line change
Expand Up @@ -31,26 +31,19 @@ function ECCOPrescribedAtmosphere(architecture = CPU(), FT = Float32;
surface_layer_height = 2, # meters
other_kw...)

ua_meta = Metadata(:eastward_wind; dataset, start_date, end_date, dir)
va_meta = Metadata(:northward_wind; dataset, start_date, end_date, dir)
Ta_meta = Metadata(:air_temperature; dataset, start_date, end_date, dir)
qa_meta = Metadata(:air_specific_humidity; dataset, start_date, end_date, dir)
pa_meta = Metadata(:sea_level_pressure; dataset, start_date, end_date, dir)
ℐꜜˡʷ_meta = Metadata(:downwelling_longwave; dataset, start_date, end_date, dir)
ℐꜜˢʷ_meta = Metadata(:downwelling_shortwave; dataset, start_date, end_date, dir)
Fr_meta = Metadata(:rain_freshwater_flux; dataset, start_date, end_date, dir)

kw = (; time_indices_in_memory, time_indexing)
kw = (; time_indexing, time_indices_in_memory)
kw = merge(kw, other_kw)

ua = FieldTimeSeries(ua_meta, architecture; kw...)
va = FieldTimeSeries(va_meta, architecture; kw...)
Ta = FieldTimeSeries(Ta_meta, architecture; kw...)
qa = FieldTimeSeries(qa_meta, architecture; kw...)
pa = FieldTimeSeries(pa_meta, architecture; kw...)
ℐꜜˡʷ = FieldTimeSeries(ℐꜜˡʷ_meta, architecture; kw...)
ℐꜜˢʷ = FieldTimeSeries(ℐꜜˢʷ_meta, architecture; kw...)
Fr = FieldTimeSeries(Fr_meta, architecture; kw...)
ECCOFieldTimeSeries(name) = FieldTimeSeries(Metadata(name; dataset, start_date, end_date, dir), architecture; kw...)

ua = ECCOFieldTimeSeries(:eastward_wind)
va = ECCOFieldTimeSeries(:northward_wind)
Ta = ECCOFieldTimeSeries(:air_temperature)
qa = ECCOFieldTimeSeries(:air_specific_humidity)
pa = ECCOFieldTimeSeries(:sea_level_pressure)
ℐꜜˡʷ = ECCOFieldTimeSeries(:downwelling_longwave)
ℐꜜˢʷ = ECCOFieldTimeSeries(:downwelling_shortwave)
Fr = ECCOFieldTimeSeries(:rain_freshwater_flux)

freshwater_flux = (; rain = Fr)

Expand Down
4 changes: 4 additions & 0 deletions src/DataWrangling/EN4/EN4.jl
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@ import NumericalEarth.DataWrangling:
z_interfaces,
longitude_interfaces,
latitude_interfaces,
longitude_name,
latitude_name,
is_three_dimensional,
reversed_vertical_axis,
inpainted_metadata_path,
Expand All @@ -70,6 +72,8 @@ reversed_vertical_axis(::EN4Monthly) = true

longitude_interfaces(::EN4Monthly) = (0.5, 360.5)
latitude_interfaces(::EN4Monthly) = (-83.5, 89.5)
longitude_name(::Metadata{<:EN4Monthly}) = "lon"
latitude_name(::Metadata{<:EN4Monthly}) = "lat"
available_variables(::EN4Monthly) = EN4_dataset_variable_names

z_interfaces(::EN4Monthly) = [
Expand Down
2 changes: 1 addition & 1 deletion src/DataWrangling/JRA55/JRA55.jl
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
module JRA55

export JRA55FieldTimeSeries, JRA55PrescribedAtmosphere, JRA55PrescribedLand, RepeatYearJRA55, MultiYearJRA55
export JRA55PrescribedAtmosphere, JRA55PrescribedLand, RepeatYearJRA55, MultiYearJRA55, JRA55FieldTimeSeries

using Oceananigans
using Oceananigans.Units
Expand Down
Loading
Loading