Add horizontal mesh operator for reconstructing edge-normal vector fields at cell centers - #480
Conversation
|
For testing, you'll need to download the CTest meshes with reconstruction variables in them. Developer quick start guide has been updated, but for reference the new |
TestingI've run CTest on frontier (both For reference using E3SM-Project/polaris#649 I'm also able to approximately reproduce the norms used the CTest in a standalone python script. Here's a sample script: Standalone python scriptfrom pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from importlib import resources
from ruamel.yaml import YAML
from polaris.constants import get_constant
from polaris.mesh.reconstruct import (
fix_out_of_bounds_indices,
build_reconstruction_weights,
cartesian_to_local_geographic,
tangential_reconstruction,
)
from polaris.mpas.area import area_for_field
from polaris.tasks.ocean.sphere_transport.resources.flow_types import (
normal_velocity_from_zonal_meridional,
)
class MeshConverter:
def __init__(self):
self.dim_map, self.var_map = self._load_mpaso_to_omega_map()
def _load_mpaso_to_omega_map(self):
text = (
resources.files('polaris.ocean.model')
.joinpath('mpaso_to_omega.yaml')
.read_text()
)
nested = YAML(typ='rt').load(text)
return nested['dimensions'], nested['variables']
def omega_to_mpas(self, ds: xr.Dataset) -> xr.Dataset:
# omega -> mpas
rename = {v: k for k, v in self.dim_map.items() if v in ds.dims}
rename.update({v: k for k, v in self.var_map.items() if v in ds})
return ds.rename(rename)
def mpas_to_omega(self, ds: xr.Dataset) -> xr.Dataset:
# mpas -> omega
rename = {k: v for k, v in self.dim_map.items() if k in ds.dims}
rename.update({k: v for k, v in self.var_map.items() if k in ds})
def detect_format(self, ds: xr.Dataset) -> Literal['mpaso', 'omega']:
dims = ds.dims
items = self.dim_map.items()
is_omega = all(v in dims for k, v in items if k in dims or v in dims)
is_mpas = all(k in dims for k, v in items if k in dims or v in dims)
if is_omega and is_mpas:
raise ValueError(
'Invalid input: dataset contains both MPASO and Omega '
'dimensions names.'
)
return 'omega' if is_omega else 'mpaso'
class TestSetupSphere:
def __init__(self, ds):
self.ds = ds
self.radius = get_constant("mean_radius")
def analytic_velocity(self, lon, lat):
clat = np.cos(lat)
clon = np.cos(lon)
slat = np.sin(lat)
slon = np.sin(lon)
u_x = -self.radius * slon**2 * clat**3
u_y = -4 * self.radius * slon * clon * clat**3 * slat
return u_x, u_y
def compute_error(ds, numerical, exact):
area = np.sqrt(area_for_field(ds, numerical))
diff = (numerical - exact).values.flatten()
frob_norm = np.linalg.norm(diff * area, ord=2)
inf_norm = np.linalg.norm(diff, ord=np.inf)
frob_norm /= np.linalg.norm(exact.values.flatten() * area, ord=2)
inf_norm /= np.linalg.norm(exact.values.flatten(), ord=np.inf)
return frob_norm, inf_norm
if __name__ == "__main__":
#ds = xr.open_dataset("OmegaMesh.nc")
ds = xr.open_dataset("OmegaSphereMesh.nc")
converter = MeshConverter()
ds = converter.omega_to_mpas(ds)
test_setup = TestSetupSphere(ds)
u_edge, v_edge = test_setup.analytic_velocity(ds.lonEdge, ds.latEdge)
normal_velocity = normal_velocity_from_zonal_meridional(
ds, u_edge, v_edge
)
stencil = ds.reconstructStencilCell
weights = ds.reconstructWeightsCell
u_x, u_y, u_z = tangential_reconstruction(
ds, normal_velocity, stencil=stencil, weights=weights
)
u_recon, v_recon, _ = cartesian_to_local_geographic(ds, u_x, u_y, u_z)
u_exact, v_exact = test_setup.analytic_velocity(ds.lonCell, ds.latCell)
magnitude_recon = np.sqrt(u_recon**2 + v_recon**2)
magnitude_exact = np.sqrt(u_exact**2 + v_exact**2)
frob_norm, inf_norm = compute_error(ds, magnitude_recon, magnitude_exact)
print(f"L2 = {frob_norm}, Linf = {inf_norm}") |
There was a problem hiding this comment.
Pull request overview
Adds a new horizontal mesh operator, VectorReconstructOnCell, to reconstruct an
edge-normal vector field onto cell centers using precomputed least-squares
weights/stencils stored in spherical mesh files. This integrates the
reconstruction inputs into Decomp/HorzMesh, updates tests to validate the
new mesh data and operator behavior, and documents the new mesh requirements and
operator availability.
Changes:
- Extend
Decomp/HorzMeshto read/redistribute reconstruction stencil arrays
and read/halo-exchangeReconstructWeightsCell(spherical meshes only). - Add
VectorReconstructOnCelloperator for reconstructing zonal/meridional
components at cell centers. - Update tests and docs to cover/describe the new reconstruction capability and
required mesh variables.
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| components/omega/test/ocn/OceanTestCommon.h | Add vecMagnitude helper for operator test comparisons. |
| components/omega/test/ocn/HorzOperatorsTest.cpp | Add spherical-only vector reconstruction test and expected error baselines; rename tangent recon test. |
| components/omega/test/base/DecompTest.cpp | Add consistency checks for reconstruction stencil arrays on spherical meshes. |
| components/omega/src/ocn/HorzOperators.h | Declare VectorReconstructOnCell functor/operator. |
| components/omega/src/ocn/HorzOperators.cpp | Implement VectorReconstructOnCell constructor and mesh gating. |
| components/omega/src/ocn/HorzMesh.h | Add mesh storage for reconstruction stencil/weights arrays. |
| components/omega/src/ocn/HorzMesh.cpp | Read/define ReconstructWeightsCell, create R3 dimension, and halo-exchange weights slices for spherical meshes. |
| components/omega/src/base/Decomp.h | Add MaxEdges2, OnSphere, and reconstruction stencil members + redistribution method declaration. |
| components/omega/src/base/Decomp.cpp | Read MaxEdges2 and OnSphere, read reconstruction stencil arrays for spherical meshes, redistribute and translate to local indices. |
| components/omega/doc/userGuide/HorzOperators.md | Document VectorReconstructOnCell availability. |
| components/omega/doc/userGuide/HorzMesh.md | Document ReconstructWeightsCell mesh variable (spherical-only). |
| components/omega/doc/userGuide/Decomp.md | Document required stencil arrays for spherical meshes. |
| components/omega/doc/devGuide/QuickStart.md | Update example mesh download links to newer files. |
| components/omega/doc/devGuide/HorzOperators.md | Document VectorReconstructOnCell availability. |
| components/omega/doc/devGuide/Decomp.md | Document new Decomp members for reconstruction stencil support. |
There was a problem hiding this comment.
I spent some time looking at this PR and I'm not sure that reading the reconstruction stencil inside Decomp is necessary. It seems to me that, with a few changes, it could be read from the HorzMeshIn stream, just like the weights. The only new functionality needed is a way to map global to local indices in ReconstructStencilCell. The global to local maps are currently created on-the-fly in Decomp. We could store them as private members and provide public member functions for translating arrays that store global indices. This way, the Decomp class stays relatively small, there is no unnecessary MPI communication for rearranging arrays, and all the code for vector reconstruction input is in one place.
What do you think @andrewdnolan ? If needed, I can help with implementing this approach.
|
I like @mwarusz's idea. I'm going to hold of on reviewing until the dust settles on that. |
|
I have a slight preference for keeping all index-space stuff in Decomp with Mesh holding the physical mesh info. And it saves carrying around the extra stored global-to-local stuff. But I'm ok with whatever you all want to do. |
|
I don't have much of an opinion. It is worth noting that the vector reconstruction at cell centers is needed for passing Omega's velocity and ssh gradient to the coupler. So without this we are limited in the coupling progress we can make. Given that, I think I'd lean to what already implemented and maybe we can come back to @mwarusz suggestion in a few weeks? |
mwarusz
left a comment
mwarusz
left a comment
There was a problem hiding this comment.
I am fine with changing where the read is done later. I ran CTests on aurora with oneapi-ifx and everything passed. I found one confusing comment and some names that don't conform to our naming convention. Otherwise this looks good. Approving.
3c77a0d to
0ef148d
Compare
philipwjones
left a comment
philipwjones
left a comment
There was a problem hiding this comment.
I just had one renaming suggestion that will impact a lot of lines. Otherwise, I approve based on code inspection and testing by others.
| @@ -86,6 +87,12 @@ described in the mesh specification above. In particular, it contains | |||
| at a vertex | |||
| - NEdgesOnCell(NCellsSize): the number of actual edges on each cell | |||
| - NEdgesOnEdge(NEdgesSize): the number of actual edges on each edge | |||
| - NCellReconstructEdges(NCellsSize): number of edges in the vector | |||
There was a problem hiding this comment.
For consistency with other names, I would start this name with NEdges, so something like NEdgesRecon (Note that recon is used elsewhere for reconstruction so think it's ok to shorten too)
There was a problem hiding this comment.
Yes! Thanks for catching that. Following the naming convention, it should be NEdgesReconstructCell (keeping the full name).
I'm on board with NEdgesReconOnCell, having the OnCell also keep with syntax. While having the OnCell might seem a little verbose, I'm inclined to keep it because the weights can also be generated to reconstruct at vertices (see Peixoto and Barros (2014)). I'm not sure if Omega will ever really need that, but nonetheless causes less confusion about how/what the variable is used for.
| @@ -37,6 +37,7 @@ Currently, the following operators are implemented: | |||
| - `GradientOnEdge` | |||
| - `CurlOnVertex` | |||
| - `TangentialReconOnEdge` | |||
| - `VectorReconstructOnCell` | |||
There was a problem hiding this comment.
Like TangetianRecon, I think you can use VectorRecon here to shorten a bit
| ///< the OnSphere attribute - only spherical meshes | ||
| ///< currently have the reconstruction stencil below) | ||
|
|
||
| Array1DI4 NCellReconstructEdges; ///< Num of edges in reconstruction stencil |
There was a problem hiding this comment.
As in an comment elsewhere, recommend change to NEdgesRecon for this and Host array to be more consistent with other NEdges arrays.
| @@ -254,6 +254,16 @@ class HorzMesh { | |||
| Array1DReal MeshScalingDel4; /// Coef to biharmonic mixing terms | |||
| HostArray1DReal MeshScalingDel4H; /// Coef to biharmonic mixing terms | |||
|
|
|||
| // Vector reconstruction | |||
| Array1DI4 NCellReconstructEdges; /// Num of edges used in stencil | |||
| @@ -67,6 +67,7 @@ described in the mesh specification above. In particular, it contains | |||
| - NCellsHalo(i): the number of owned+halo cells for each halo layer | |||
| - Analogous size variables for Edges and Vertices | |||
| - MaxEdges: the max number of edges on a cell (and array size) | |||
| - MaxEdges2: the max number of edges on a edge | |||
There was a problem hiding this comment.
To respond to the CoPilot comment, I would just add in parentheses (2*MaxEdges) at the end of this line
xylar
left a comment
xylar
left a comment
There was a problem hiding this comment.
@andrewdnolan, I think this looks good. Thanks for putting in the work both here and in Polaris to make this happen!
|
@sbrus89, I'm assigning this to you. Please merge when the time is right, once @andrewdnolan has had a chance to address the remaining comments from @philipwjones. |
0ef148d to
48e088e
Compare
|
@philipwjones Thanks for catching the variable naming issue. In addition to you suggested shortening of
|
|
|
@andrewdnolan, I'm seeing errors (seg faults) on the |
Read, redistribute, and translate NCellReconstructEdges and ReconstructStencilCell in Decomp like other mesh connectivity arrays. Ensures downstream code consumes already-local, halo-consistentindices Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Currently only supporting single level fields and returning local geographic vector components. Further support will be added as needed.
All meshes now have `SurfacePressure` in them and spherical meshes now also have the reconstruction realted variable
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Maciej Waruszewski <mwarusz@igf.fuw.edu.pl>
48e088e to
085d04c
Compare
|
The issues with In my initial implementation of reconstruction weight generation (E3SM-Project/polaris#649) in did not wire the weight generation anywhere in the mesh steps. They were just used for visualization. In addition to addressing the name issue I'm wired in the weight generation in E3SM-Project/polaris#687, which should fix the fails Steven was seeing. Testing this PR will require using E3SM-Project/polaris#687. |
|
For convenience here are the cd test
wget -O OmegaMesh.nc https://web.lcrc.anl.gov/public/e3sm/polaris/ocean/omega_ctest/ocean.QU.240km.omega_vars.260807.nc
wget -O OmegaSphereMesh.nc https://web.lcrc.anl.gov/public/e3sm/polaris/ocean/omega_ctest/cosine_bell_icos480.omega_vars.260807.nc
wget -O OmegaPlanarMesh.nc https://web.lcrc.anl.gov/public/e3sm/polaris/ocean/omega_ctest/PlanarPeriodic48x48.omega_vars.260720.nc
cd .. |
|
Point to mesh with correct reconstruction variables names, which were changed as part of E3SM-Project#480
|
@andrewdnolan Have you made all the changes you intended to the code? Is this ready for both ctest-ing and polaris testing? I can do the aurora testing. |
TestingI've rebased this branch onto #461 and have been running thermodynamically active C-Cases. Here's an animation of the surface speed as provided to the coupler, which uses the vector reconstruction from this PR. SfcSpeed.mp4 |
|
@cbegeman Yes! All code updates have been made. With E3SM-Project/polaris#687 merged polaris should be ready for testing this. |
|
Local merge passes Polaris
|
|
Passes Polaris
|
Point to mesh with correct reconstruction variables names, which were changed as part of #480
Point to mesh with correct reconstruction variables names, which were changed as part of E3SM-Project#480
TestingAurora, gpu ctesting: All fails are related to #368. Two tests timed out for unknown reasons at 1500s. |
Adds
VectorReconstructOnCell: an operator for reconstructing an edge normal vector field at cell-centers. The operator uses pre-computed weights and stencil to perform the reconstuciton following section 2.5 from Peixoto and Barros (2014). This PR requires the reconstruction variables (NCellReconstructEdges,ReconstructStencilCell, andReconstructWeightsCell) to be present in spherical meshes. Support for reconstruction on planar meshes will come in a follow up PR.The "offline" computation of the weights and stencil is a deliberate design choice, because these variables only depend of the mesh. This way the reconstruction variables are present in mesh from it's definition and downstream analysis tool can do reconstruction without needing to run the forward model and write the weights to disk (i.e. what is currently required in MPAS-O).
Polaris PR for computing weights/stencils and adding them to an existing mesh: E3SM-Project/polaris#649.
The
VectorReconstructOnCellcurrently only has a single overload, which returns the local geographic (i.e. zonal and meridional) components of the vector field. We need this functionality to pass theSo_uandSo_vfields to the coupler, which expect values at cell centers. As more uses for vector reconstruction come up, additional overloads can be set up (e.g. returning the local Cartesian components of the vector field).Changes made:
Decomp: reads and redistributes the integer stencil arrays (NCellReconstructEdges,ReconstructStencilCell) andMaxEdges2, mirroring the existingCellsOnCell/EdgesOnCellhandling. Adds anOnSphereflag (read from the mesh file) to gate this since planar weights/stencils are not yet supported.HorzMesh: reads the real-valuedReconstructWeightsCellfield and halo-exchanges it component-by-component, since Halo's generic rank-3 exchange convention doesn't match this field's file-driven dimension order.HorzOperators: adds theVectorReconstructOnCellfunctor.DecompTestchecks stencil array consistency;HorzOperatorsTestadds
testVectorReconstruction, renamestestRecontotestTangentReconfor clarity, and adds avecMagnitudehelper toOceanTestCommon.devGuide/userGuideentries forDecomp,HorzMesh, andHorzOperators.Checklist:
Documentation:
Linting
Building
Testing
aurora, oneapi-ifx, mpich
chrysalis, oneapi-ifx, openmpi
frontier, craygnu, mpich
frontier, craygnu-mphipcc, mpich
pm-cpu, gnu, mpich
pm-gpu, gnugpu, mpich
Provide relevant details in a comment to the PR titled
Testingwith the following:have been run on and indicate that are all passing.
has passed, using the Polaris
e3sm_submodules/Omegabaseline-pfor both the baseline (Polarise3sm_submodules/Omega) and the PR buildPerformance related PRs: Please include a relevant PACE experiment link documenting performance before and after.
New tests:
Stealth Features