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