diff --git a/windspharm/iris.py b/windspharm/iris.py index 03670c0..9cb1cb5 100644 --- a/windspharm/iris.py +++ b/windspharm/iris.py @@ -756,6 +756,62 @@ def truncate(self, field, truncation=None): field.transpose(reorder) return field + def getuv(self, vorticity, divergence): + """Compute vector winds from vorticity and divergence fields. + + **Argument:** + + *vorticity* + A scalar field of vorticity. It must be a `~iris.cube.Cube` + with the same latitude and longitude dimensions as the + vector wind components that initialized the `VectorWind` + instance. + + *divergence* + A scalar field of divergence. It must be a `~iris.cube.Cube` + with the same latitude and longitude dimensions as the + vector wind components that initialized the `VectorWind` + instance. + + **Returns:** + + *u*, *v* + Zonal and meridional wind components respectively. Their types will + match input types to passed to `VectorWind` instance. + """ + + def clean_array(field): + if type(field) is not Cube: + raise TypeError('scalar field must be an iris cube') + name = field.name() + lat, lat_dim = _dim_coord_and_dim(field, 'latitude') + lon, lon_dim = _dim_coord_and_dim(field, 'longitude') + if (lat.points[0] < lat.points[1]): + # need to reverse latitude dimension + field = reverse(field, lat_dim) + lat, lat_dim = _dim_coord_and_dim(field, 'latitude') + apiorder, reorder = get_apiorder(field.ndim, lat_dim, lon_dim) + field = field.copy() + field.transpose(apiorder) + ishape = field.shape + coords = field.dim_coords + field = to3d(field.data) + return field + + vortic = clean_array(vorticity) + diverg = clean_array(divergence) + + ugrd, vgrd = self._api.getuv(vortic, diverg) + ugrd = self._metadata(ugrd, + units='m s**-1', + standard_name='eastward_wind', + long_name='eastward_component_of_wind') + vgrd = self._metadata(vgrd, + units='m s**-1', + standard_name='northward_wind', + long_name='northward_component_of_wind') + return ugrd, vgrd + def _dim_coord_and_dim(cube, coord): """ diff --git a/windspharm/standard.py b/windspharm/standard.py index 8215fee..3d652f9 100644 --- a/windspharm/standard.py +++ b/windspharm/standard.py @@ -657,3 +657,33 @@ def truncate(self, field, truncation=None): raise ValueError('field is not compatible') fieldtrunc = self.s.spectogrd(fieldspec) return fieldtrunc + + def getuv(self, vorticity, divergence): + """Compute vector winds from vorticity and divergence fields. + + **Argument:** + + *vorticity* + A scalar field of vorticity. Its shape must be either (nlat, nlon) or + (nlat, nlon, nfields) where nlat and nlon are the same + as those for the vector wind components that initialized the + `VectorWind` instance. + + *divergence* + A scalar field of divergence. Its shape must be either (nlat, nlon) or + (nlat, nlon, nfields) where nlat and nlon are the same + as those for the vector wind components that initialized the + `VectorWind` instance. + + **Optional argument:** + + **Returns:** + + *u*, *v* + Zonal and meridional wind components respectively. Their types will + match input types to passed to `VectorWind` instance. + """ + vrspec = self.s.grdtospec(vorticity) + dvspec = self.s.grdtospec(divergence) + ugrd, vgrid = self.s.getuv(vrspec, dvspec) + return ugrd, vgrid diff --git a/windspharm/tests/test_solution.py b/windspharm/tests/test_solution.py index 7bb5e5c..7fa7ba6 100644 --- a/windspharm/tests/test_solution.py +++ b/windspharm/tests/test_solution.py @@ -148,6 +148,16 @@ def test_truncate(self): vrt_trunc = self.vw.truncate(self.solution['vrt'], truncation=21) self.assert_error_is_zero(vrt_trunc, self.solution['vrt_trunc']) + def test_getuv(self): + from .utils import error + + div1 = self.vw.divergence() + vrt1 = self.vw.vorticity() + + ugrd, vgrd = self.vw.getuv(vrt1, div1) + # this isn't a perfect 1-1 function so some error is introduced + assert error(ugrd, self.solution['uwnd']) == pytest.approx(0., abs=1e-4) + assert error(vgrd, self.solution['vwnd']) == pytest.approx(0., abs=1e-4) # ---------------------------------------------------------------------------- # Tests for the standard interface diff --git a/windspharm/xarray.py b/windspharm/xarray.py index b8cb69e..ed39fc2 100644 --- a/windspharm/xarray.py +++ b/windspharm/xarray.py @@ -747,6 +747,61 @@ def truncate(self, field, truncation=None): field = field.transpose(*reorder) return field + def getuv(self, vorticity, divergence): + """Compute vector winds from vorticity and divergence fields. + + **Argument:** + + *vorticity* + A scalar field of vorticity. Its shape must be either (nlat, nlon) or + (nlat, nlon, nfields) where nlat and nlon are the same + as those for the vector wind components that initialized the + `VectorWind` instance. + + *divergence* + A scalar field of divergence. Its shape must be either (nlat, nlon) or + (nlat, nlon, nfields) where nlat and nlon are the same + as those for the vector wind components that initialized the + `VectorWind` instance. + + **Returns:** + + *u*, *v* + Zonal and meridional wind components respectively. Their types will + match input types to passed to `VectorWind` instance. + """ + def clean_array(field): + if not isinstance(field, xr.DataArray): + raise TypeError('scalar field must be an xarray.DataArray') + name = field.name + lat, lat_dim = _find_latitude_coordinate(field) + lon, lon_dim = _find_longitude_coordinate(field) + if (lat.values[0] < lat.values[1]): + # need to reverse latitude dimension + field = _reverse(field, lat_dim) + lat, lat_dim = _find_latitude_coordinate(field) + apiorder, _ = get_apiorder(field.ndim, lat_dim, lon_dim) + apiorder = [field.dims[i] for i in apiorder] + reorder = field.dims + field = field.copy().transpose(*apiorder) + ishape = field.shape + coords = [field.coords[n] for n in field.dims] + field = to3d(field.values) + return field + + vortic = clean_array(vorticity) + diverg = clean_array(divergence) + + ugrd, vgrd = self._api.getuv(vortic, diverg) + ugrd = self._metadata(ugrd, 'u', + units='m s**-1', + standard_name='eastward_wind', + long_name='eastward_component_of_wind') + vgrd = self._metadata(vgrd, 'v', + units='m s**-1', + standard_name='northward_wind', + long_name='northward_component_of_wind') + return ugrd, vgrd def _reverse(array, dim): """Reverse an `xarray.DataArray` along a given dimension."""