use TripolarGrid for Oceananigans sim #1409
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,11 +1,19 @@ | ||
| import Oceananigans as OC | ||
| import ClimaOcean as CO | ||
| import ClimaAtmos as CA | ||
| import ClimaCoupler: Checkpointer, FieldExchanger, FluxCalculator, Interfacer, Utilities | ||
| import ClimaComms | ||
| import ClimaCore as CC | ||
| import Thermodynamics as TD | ||
| import ClimaOcean.EN4: download_dataset | ||
| using KernelAbstractions: @kernel, @index, @inbounds | ||
| using XESMF # to load Oceananigans regridding extension | ||
|
|
||
| OceananigansXESMFExt = | ||
| Base.get_extension( | ||
| OC, | ||
| :OceananigansXESMFExt, | ||
| ).OceananigansXESMFExt; | ||
|
|
||
| """ | ||
| OceananigansSimulation{SIM, A, OPROP, REMAP} | ||
|
|
@@ -45,6 +53,9 @@ function OceananigansSimulation( | |
| output_dir, | ||
| comms_ctx = ClimaComms.context(), | ||
| ) | ||
| # using Dates | ||
| # start_date = Dates.DateTime(2008) | ||
| # stop_date = Dates.DateTime(2008, 1, 2) | ||
| arch = comms_ctx.device isa ClimaComms.CUDADevice ? OC.GPU() : OC.CPU() | ||
|
|
||
| # Retrieve EN4 data (monthly) | ||
|
|
@@ -55,33 +66,20 @@ function OceananigansSimulation( | |
| download_dataset(en4_temperature) | ||
| download_dataset(en4_salinity) | ||
|
|
||
| # Set up tripolar ocean grid (1 degree) | ||
| Nx = 360 | ||
| Ny = 180 | ||
| Nz = 40 | ||
| depth = 4000 # meters | ||
| z = OC.ExponentialDiscretization(Nz, -depth, 0; scale = 0.85 * depth) | ||
|
|
||
| underlying_grid = OC.TripolarGrid(arch; size = (Nx, Ny, Nz), halo = (7, 7, 4), z) | ||
| bottom_height = CO.regrid_bathymetry( | ||
| underlying_grid; | ||
| minimum_depth = 30, | ||
| interpolation_passes = 20, | ||
| major_basins = 1, | ||
| ) | ||
| grid = OC.ImmersedBoundaryGrid( | ||
| underlying_grid, | ||
| OC.GridFittedBottom(bottom_height); | ||
|
|
@@ -113,17 +111,99 @@ function OceananigansSimulation( | |
| # Set initial condition to EN4 state estimate at start_date | ||
| OC.set!(ocean.model, T = en4_temperature[1], S = en4_salinity[1]) | ||
|
|
||
| # Construct a remapper from the exchange grid to `Center, Center` fields | ||
|
|
||
| # Tripolar to cubed sphere (fails currently) | ||
| # long_oc = OC.λnodes(grid, OC.Center(), OC.Center(), OC.Center()) | ||
| # lat_oc = OC.φnodes(grid, OC.Center(), OC.Center(), OC.Center()) | ||
|
|
||
| T = OC.CenterField(grid) # field on tripolar grid | ||
| OC.set!(T, en4_temperature[1]) | ||
| coords_tripolar = OceananigansXESMFExt.xesmf_coordinates(T) | ||
|
|
||
| # Put everything on CPU | ||
| coords_tripolar = Dict(k => Array(v) for (k, v) in coords_tripolar) | ||
| # Create a 2D matrix containing each lat/long combination as a LatLongPoint | ||
| # Note this must be done on CPU since the CC.Remapper module is not GPU-compatible | ||
| # target_points_oc = Array(CC.Geometry.LatLongPoint.(lat_oc, long_oc)) | ||
|
|
||
| # Get the latitude and longitude of each node on the boundary space | ||
| boundary_space = axes(area_fraction) | ||
| # CC.CommonSpaces.CubedSphereSpace( | ||
| # Float32; | ||
| # radius = 1.0, | ||
| # n_quad_points = 2, | ||
| # h_elem = 10, | ||
| # ) | ||
| cubedsphere_coords = CC.Fields.coordinate_field(boundary_space) | ||
| # Put everything on CPU | ||
| cubedsphere_lat_arr = Array(CC.Fields.field2array(cubedsphere_coords.lat)) | ||
| cubedsphere_long_arr = Array(CC.Fields.field2array(cubedsphere_coords.long)) | ||
|
|
||
| # lon2d varies along cols (dim 2) — x direction | ||
| cubedsphere_lon2d = repeat(cubedsphere_long_arr', length(cubedsphere_lat_arr), 1) # transpose lon to make it row vector | ||
| # lat2d varies along rows (dim 1) — y direction | ||
| cubedsphere_lat2d = repeat(cubedsphere_lat_arr, 1, length(cubedsphere_long_arr)) # column vector repeated across | ||
|
|
||
| coords_cubedsphere = Dict("lat" => cubedsphere_lat2d, "lon" => cubedsphere_lon2d) | ||
|
|
||
| regridder_tripolar_to_cubedsphere = | ||
| XESMF.Regridder(coords_cubedsphere, coords_tripolar; method = "bilinear") | ||
| regridder_cubedsphere_to_tripolar = | ||
| XESMF.Regridder(coords_tripolar, coords_cubedsphere; method = "bilinear") | ||
|
|
||
| # Tripolar to LatLon | ||
| T = OC.CenterField(grid) # field on tripolar grid | ||
| OC.set!(T, en4_temperature[1]) | ||
|
|
||
| grid_latlon = OC.LatitudeLongitudeGrid( | ||
| arch; | ||
| size = (Nx, Ny, Nz), | ||
| longitude = (0, 360), | ||
| latitude = (-81, 90), | ||
| z, | ||
| ) | ||
| field_latlon = OC.CenterField(grid_latlon) | ||
| remapper_tripolar_to_latlon = XESMF.Regridder(T, field_latlon; method = "conservative") | ||
| remapper_tripolar_to_latlon( | ||
| vec(OC.interior(field_latlon, :, :, Nz)), | ||
| vec(OC.interior(T, :, :, Nz)), | ||
| ) | ||
| OC.fill_halo_regions!(field_latlon) | ||
| # heatmap(field_latlon) # using GeoMakie, using CairoMakie | ||
|
|
||
| # LatLon to cubed sphere | ||
| boundary_space = CC.CommonSpaces.CubedSphereSpace( | ||
| Float32; | ||
| radius = 1.0, | ||
| n_quad_points = 2, | ||
| h_elem = 10, | ||
| ) | ||
| field_cubedsphere = Interfacer.remap(field_latlon, boundary_space) | ||
| # fieldheatmap(field_cubedsphere) # using ClimaCoreMakie | ||
|
|
||
| # Cubed sphere to LatLon | ||
| long_oc = OC.λnodes(grid_latlon, OC.Center(), OC.Center(), OC.Center()) | ||
| lat_oc = OC.φnodes(grid_latlon, OC.Center(), OC.Center(), OC.Center()) | ||
| long_oc = reshape(long_oc, length(long_oc), 1) | ||
| lat_oc = reshape(lat_oc, 1, length(lat_oc)) | ||
| target_points_oc = @. CC.Geometry.LatLongPoint(lat_oc, long_oc) | ||
|
|
||
| remapper_cubedsphere_to_latlon = CC.Remapping.Remapper(boundary_space, target_points_oc) | ||
| field_cubedsphere = CC.Fields.zeros(boundary_space) | ||
| CC.Remapping.interpolate!( | ||
| field_cubedsphere, | ||
| remapper_cubedsphere_to_latlon, | ||
| field_latlon, | ||
| ) | ||
|
|
||
|
|
||
| # Previous remapper for LatLon | ||
| if pkgversion(CC) >= v"0.14.34" | ||
| remapper_cc = CC.Remapping.Remapper(axes(area_fraction), target_points_cc) | ||
| else | ||
| remapper_cc = CC.Remapping.Remapper(axes(area_fraction), target_points_cc, nothing) | ||
| end | ||
|
|
||
| # Construct two 2D Center/Center fields to use as scratch space while remapping | ||
| scratch_cc1 = OC.Field{OC.Center, OC.Center, Nothing}(grid) | ||
|
|
@@ -145,23 +225,18 @@ function OceananigansSimulation( | |
| ocean_fresh_water_density = 999.8, | ||
| ) | ||
|
|
||
|
There was a problem hiding this comment. Tracked in this issue now instead of a TODO: #1478 |
||
| # Save all tracers and velocities to a JLD2 file at daily frequency | ||
| outputs = merge(ocean.model.tracers, ocean.model.velocities) | ||
| jld2_writer = OC.JLD2Writer( | ||
| ocean.model, | ||
| outputs; | ||
| schedule = OC.TimeInterval(86400), # Daily output | ||
| filename = joinpath(output_dir, "ocean_diagnostics"), | ||
| indices = (:, :, grid.Nz), | ||
| overwrite_existing = true, | ||
| array_type = Array{Float32}, | ||
| ) | ||
| ocean.output_writers[:diagnostics] = jld2_writer | ||
|
|
||
| sim = OceananigansSimulation(ocean, area_fraction, ocean_properties, remapping) | ||
| return sim | ||
|
|
@@ -209,14 +284,14 @@ Interfacer.step!(sim::OceananigansSimulation, t) = | |
|
|
||
| # We always want the surface, so we always set zero(pt.lat) for z | ||
| """ | ||
| to_node(pt::CCGeometry.LatLongPoint) | ||
|
|
||
| Transform `LatLongPoint` into a tuple (long, lat, 0), where the 0 is needed because we only | ||
| care about the surface. | ||
| """ | ||
| @inline to_node(pt::CC.Geometry.LatLongPoint) = pt.long, pt.lat, zero(pt.lat) | ||
| # This next one is needed if we have "LevelGrid" | ||
| @inline to_node(pt::CC.Geometry.LatLongZPoint) = pt.long, pt.lat, zero(pt.lat) | ||
|
|
||
| """ | ||
| map_interpolate(points, oc_field::OC.Field) | ||
|
|
||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.