def _main(args): """Run the command line program.""" temperature_cube, temperature_history = gio.combine_files(args.temperature_file, args.temperature_var, checks=True) salinity_cube, salinity_history = gio.combine_files(args.salinity_file, 'sea_water_salinity', checks=True) assert 'c' in str(temperature_cube.units).lower(), "Input temperature units must be in celsius" # if not 'C' in str(bigthetao_cube.units): # bigthetao_cube.data = bigthetao_cube.data - 273.15 # data_median = np.ma.median(bigthetao_cube.data) # assert data_median < 100 # assert data_median > -10 # bigthetao_cube.units = 'C' target_shape = temperature_cube.shape[1:] depth = temperature_cube.coord('depth').points * -1 broadcast_depth = uconv.broadcast_array(depth, 0, target_shape) broadcast_longitude = uconv.broadcast_array(temperature_cube.coord('longitude').points, [1, 2], target_shape) broadcast_latitude = uconv.broadcast_array(temperature_cube.coord('latitude').points, [1, 2], target_shape) pressure = gsw.p_from_z(broadcast_depth, broadcast_latitude) absolute_salinity = gsw.SA_from_SP(salinity_cube.data, pressure, broadcast_longitude, broadcast_latitude) if args.temperature_var == 'sea_water_conservative_temperature': conservative_temperature = temperature_cube.data elif args.temperature_var == 'sea_water_potential_temperature': conservative_temperature = gsw.CT_from_pt(absolute_salinity, temperature_cube.data) else: raise ValueError('Invalid temperature variable') if args.coefficient == 'alpha': coefficient_data = gsw.alpha(absolute_salinity, conservative_temperature, pressure) var_name = 'alpha' standard_name = 'thermal_expansion_coefficient' long_name = 'thermal expansion coefficient' units = '1/K' elif args.coefficient == 'beta': coefficient_data = gsw.beta(absolute_salinity, conservative_temperature, pressure) var_name = 'beta' standard_name = 'saline_contraction_coefficient' long_name = 'saline contraction coefficient' units = 'kg/g' else: raise ValueError('Coefficient must be alpha or beta') iris.std_names.STD_NAMES[standard_name] = {'canonical_units': units} coefficient_cube = temperature_cube.copy() coefficient_cube.data = coefficient_data coefficient_cube.standard_name = standard_name coefficient_cube.long_name = long_name coefficient_cube.var_name = var_name coefficient_cube.units = units coefficient_cube.attributes['history'] = cmdprov.new_log(git_repo=repo_dir) iris.save(coefficient_cube, args.outfile)
def create_df(dcube, variable_name, vcube, bcube, basin, abort=True): """Create DataFrame""" if 'temperature' in variable_name: dcube = gio.temperature_unit_check(dcube, convert_to_celsius=True, abort=abort) elif 'salinity' in variable_name: dcube = gio.salinity_unit_check(dcube, abort=abort) assert dcube.ndim == 3 coord_names = [coord.name() for coord in dcube.dim_coords] assert coord_names[0] == 'depth' if dcube.coord('latitude').ndim == 1: lat_loc = coord_names.index('latitude') lon_loc = coord_names.index('longitude') else: lat_loc = lon_loc = [1, 2] lats = uconv.broadcast_array( dcube.coord('latitude').points, lat_loc, dcube.shape) lons = uconv.broadcast_array( dcube.coord('longitude').points, lon_loc, dcube.shape) levs = uconv.broadcast_array(dcube.coord('depth').points, 0, dcube.shape) ddata = dcube.data.flatten() vdata = vcube.data.flatten() bdata = bcube.data.flatten() lat_data = lats.flatten() lon_data = lons.flatten() depth_data = levs.flatten() df = pandas.DataFrame(index=range(ddata.shape[0])) df[variable_name] = ddata.filled(fill_value=5000) df['volume'] = vdata.filled(fill_value=5000) df['basin'] = bdata.filled(fill_value=5000) df['latitude'] = lat_data df['longitude'] = lon_data df['depth'] = depth_data df = df[df[variable_name] != 5000] df = df[df[variable_name] != -273.15] if basin: df = select_basin(df, basin) return df
def calc_zonal_weights(cube, coord_names): """Calculate zonal weights. Defined as the zonal distance (m) spanned by each grid box. The length of a degree of longitude is a function of latitude. The formula for the length of one degree of longitude is (pi/180) a cos(lat), where a is the radius of the earth. Args: cube (iris.cube.Cube): Data cube coord_names (list): Names of each data coordinate Returns: iris.cube: Array of weights with shape matching data_shape """ lon_coord = cube.coord('longitude') lon_index = coord_names.index('longitude') if not lon_coord.has_bounds(): lon_coord.guess_bounds() coslat = cosine_latitude_weights(cube) radius = iris.analysis.cartography.DEFAULT_SPHERICAL_EARTH_RADIUS lon_diffs = numpy.apply_along_axis(lambda x: x[1] - x[0], 1, lon_coord.bounds) lon_diffs = uconv.broadcast_array(lon_diffs, lon_index, cube.shape) lon_extents = (math.pi / 180.) * radius * coslat * lon_diffs return lon_extents
def calc_meridional_weights(lat_coord, coord_names, data_shape): """Calculate meridional weights. Defined as the zonal distance (m) spanned by each grid box. The length of a degree of latitude is (pi/180)*a, where a is the radius of the earth. Args: lat_coord (iris.coords.DimCoord): One-dimensional latitude coordinate coord_names (list): Names of each data coordinate data_shape (tuple): Shape of data Returns: iris.cube: Array of weights with shape matching data_shape """ lat_index = coord_names.index('latitude') if not lat_coord.has_bounds(): lat_coord.guess_bounds() lat_diffs = numpy.apply_along_axis(lambda x: x[1] - x[0], 1, lat_coord.bounds) lat_diffs = uconv.broadcast_array(lat_diffs, lat_index, data_shape) lat_extents = ( math.pi / 180. ) * iris.analysis.cartography.DEFAULT_SPHERICAL_EARTH_RADIUS * lat_diffs return lat_extents
def multiply_by_volume(cube, volume_cube=None): """Multiply by cell volume.""" if volume_cube: cube_ndim = cube.ndim volume_ndim = volume_cube.ndim if not cube_ndim == volume_ndim: diff = cube_ndim - volume_ndim assert cube.shape[diff:] == volume_cube.shape volume_data = uconv.broadcast_array( volume_cube.data, [cube_ndim - volume_ndim, cube_ndim - 1], cube.shape) else: volume_data = volume_cube.data else: volume_data = volume_array(cube) units = str(cube.units) cube.data = cube.data * volume_data if 'm-3' in units: cube.units = units.replace('m-3', '').replace(" ", " ") else: cube.units = units + ' m3' return cube
def multiply_by_area(cube, area_cube=None): """Multiply by cell area.""" if area_cube: cube_ndim = cube.ndim area_ndim = area_cube.ndim if not cube_ndim == area_ndim: diff = cube_ndim - area_ndim assert cube.shape[diff:] == area_cube.shape area_data = uconv.broadcast_array( area_cube.data, [cube_ndim - area_ndim, cube_ndim - 1], cube.shape) else: area_data = area_cube.data else: area_data = area_array(cube) units = str(cube.units) cube.data = cube.data * area_data if 'm-2' in units: cube.units = units.replace('m-2', '').replace(" ", " ") else: cube.units = units + ' m2' return cube
def read_data(infiles, var, area_cube, annual=False, multiply_by_area=False, chunk_annual=False): """Read the input data.""" cube, history = gio.combine_files(infiles, var) if annual: cube = timeseries.convert_to_annual(cube, days_in_month=True, chunk=chunk_annual) cube = uconv.flux_to_magnitude(cube) if multiply_by_area: cube = spatial_weights.multiply_by_area(cube, area_cube=area_cube) coord_names = [coord.name() for coord in cube.coords(dim_coords=True)] assert cube.ndim == 3 lats = cube.coord('latitude').points if lats.ndim == 1: lat_pos = coord_names.index('latitude') lats = uconv.broadcast_array(lats, lat_pos - 1, cube.shape[1:]) else: assert lats.shape == cube.shape[1:] return cube, lats, history
def calc_vertical_weights_1D(depth_coord, coord_names, data_shape): """Calculate vertical weights for a 1D depth axis with units = m. Args: depth_coord (iris.coords.DimCoord): One-dimensional depth coordinate coord_names (list): Names of each data coordinate data_shape (tuple): Shape of data Returns: iris.cube: Array of weights with shape matching data_shape """ assert depth_coord.units == 'm' # Calculate weights if not depth_coord.has_bounds(): depth_coord.guess_bounds() level_bounds = depth_coord.bounds level_diffs = numpy.apply_along_axis(lambda x: x[1] - x[0], 1, level_bounds) #guess_bounds can produce negative bound at surface if level_bounds[0][0] < 0.0: level_diffs[0] = level_diffs[0] + level_bounds[0][0] # Broadcast to size of data depth_index = coord_names.index('depth') level_diffs = uconv.broadcast_array(level_diffs, depth_index, data_shape) assert level_diffs.min() > 0.0 return level_diffs
def calc_zonal_vertical_mean(vertical_mean_cube, depth_cube, basin_array, basin_name, layer, atts, original_standard_name, original_var_name): """Calculate the zonal mean of the vertical mean field.""" assert layer in ['surface', 'argo'] if not basin_name == 'globe': vertical_mean_cube.data.mask = numpy.where((vertical_mean_cube.data.mask == False) & (basin_array == basins[basin_name]), False, True) if depth_cube: ndim = vertical_mean_cube.ndim depth_array = uconv.broadcast_array(depth_cube.data, [ndim - 2, ndim - 1], vertical_mean_cube.shape) else: depth_array = create_depth_array(vertical_mean_cube) max_depth = vertical_layers[layer][-1] depth_weights = numpy.ma.where(depth_array > max_depth, max_depth, depth_array) zonal_vertical_mean_cube = vertical_mean_cube.collapsed(['longitude'], iris.analysis.MEAN, weights=depth_weights.astype(numpy.float32)) zonal_vertical_mean_cube.remove_coord('longitude') zonal_vertical_mean_cube.data = zonal_vertical_mean_cube.data.astype(numpy.float32) units = str(vertical_mean_cube.units) standard_name = 'zonal_vertical_mean_%s_%s_%s' %(basin_name, layer, original_standard_name) var_name = '%s_zvm_%s_%s' %(original_var_name, basin_name, layer) zonal_vertical_mean_cube = add_metadata(atts, zonal_vertical_mean_cube, standard_name, var_name, units) return zonal_vertical_mean_cube
def main(inargs): """Run the program.""" standard_names = { 'thetao': 'sea_water_potential_temperature', 'so': 'sea_water_salinity' } volume_cube = gio.get_ocean_weights(inargs.volfile) output_cubelist = iris.cube.CubeList([]) for infile in inargs.infiles: cube = iris.load_cube(infile, standard_names[inargs.invar]) weights = uconv.broadcast_array(volume_cube.data, [1, 3], cube.shape) coord_names = [coord.name() for coord in cube.dim_coords] aux_coord_names = [coord.name() for coord in cube.aux_coords] ga = cube.collapsed(coord_names[1:], iris.analysis.MEAN, weights=weights) for coord in coord_names[1:] + aux_coord_names: ga.remove_coord(coord) ga.var_name = inargs.invar + 'ga' output_cubelist.append(ga) print(infile) outcube = gio.combine_cubes(output_cubelist) metadata_dict = {} metadata_dict[infile] = cube.attributes['history'] metadata_dict[inargs.volfile] = volume_cube.attributes['history'] outcube.attributes['history'] = cmdprov.new_log( infile_history=metadata_dict, git_repo=repo_dir) iris.save(outcube, inargs.outfile)
def main(inargs): """Run the program.""" data_dict = {} for experiment in list(experiment_colors.keys()): data_dict[(experiment, 'x_data')] = numpy.ma.array([]) data_dict[(experiment, 'y_data')] = numpy.ma.array([]) metadata_dict = {} for data_file, basin_file in inargs.file_pair: try: time_constraint = gio.get_time_constraint(inargs.time) except AttributeError: time_constraint = iris.Constraint() with iris.FUTURE.context(cell_datetime_objects=True): cube = iris.load_cube(data_file, 'sea_surface_salinity' & time_constraint) basin_cube = read_basin(basin_file) ndim = cube.ndim basin_array = uconv.broadcast_array(basin_cube.data, [ndim - 2, ndim - 1], cube.shape) metadata_dict[data_file] = cube.attributes['history'] metadata_dict[basin_file] = basin_cube.attributes['history'] model, experiment = get_experiment_details(cube) for basin in list(basins.keys()): zonal_climatology, zonal_trends = calc_zonal_stats( cube.copy(), basin_array, basin) data_dict[(experiment, 'x_data')] = numpy.ma.append( data_dict[(experiment, 'x_data')], zonal_climatology) data_dict[(experiment, 'y_data')] = numpy.ma.append( data_dict[(experiment, 'y_data')], zonal_trends) fig = plt.figure(figsize=(12, 8)) for experiment, color in experiment_colors.items(): x_data = data_dict[(experiment, 'x_data')] y_data = data_dict[(experiment, 'y_data')] if numpy.any(x_data): plt.scatter(x_data[::inargs.thin], y_data[::inargs.thin], facecolors='none', edgecolors=color, label=experiment) if experiment in ['AA', 'noAA']: x_trend, y_trend = calc_trend(x_data, y_data, experiment) plt.plot(x_trend, y_trend, color=color) plt.legend(loc=4) plt.xlabel('Climatological mean salinity') plt.ylabel('Salinity trend (per 50 years)') plt.title(model) # Write output plt.savefig(inargs.outfile, bbox_inches='tight') gio.write_metadata(inargs.outfile, file_info=metadata_dict)
def volume_from_volume(target_cube, volume_cube): """Create volume array from volume data cube.""" target_coord_names = [coord.name() for coord in target_cube.dim_coords] volume_coord_names = [coord.name() for coord in volume_cube.dim_coords] assert target_coord_names[0:2] == ['time', 'depth' ] and len(target_coord_names) == 4 assert volume_coord_names[0] == 'depth' depth_match = numpy.array_equal( volume_cube.coord('depth').points, target_cube.coord('depth').points) if depth_match: volume_data = volume_cube.data else: depth_match = numpy.array_equal( volume_cube[::-1, ::].coord('depth').points, target_cube.coord('depth').points) assert depth_match volume_data = volume_cube[::-1, ::].data volume_data = uconv.broadcast_array(volume_data, [1, 3], target_cube.shape) return volume_data
def volume_from_area(target_cube, area_cube=None): """Create volume array from area data.""" target_coord_names = [coord.name() for coord in target_cube.dim_coords] assert target_coord_names[0:2] == ['time', 'depth' ] and len(target_coord_names) == 4 if area_cube: area_data = uconv.broadcast_array(area_cube.data, [2, 3], target_cube.shape) else: area_data = spatial_weights.area_array(target_cube) depth_axis = target_cube.coord('depth') assert depth_axis.units in ['m', 'dbar'], "Unrecognised depth axis units" if depth_axis.units == 'm': depth_data = calc_vertical_weights_1D(depth_axis, target_coord_names, target_cube.shape) elif depth_axis.units == 'dbar': assert coord_names == [ 'depth', 'latitude', 'longitude' ], "2D weights will not work for curvilinear grid" depth_data = calc_vertical_weights_2D(depth_axis, target_cube.coord('latitude'), target_coord_names, target_cube.shape) volume_data = area_data * depth_data assert volume_data.shape == target_cube.shape return volume_data
def main(inargs): """Run the program.""" sfc_tbin_cube = iris.load_cube( inargs.sfc_file, 'total surface forcing binned by temperature') wfo_tbin_cube = iris.load_cube( inargs.wfo_file, 'Water Flux into Sea Water binned by temperature') cp = 3992.10322329649 #J kg-1 degC-1 lower_tos_bounds = sfc_tbin_cube.coord('sea_surface_temperature').bounds[:, 0] coord_names_tbin = [coord.name() for coord in sfc_tbin_cube.dim_coords] theta = uconv.broadcast_array( lower_tos_bounds, coord_names_tbin.index('sea_surface_temperature'), sfc_tbin_cube.shape) sfci_tbin_cube = sfc_tbin_cube.copy() sfci_tbin_cube.data = sfc_tbin_cube.data - (cp * theta * wfo_tbin_cube.data ) # SFCI = SFC - Cp*THETA*SVF sfci_tbin_cube.var_name = 'sfci_tbin' sfci_tbin_cube.long_name = 'total internal surface forcing binned by temperature' metadata = { inargs.sfc_file: sfc_tbin_cube.attributes['history'], inargs.wfo_file: wfo_tbin_cube.attributes['history'] } log = cmdprov.new_log(infile_history=metadata, git_repo=repo_dir) sfci_tbin_cube.attributes['history'] = log sfc_tsbin_cube = iris.load_cube( inargs.sfc_file, 'total surface forcing binned by temperature and salinity') wfo_tsbin_cube = iris.load_cube( inargs.wfo_file, 'Water Flux into Sea Water binned by temperature and salinity') coord_names_tsbin = [coord.name() for coord in sfc_tsbin_cube.dim_coords] theta = uconv.broadcast_array( lower_tos_bounds, coord_names_tsbin.index('sea_surface_temperature'), sfc_tsbin_cube.shape) sfci_tsbin_cube = sfc_tsbin_cube.copy() sfci_tsbin_cube.data = sfc_tsbin_cube.data - ( cp * theta * wfo_tsbin_cube.data) # SFCI = SFC - Cp*THETA*SVF sfci_tsbin_cube.var_name = 'sfci_tsbin' sfci_tsbin_cube.long_name = 'total internal surface forcing binned by temperature and salinity' sfci_tsbin_cube.attributes['history'] = log cube_list = iris.cube.CubeList([sfci_tbin_cube, sfci_tsbin_cube]) iris.save(cube_list, inargs.sfci_file)
def calc_spatial_agg(cube, coord_names, aux_coord_names, grid_type, aggregation_method, area_cube, lat_bounds=None, chunk=False): """Load the infiles and calculate the spatial aggregate (sum or mean).""" cube = cube.copy() coord_names = coord_names.copy() # Extract region if lat_bounds: if grid_type == 'curvilinear': cube = extract_region_curvilinear(cube, lat_bounds) else: cube = extract_region_latlon(cube, lat_bounds) # Get area weights if type(area_cube) == iris.cube.Cube: if grid_type == 'latlon' and lat_bounds: area_cube = extract_region_latlon(area_cube, lat_bounds) area_weights = uconv.broadcast_array(area_cube.data, [1, 2], cube.shape) elif type(area_cube) == str: area_weights = spatial_weights.area_array(cube) else: area_weights = None # Calculate spatial aggregate coord_names.remove('time') if chunk: spatial_agg = uconv.chunked_collapse_by_time(cube, coord_names, aggregation_method, weights=area_weights) else: spatial_agg = cube.collapsed(coord_names, aggregation_method, weights=area_weights) if area_cube and (aggregation_method == iris.analysis.SUM): units = str(spatial_agg.units) spatial_agg.units = units.replace('m-2', '') spatial_agg.remove_coord('latitude') spatial_agg.remove_coord('longitude') if grid_type == 'curvilinear': spatial_agg.remove_coord(coord_names[0]) spatial_agg.remove_coord(coord_names[1]) return spatial_agg
def calc_mean_anomaly(data_cube, clim_cube, sign, grid_areas): """Calculate the mean of all the positive or negative anomalies.""" clim_data = uconv.broadcast_array(clim_cube.data, [1, 2], data_cube.shape) grid_areas = uconv.broadcast_array(grid_areas, [1, 2], data_cube.shape) if sign == 'positive': new_mask = numpy.where( (data_cube.data.mask == False) & (clim_data > 0.0), False, True) elif sign == 'negative': new_mask = numpy.where( (data_cube.data.mask == False) & (clim_data < 0.0), False, True) data_cube.data.mask = new_mask data_cube = data_cube.collapsed(['longitude', 'latitude'], iris.analysis.MEAN, weights=grid_areas) data_cube.remove_coord('longitude') data_cube.remove_coord('latitude') return data_cube
def main(inargs): """Run the program.""" # Depth data data_cube = iris.load_cube(inargs.dummy_file, inargs.dummy_var) dim_coord_names = [coord.name() for coord in data_cube.dim_coords] aux_coord_names = [coord.name() for coord in data_cube.aux_coords] assert dim_coord_names[0] == 'time' depth_name = dim_coord_names[1] data_cube = data_cube[0, ::] data_cube.remove_coord('time') depth_data = spatial_weights.get_depth_array(data_cube, depth_name) # Area data if inargs.area_file: area_cube = iris.load_cube(inargs.area_file, 'cell_area') gio.check_global_ocean_area(area_cube.data.sum()) area_data = uconv.broadcast_array(area_cube.data, [1, 2], depth_data.shape) else: area_data = spatial_weights.area_array(data_cube) volume_data = depth_data * area_data if inargs.sftof_file: sftof_cube = iris.load_cube(inargs.sftof_file) assert sftof_cube.data.max() == 100 sftof_data = uconv.broadcast_array(sftof_cube.data, [1, 2], depth_data.shape) volume_data = volume_data * (sftof_data / 100.0) volume_data = numpy.ma.asarray(volume_data) data = numpy.ma.masked_invalid(data_cube.data) volume_data.mask = data.mask global_atts = area_cube.attributes if inargs.area_file else None volume_cube = construct_volume_cube(volume_data, data_cube, global_atts) volume_cube.attributes['history'] = gio.write_metadata() gio.check_global_ocean_volume(volume_cube.data.sum()) iris.save(volume_cube, inargs.outfile)
def create_basin_array(cube, ref_is_basin): """Create an ocean basin array. Args: cube - reference cube (iris.cube.Cube) ref_is_basin - flag to indicate if ref cube is a CMIP basin cube (bool) """ # Build latitude and longitude reference arrays lat_coord = cube.coord('latitude') lon_coord = cube.coord('longitude') if lat_coord.ndim == 1: # rectilinear grid lat_axis = lat_coord.points lon_axis = uconv.adjust_lon_range(lon_coord.points, radians=False) coord_names = [coord.name() for coord in cube.dim_coords] lat_index = coord_names.index('latitude') lon_index = coord_names.index('longitude') lat_array = uconv.broadcast_array(lat_axis, lat_index, cube.shape) lon_array = uconv.broadcast_array(lon_axis, lon_index, cube.shape) else: # curvilinear grid lon_coord = check_lon_coord(lon_coord) assert cube.ndim == 2 lat_array = lat_coord.points lon_array = lon_coord.points # Create basin array pacific_lon_bounds = [147, 294] indian_lon_bounds = [23, 147] if ref_is_basin: basin_array = basins_from_cmip(cube, lat_array, lon_array, pacific_lon_bounds, indian_lon_bounds) else: basin_array = basins_manually(cube, lat_array, lon_array, pacific_lon_bounds, indian_lon_bounds) return basin_array
def create_mask(mask_cube, target_shape): """Create mask from an sftlf (land surface fraction) file. There is no land when cell value == 0 """ target_ndim = len(target_shape) land_array = numpy.where(mask_cube.data < 0.001, False, True) mask = uconv.broadcast_array(land_array, [target_ndim - 2, target_ndim - 1], target_shape) assert mask.shape == target_shape return mask
def get_area_weights(cube, area_cube): """Get area weights for averaging""" if area_cube: area_weights = uconv.broadcast_array(area_cube.data, [1, 2], cube.shape) else: if not cube.coord('latitude').has_bounds(): cube.coord('latitude').guess_bounds() if not cube.coord('longitude').has_bounds(): cube.coord('longitude').guess_bounds() area_weights = iris.analysis.cartography.area_weights(cube) return area_weights
def _days_in_month_annual_mean(cube): """Calculate the annual mean timeseries accounting for days in month.""" df, days_in_year = get_days_in_year(cube, return_df=True) df['weight'] = df.apply(lambda row: row['days_in_month'] / days_in_year.loc[row['year']], axis=1) np.testing.assert_allclose(df.groupby('year').sum()['weight'].min(), 1.0) np.testing.assert_allclose(df.groupby('year').sum()['weight'].max(), 1.0) weights_array = uconv.broadcast_array(df['weight'].values, 0, cube.shape) cube.data = cube.data * weights_array cube = cube.aggregated_by(['year'], iris.analysis.SUM) return cube
def convert_to_joules(cube): """Convert units to Joules""" assert 'W' in str(cube.units) assert 'days' in str(cube.coord('time').units) time_span_days = cube.coord('time').bounds[:, 1] - cube.coord( 'time').bounds[:, 0] time_span_seconds = time_span_days * 60 * 60 * 24 cube.data = cube.data * uconv.broadcast_array(time_span_seconds, 0, cube.shape) cube.units = str(cube.units).replace('W', 'J') return cube
def broadcast_data(cube, target_shape): """Broadcast a static cube (e.g. basin, area, volume).""" target_ndim = len(target_shape) if not cube.ndim == target_ndim: diff = target_ndim - cube.ndim assert target_shape[diff:] == cube.shape data = uconv.broadcast_array(cube.data, [diff, target_ndim - 1], target_shape) else: data = cube.data assert target_shape == data.shape return data
def get_area_weights(cube, area_file, lat_constraint): """Get area weights for averaging""" if area_file: area_cube = iris.load_cube(inargs.area_file, lat_constraint) else: area_cube = None if area_cube: area_weights = uconv.broadcast_array(area_cube.data, [1, 2], cube.shape) else: area_weights = spatial_weights.area_array(cube) return area_weights
def select_basin(cube, basin_cube, basin_name): """Select an ocean basin.""" basins = {'atlantic': 2, 'pacific': 3, 'indian': 5} assert basin_name in basins.keys() assert cube.ndim == 4 basin_array = uconv.broadcast_array(basin_cube.data, [1, 3], cube.shape) cube.data.mask = numpy.where((cube.data.mask == False) & (basin_array == basins[basin_name]), False, True) return cube
def create_region_mask(latitude_array, target_shape, lat_bounds): """Create mask from the latitude auxillary coordinate""" target_ndim = len(target_shape) southern_lat, northern_lat = lat_bounds mask_array = numpy.where( (latitude_array >= southern_lat) & (latitude_array < northern_lat), False, True) mask = uconv.broadcast_array(mask_array, [target_ndim - 2, target_ndim - 1], target_shape) assert mask.shape == target_shape return mask
def _create_region_mask(horiz_array, target_shape, horiz_bounds): """Create mask from a horizontal (i.e. latitude or longitude) auxillary coordinate""" target_ndim = len(target_shape) lower_bound, upper_bound = horiz_bounds mask_array = numpy.where( (horiz_array >= lower_bound) & (horiz_array < upper_bound), False, True) mask = uconv.broadcast_array(mask_array, [target_ndim - 2, target_ndim - 1], target_shape) assert mask.shape == target_shape return mask
def select_basin(cube, basin_cube, basin_name): """Select an ocean basin.""" assert basin_cube.shape == cube.shape[-2:] basin_array = uconv.broadcast_array(basin_cube.data, [cube.ndim - 2, cube.ndim - 1], cube.shape) assert basin_array.min() == 11 assert basin_array.max() == 18 basins = {'atlantic': basin_array <= 12, 'indo-pacific': (basin_array >= 13) & (basin_array <= 15), 'globe': basin_array <= 16} assert basin_name in basins.keys() cube.data.mask = numpy.where((cube.data.mask == False) & basins[basin_name], False, True) return cube
def calc_region_sum(cube, coord_names, aux_coord_names, grid_type, area_cube, region): """Calculate the spatial sum.""" if grid_type == 'curvilinear': assert area_cube, "Must provide an area cube of curvilinear data" cube = cube.copy() coord_names = coord_names.copy() lat_bounds = region_bounds[region] # Extract region if lat_bounds: if grid_type == 'curvilinear': cube = extract_region_curvilinear(cube, lat_bounds) else: cube = extract_region_latlon(cube, lat_bounds) if 'm-2' in str(cube.units): # Get area weights if area_cube: if grid_type == 'latlon' and lat_bounds: area_cube = extract_region_latlon(area_cube, lat_bounds) area_data = uconv.broadcast_array(area_cube.data, [1, 2], cube.shape) else: area_data = spatial_weights.area_array(cube) # Multiply by area cube.data = cube.data * area_data units = str(cube.units) cube.units = units.replace('m-2', '') assert cube.units == 'J' coord_names.remove('time') spatial_agg = cube.collapsed(coord_names, iris.analysis.SUM) spatial_agg.remove_coord('latitude') spatial_agg.remove_coord('longitude') if grid_type == 'curvilinear': spatial_agg.remove_coord(coord_names[0]) spatial_agg.remove_coord(coord_names[1]) return spatial_agg
def area_ajustment(data_cube, area_file, metadata_dict): """Multipy a data cube by its cell area.""" if area_file: area_cube = iris.load_cube(area_file[0]) area_data = uconv.broadcast_array(area_cube.data, [1, 2], data_cube.shape) metadata_dict[area_file[0]] = area_cube.attributes['history'] else: area_data = get_areacello_data(data_cube) data_cube.data = data_cube.data * area_data if 'm-2' in str(data_cube.units): units = str(data_cube.units).replace('m-2', "") else: units = str(data_cube.units) + ' m2' return data_cube, units, metadata_dict