diff --git a/conda_package/docs/vector.rst b/conda_package/docs/vector.rst index 11520613b..53e02d2f9 100644 --- a/conda_package/docs/vector.rst +++ b/conda_package/docs/vector.rst @@ -28,6 +28,14 @@ The command-line tool ``vector_reconstruct`` and the function reconstruct Cartesian (X, Y, Z), zonal and meridional components of an MPAS vector field at cell centers, given the field on edge normals. +On spherical meshes, the zonal and meridional components are found by rotating +the Cartesian components into the local geographic frame at each cell center, +using ``latCell`` and ``lonCell``. On planar meshes (indicated by the +``on_a_sphere`` attribute of the mesh being ``NO``), no rotation is performed: +the zonal and meridional components are the x and y components, respectively. +This matches the convention in MPAS' own ``mpas_reconstruct_*`` routines. A +mesh without the ``on_a_sphere`` attribute is assumed to be spherical. + This tool requires that the field ``coeffs_reconstruct`` has been saved to a NetCDF file. The simplest way to do this is to include the following stream in a forward run: diff --git a/conda_package/mpas_tools/vector/reconstruct.py b/conda_package/mpas_tools/vector/reconstruct.py index f3fcfaf2a..794b6826f 100755 --- a/conda_package/mpas_tools/vector/reconstruct.py +++ b/conda_package/mpas_tools/vector/reconstruct.py @@ -2,6 +2,12 @@ Extract Cartesian (X, Y, Z), zonal and meridional components of an MPAS vector field, given the field on edge normals. +On spherical meshes, the zonal and meridional components are found by +rotating the Cartesian components into the local geographic frame at each +cell center. On planar meshes (``on_a_sphere = 'NO'``), no rotation is +performed: the zonal and meridional components are the x and y components, +respectively, matching the convention in MPAS' own reconstruction routines. + This tool requires that the field 'coeffs_reconstruct' has been saved to a NetCDF file. The simplest way to do this is to include the following stream in a forward run: @@ -17,6 +23,7 @@ and run the model for one time step. """ + import argparse import sys from datetime import datetime @@ -28,13 +35,27 @@ from mpas_tools.io import write_netcdf -def reconstruct_variable(out_var_name, variable_on_edges, ds_mesh, - coeffs_reconstruct, ds_out, chunk_size=32768, - quiet=False): +def reconstruct_variable( + out_var_name, + variable_on_edges, + ds_mesh, + coeffs_reconstruct, + ds_out, + chunk_size=32768, + quiet=False, +): """ Extract Cartesian (X, Y, Z), zonal and meridional components of an MPAS vector field, given the field on edge normals. + On spherical meshes, the zonal and meridional components are found by + rotating the Cartesian components into the local geographic frame using + ``latCell`` and ``lonCell``. On planar meshes (indicated by the + ``on_a_sphere`` attribute of ``ds_mesh`` being ``'NO'``), the zonal and + meridional components are the x and y components, respectively, as in + MPAS' own reconstruction routines. A mesh without the ``on_a_sphere`` + attribute is assumed to be spherical. + Parameters ---------- out_var_name : str @@ -44,7 +65,9 @@ def reconstruct_variable(out_var_name, variable_on_edges, ds_mesh, The variable at edge normals ds_mesh : xarray.Dataset - A dataset with the mesh variables + A dataset with the mesh variables (``edgesOnCell``, along with + ``latCell`` and ``lonCell`` if the mesh is spherical) and the + ``on_a_sphere`` attribute coeffs_reconstruct : xarray.DataArray A data array with the reconstruction coefficients @@ -94,8 +117,11 @@ def reconstruct_variable(out_var_name, variable_on_edges, ds_mesh, if not quiet: print('Computing Cartesian components:') for index, component in enumerate(['X', 'Y', 'Z']): - var = (coeffs_reconstruct.isel(R3=index)*variable).sum( - dim='maxEdges').transpose(*dims) + var = ( + (coeffs_reconstruct.isel(R3=index) * variable) + .sum(dim='maxEdges') + .transpose(*dims) + ) out_name = f'{out_var_name}{component}' if quiet: var.compute() @@ -106,21 +132,33 @@ def reconstruct_variable(out_var_name, variable_on_edges, ds_mesh, ds_out[out_name] = var var_cart.append(var) - lat_cell = ds_mesh.latCell - lon_cell = ds_mesh.lonCell - lat_cell.load() - lon_cell.load() - - clat = np.cos(lat_cell) - slat = np.sin(lat_cell) - clon = np.cos(lon_cell) - slon = np.sin(lon_cell) - if not quiet: print('Computing zonal and meridional components:') + if _on_a_sphere(ds_mesh): + lat_cell = ds_mesh.latCell + lon_cell = ds_mesh.lonCell + lat_cell.load() + lon_cell.load() + + clat = np.cos(lat_cell) + slat = np.sin(lat_cell) + clon = np.cos(lon_cell) + slon = np.sin(lon_cell) + + zonal = -var_cart[0] * slon + var_cart[1] * clon + merid = ( + -(var_cart[0] * clon + var_cart[1] * slon) * slat + + var_cart[2] * clat + ) + else: + # On a planar mesh, there is no rotation to perform: the x and y + # axes are the "zonal" and "meridional" directions, matching the + # convention in MPAS' own mpas_reconstruct_* routines. + zonal = var_cart[0] + merid = var_cart[1] + out_name = f'{out_var_name}Zonal' - zonal = -var_cart[0] * slon + var_cart[1] * clon if quiet: zonal.compute() else: @@ -130,8 +168,6 @@ def reconstruct_variable(out_var_name, variable_on_edges, ds_mesh, ds_out[out_name] = zonal out_name = f'{out_var_name}Meridional' - merid = (-(var_cart[0] * clon + var_cart[1] * slon) * slat + - var_cart[2] * clat) if quiet: merid.compute() else: @@ -141,34 +177,75 @@ def reconstruct_variable(out_var_name, variable_on_edges, ds_mesh, ds_out[out_name] = merid +def _on_a_sphere(ds_mesh): + """ + Whether ``ds_mesh`` is a spherical mesh, based on its ``on_a_sphere`` + attribute. A mesh without the attribute is assumed to be spherical. + """ + if 'on_a_sphere' not in ds_mesh.attrs: + return True + return str(ds_mesh.attrs['on_a_sphere']).strip().upper() != 'NO' + + def main(): # client = Client(n_workers=1, threads_per_worker=4, memory_limit='10GB') parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawTextHelpFormatter) - parser.add_argument("-m", "--mesh_filename", dest="mesh_filename", - type=str, required=False, - help="An MPAS file with mesh data (edgesOnCell, etc.) " - "if not from in_filename") - parser.add_argument("-w", "--weights_filename", dest="weights_filename", - type=str, required=False, - help="An MPAS file with coeffs_reconstruct if not " - "from in_filename") - parser.add_argument("-i", "--in_filename", dest="in_filename", type=str, - required=True, - help="An MPAS file with one or more fields on edges " - "to be reconstructed at cell centers. Used for " - "mesh data and/or weights if a separate files " - "are not provided.") - parser.add_argument("-v", "--variables", nargs='+', dest="variables", - type=str, required=True, - help="variables on edges to reconstruct") - parser.add_argument("--out_variables", nargs='+', dest="out_variables", - type=str, required=False, - help="prefixes for output variable names") - parser.add_argument("-o", "--out_filename", dest="out_filename", type=str, - required=True, - help="An output MPAS file with the reconstructed " - "X, Y, Z, zonal and meridional fields") + description=__doc__, formatter_class=argparse.RawTextHelpFormatter + ) + parser.add_argument( + '-m', + '--mesh_filename', + dest='mesh_filename', + type=str, + required=False, + help='An MPAS file with mesh data (edgesOnCell, etc.) ' + 'if not from in_filename', + ) + parser.add_argument( + '-w', + '--weights_filename', + dest='weights_filename', + type=str, + required=False, + help='An MPAS file with coeffs_reconstruct if not from in_filename', + ) + parser.add_argument( + '-i', + '--in_filename', + dest='in_filename', + type=str, + required=True, + help='An MPAS file with one or more fields on edges ' + 'to be reconstructed at cell centers. Used for ' + 'mesh data and/or weights if a separate files ' + 'are not provided.', + ) + parser.add_argument( + '-v', + '--variables', + nargs='+', + dest='variables', + type=str, + required=True, + help='variables on edges to reconstruct', + ) + parser.add_argument( + '--out_variables', + nargs='+', + dest='out_variables', + type=str, + required=False, + help='prefixes for output variable names', + ) + parser.add_argument( + '-o', + '--out_filename', + dest='out_filename', + type=str, + required=True, + help='An output MPAS file with the reconstructed ' + 'X, Y, Z, zonal and meridional fields', + ) args = parser.parse_args() if args.mesh_filename: @@ -192,9 +269,16 @@ def main(): coeffs_reconstruct = ds_weights.coeffs_reconstruct ds_out = xr.Dataset() - for in_var_name, out_var_name in zip(args.variables, out_variables): - reconstruct_variable(out_var_name, ds_in[in_var_name], ds_mesh, - coeffs_reconstruct, ds_out) + for in_var_name, out_var_name in zip( + args.variables, out_variables, strict=False + ): + reconstruct_variable( + out_var_name, + ds_in[in_var_name], + ds_mesh, + coeffs_reconstruct, + ds_out, + ) for attr_name in ds_in.attrs: ds_out.attrs[attr_name] = ds_in.attrs[attr_name] diff --git a/conda_package/tests/test_vector_reconstruct.py b/conda_package/tests/test_vector_reconstruct.py new file mode 100644 index 000000000..b1cce4493 --- /dev/null +++ b/conda_package/tests/test_vector_reconstruct.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python + +import numpy as np +import xarray as xr + +from mpas_tools.vector.reconstruct import reconstruct_variable + + +def test_reconstruct_planar(): + """ + On a planar mesh, the zonal and meridional components are the x and y + components, with no rotation applied. + """ + ds_mesh, coeffs, var_on_edges = _get_mesh(on_a_sphere='NO') + + ds_out = xr.Dataset() + reconstruct_variable( + 'velocity', var_on_edges, ds_mesh, coeffs, ds_out, quiet=True + ) + + assert np.allclose(ds_out.velocityZonal, ds_out.velocityX) + assert np.allclose(ds_out.velocityMeridional, ds_out.velocityY) + + +def test_reconstruct_planar_without_lat_lon(): + """ + A planar mesh need not have ``latCell`` and ``lonCell`` at all. + """ + ds_mesh, coeffs, var_on_edges = _get_mesh(on_a_sphere='NO') + ds_mesh = ds_mesh.drop_vars(['latCell', 'lonCell']) + + ds_out = xr.Dataset() + reconstruct_variable( + 'velocity', var_on_edges, ds_mesh, coeffs, ds_out, quiet=True + ) + + assert np.allclose(ds_out.velocityZonal, ds_out.velocityX) + assert np.allclose(ds_out.velocityMeridional, ds_out.velocityY) + + +def test_reconstruct_spherical(): + """ + On a spherical mesh, the Cartesian components are rotated into the local + geographic frame at each cell center. + """ + ds_mesh, coeffs, var_on_edges = _get_mesh(on_a_sphere='YES') + + ds_out = xr.Dataset() + reconstruct_variable( + 'velocity', var_on_edges, ds_mesh, coeffs, ds_out, quiet=True + ) + + lat = ds_mesh.latCell + lon = ds_mesh.lonCell + u_x = ds_out.velocityX + u_y = ds_out.velocityY + u_z = ds_out.velocityZ + zonal = -u_x * np.sin(lon) + u_y * np.cos(lon) + merid = -(u_x * np.cos(lon) + u_y * np.sin(lon)) * np.sin( + lat + ) + u_z * np.cos(lat) + + assert np.allclose(ds_out.velocityZonal, zonal) + assert np.allclose(ds_out.velocityMeridional, merid) + + +def test_reconstruct_missing_on_a_sphere(): + """ + A mesh without the ``on_a_sphere`` attribute is treated as spherical, as + it was before planar meshes were handled. + """ + ds_mesh, coeffs, var_on_edges = _get_mesh(on_a_sphere='YES') + ds_with_attr = ds_mesh + ds_without_attr = ds_mesh.copy() + del ds_without_attr.attrs['on_a_sphere'] + + ds_ref = xr.Dataset() + reconstruct_variable( + 'velocity', var_on_edges, ds_with_attr, coeffs, ds_ref, quiet=True + ) + ds_out = xr.Dataset() + reconstruct_variable( + 'velocity', var_on_edges, ds_without_attr, coeffs, ds_out, quiet=True + ) + + assert np.allclose(ds_out.velocityZonal, ds_ref.velocityZonal) + assert np.allclose(ds_out.velocityMeridional, ds_ref.velocityMeridional) + + +def test_reconstruct_on_a_sphere_padded(): + """ + MPAS files often pad the ``on_a_sphere`` attribute with spaces. + """ + ds_mesh, coeffs, var_on_edges = _get_mesh(on_a_sphere='NO ') + + ds_out = xr.Dataset() + reconstruct_variable( + 'velocity', var_on_edges, ds_mesh, coeffs, ds_out, quiet=True + ) + + assert np.allclose(ds_out.velocityZonal, ds_out.velocityX) + assert np.allclose(ds_out.velocityMeridional, ds_out.velocityY) + + +def _get_mesh(on_a_sphere): + """ + A tiny synthetic mesh, along with reconstruction coefficients and a + variable on edges, for testing ``reconstruct_variable()`` + """ + n_cells = 4 + max_edges = 6 + n_edges = n_cells * max_edges + rng = np.random.default_rng(seed=0) + + ds_mesh = xr.Dataset() + ds_mesh['edgesOnCell'] = ( + ('nCells', 'maxEdges'), + np.arange(1, n_edges + 1).reshape(n_cells, max_edges), + ) + if on_a_sphere.strip() == 'NO': + # as in meshes from ``mpas_tools.planar_hex`` + lat = np.zeros(n_cells) + lon = np.zeros(n_cells) + else: + lat = np.deg2rad(np.array([-80.0, -20.0, 15.0, 70.0])) + lon = np.deg2rad(np.array([10.0, 135.0, 200.0, 350.0])) + ds_mesh['latCell'] = (('nCells',), lat) + ds_mesh['lonCell'] = (('nCells',), lon) + ds_mesh.attrs['on_a_sphere'] = on_a_sphere + + coeffs = xr.DataArray( + rng.random((n_cells, max_edges, 3)), + dims=('nCells', 'maxEdges', 'R3'), + ) + var_on_edges = xr.DataArray(rng.random(n_edges), dims=('nEdges',)) + + return ds_mesh, coeffs, var_on_edges