def lon_lat_to_cartesian(lon, lat, radius=1): """ calculates lon, lat coordinates of a point on a sphere with radius radius """ # Unpack xarray object into plane arrays if hasattr(lon, 'data'): lon = lon.data if hasattr(lat, 'data'): lat = lat.data if lon.ndim != lat.ndim: raise ValueError('coordinate must share the same number of dimensions') if lon.ndim == 1: lon, lat = np.meshgrid(lon, lat) lon_r = xu.radians(lon) lat_r = xu.radians(lat) x = radius * xu.cos(lat_r) * xu.cos(lon_r) y = radius * xu.cos(lat_r) * xu.sin(lon_r) z = radius * xu.sin(lat_r) return x.flatten(), y.flatten(), z.flatten()
def angle2xyz(azi, zen): """Convert azimuth and zenith to cartesian.""" azi = xu.deg2rad(azi) zen = xu.deg2rad(zen) x = xu.sin(zen) * xu.sin(azi) y = xu.sin(zen) * xu.cos(azi) z = xu.cos(zen) return x, y, z
def lonlat2xyz(lon, lat): """Convert lon lat to cartesian.""" lat = xu.deg2rad(lat) lon = xu.deg2rad(lon) x = xu.cos(lat) * xu.cos(lon) y = xu.cos(lat) * xu.sin(lon) z = xu.sin(lat) return x, y, z
def distance_to_point(self, lat, lon): """ Use Haversine formula to estimate distances from all gridpoints to a given location (lat, lon) """ R = 6371. # Radius of earth in km lat = np.radians(lat) lon = np.radians(lon) dlat = lat - xu.radians(self['lat'].values) dlon = lon - xu.radians(self['lon'].values) a = xu.sin(dlat/2)**2 + xu.cos(lat) * xu.cos(xu.radians(self['lat'].values)) * \ xu.sin(dlon/2)**2 c = 2 * xu.arctan2(xu.sqrt(a), xu.sqrt(1.0 - a)) return R * c
def distance_to_point(self, lat, lon): """ Use Haversine formula to estimate distances from all gridpoints to a given location (lat, lon) """ R = 6371. # Radius of earth in km lat = np.radians(lat) lon = np.radians(lon) dlat = lat - xu.radians(self['lat'].values) dlon = lon - xu.radians(self['lon'].values) a = xu.sin(dlat/2)**2 + xu.cos(lat) * xu.cos(xu.radians(self['lat'].values)) * \ xu.sin(dlon/2)**2 c = 2 * xu.arctan2(xu.sqrt(a), xu.sqrt(1.0-a)) return R*c
def rotate_vector(u, v, angle): '''Rotate the vector (u, v) to (u_lon, v_lat). ** Input ** u: velocity in x direction (shape: nt,nj,ni) v: velocity in y direction (shape: nt,nj,ni) angle: the angle between x and longitude (in radians, shape: nj,ni) ** Returns ** u_lon: velocity in zonal direction (shape: nt,nj,ni) v_lat: velocity in meridional direction (shape: nt,nj,ni) ''' u_lon = u * cos(angle) - v * sin(angle) v_lat = u * sin(angle) + v * cos(angle) return u_lon, v_lat
def test_horizontal_gradient(self): tols = {'rtol': 1e-3, 'atol': 1e-3} gradvar = self.grd.horizontal_gradient(self.tvar1) gradx = self.grd.horizontal_gradient(self.xt) #dxdy = gradx.y_component s = self.scale actual_gx = gradvar.x_component.to_masked_array() actual_gy = gradvar.y_component.to_masked_array() expected_gx = (xu.cos(self.xu * s) * s).to_masked_array() expected_gy = ( -xu.sin(self.yv * s) * s # + xu.cos(self.xv * s) * s * dxdy ).to_masked_array() self.assertArray2dCloseInside(actual_gx / s, expected_gx / s, depth=2, **tols) #print dxdy[:,20].values #print self.grd._arrays['cell_y_size_at_t_location'][:,20].values #print self.grd._arrays['cell_y_size_at_v_location'][:,20].values #print_array_around(expected=expected_gy/s,actual=actual_gy/s) self.assertArray2dCloseInside(actual_gy / s, expected_gy / s, depth=2, **tols)
def setUp(self): self.x = np.arange(start=0, stop=101, step=10, dtype=float) self.y = np.arange(start=0, stop=101, step=10, dtype=float) self.xx, self.yy = np.meshgrid(self.x, self.y) self.xrx = xr.DataArray(self.xx, dims=['y', 'x']) self.xry = xr.DataArray(self.yy, dims=['y', 'x']) self.xrt = xu.sin(self.xrx / 30.) + xu.cos(self.xry / 30.) self.t = np.sin(self.xx / 30.) + np.cos(self.yy / 30.)
def setUp(self): self.x = np.arange(start=0, stop=101, step=10,dtype=float) self.y = np.arange(start=0, stop=101, step=10,dtype=float) self.xx,self.yy = np.meshgrid(self.x, self.y) self.xrx = xr.DataArray(self.xx,dims=['y','x']) self.xry = xr.DataArray(self.yy,dims=['y','x']) self.xrt = xu.sin(self.xrx/30.) + xu.cos(self.xry/30.) self.t = np.sin(self.xx/30.) + np.cos(self.yy/30.)
def nearest_points(self, lat, lon, npt=1): """ Use the lat-lon arrays to return a list of indices of the nearest npt points to the given lat-lon """ # Use sin of lat lon to handle periodic # and not worry about if we are in negative # degrees dist = xu.hypot(xu.sin(xu.radians(self['lat'].values)) - xu.sin(xu.radians(lat)),\ xu.cos(xu.radians(self['lon'].values)) - xu.cos(xu.radians(lon))) # Get indices of the flattened array nearest_raw = dist.argsort(axis=None)[:npt] # Convert back to 2-d coords nearest = np.unravel_index(nearest_raw, self['lat'].shape) return nearest
def setUp(self): # only testing on a regular grid x = np.arange(start=0, stop=1.e7, step=1.e5,dtype=float) y = np.arange(start=0, stop=1.2e7, step=1.e5,dtype=float) x,y = np.meshgrid(x,y) self.grd = agrids.plane_2d_grid(ycoord=y,xcoord=x) self.scale = 2. * 3.14159 / 1.e7 self.xt = self.grd._arrays['plane_x_coordinate_at_t_location'] self.yt = self.grd._arrays['plane_y_coordinate_at_t_location'] self.xu = self.grd.change_grid_location_t_to_u(self.xt) self.yu = self.grd.change_grid_location_t_to_u(self.yt) self.xv = self.grd.change_grid_location_t_to_v(self.xt) self.yv = self.grd.change_grid_location_t_to_v(self.yt) self.tvar1 = xu.sin(self.xt * self.scale) + xu.cos(self.yt * self.scale) self.uvar1 = xu.sin(self.xu * self.scale) + xu.cos(self.yu * self.scale) self.vvar1 = xu.sin(self.xv * self.scale) + xu.cos(self.yv * self.scale) self.vector = grids.VectorField2d(self.uvar1 ,self.vvar1, x_component_grid_location = 'u', y_component_grid_location = 'v') self.tvar2 = xu.cos(self.yt * self.scale)
def test_horizontal_divergence(self): # TODO : in 3D test divergence of curl is zero. tols = {'rtol':1e-3,'atol':1e-3} s = self.scale divvar = self.grd.horizontal_divergence(self.vector) dxdy = grids._dj(self.xv).shift(y=1) \ / self.grd._arrays["cell_y_size_at_t_location"] # custom derivative actual_div = divvar.to_masked_array() expected_div = (xu.cos(self.xt * s) * s - xu.sin(self.yt * s) * s + xu.cos(self.xt * s) * s * dxdy ).to_masked_array() #print_array_around(expected=expected_div/s,actual=actual_div/s) self.assertArray2dCloseInside(actual_div/s,expected_div/s, depth=2,**tols)
def setUp(self): # only testing on a regular grid x = np.arange(start=0, stop=1.e7, step=1.e5, dtype=float) y = np.arange(start=0, stop=1.2e7, step=1.e5, dtype=float) x, y = np.meshgrid(x, y) self.grd = agrids.plane_2d_grid(ycoord=y, xcoord=x) self.scale = 2. * 3.14159 / 1.e7 self.xt = self.grd._arrays['plane_x_coordinate_at_t_location'] self.yt = self.grd._arrays['plane_y_coordinate_at_t_location'] self.xu = self.grd.change_grid_location_t_to_u(self.xt) self.yu = self.grd.change_grid_location_t_to_u(self.yt) self.xv = self.grd.change_grid_location_t_to_v(self.xt) self.yv = self.grd.change_grid_location_t_to_v(self.yt) self.tvar1 = xu.sin(self.xt * self.scale) + xu.cos( self.yt * self.scale) self.uvar1 = xu.sin(self.xu * self.scale) + xu.cos( self.yu * self.scale) self.vvar1 = xu.sin(self.xv * self.scale) + xu.cos( self.yv * self.scale) self.vector = grids.VectorField2d(self.uvar1, self.vvar1, x_component_grid_location='u', y_component_grid_location='v') self.tvar2 = xu.cos(self.yt * self.scale)
def test_horizontal_divergence(self): # TODO : in 3D test divergence of curl is zero. tols = {'rtol': 1e-3, 'atol': 1e-3} s = self.scale divvar = self.grd.horizontal_divergence(self.vector) dxdy = grids._dj(self.xv).shift(y=1) \ / self.grd._arrays["cell_y_size_at_t_location"] # custom derivative actual_div = divvar.to_masked_array() expected_div = (xu.cos(self.xt * s) * s - xu.sin(self.yt * s) * s + xu.cos(self.xt * s) * s * dxdy).to_masked_array() #print_array_around(expected=expected_div/s,actual=actual_div/s) self.assertArray2dCloseInside(actual_div / s, expected_div / s, depth=2, **tols)
def test_horizontal_gradient(self): tols = {'rtol':1e-3,'atol':1e-3} gradvar = self.grd.horizontal_gradient(self.tvar1) gradx = self.grd.horizontal_gradient(self.xt) #dxdy = gradx.y_component s = self.scale actual_gx = gradvar.x_component.to_masked_array() actual_gy = gradvar.y_component.to_masked_array() expected_gx = (xu.cos(self.xu * s) * s).to_masked_array() expected_gy = (- xu.sin(self.yv * s) * s # + xu.cos(self.xv * s) * s * dxdy ).to_masked_array() self.assertArray2dCloseInside(actual_gx / s,expected_gx / s , depth=2,**tols) #print dxdy[:,20].values #print self.grd._arrays['cell_y_size_at_t_location'][:,20].values #print self.grd._arrays['cell_y_size_at_v_location'][:,20].values #print_array_around(expected=expected_gy/s,actual=actual_gy/s) self.assertArray2dCloseInside(actual_gy / s,expected_gy /s , depth=2, **tols)
def eady_growth_rate(data): """Calculate the local Eady Growth rate. Following Vallis (2017) p.354. EGR = 0.31*du/dz*f/N Parameters ---------- data : xarray.DataSet The Isca dataset. Requires fields 'temp', 'ps', 'pfull' and 'phalf' Returns a new xarray.DataArray of growth rate values on phalf levels, in s^-1. """ N2 = ixr.brunt_vaisala(data) f = 2.0 * omega * xruf.sin(xruf.deg2rad(data.lat)) dz = ixr.domain.calculate_dz(data) du = ixr.domain.diff_pfull(data.ucomp, data) N = xruf.sqrt(N2.where(N2 > 0)) egr = 0.31 * du / dz * f / N return np.abs(egr)
def test_ufuncs(self): u = self.eager_array v = self.lazy_array self.assertLazyAndAllClose(np.sin(u), xu.sin(v))
def test_univariate_ufunc(self): u = self.eager_var v = self.lazy_var self.assertLazyAndAllClose(np.sin(u), xu.sin(v))
def test_ufuncs(self): x = self.sp_xr assert_equal(np.sin(x), xu.sin(x))
def test_univariate_ufunc(self): assert_sparse_equal(np.sin(self.data), xu.sin(self.var).data)
def test_univariate_ufunc(self): sparse.utils.assert_eq(np.sin(self.data), xu.sin(self.var).data)