diff --git a/Project.toml b/Project.toml index 9faa65131..b13f98844 100644 --- a/Project.toml +++ b/Project.toml @@ -8,8 +8,10 @@ authors = ["NumericalEarth contributors"] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" CFTime = "179af706-886a-5703-950a-314cd64e0468" ClimaSeaIce = "6ba0ff68-24e6-4315-936c-2e99227c95a4" +CubedSphere = "7445602f-e544-4518-8976-18f8e8ae6cdb" DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" +Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" @@ -56,10 +58,12 @@ Breeze = "0.4" CDSAPI = "2.2.1" CFTime = "0.1, 0.2" ClimaSeaIce = "0.4.4, 0.5" +CubedSphere = "0.3.4" CondaPkg = "0.2.33" CopernicusMarine = "0.1.1" DataDeps = "0.7" Dates = "<0.0.1, 1" +Distances = "0.10" DocStringExtensions = "0.9" Downloads = "<0.0.1, 1" GPUArraysCore = "0.2.0" diff --git a/src/Bathymetry/Bathymetry.jl b/src/Bathymetry/Bathymetry.jl index 40eed858d..9801b75df 100644 --- a/src/Bathymetry/Bathymetry.jl +++ b/src/Bathymetry/Bathymetry.jl @@ -1,6 +1,6 @@ module Bathymetry -export regrid_bathymetry, ORCAGrid +export regrid_bathymetry, ORCAGrid, resample_grid using Downloads using ImageMorphology diff --git a/src/Bathymetry/orca_grid.jl b/src/Bathymetry/orca_grid.jl index e214a3fed..6896ad53b 100644 --- a/src/Bathymetry/orca_grid.jl +++ b/src/Bathymetry/orca_grid.jl @@ -1,12 +1,15 @@ using Oceananigans.BoundaryConditions: fill_halo_regions!, FPivotZipperBoundaryCondition, NoFluxBoundaryCondition, FieldBoundaryConditions -using Oceananigans.Fields: set! +using Oceananigans.Fields: set!, convert_to_0_360 using Oceananigans.Grids: RightFaceFolded, generate_coordinate using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom using Oceananigans.OrthogonalSphericalShellGrids: Tripolar, continue_south! +using CubedSphere.SphericalGeometry: lat_lon_to_cartesian, cartesian_to_lat_lon, + spherical_area_quadrilateral +using Distances: haversine -using ..DataWrangling: dataset_variable_name -using ..DataWrangling.ORCA: ORCA1, default_south_rows_to_remove +using ..DataWrangling: dataset_variable_name, default_download_directory +using ..DataWrangling.ORCA: ORCA1, ORCA12, default_south_rows_to_remove """ read_2d_nemo_variable(ds, name) @@ -16,16 +19,277 @@ dimension layouts: `(x, y)`, `(x, y, z)`, or `(x, y, z, t)`. """ function read_2d_nemo_variable(ds, name) var = ds[name] - nd = ndims(var) - if nd == 2 - return Array(var[:, :]) - elseif nd == 3 - return Array(var[:, :, 1]) + # NOTE: `var[:]` does linear indexing and flattens to 1D. + # We need the full shaped array here. + data = Array(var) + + if ndims(data) < 2 + throw(ArgumentError("Variable $name could not be reduced to 2D. Size after slicing: $(size(data))")) + end + + if ndims(data) > 2 + # Keep the two largest dimensions as horizontal (x, y), and pick + # the first index for all other dimensions (for example t=1, z=1). + # This handles layouts like (t, x, y), (x, y, t), (t, z, y, x), etc. + sizes = collect(size(data)) + keep = sort(sortperm(sizes; rev = true)[1:2]) + indices = ntuple(d -> (d in keep ? Colon() : 1), ndims(data)) + data = @view data[indices...] + end + + if ndims(data) != 2 + throw(ArgumentError("Variable $name could not be reduced to 2D. Size after slicing: $(size(data))")) + end + + return Array(data) +end + +has_all_variables(ds, names) = all(name -> name in keys(ds), names) + +function orient_xy(data, Nx, Ny; name = "variable") + sx, sy = size(data) + if (sx, sy) == (Nx, Ny) + return data + elseif (sx, sy) == (Ny, Nx) + return permutedims(data, (2, 1)) else - return Array(var[:, :, 1, 1]) + throw(ArgumentError("Cannot orient $name with size $(size(data)) to (Nx, Ny)=($Nx, $Ny).")) end end +function orient_orca_xy(data; name = "variable") + sx, sy = size(data) + if sx >= sy + return data + else + # ORCA global grids should have Nx > Ny. If not, we assume (y, x) layout. + return permutedims(data, (2, 1)) + end +end + +@inline wrap_longitude(λ) = convert_to_0_360(λ + 180) - 180 + +@inline function midpoint_longitude(λ₁, λ₂) + Δλ = λ₂ - λ₁ + Δλ = ifelse(Δλ > 180, Δλ - 360, Δλ) + Δλ = ifelse(Δλ < -180, Δλ + 360, Δλ) + return wrap_longitude(λ₁ + Δλ / 2) +end + +@inline function spherical_midpoint(λ₁, φ₁, λ₂, φ₂) + x₁, y₁, z₁ = lat_lon_to_cartesian(φ₁, λ₁; radius = 1, check_latitude_bounds = false) + x₂, y₂, z₂ = lat_lon_to_cartesian(φ₂, λ₂; radius = 1, check_latitude_bounds = false) + x = x₁ + x₂ + y = y₁ + y₂ + z = z₁ + z₂ + n = sqrt(x^2 + y^2 + z^2) + + if n < 1e-12 + λm = midpoint_longitude(λ₁, λ₂) + φm = (φ₁ + φ₂) / 2 + return λm, φm + end + + x /= n + y /= n + z /= n + + φm, λm = cartesian_to_lat_lon(x, y, z) + λm = wrap_longitude(λm) + return λm, φm +end + +@inline function spherical_quadrilateral_area_unit(λ₁, φ₁, λ₂, φ₂, λ₃, φ₃, λ₄, φ₄) + a = lat_lon_to_cartesian(φ₁, λ₁; radius = 1, check_latitude_bounds = false) + b = lat_lon_to_cartesian(φ₂, λ₂; radius = 1, check_latitude_bounds = false) + c = lat_lon_to_cartesian(φ₃, λ₃; radius = 1, check_latitude_bounds = false) + d = lat_lon_to_cartesian(φ₄, λ₄; radius = 1, check_latitude_bounds = false) + return spherical_area_quadrilateral(a, b, c, d; radius = 1) +end + +function reconstruct_orca_staggered_mesh_from_t_f_points(λCC, φCC, λFF, φFF; radius) + size(λCC) == size(φCC) || throw(ArgumentError("glamt and gphit size mismatch: $(size(λCC)) vs $(size(φCC)).")) + size(λFF) == size(φFF) || throw(ArgumentError("glamf and gphif size mismatch: $(size(λFF)) vs $(size(φFF)).")) + size(λCC) == size(λFF) || throw(ArgumentError("T-point and F-point grids must have matching size, got $(size(λCC)) and $(size(λFF)).")) + + Nx, Ny = size(λCC) + AFT = promote_type(eltype(λCC), eltype(φCC), eltype(λFF), eltype(φFF), typeof(radius)) + + λFC = similar(λCC, AFT) + φFC = similar(φCC, AFT) + λCF = similar(λCC, AFT) + φCF = similar(φCC, AFT) + + @inbounds for j in 1:Ny, i in 1:Nx + iE = mod1(i + 1, Nx) + λm, φm = spherical_midpoint(λCC[i, j], φCC[i, j], λCC[iE, j], φCC[iE, j]) + λFC[i, j] = λm + φFC[i, j] = φm + end + + if Ny > 1 + @inbounds for j in 1:Ny-1, i in 1:Nx + λm, φm = spherical_midpoint(λCC[i, j], φCC[i, j], λCC[i, j+1], φCC[i, j+1]) + λCF[i, j] = λm + φCF[i, j] = φm + end + end + + # Northern V-points are inferred from the northern F-point edge. + @inbounds for i in 1:Nx + iE = mod1(i + 1, Nx) + λm, φm = spherical_midpoint(λFF[i, Ny], φFF[i, Ny], λFF[iE, Ny], φFF[iE, Ny]) + λCF[i, Ny] = λm + φCF[i, Ny] = φm + end + + e1u = similar(λCC, AFT) + e2u = similar(λCC, AFT) + e1v = similar(λCC, AFT) + e2v = similar(λCC, AFT) + e1f = similar(λCC, AFT) + e2f = similar(λCC, AFT) + e1t = similar(λCC, AFT) + e2t = similar(λCC, AFT) + + @inbounds for j in 1:Ny, i in 1:Nx + iE = mod1(i + 1, Nx) + e1u[i, j] = haversine((λCC[i, j], φCC[i, j]), (λCC[iE, j], φCC[iE, j]), radius) + e1v[i, j] = haversine((λFF[i, j], φFF[i, j]), (λFF[iE, j], φFF[iE, j]), radius) + e1f[i, j] = haversine((λCF[i, j], φCF[i, j]), (λCF[iE, j], φCF[iE, j]), radius) + end + + @inbounds for j in 1:Ny-1, i in 1:Nx + e2u[i, j] = haversine((λFC[i, j], φFC[i, j]), (λFC[i, j+1], φFC[i, j+1]), radius) + e2v[i, j] = haversine((λCC[i, j], φCC[i, j]), (λCC[i, j+1], φCC[i, j+1]), radius) + e2f[i, j] = haversine((λFC[i, j], φFC[i, j]), (λFC[i, j+1], φFC[i, j+1]), radius) + end + + if Ny > 1 + @inbounds for i in 1:Nx + e2u[i, Ny] = e2u[i, Ny-1] + e2v[i, Ny] = e2v[i, Ny-1] + e2f[i, Ny] = e2f[i, Ny-1] + end + else + @inbounds for i in 1:Nx + e2u[i, 1] = e1u[i, 1] + e2v[i, 1] = e1v[i, 1] + e2f[i, 1] = e1f[i, 1] + end + end + + @inbounds for j in 1:Ny, i in 1:Nx + iW = mod1(i - 1, Nx) + e1t[i, j] = haversine((λFC[iW, j], φFC[iW, j]), (λFC[i, j], φFC[i, j]), radius) + end + + if Ny > 1 + @inbounds for i in 1:Nx + e2t[i, 1] = e2v[i, 1] + for j in 2:Ny + e2t[i, j] = (e2v[i, j-1] + e2v[i, j]) / 2 + end + end + else + @inbounds for i in 1:Nx + e2t[i, 1] = e2v[i, 1] + end + end + + AzCC = similar(λCC, AFT) + AzFC = e1u .* e2u + AzCF = e1v .* e2v + AzFF = similar(λCC, AFT) + + if Ny > 1 + @inbounds for j in 1:Ny-1, i in 1:Nx + iE = mod1(i + 1, Nx) + A = spherical_quadrilateral_area_unit(λFF[i, j], φFF[i, j], + λFF[iE, j], φFF[iE, j], + λFF[iE, j+1], φFF[iE, j+1], + λFF[i, j+1], φFF[i, j+1]) + AzCC[i, j] = A * radius^2 + end + @inbounds for i in 1:Nx + AzCC[i, Ny] = AzCC[i, Ny-1] + end + + @inbounds for j in 2:Ny, i in 1:Nx + iW = mod1(i - 1, Nx) + A = spherical_quadrilateral_area_unit(λCC[iW, j-1], φCC[iW, j-1], + λCC[i, j-1], φCC[i, j-1], + λCC[i, j], φCC[i, j], + λCC[iW, j], φCC[iW, j]) + AzFF[i, j] = A * radius^2 + end + @inbounds for i in 1:Nx + AzFF[i, 1] = AzFF[i, 2] + end + else + AzCC .= e1t .* e2t + AzFF .= AzCC + end + + return (; λCC, λFC, λCF, λFF, φCC, φFC, φCF, φFF, + e1t, e1u, e1v, e1f, e2t, e2u, e2v, e2f, + AzCC, AzFC, AzCF, AzFF) +end + +""" + read_orca_staggered_mesh(ds) + +Read ORCA horizontal coordinates and metrics. + +Supports: +- full NEMO staggered mesh variables (`glamt/gphit/e1u/...`), and +- approximate reconstruction from T/F coordinates only (`glamt/gphit/glamf/gphif`) + using Tripolar-style spherical metric assumptions. +""" +function read_orca_staggered_mesh(ds; radius = Oceananigans.defaults.planet_radius) + full_stagger_vars = ("glamt", "glamu", "glamv", "glamf", + "gphit", "gphiu", "gphiv", "gphif", + "e1t", "e1u", "e1v", "e1f", + "e2t", "e2u", "e2v", "e2f") + + if has_all_variables(ds, full_stagger_vars) + read_2d = read_2d_nemo_variable + λCC = orient_orca_xy(read_2d(ds, "glamt"); name = "glamt") + Nx, Ny = size(λCC) + + orient(data, name) = orient_xy(data, Nx, Ny; name) + + λFC, λCF, λFF = orient(read_2d(ds, "glamu"), "glamu"), orient(read_2d(ds, "glamv"), "glamv"), orient(read_2d(ds, "glamf"), "glamf") + φCC, φFC, φCF, φFF = orient(read_2d(ds, "gphit"), "gphit"), orient(read_2d(ds, "gphiu"), "gphiu"), orient(read_2d(ds, "gphiv"), "gphiv"), orient(read_2d(ds, "gphif"), "gphif") + e1t, e1u, e1v, e1f = orient(read_2d(ds, "e1t"), "e1t"), orient(read_2d(ds, "e1u"), "e1u"), orient(read_2d(ds, "e1v"), "e1v"), orient(read_2d(ds, "e1f"), "e1f") + e2t, e2u, e2v, e2f = orient(read_2d(ds, "e2t"), "e2t"), orient(read_2d(ds, "e2u"), "e2u"), orient(read_2d(ds, "e2v"), "e2v"), orient(read_2d(ds, "e2f"), "e2f") + + if "e1e2t" in keys(ds) + AzCC, AzFC = orient(read_2d(ds, "e1e2t"), "e1e2t"), orient(read_2d(ds, "e1e2u"), "e1e2u") + AzCF, AzFF = orient(read_2d(ds, "e1e2v"), "e1e2v"), orient(read_2d(ds, "e1e2f"), "e1e2f") + else + AzCC, AzFC, AzCF, AzFF = e1t .* e2t, e1u .* e2u, e1v .* e2v, e1f .* e2f + end + + return (; λCC, λFC, λCF, λFF, φCC, φFC, φCF, φFF, + e1t, e1u, e1v, e1f, e2t, e2u, e2v, e2f, + AzCC, AzFC, AzCF, AzFF) + end + + tf_vars = ("glamt", "gphit", "glamf", "gphif") + if has_all_variables(ds, tf_vars) + read_2d = read_2d_nemo_variable + λCC = orient_orca_xy(read_2d(ds, "glamt"); name = "glamt") + Nx, Ny = size(λCC) + λFF = orient_xy(read_2d(ds, "glamf"), Nx, Ny; name = "glamf") + φCC = orient_xy(read_2d(ds, "gphit"), Nx, Ny; name = "gphit") + φFF = orient_xy(read_2d(ds, "gphif"), Nx, Ny; name = "gphif") + return reconstruct_orca_staggered_mesh_from_t_f_points(λCC, φCC, λFF, φFF; radius) + end + + throw(ArgumentError("Unsupported ORCA mesh format. Missing either full staggered variables $(full_stagger_vars) or T/F variables $(tf_vars).")) +end + # Detect periodic overlap columns in NEMO data. # The eORCA grid has `n` trailing columns that are copies of the first `n` columns # (e.g., columns 361:362 repeat columns 1:2 for eORCA1). @@ -94,19 +358,21 @@ end radius = Oceananigans.defaults.planet_radius, with_bathymetry = true, active_cells_map = true, - south_rows_to_remove = default_south_rows_to_remove(dataset)) + south_rows_to_remove = default_south_rows_to_remove(dataset), + dir = default_download_directory(dataset)) Construct an `OrthogonalSphericalShellGrid` with `(Periodic, RightFaceFolded, Bounded)` topology using coordinate and metric data from a NEMO eORCA `mesh_mask` file. -The `dataset` keyword argument specifies which ORCA configuration to use (e.g., `ORCA1()`). +The `dataset` keyword argument specifies which ORCA configuration to use (e.g., `ORCA1() or ORCA12()`). The mesh mask and bathymetry files are downloaded automatically via the `DataWrangling.ORCA` metadata interface. The horizontal grid (including coordinates, scale factors, and areas) is loaded -directly from the `mesh_mask` NetCDF file, which contains data at all four staggered -locations (T, U, V, F points). The user provides the vertical discretization via the `z` -keyword argument. +directly from the `mesh_mask` NetCDF file. If all staggered NEMO fields are present +(`T`, `U`, `V`, `F` points), they are used directly. If only `T` and `F` +coordinates are available (`glamt/gphit/glamf/gphif`), staggered coordinates and +metrics are reconstructed approximately using Tripolar-style spherical assumptions. When `with_bathymetry = true` (the default), the bathymetry is also downloaded and the grid is returned as an `ImmersedBoundaryGrid` with a `GridFittedBottom`. @@ -120,7 +386,7 @@ Positional Arguments Keyword Arguments ================= -- `dataset`: The ORCA dataset to use. Default: `ORCA1()` (from Zenodo; ). +- `dataset`: The ORCA dataset to use. Default: `ORCA1()`. `ORCA12()` is also supported (ORCA1 data from Zenodo; ). - `halo`: Halo size tuple `(Hx, Hy, Hz)`. Default: `(4, 4, 4)`. - `z`: Vertical coordinate specification. Can be a 2-tuple `(z_bottom, z_top)`, an array of z-interfaces, or, e.g., an `ExponentialDiscretization`. Default: `(-6000, 0)`. @@ -133,6 +399,8 @@ Keyword Arguments - `south_rows_to_remove`: Number of southern rows to remove from the eORCA grid. The "extended" eORCA grid contains degenerate padding rows near Antarctica that are entirely land. Removing them reduces memory usage and computation. +- `dir`: Directory to store and look up ORCA files (`mesh_mask` and bathymetry). + Defaults to the dataset scratch cache via `default_download_directory(dataset)`. """ function ORCAGrid(arch = CPU(), FT::DataType = Float64; dataset = ORCA1(), @@ -142,41 +410,31 @@ function ORCAGrid(arch = CPU(), FT::DataType = Float64; radius = Oceananigans.defaults.planet_radius, with_bathymetry = true, active_cells_map = true, - south_rows_to_remove = default_south_rows_to_remove(dataset)) + south_rows_to_remove = default_south_rows_to_remove(dataset), + dir = default_download_directory(dataset)) # Download mesh_mask via the metadata interface - mesh_meta = Metadatum(:mesh_mask; dataset) + mesh_meta = Metadatum(:mesh_mask; dataset, dir) mesh_mask_path = download_dataset(mesh_meta) ds = Dataset(mesh_mask_path) - - # Read 2D arrays at all four NEMO stagger locations: - # T → (Center, Center), U → (Face, Center), - # V → (Center, Face), F → (Face, Face) - read_2d = read_2d_nemo_variable - - λCC, λFC, λCF, λFF = read_2d(ds, "glamt"), read_2d(ds, "glamu"), read_2d(ds, "glamv"), read_2d(ds, "glamf") - φCC, φFC, φCF, φFF = read_2d(ds, "gphit"), read_2d(ds, "gphiu"), read_2d(ds, "gphiv"), read_2d(ds, "gphif") - e1t, e1u, e1v, e1f = read_2d(ds, "e1t"), read_2d(ds, "e1u"), read_2d(ds, "e1v"), read_2d(ds, "e1f") - e2t, e2u, e2v, e2f = read_2d(ds, "e2t"), read_2d(ds, "e2u"), read_2d(ds, "e2v"), read_2d(ds, "e2f") - - # Areas: read pre-computed if available, otherwise compute from scale factors - if "e1e2t" in keys(ds) - AzCC, AzFC = read_2d(ds, "e1e2t"), read_2d(ds, "e1e2u") - AzCF, AzFF = read_2d(ds, "e1e2v"), read_2d(ds, "e1e2f") - else - AzCC, AzFC, AzCF, AzFF = e1t .* e2t, e1u .* e2u, e1v .* e2v, e1f .* e2f - end - + mesh = read_orca_staggered_mesh(ds; radius) close(ds) + λCC, λFC, λCF, λFF = mesh.λCC, mesh.λFC, mesh.λCF, mesh.λFF + φCC, φFC, φCF, φFF = mesh.φCC, mesh.φFC, mesh.φCF, mesh.φFF + e1t, e1u, e1v, e1f = mesh.e1t, mesh.e1u, mesh.e1v, mesh.e1f + e2t, e2u, e2v, e2f = mesh.e2t, mesh.e2u, mesh.e2v, mesh.e2f + AzCC, AzFC, AzCF, AzFF = mesh.AzCC, mesh.AzFC, mesh.AzCF, mesh.AzFF + # Extract tripolar pole parameters from F-point coordinates last_row_φ = φFF[:, end] pole_idx = argmax(last_row_φ) - north_poles_latitude = Float64(last_row_φ[pole_idx]) - first_pole_longitude = Float64(λFF[pole_idx, end]) + north_poles_latitude = min(Float64(last_row_φ[pole_idx]), 89.999) + first_pole_longitude = Float64(λFF[pole_idx, end]) Nx, Ny = size(λCC) + Ny_full = Ny # Detect periodic overlap columns (e.g., eORCA1 has 2 trailing overlap columns) overlap = periodic_overlap_index(λCC) @@ -258,22 +516,224 @@ function ORCAGrid(arch = CPU(), FT::DataType = Float64; with_bathymetry || return underlying_grid # Load bathymetry - bathy_meta = Metadatum(:bottom_height; dataset) + bathy_meta = Metadatum(:bottom_height; dataset, dir) bathymetry_path = download_dataset(bathy_meta) bathy_ds = Dataset(bathymetry_path) - bathy_data = Array(bathy_ds[dataset_variable_name(bathy_meta)][:, :]) + bathy_name = dataset_variable_name(bathy_meta) + bathy_data = read_2d_nemo_variable(bathy_ds, bathy_name) close(bathy_ds) + bathy_data = orient_xy(bathy_data, Nx, Ny_full; name = string(bathy_name)) + if jr > 0 bathy_data = chop(bathy_data) end # NEMO bathymetry is positive depth; convert to negative bottom height. # Land (bathymetry == 0) gets mapped to +100 so GridFittedBottom masks it. - bottom_height = convert.(FT, bathy_data) - bottom_height .= ifelse.(bottom_height .> 0, .-bottom_height, FT(100)) + bottom_height = FT.(coalesce.(bathy_data, FT(0))) + bottom_height .= ifelse.(isfinite.(bottom_height) .& (bottom_height .> 0), .-bottom_height, FT(100)) bottom_height = on_architecture(arch, bottom_height) return ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height); active_cells_map) end + +function _extract_t_and_f_coordinates(grid) + Nx, Ny, _ = size(grid) + λCC = Array(on_architecture(CPU(), grid.λᶜᶜᵃ[1:Nx, 1:Ny])) + φCC = Array(on_architecture(CPU(), grid.φᶜᶜᵃ[1:Nx, 1:Ny])) + λFF = Array(on_architecture(CPU(), grid.λᶠᶠᵃ[1:Nx, 1:Ny])) + φFF = Array(on_architecture(CPU(), grid.φᶠᶠᵃ[1:Nx, 1:Ny])) + return λCC, φCC, λFF, φFF +end + +@inline function _resampled_point(λsrc, φsrc, x, y) + Nx, Ny = size(λsrc) + + i0_raw = floor(Int, x) + i0 = mod1(i0_raw, Nx) + i1 = mod1(i0_raw + 1, Nx) + fx = x - floor(x) + + j0 = clamp(floor(Int, y), 1, Ny) + j1 = clamp(j0 + 1, 1, Ny) + fy = clamp(y - floor(y), 0.0, 1.0) + + w00 = (1 - fx) * (1 - fy) + w10 = fx * (1 - fy) + w01 = (1 - fx) * fy + w11 = fx * fy + + x00, y00, z00 = lat_lon_to_cartesian(φsrc[i0, j0], λsrc[i0, j0]; radius = 1, check_latitude_bounds = false) + x10, y10, z10 = lat_lon_to_cartesian(φsrc[i1, j0], λsrc[i1, j0]; radius = 1, check_latitude_bounds = false) + x01, y01, z01 = lat_lon_to_cartesian(φsrc[i0, j1], λsrc[i0, j1]; radius = 1, check_latitude_bounds = false) + x11, y11, z11 = lat_lon_to_cartesian(φsrc[i1, j1], λsrc[i1, j1]; radius = 1, check_latitude_bounds = false) + + xc = w00 * x00 + w10 * x10 + w01 * x01 + w11 * x11 + yc = w00 * y00 + w10 * y10 + w01 * y01 + w11 * y11 + zc = w00 * z00 + w10 * z10 + w01 * z01 + w11 * z11 + + nrm = sqrt(xc^2 + yc^2 + zc^2) + if nrm < 1e-12 + return λsrc[i0, j0], φsrc[i0, j0] + end + + xc /= nrm + yc /= nrm + zc /= nrm + + φr, λr = cartesian_to_lat_lon(xc, yc, zc) + return wrap_longitude(λr), φr +end + +function _resample_lon_lat(λsrc, φsrc, Nx_target, Ny_target) + Nx_src, Ny_src = size(λsrc) + T = promote_type(eltype(λsrc), eltype(φsrc), Float64) + + λdst = Matrix{T}(undef, Nx_target, Ny_target) + φdst = Matrix{T}(undef, Nx_target, Ny_target) + + xq = 1 .+ (0:Nx_target-1) .* (Nx_src / Nx_target) + yq = Ny_target == 1 ? [1.0] : collect(range(1, Ny_src, length = Ny_target)) + + @inbounds for j in 1:Ny_target, i in 1:Nx_target + λp, φp = _resampled_point(λsrc, φsrc, xq[i], yq[j]) + λdst[i, j] = λp + φdst[i, j] = φp + end + + return λdst, φdst +end + +function _build_grid_from_resampled_orca_mesh(mesh; + arch, + FT::DataType, + Nz, + z, + halo, + radius) + λCC, λFC, λCF, λFF = mesh.λCC, mesh.λFC, mesh.λCF, mesh.λFF + φCC, φFC, φCF, φFF = mesh.φCC, mesh.φFC, mesh.φCF, mesh.φFF + e1t, e1u, e1v, e1f = mesh.e1t, mesh.e1u, mesh.e1v, mesh.e1f + e2t, e2u, e2v, e2f = mesh.e2t, mesh.e2u, mesh.e2v, mesh.e2f + AzCC, AzFC, AzCF, AzFF = mesh.AzCC, mesh.AzFC, mesh.AzCF, mesh.AzFF + + Nx, Ny = size(λCC) + + overlap = periodic_overlap_index(λCC) + southernmost_latitude = Float64(minimum(φCC)) + + last_row_φ = φFF[:, end] + pole_idx = argmax(last_row_φ) + north_poles_latitude = min(Float64(last_row_φ[pole_idx]), 89.999) + first_pole_longitude = Float64(λFF[pole_idx, end]) + + Hx, Hy, Hz = halo + topo = (Periodic, RightFaceFolded, Bounded) + Lz, z_coord = generate_coordinate(FT, topo, (Nx, Ny, Nz), halo, z, :z, 3, CPU()) + + helper_grid = RectilinearGrid(; size = (Nx, Ny), halo = (Hx, Hy), + x = (0, 1), y = (0, 1), + topology = (Periodic, RightFaceFolded, Flat)) + + bcs = FieldBoundaryConditions(north = FPivotZipperBoundaryCondition(), + south = NoFluxBoundaryCondition(), + west = Oceananigans.PeriodicBoundaryCondition(), + east = Oceananigans.PeriodicBoundaryCondition(), + top = nothing, + bottom = nothing) + + λᶜᶜᵃ, λᶠᶜᵃ, λᶜᶠᵃ, λᶠᶠᵃ = halo_fill_stagger(λCC, λFC, λCF, λFF, helper_grid, bcs, overlap) + φᶜᶜᵃ, φᶠᶜᵃ, φᶜᶠᵃ, φᶠᶠᵃ = halo_fill_stagger(φCC, φFC, φCF, φFF, helper_grid, bcs, overlap) + Δxᶜᶜᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ = halo_fill_stagger(e1t, e1u, e1v, e1f, helper_grid, bcs, overlap) + Δyᶜᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶠᵃ = halo_fill_stagger(e2t, e2u, e2v, e2f, helper_grid, bcs, overlap) + Azᶜᶜᵃ, Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ = halo_fill_stagger(AzCC, AzFC, AzCF, AzFF, helper_grid, bcs, overlap) + + ref_grid = LatitudeLongitudeGrid(; size = (Nx, Ny, Nz), + latitude = (southernmost_latitude, 90), + longitude = (-180, 180), + halo, z = (0, 1), radius) + + for (field, ref_name) in ((Δxᶜᶜᵃ, :Δxᶜᶜᵃ), (Δxᶠᶜᵃ, :Δxᶠᶜᵃ), (Δxᶜᶠᵃ, :Δxᶜᶠᵃ), (Δxᶠᶠᵃ, :Δxᶠᶠᵃ), + (Δyᶜᶜᵃ, :Δyᶜᶠᵃ), (Δyᶠᶜᵃ, :Δyᶠᶜᵃ), (Δyᶜᶠᵃ, :Δyᶜᶠᵃ), (Δyᶠᶠᵃ, :Δyᶠᶜᵃ), + (Azᶜᶜᵃ, :Azᶜᶜᵃ), (Azᶠᶜᵃ, :Azᶠᶜᵃ), (Azᶜᶠᵃ, :Azᶜᶠᵃ), (Azᶠᶠᵃ, :Azᶠᶠᵃ)) + continue_south!(field, getproperty(ref_grid, ref_name)) + end + + to_arch(data) = on_architecture(arch, map(FT, data)) + + return OrthogonalSphericalShellGrid{Periodic, RightFaceFolded, Bounded}( + arch, + Nx, Ny, Nz, + Hx, Hy, Hz, + convert(FT, Lz), + to_arch(λᶜᶜᵃ), to_arch(λᶠᶜᵃ), to_arch(λᶜᶠᵃ), to_arch(λᶠᶠᵃ), + to_arch(φᶜᶜᵃ), to_arch(φᶠᶜᵃ), to_arch(φᶜᶠᵃ), to_arch(φᶠᶠᵃ), + on_architecture(arch, z_coord), + to_arch(Δxᶜᶜᵃ), to_arch(Δxᶠᶜᵃ), to_arch(Δxᶜᶠᵃ), to_arch(Δxᶠᶠᵃ), + to_arch(Δyᶜᶜᵃ), to_arch(Δyᶠᶜᵃ), to_arch(Δyᶜᶠᵃ), to_arch(Δyᶠᶠᵃ), + to_arch(Azᶜᶜᵃ), to_arch(Azᶠᶜᵃ), to_arch(Azᶜᶠᵃ), to_arch(Azᶠᶠᵃ), + convert(FT, radius), + Tripolar(north_poles_latitude, first_pole_longitude, southernmost_latitude)) +end + +""" + resample_grid(source_grid; size, halo = nothing, Nz_target = nothing, z = nothing) + +Resample ORCA horizontal grid coordinates/metrics to a new `(Nx, Ny)` horizontal size. + +Only horizontal resampling is supported. Vertical resampling is intentionally unsupported. +If `source_grid` is an `ImmersedBoundaryGrid`, this function warns and returns a non-immersed +resampled underlying grid (bathymetry is not resampled). +""" +function resample_grid(source_grid; size, halo = nothing, Nz_target = nothing, z = nothing) + if !(source_grid isa ImmersedBoundaryGrid) && + !(source_grid isa Oceananigans.Grids.OrthogonalSphericalShellGrid) + throw(ArgumentError("`resample_grid` expects an ORCA underlying grid or ImmersedBoundaryGrid returned by ORCAGrid. Received $(typeof(source_grid)).")) + end + + base_grid = source_grid + if source_grid isa ImmersedBoundaryGrid + @warn "Bathymetry resampling is not supported. The resampled grid will not be an ImmersedBoundaryGrid even if the input grid has attached bathymetry. Recompute bathymetry with regrid_bathymetry(...) after resampling." + base_grid = source_grid.underlying_grid + end + + if !isnothing(Nz_target) + throw(ArgumentError("Vertical resampling is not supported in `resample_grid`. Received `Nz_target = $(Nz_target)`. Only horizontal resampling via `size=(Nx, Ny)` is supported.")) + end + + if !isnothing(z) + throw(ArgumentError("Vertical coordinate overrides are not supported in `resample_grid`. Received `z = $(z)`. The source grid vertical coordinate is always preserved.")) + end + + Nx_target, Ny_target = if size isa Tuple + length(size) == 2 || throw(ArgumentError("`size` must have exactly two entries `(Nx, Ny)` for horizontal resampling. Received size = $(size).")) + Int(size[1]), Int(size[2]) + elseif size isa AbstractVector + length(size) == 2 || throw(ArgumentError("`size` must have exactly two entries `[Nx, Ny]` for horizontal resampling. Received size = $(size).")) + Int(size[1]), Int(size[2]) + else + throw(ArgumentError("`size` must be a Tuple or Vector with two entries `(Nx, Ny)`. Received $(typeof(size)).")) + end + + Nx_target > 0 || throw(ArgumentError("Target Nx must be > 0, got Nx = $(Nx_target).")) + Ny_target > 0 || throw(ArgumentError("Target Ny must be > 0, got Ny = $(Ny_target).")) + + arch = Oceananigans.Architectures.architecture(base_grid) + FT = eltype(base_grid) + Nz = base_grid.Nz + radius = hasproperty(base_grid, :radius) ? base_grid.radius : Oceananigans.defaults.planet_radius + halo_out = isnothing(halo) ? (base_grid.Hx, base_grid.Hy, base_grid.Hz) : Tuple(halo) + + length(halo_out) == 3 || throw(ArgumentError("`halo` must be a 3-tuple `(Hx, Hy, Hz)`. Received halo = $(halo_out).")) + + λCC, φCC, λFF, φFF = _extract_t_and_f_coordinates(base_grid) + λCCr, φCCr = _resample_lon_lat(λCC, φCC, Nx_target, Ny_target) + λFFr, φFFr = _resample_lon_lat(λFF, φFF, Nx_target, Ny_target) + + mesh = reconstruct_orca_staggered_mesh_from_t_f_points(λCCr, φCCr, λFFr, φFFr; radius) + z_interfaces = Array(on_architecture(CPU(), base_grid.z.cᵃᵃᶠ[1:Nz+1])) + + return _build_grid_from_resampled_orca_mesh(mesh; arch, FT, Nz, z = z_interfaces, halo = halo_out, radius) +end diff --git a/src/Bathymetry/regrid_bathymetry.jl b/src/Bathymetry/regrid_bathymetry.jl index 430cad426..3911f4a40 100644 --- a/src/Bathymetry/regrid_bathymetry.jl +++ b/src/Bathymetry/regrid_bathymetry.jl @@ -352,6 +352,13 @@ function interpolate_bathymetry_in_passes(native_z, target_grid; Nλt, Nφt = Nt = size(target_grid) Nλn, Nφn = Nn = size(native_z) + source_cells = Nλn * Nφn + target_cells = Nλt * Nφt + + if target_cells > source_cells + upsampling_ratio = target_cells / source_cells + @warn "Interpolating bathymetry to higher horizontal resolution (refining, not coarsening)." source_size = (Nλn, Nφn) target_size = (Nλt, Nφt) source_cells target_cells upsampling_ratio + end # Interpolate in passes latitude = y_domain(native_z.grid) diff --git a/src/DataWrangling/ORCA/ORCA.jl b/src/DataWrangling/ORCA/ORCA.jl index 90db416e0..b6a8a89bb 100644 --- a/src/DataWrangling/ORCA/ORCA.jl +++ b/src/DataWrangling/ORCA/ORCA.jl @@ -1,6 +1,6 @@ module ORCA -export ORCA1 +export ORCA1, ORCA12 using Downloads using Oceananigans @@ -28,28 +28,52 @@ function __init__() end struct ORCA1 end +struct ORCA12 end default_download_directory(::ORCA1) = download_ORCA_cache +default_download_directory(::ORCA12) = download_ORCA_cache + reversed_vertical_axis(::ORCA1) = false +reversed_vertical_axis(::ORCA12) = false + longitude_interfaces(::ORCA1) = (-180, 180) +longitude_interfaces(::ORCA12) = (-180, 180) + latitude_interfaces(::ORCA1) = (-80, 90) +latitude_interfaces(::ORCA12) = (-80, 90) all_dates(::ORCA1, args...) = nothing +all_dates(::ORCA12, args...) = nothing + first_date(::ORCA1, args...) = nothing +first_date(::ORCA12, args...) = nothing + last_date(::ORCA1, args...) = nothing +last_date(::ORCA12, args...) = nothing const ORCA1Metadatum = Metadatum{<:ORCA1} +const ORCA12Metadatum = Metadatum{<:ORCA12} ORCA1_variable_names = Dict( :bottom_height => "Bathymetry", :mesh_mask => "glamt", ) +ORCA12_variable_names = Dict( + :bottom_height => "Bathymetry", + :mesh_mask => "e1t", +) + dataset_variable_name(data::ORCA1Metadatum) = ORCA1_variable_names[data.name] +dataset_variable_name(data::ORCA12Metadatum) = ORCA12_variable_names[data.name] # Zenodo record 4436658: eORCA1 mesh_mask and bathymetry -const ORCA1_mesh_mask_url = "https://zenodo.org/records/4436658/files/eORCA1.2_mesh_mask.nc" -const ORCA1_bathymetry_url = "https://zenodo.org/records/4436658/files/eORCA_R1_bathy_meter_v2.2.nc" +const ORCA1_mesh_mask_url = "https://zenodo.org/records/4436658/files/eORCA1.2_mesh_mask.nc" +const ORCA1_bathymetry_url = "https://zenodo.org/records/4436658/files/eORCA_R1_bathy_meter_v2.2.nc" + +# Google Drive records for GLORYS12 (ORCA12) mesh mask and bathymetry +const ORCA12_mesh_mask_url = "https://zenodo.org/records/15495870/files/grid_mask_eORCA12-GO6.nc" +const ORCA12_bathymetry_url = "https://zenodo.org/records/15495870/files/bathy_eORCA12_noclosea_from_GEBCO2021_FillZero_S21TT_CloseaCopy.nc" function metadata_url(metadatum::ORCA1Metadatum) if metadatum.name == :mesh_mask @@ -61,6 +85,16 @@ function metadata_url(metadatum::ORCA1Metadatum) end end +function metadata_url(metadatum::ORCA12Metadatum) + if metadatum.name == :mesh_mask + return ORCA12_mesh_mask_url + elseif metadatum.name == :bottom_height + return ORCA12_bathymetry_url + else + error("Unknown ORCA12 variable: $(metadatum.name)") + end +end + function metadata_filename(::ORCA1, name, date, bounding_box) if name == :mesh_mask return "eORCA1.2_mesh_mask.nc" @@ -71,14 +105,26 @@ function metadata_filename(::ORCA1, name, date, bounding_box) end end +function metadata_filename(::ORCA12, name, date, bounding_box) + if name == :mesh_mask + return "grid_mask_eORCA12-GO6.nc" + elseif name == :bottom_height + return "bathy_eORCA12_noclosea_from_GEBCO2021_FillZero_S21TT_CloseaCopy.nc" + else + error("Unknown ORCA12 variable: $name") + end +end + z_interfaces(::ORCA1Metadatum) = nothing +z_interfaces(::ORCA12Metadatum) = nothing -function download_dataset(metadatum::ORCA1Metadatum) +function download_dataset(metadatum::Union{ORCA1Metadatum, ORCA12Metadatum}) fileurl = metadata_url(metadatum) filepath = metadata_path(metadatum) @root if !isfile(filepath) - @info "Downloading ORCA1 data: $(metadatum.name) to $(metadatum.dir)..." + dataset_name = nameof(typeof(metadatum.dataset)) + @info "Downloading $(dataset_name) data: $(metadatum.name) to $(metadatum.dir)..." Downloads.download(fileurl, filepath; progress=download_progress) end @@ -86,5 +132,7 @@ function download_dataset(metadatum::ORCA1Metadatum) end default_south_rows_to_remove(::ORCA1) = 35 +# Conservative default: keep all rows unless user requests trimming. +default_south_rows_to_remove(::ORCA12) = 0 end # module diff --git a/src/NumericalEarth.jl b/src/NumericalEarth.jl index 9dd134f19..99831004a 100644 --- a/src/NumericalEarth.jl +++ b/src/NumericalEarth.jl @@ -42,7 +42,7 @@ export EN4Monthly, WOAClimatology, WOAAnnual, WOAMonthly, GLORYSDaily, GLORYSMonthly, GLORYSStatic, - ORCA1, + ORCA1, ORCA12, RepeatYearJRA55, MultiYearJRA55, first_date, last_date, @@ -52,6 +52,7 @@ export DatasetRestoring, ocean_simulation, ORCAGrid, + resample_grid, sea_ice_simulation, atmosphere_simulation, sea_ice_dynamics, diff --git a/test/test_bathymetry.jl b/test/test_bathymetry.jl index 73faf02c2..9a92aba11 100644 --- a/test/test_bathymetry.jl +++ b/test/test_bathymetry.jl @@ -1,11 +1,13 @@ include("runtests_setup.jl") using JLD2 +using Logging using NumericalEarth.Bathymetry: remove_minor_basins!, BathymetryRegridding, cache_filename, load_bathymetry_cache, - save_bathymetry_cache + save_bathymetry_cache, + interpolate_bathymetry_in_passes using NumericalEarth.DataWrangling.ETOPO using Statistics @@ -72,6 +74,38 @@ using Statistics end end +@testset "Bathymetry upsampling warning" begin + source_grid = LatitudeLongitudeGrid(CPU(); + size = (20, 20, 1), + longitude = (0, 100), + latitude = (0, 50), + z = (-6000, 0)) + + native_z = Field{Center, Center, Nothing}(source_grid) + set!(native_z, -1000.0) + fill_halo_regions!(native_z) + + finer_grid = LatitudeLongitudeGrid(CPU(); + size = (40, 40, 1), + longitude = (0, 100), + latitude = (0, 50), + z = (-6000, 0)) + + @test_logs (:warn, r"higher horizontal resolution") match_mode=:any begin + interpolate_bathymetry_in_passes(native_z, finer_grid; passes = 2) + end + + coarser_grid = LatitudeLongitudeGrid(CPU(); + size = (10, 10, 1), + longitude = (0, 100), + latitude = (0, 50), + z = (-6000, 0)) + + @test_logs min_level=Logging.Warn begin + interpolate_bathymetry_in_passes(native_z, coarser_grid; passes = 2) + end +end + @testset "BathymetryRegridding configuration" begin @info "Testing BathymetryRegridding configuration..." diff --git a/test/test_orca_grid.jl b/test/test_orca_grid.jl index 823e7f125..470d23c37 100644 --- a/test/test_orca_grid.jl +++ b/test/test_orca_grid.jl @@ -5,7 +5,7 @@ using NumericalEarth.DataWrangling: download_dataset, metadata_path using NumericalEarth.DataWrangling.ORCA: default_south_rows_to_remove using Oceananigans using Oceananigans.OrthogonalSphericalShellGrids: TripolarGrid -using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid +using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom using NCDatasets using Statistics using Test @@ -20,6 +20,20 @@ using Test @test mesh_meta.dataset isa ORCA1 end + +@testset "ORCA12 Metadatum construction" begin + bathy_meta = Metadatum(:bottom_height; dataset=ORCA12()) + @test bathy_meta.name == :bottom_height + @test bathy_meta.dataset isa ORCA12 + + mesh_meta = Metadatum(:mesh_mask; dataset=ORCA12()) + @test mesh_meta.name == :mesh_mask + @test mesh_meta.dataset isa ORCA12 + + @test default_south_rows_to_remove(ORCA12()) == 0 + @test occursin("eORCA12", metadata_path(mesh_meta)) + @test occursin("eORCA12", metadata_path(bathy_meta)) +end @testset "ORCAGrid with ORCA1 dataset on $(arch)" for arch in test_architectures south_rows_to_remove = 43 grid = ORCAGrid(arch; dataset=ORCA1(), Nz=5, z=(-5000, 0), halo=(4, 4, 4), south_rows_to_remove) @@ -144,3 +158,31 @@ end @test size(bathy) == (362, 332) @test maximum(bathy) > 5000 # Deep ocean end + +@testset "resample_grid horizontal-only behavior" begin + source = ORCAGrid(CPU(); dataset=ORCA1(), Nz=5, z=(-5000, 0), halo=(4, 4, 4), + with_bathymetry=false) + + Nx_target, Ny_target = 180, 160 + resampled = resample_grid(source; size=(Nx_target, Ny_target)) + + @test resampled isa Oceananigans.Grids.OrthogonalSphericalShellGrid + @test !(resampled isa ImmersedBoundaryGrid) + @test resampled.Nx == Nx_target + @test resampled.Ny == Ny_target + @test resampled.Nz == source.Nz + + @test_throws ArgumentError resample_grid(source; size=(Nx_target, Ny_target), Nz_target = 3) + @test_throws ArgumentError resample_grid(source; size=(Nx_target, Ny_target), z = (-3000, 0)) + + immersed_source = ImmersedBoundaryGrid(source, GridFittedBottom(fill(-100.0, source.Nx, source.Ny)); + active_cells_map = false) + + @test_logs (:warn, r"Bathymetry resampling is not supported") match_mode=:any begin + rs = resample_grid(immersed_source; size=(Nx_target, Ny_target)) + @test !(rs isa ImmersedBoundaryGrid) + @test rs.Nx == Nx_target + @test rs.Ny == Ny_target + @test rs.Nz == source.Nz + end +end