Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
8 changes: 8 additions & 0 deletions conda_package/docs/vector.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
178 changes: 131 additions & 47 deletions conda_package/mpas_tools/vector/reconstruct.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -17,6 +23,7 @@

and run the model for one time step.
"""

import argparse
import sys
from datetime import datetime
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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()
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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]
Expand Down
Loading
Loading