def test_no_transpose(): data = avg_center_data() avg_result = avg_to_cell_center(data, 1) expected = np.array([[4.5, 7, 9.5], [23.5, 40, 30.5], [17, 23.5, 44.5]]) np.testing.assert_almost_equal(avg_result, expected, decimal=3)
def getmap(self, layer, request): time_index, time_value = self.nearest_time(layer, request.GET['time']) wgs84_bbox = request.GET['wgs84_bbox'] with self.dataset() as nc: cached_sg = from_ncfile(self.topology_file) lon_name, lat_name = cached_sg.face_coordinates lon_obj = getattr(cached_sg, lon_name) lat_obj = getattr(cached_sg, lat_name) centers = cached_sg.centers lon = centers[..., 0][lon_obj.center_slicing] lat = centers[..., 1][lat_obj.center_slicing] if isinstance(layer, Layer): data_obj = getattr(cached_sg, layer.access_name) raw_var = nc.variables[layer.access_name] if len(raw_var.shape) == 4: z_index, z_value = self.nearest_z(layer, request.GET['elevation']) raw_data = raw_var[time_index, z_index, data_obj.center_slicing[-2], data_obj.center_slicing[-1]] elif len(raw_var.shape) == 3: raw_data = raw_var[time_index, data_obj.center_slicing[-2], data_obj.center_slicing[-1]] elif len(raw_var.shape) == 2: raw_data = raw_var[data_obj.center_slicing] else: raise BaseException('Unable to trim variable {0} data.'.format(layer.access_name)) # handle edge variables if data_obj.location is not None and 'edge' in data_obj.location: raw_data = avg_to_cell_center(raw_data, data_obj.center_axis) if request.GET['image_type'] == 'pcolor': return mpl_handler.pcolormesh_response(lon, lat, data=raw_data, request=request) elif request.GET['image_type'] == 'filledcontours': return mpl_handler.contourf_response(lon, lat, data=raw_data, request=request) else: raise NotImplementedError('Image type "{}" is not supported.'.format(request.GET['image_type'])) elif isinstance(layer, VirtualLayer): x_var = None y_var = None raw_vars = [] for l in layer.layers: data_obj = getattr(cached_sg, l.access_name) raw_var = nc.variables[l.access_name] raw_vars.append(raw_var) if len(raw_var.shape) == 4: z_index, z_value = self.nearest_z(layer, request.GET['elevation']) raw_data = raw_var[time_index, z_index, data_obj.center_slicing[-2], data_obj.center_slicing[-1]] elif len(raw_var.shape) == 3: raw_data = raw_var[time_index, data_obj.center_slicing[-2], data_obj.center_slicing[-1]] elif len(raw_var.shape) == 2: raw_data = raw_var[data_obj.center_slicing] else: raise BaseException('Unable to trim variable {0} data.'.format(l.access_name)) raw_data = avg_to_cell_center(raw_data, data_obj.center_axis) if x_var is None: if data_obj.vector_axis and data_obj.vector_axis.lower() == 'x': x_var = raw_data elif data_obj.center_axis == 1: x_var = raw_data if y_var is None: if data_obj.vector_axis and data_obj.vector_axis.lower() == 'y': y_var = raw_data elif data_obj.center_axis == 0: y_var = raw_data if x_var is None or y_var is None: raise BaseException('Unable to determine x and y variables.') dim_lengths = [ len(v.dimensions) for v in raw_vars ] if len(list(set(dim_lengths))) != 1: raise AttributeError('One or both of the specified variables has screwed up dimensions.') if request.GET['image_type'] == 'vectors': angles = cached_sg.angles[lon_obj.center_slicing] vectorstep = request.GET['vectorstep'] # don't do this if the vectorstep is 1; let's save a microsecond or two # it's identical to getting all the data if vectorstep > 1: data_dim = len(lon.shape) step_slice = (np.s_[::vectorstep],) * data_dim # make sure the vector step is used for all applicable dimensions lon = lon[step_slice] lat = lat[step_slice] x_var = x_var[step_slice] y_var = y_var[step_slice] angles = angles[step_slice] vectorscale = request.GET['vectorscale'] padding_factor = calc_safety_factor(vectorscale) # figure out the average distance between lat/lon points # do the math after taking into the vectorstep if specified spatial_idx_padding = calc_lon_lat_padding(lon, lat, padding_factor) spatial_idx = data_handler.lat_lon_subset_idx(lon, lat, lonmin=wgs84_bbox.minx, latmin=wgs84_bbox.miny, lonmax=wgs84_bbox.maxx, latmax=wgs84_bbox.maxy, padding=spatial_idx_padding ) subset_lon = self._spatial_data_subset(lon, spatial_idx) subset_lat = self._spatial_data_subset(lat, spatial_idx) # rotate vectors x_rot, y_rot = rotate_vectors(x_var, y_var, angles) spatial_subset_x_rot = self._spatial_data_subset(x_rot, spatial_idx) spatial_subset_y_rot = self._spatial_data_subset(y_rot, spatial_idx) return mpl_handler.quiver_response(subset_lon, subset_lat, spatial_subset_x_rot, spatial_subset_y_rot, request, vectorscale ) else: raise NotImplementedError('Image type "{}" is not supported.'.format(request.GET['image_type']))
def minmax(self, layer, request): time_index, time_value = self.nearest_time(layer, request.GET['time']) wgs84_bbox = request.GET['wgs84_bbox'] with self.dataset() as nc: cached_sg = from_ncfile(self.topology_file) lon_name, lat_name = cached_sg.face_coordinates lon_obj = getattr(cached_sg, lon_name) lat_obj = getattr(cached_sg, lat_name) centers = cached_sg.centers lon = centers[..., 0][lon_obj.center_slicing] lat = centers[..., 1][lat_obj.center_slicing] spatial_idx = data_handler.lat_lon_subset_idx(lon, lat, lonmin=wgs84_bbox.minx, latmin=wgs84_bbox.miny, lonmax=wgs84_bbox.maxx, latmax=wgs84_bbox.maxy) subset_lon = np.unique(spatial_idx[0]) subset_lat = np.unique(spatial_idx[1]) grid_variables = cached_sg.grid_variables vmin = None vmax = None raw_data = None if isinstance(layer, Layer): data_obj = getattr(cached_sg, layer.access_name) raw_var = nc.variables[layer.access_name] if len(raw_var.shape) == 4: z_index, z_value = self.nearest_z(layer, request.GET['elevation']) raw_data = raw_var[time_index, z_index, subset_lon, subset_lat] elif len(raw_var.shape) == 3: raw_data = raw_var[time_index, subset_lon, subset_lat] elif len(raw_var.shape) == 2: raw_data = raw_var[subset_lon, subset_lat] else: raise BaseException('Unable to trim variable {0} data.'.format(layer.access_name)) # handle grid variables if set([layer.access_name]).issubset(grid_variables): raw_data = avg_to_cell_center(raw_data, data_obj.center_axis) vmin = np.nanmin(raw_data).item() vmax = np.nanmax(raw_data).item() elif isinstance(layer, VirtualLayer): x_var = None y_var = None raw_vars = [] for l in layer.layers: data_obj = getattr(cached_sg, l.access_name) raw_var = nc.variables[l.access_name] raw_vars.append(raw_var) if len(raw_var.shape) == 4: z_index, z_value = self.nearest_z(layer, request.GET['elevation']) raw_data = raw_var[time_index, z_index, subset_lon, subset_lat] elif len(raw_var.shape) == 3: raw_data = raw_var[time_index, subset_lon, subset_lat] elif len(raw_var.shape) == 2: raw_data = raw_var[subset_lon, subset_lat] else: raise BaseException('Unable to trim variable {0} data.'.format(l.access_name)) if x_var is None: if data_obj.vector_axis and data_obj.vector_axis.lower() == 'x': x_var = raw_data elif data_obj.center_axis == 1: x_var = raw_data if y_var is None: if data_obj.vector_axis and data_obj.vector_axis.lower() == 'y': y_var = raw_data elif data_obj.center_axis == 0: y_var = raw_data if ',' in layer.var_name and raw_data is not None: # Vectors, so return magnitude data = [ sqrt((u*u) + (v*v)) for (u, v,) in zip(x_var.flatten(), y_var.flatten()) if u != np.nan and v != np.nan] vmin = min(data) vmax = max(data) return gmd_handler.from_dict(dict(min=vmin, max=vmax))
# The raw API is just that: an interpretation of the SGRID conventions that translate to a center slicing object. # It does not look like much, but anyone familiar with staggered grids models know the amount of code that goes into performing that slice. # # The API is meant to be used by higher level tools like `iris` and `xray` where plotting/slicing would be done automatically. # # However, `pysgrid` **does** bring some convenience processing methods like: # # - avg_to_cell_center # In[6]: from pysgrid.processing_2d import avg_to_cell_center u = avg_to_cell_center(u, sgrid.u.center_axis) v = avg_to_cell_center(v, sgrid.v.center_axis) # - rotate_vectors # In[7]: from pysgrid.processing_2d import rotate_vectors angles = nc.variables[sgrid.angle.variable][sgrid.angle.center_slicing] u, v = rotate_vectors(u, v, angles) # - vector_sum
def getmap(self, layer, request): time_index, time_value = self.nearest_time(layer, request.GET['time']) wgs84_bbox = request.GET['wgs84_bbox'] with self.dataset() as nc: cached_sg = from_ncfile(self.topology_file) lon_name, lat_name = cached_sg.face_coordinates lon_obj = getattr(cached_sg, lon_name) lat_obj = getattr(cached_sg, lat_name) centers = cached_sg.centers lon = centers[..., 0][lon_obj.center_slicing] lat = centers[..., 1][lat_obj.center_slicing] if isinstance(layer, Layer): data_obj = getattr(cached_sg, layer.access_name) raw_var = nc.variables[layer.access_name] if len(raw_var.shape) == 4: z_index, z_value = self.nearest_z(layer, request.GET['elevation']) raw_data = raw_var[time_index, z_index, data_obj.center_slicing[-2], data_obj.center_slicing[-1]] elif len(raw_var.shape) == 3: raw_data = raw_var[time_index, data_obj.center_slicing[-2], data_obj.center_slicing[-1]] elif len(raw_var.shape) == 2: raw_data = raw_var[data_obj.center_slicing] else: raise BaseException( 'Unable to trim variable {0} data.'.format( layer.access_name)) # handle edge variables if data_obj.location is not None and 'edge' in data_obj.location: raw_data = avg_to_cell_center(raw_data, data_obj.center_axis) if request.GET['image_type'] == 'pcolor': return mpl_handler.pcolormesh_response(lon, lat, data=raw_data, request=request) elif request.GET['image_type'] in [ 'filledhatches', 'hatches', 'filledcontours', 'contours' ]: return mpl_handler.contouring_response(lon, lat, data=raw_data, request=request) else: raise NotImplementedError( 'Image type "{}" is not supported.'.format( request.GET['image_type'])) elif isinstance(layer, VirtualLayer): x_var = None y_var = None raw_vars = [] for l in layer.layers: data_obj = getattr(cached_sg, l.access_name) raw_var = nc.variables[l.access_name] raw_vars.append(raw_var) if len(raw_var.shape) == 4: z_index, z_value = self.nearest_z( layer, request.GET['elevation']) raw_data = raw_var[time_index, z_index, data_obj.center_slicing[-2], data_obj.center_slicing[-1]] elif len(raw_var.shape) == 3: raw_data = raw_var[time_index, data_obj.center_slicing[-2], data_obj.center_slicing[-1]] elif len(raw_var.shape) == 2: raw_data = raw_var[data_obj.center_slicing] else: raise BaseException( 'Unable to trim variable {0} data.'.format( l.access_name)) raw_data = avg_to_cell_center(raw_data, data_obj.center_axis) if x_var is None: if data_obj.vector_axis and data_obj.vector_axis.lower( ) == 'x': x_var = raw_data elif data_obj.center_axis == 1: x_var = raw_data if y_var is None: if data_obj.vector_axis and data_obj.vector_axis.lower( ) == 'y': y_var = raw_data elif data_obj.center_axis == 0: y_var = raw_data if x_var is None or y_var is None: raise BaseException( 'Unable to determine x and y variables.') dim_lengths = [len(v.dimensions) for v in raw_vars] if len(list(set(dim_lengths))) != 1: raise AttributeError( 'One or both of the specified variables has screwed up dimensions.' ) if request.GET['image_type'] == 'vectors': angles = cached_sg.angles[lon_obj.center_slicing] vectorstep = request.GET['vectorstep'] # don't do this if the vectorstep is 1; let's save a microsecond or two # it's identical to getting all the data if vectorstep > 1: data_dim = len(lon.shape) step_slice = ( np.s_[::vectorstep], ) * data_dim # make sure the vector step is used for all applicable dimensions lon = lon[step_slice] lat = lat[step_slice] x_var = x_var[step_slice] y_var = y_var[step_slice] angles = angles[step_slice] vectorscale = request.GET['vectorscale'] padding_factor = calc_safety_factor(vectorscale) # figure out the average distance between lat/lon points # do the math after taking into the vectorstep if specified spatial_idx_padding = calc_lon_lat_padding( lon, lat, padding_factor) spatial_idx = data_handler.lat_lon_subset_idx( lon, lat, lonmin=wgs84_bbox.minx, latmin=wgs84_bbox.miny, lonmax=wgs84_bbox.maxx, latmax=wgs84_bbox.maxy, padding=spatial_idx_padding) subset_lon = self._spatial_data_subset(lon, spatial_idx) subset_lat = self._spatial_data_subset(lat, spatial_idx) # rotate vectors x_rot, y_rot = rotate_vectors(x_var, y_var, angles) spatial_subset_x_rot = self._spatial_data_subset( x_rot, spatial_idx) spatial_subset_y_rot = self._spatial_data_subset( y_rot, spatial_idx) return mpl_handler.quiver_response(subset_lon, subset_lat, spatial_subset_x_rot, spatial_subset_y_rot, request, vectorscale) else: raise NotImplementedError( 'Image type "{}" is not supported.'.format( request.GET['image_type']))
def minmax(self, layer, request): time_index, time_value = self.nearest_time(layer, request.GET['time']) wgs84_bbox = request.GET['wgs84_bbox'] with self.dataset() as nc: cached_sg = from_ncfile(self.topology_file) lon_name, lat_name = cached_sg.face_coordinates lon_obj = getattr(cached_sg, lon_name) lat_obj = getattr(cached_sg, lat_name) centers = cached_sg.centers lon = centers[..., 0][lon_obj.center_slicing] lat = centers[..., 1][lat_obj.center_slicing] spatial_idx = data_handler.lat_lon_subset_idx( lon, lat, lonmin=wgs84_bbox.minx, latmin=wgs84_bbox.miny, lonmax=wgs84_bbox.maxx, latmax=wgs84_bbox.maxy) subset_lon = np.unique(spatial_idx[0]) subset_lat = np.unique(spatial_idx[1]) grid_variables = cached_sg.grid_variables vmin = None vmax = None raw_data = None if isinstance(layer, Layer): data_obj = getattr(cached_sg, layer.access_name) raw_var = nc.variables[layer.access_name] if len(raw_var.shape) == 4: z_index, z_value = self.nearest_z(layer, request.GET['elevation']) raw_data = raw_var[time_index, z_index, subset_lon, subset_lat] elif len(raw_var.shape) == 3: raw_data = raw_var[time_index, subset_lon, subset_lat] elif len(raw_var.shape) == 2: raw_data = raw_var[subset_lon, subset_lat] else: raise BaseException( 'Unable to trim variable {0} data.'.format( layer.access_name)) # handle grid variables if set([layer.access_name]).issubset(grid_variables): raw_data = avg_to_cell_center(raw_data, data_obj.center_axis) vmin = np.nanmin(raw_data).item() vmax = np.nanmax(raw_data).item() elif isinstance(layer, VirtualLayer): x_var = None y_var = None raw_vars = [] for l in layer.layers: data_obj = getattr(cached_sg, l.access_name) raw_var = nc.variables[l.access_name] raw_vars.append(raw_var) if len(raw_var.shape) == 4: z_index, z_value = self.nearest_z( layer, request.GET['elevation']) raw_data = raw_var[time_index, z_index, subset_lon, subset_lat] elif len(raw_var.shape) == 3: raw_data = raw_var[time_index, subset_lon, subset_lat] elif len(raw_var.shape) == 2: raw_data = raw_var[subset_lon, subset_lat] else: raise BaseException( 'Unable to trim variable {0} data.'.format( l.access_name)) if x_var is None: if data_obj.vector_axis and data_obj.vector_axis.lower( ) == 'x': x_var = raw_data elif data_obj.center_axis == 1: x_var = raw_data if y_var is None: if data_obj.vector_axis and data_obj.vector_axis.lower( ) == 'y': y_var = raw_data elif data_obj.center_axis == 0: y_var = raw_data if ',' in layer.var_name and raw_data is not None: # Vectors, so return magnitude data = [ sqrt((u * u) + (v * v)) for ( u, v, ) in zip(x_var.flatten(), y_var.flatten()) if u != np.nan and v != np.nan ] vmin = min(data) vmax = max(data) return gmd_handler.from_dict(dict(min=vmin, max=vmax))
def test_with_transpose(): data = avg_center_data() avg_result = avg_to_cell_center(data, 0) expected = np.array([[6, 22, 25, 15], [6.5, 34, 29.5, 45.5]]) np.testing.assert_almost_equal(avg_result, expected, decimal=3)
# <codecell> angle = sgrid.angle angles = nc.variables[angle.variable][angle.center_slicing] # <markdowncell> # ### Average velocity vectors to cell centers # <codecell> from pysgrid.processing_2d import avg_to_cell_center u_avg = avg_to_cell_center(u_data, u_var.center_axis) v_avg = avg_to_cell_center(v_data, v_var.center_axis) # <markdowncell> # ### Rotate vectors by angles # <codecell> from pysgrid.processing_2d import rotate_vectors u_rot, v_rot = rotate_vectors(u_avg, v_avg, angles) # <markdowncell> # ### Speed
def getmap(self, layer, request): time_index, time_value = self.nearest_time(layer, request.GET['time']) wgs84_bbox = request.GET['wgs84_bbox'] with self.dataset() as nc: cached_sg = from_ncfile(self.topology_file) lon_name, lat_name = cached_sg.face_coordinates lon_obj = getattr(cached_sg, lon_name) lat_obj = getattr(cached_sg, lat_name) centers = cached_sg.centers lon = centers[..., 0][lon_obj.center_slicing] lat = centers[..., 1][lat_obj.center_slicing] if request.GET['image_type'] == 'vectors': vectorstep = request.GET['vectorstep'] vectorscale = request.GET['vectorscale'] padding_factor = calc_safety_factor(vectorscale) spatial_idx_padding = calc_lon_lat_padding(lon, lat, padding_factor) else: spatial_idx_padding = 0.18 vectorstep = None spatial_idx = data_handler.lat_lon_subset_idx(lon, lat, lonmin=wgs84_bbox.minx, latmin=wgs84_bbox.miny, lonmax=wgs84_bbox.maxx, latmax=wgs84_bbox.maxy, padding=spatial_idx_padding ) subset_x = np.unique(spatial_idx[0]) subset_y = np.unique(spatial_idx[1]) if subset_x.shape == (0, ) and subset_y.shape == (0, ): return mpl_handler.empty_response() # return an empty tile if subset contains no data else: x_min_idx = subset_x.min() x_max_idx = subset_x.max() + 1 y_min_idx = subset_y.min() y_max_idx = subset_y.max() + 1 lonlat_mask = np.ones(lon.shape) lonlat_mask[spatial_idx[0], spatial_idx[1]] = 0 trimmed_lon = ma.masked_array(lon, mask=lonlat_mask).data[x_min_idx:x_max_idx:vectorstep, y_min_idx:y_max_idx:vectorstep] trimmed_lat = ma.masked_array(lat, mask=lonlat_mask).data[x_min_idx:x_max_idx:vectorstep, y_min_idx:y_max_idx:vectorstep] if isinstance(layer, Layer): data_obj = getattr(cached_sg, layer.access_name) raw_var = nc.variables[layer.access_name] raw_data = self._retrieve_data(request=request, nc_variable=raw_var, sg_variable=data_obj, layer=layer, subset_x=subset_x, subset_y=subset_y, time_index=time_index, vectorstep=1 ) # handle edge variables if data_obj.location is not None and 'edge' in data_obj.location: raw_data = avg_to_cell_center(raw_data, data_obj.center_axis) if request.GET['image_type'] == 'pcolor': return mpl_handler.pcolormesh_response(trimmed_lon, trimmed_lat, data=raw_data, request=request) elif request.GET['image_type'] == 'filledcontours': return mpl_handler.contourf_response(trimmed_lon, trimmed_lat, data=raw_data, request=request) else: raise NotImplementedError('Image type "{}" is not supported.'.format(request.GET['image_type'])) elif isinstance(layer, VirtualLayer): x_var = None y_var = None raw_vars = [] for l in layer.layers: data_obj = getattr(cached_sg, l.access_name) raw_var = nc.variables[l.access_name] raw_vars.append(raw_var) raw_data = self._retrieve_data(request=request, nc_variable=raw_var, sg_variable=data_obj, layer=layer, subset_x=subset_x, subset_y=subset_y, time_index=time_index, vectorstep=vectorstep ) raw_data = self._avg_to_cell_center(raw_data, data_obj.center_axis, vectorstep) if x_var is None: if data_obj.vector_axis and data_obj.vector_axis.lower() == 'x': x_var = raw_data elif data_obj.center_axis == 1: x_var = raw_data if y_var is None: if data_obj.vector_axis and data_obj.vector_axis.lower() == 'y': y_var = raw_data elif data_obj.center_axis == 0: y_var = raw_data if x_var is None or y_var is None: raise BaseException('Unable to determine x and y variables.') dim_lengths = [ len(v.dimensions) for v in raw_vars ] if len(list(set(dim_lengths))) != 1: raise AttributeError('One or both of the specified variables has screwed up dimensions.') if request.GET['image_type'] == 'vectors': angles = cached_sg.angles[lon_obj.center_slicing] trimmed_angles = ma.masked_array(angles, mask=lonlat_mask).data[x_min_idx:x_max_idx:vectorstep, y_min_idx:y_max_idx:vectorstep] # rotate vectors x_rot, y_rot = rotate_vectors(x_var, y_var, trimmed_angles) return mpl_handler.quiver_response(trimmed_lon, trimmed_lat, x_rot, y_rot, request, vectorscale ) else: raise NotImplementedError('Image type "{}" is not supported.'.format(request.GET['image_type']))
# **Rotate the grid. # Average at the center. from pysgrid.processing_2d import avg_to_cell_center, rotate_vectors, vector_sum nc = Dataset(url) u_velocity = nc.variables[u_var.variable] v_velocity = nc.variables[v_var.variable] v_idx = 0 # surface time_idx = -1 # Last time step. u = u_velocity[time_idx, v_idx, u_var.center_slicing[-2], u_var.center_slicing[-1]] v = v_velocity[time_idx, v_idx, v_var.center_slicing[-2], v_var.center_slicing[-1]] u = avg_to_cell_center(u, u_var.center_axis) v = avg_to_cell_center(v, v_var.center_axis) angles = nc.variables[sgrid.angle.variable][sgrid.angle.center_slicing] u, v = rotate_vectors(u, v, angles) speed = vector_sum(u, v) \*\* CF convention does describe the angle variable for grids that needs rotation, but there is no action expected. For example, in the formula_terms, pysgrid must be improved to abstract that action when needed via a simpler method. ```xml <entry id="angle_of_rotation_from_east_to_x"> <canonical_units>degree</canonical_units> <grib></grib>
subset_lon = subset_data(longitudes, subset_idx) subset_lat = subset_data(latitudes, subset_idx) layer = LAYER variables = layer.split(',') var1_name = variables[0] var2_name = variables[1] cached_var1 = getattr(cached_sg, var1_name) cached_var2 = getattr(cached_sg, var2_name) raw_var1 = canon_dataset.variables[var1_name] raw_var2 = canon_dataset.variables[var2_name] vertical_index, vertical = nearest_z(raw_var1, canon_dataset, z) print('Vertical: {0}'.format(vertical)) var1_trimmed = raw_var1[time_idx, vertical_index, cached_var1.center_slicing[2], cached_var1.center_slicing[3]] var2_trimmed = raw_var2[time_idx, vertical_index, cached_var2.center_slicing[2], cached_var2.center_slicing[3]] var1_avg = avg_to_cell_center(var1_trimmed, cached_var1.center_axis) var2_avg = avg_to_cell_center(var2_trimmed, cached_var2.center_axis) if cached_var1.center_axis == 1 and cached_var2.center_axis == 0: x_var = var1_avg y_var = var2_avg else: x_var = var2_avg y_var = var1_avg x_rot, y_rot = rotate_vectors(x_var, y_var, angles) subset_x_rot = subset_data(x_rot, subset_idx) subset_y_rot = subset_data(y_rot, subset_idx) xy_vector_sum = vector_sum(subset_x_rot, subset_y_rot) # start experimental quiver_response section # end experimental quiver_response section
sg_u = sg.U sg_v = sg.V cell_centers = sg.centers lon_var_name, lat_var_name = sg.face_coordinates lon_data = cell_centers[:, :, 0] lat_data = cell_centers[:, :, 1] lon_var_obj = getattr(sg, lon_var_name) lat_var_obj = getattr(sg, lat_var_name) lon_subset = lon_data[lon_var_obj.center_slicing] lat_subset = lat_data[lat_var_obj.center_slicing] angles = sg.angles[lon_var_obj.center_slicing] u_data_trimmed = u_var[TIME_INDEX, VERTICAL_INDEX, sg_u.center_slicing[2], sg_u.center_slicing[3]] v_data_trimmed = v_var[TIME_INDEX, VERTICAL_INDEX, sg_v.center_slicing[2], sg_v.center_slicing[3]] u_data_avg = avg_to_cell_center(u_data_trimmed, sg_u.center_axis) v_data_avg = avg_to_cell_center(v_data_trimmed, sg_v.center_axis) u_rotated, v_rotated = rotate_vectors(u_data_avg, v_data_avg, angles) uv_vector_sum = vector_sum(u_rotated, v_rotated) fig = plt.figure(figsize=(12, 12)) plt.subplot(111, aspect=(1.0/np.cos(np.mean(lat_subset)*np.pi/180.0))) q = plt.quiver(lon_subset[::SUB, ::SUB], lat_subset[::SUB, ::SUB], u_rotated[::SUB, ::SUB], v_rotated[::SUB, ::SUB], uv_vector_sum[::SUB, ::SUB], scale=1.0/SCALE, pivot='middle', zorder=1e35, width=0.003