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 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 read_data(infile_list, var, basin_cube, region):
"""Read the data files.
The CSIRO-Mk3-6-0 model seems to be formatted incorrectly
and you can't select the regioins by name.
"""
cube, history = gio.combine_files(infile_list, var)
cube.attributes = atts
cube.attributes['history'] = history[0]
model = atts['model_id']
if var == 'northward_ocean_heat_transport':
region_index = {}
region_index['atlantic_arctic_ocean'] = 0
region_index['indian_pacific_ocean'] = 1
region_index['global_ocean'] = 2
if model == 'CSIRO-Mk3-6-0':
cube = cube[:, region_index[region], :]
else:
cube = cube.extract(iris.Constraint(region=region))
cube = timeseries.convert_to_annual(cube, full_months=True)
if basin_cube:
cube = uconv.mask_marginal_seas(cube, basin_cube)
if region != 'global_ocean':
basin_numbers = {}
basin_numbers['atlantic_arctic_ocean'] = [2, 4]
basin_numbers['indian_pacific_ocean'] = [3, 5]
cube = uconv.mask_unwanted_seas(cube, basin_cube, basin_numbers[region])
return 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([])
cube, history = gio.combine_files(inargs.infiles,
standard_names[inargs.invar],
checks=True)
ntsteps = cube.shape[0]
for tstep, cube_slice in enumerate(cube.slices_over('time')):
print(f'time step {tstep + 1} of {ntsteps}')
coord_names = [coord.name() for coord in cube.dim_coords]
aux_coord_names = [coord.name() for coord in cube.aux_coords]
ga = cube_slice.collapsed(coord_names[1:],
iris.analysis.MEAN,
weights=volume_cube.data)
for coord in coord_names[1:] + aux_coord_names:
ga.remove_coord(coord)
ga.var_name = inargs.invar + 'ga'
output_cubelist.append(ga)
outcube = output_cubelist.merge()[0]
metadata_dict = {}
metadata_dict[inargs.infiles[0]] = 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."""
pe_cube, pe_history = gio.combine_files(
inargs.pe_files, 'precipitation minus evaporation flux')
lat_coord = pe_cube.coord('latitude').points
region_data = np.apply_along_axis(get_regional_totals, 1, pe_cube.data,
lat_coord)
if inargs.cumsum:
region_data = np.cumsum(region_data, axis=0)
region_coord = create_region_coord()
time_coord = pe_cube.coord('time')
iris.std_names.STD_NAMES['precipitation_minus_evaporation_flux'] = {
'canonical_units': pe_cube.units
}
dim_coords_list = [(time_coord, 0), (region_coord, 1)]
out_cube = iris.cube.Cube(
region_data,
standard_name='precipitation_minus_evaporation_flux',
long_name=pe_cube.long_name,
var_name=pe_cube.var_name,
units=pe_cube.units,
attributes=pe_cube.attributes,
dim_coords_and_dims=dim_coords_list)
out_cube.attributes['history'] = cmdprov.new_log(
infile_history={inargs.pe_files[0]: pe_history[0]}, git_repo=repo_dir)
iris.save(out_cube, inargs.outfile)
def get_bin_data(files, var, w_cube):
"""Get binning variable data."""
cube, history = gio.combine_files(files, var, checks=True)
w_coord_names = [coord.name() for coord in w_cube.dim_coords]
coord_names = [coord.name() for coord in cube.dim_coords]
assert w_cube.shape[-2:] == cube.shape[-2:]
if not w_cube.shape == cube.shape:
if (w_cube.ndim == 3) and (cube.ndim == 4) and (w_coord_names[0]
== coord_names[0]):
#e.g. w_cube is surface flux (time, i, j),
#cube is full depth temperature (time, depth, i, j)
cube = cube[:, 0, ::]
cube.remove_coord(coord_names[1])
assert w_cube.shape == cube.shape
elif (w_cube.ndim == 2) and (cube.ndim == 4):
#e.g. w_cube is area (i, j),
#cube is full depth temperature (time, depth, i, j)
cube = cube[:, 0, ::]
cube.remove_coord(coord_names[1])
assert w_cube.shape == cube.shape[1:]
else:
#e.g. w_cube is area (i, j),
#cube is surface temperature (time, i, j)
#e.g. w_cube is volume (depth, i, j),
#cube is temperature (time, depth, i, j)
assert w_cube.shape == cube.shape[1:]
return cube, history
def get_data(infile, var, time_constraint, operation, ref_model=None):
"""Get the data for a particular component"""
assert operation in ['anomaly', 'mean']
cube, history = gio.combine_files(infile, var, new_calendar='365_day')
cube = cube[:, :, -1]
cube.remove_coord('region')
if time_constraint:
cube = cube.extract(time_constraint)
if operation == 'mean':
cube = cube.collapsed('time', iris.analysis.MEAN)
else:
cube.data = cube.data - cube.data[0, ::]
cube = cube[-1, ::]
cube.remove_coord('time')
if ref_model:
try:
model = cube.attributes['model_id']
except KeyError:
model = cube.attributes['source_id']
assert model == ref_model, f"Model mismatch: {ref_model}, {model}"
return cube, history
def main(inargs):
"""Run the program."""
coefficient_a_cube = iris.load_cube(inargs.coefficient_file, 'coefficient a')
coefficient_b_cube = iris.load_cube(inargs.coefficient_file, 'coefficient b')
coefficient_c_cube = iris.load_cube(inargs.coefficient_file, 'coefficient c')
coefficient_d_cube = iris.load_cube(inargs.coefficient_file, 'coefficient d')
data_cube, data_history = gio.combine_files(inargs.data_files, inargs.var)
coord_names = [coord.name() for coord in data_cube.coords(dim_coords=True)]
assert coord_names[0] == 'year'
if not inargs.branch_year == None:
branch_year = inargs.branch_year
else:
branch_year = get_branch_year(data_cube, inargs.control_time_units)
time_values = numpy.arange(branch_year, branch_year + data_cube.shape[0])
drift_signal, start_polynomial = remove_drift.apply_polynomial(time_values, coefficient_a_cube.data,
coefficient_b_cube.data, coefficient_c_cube.data,
coefficient_d_cube.data, poly_start=None)
new_cube = data_cube - drift_signal
#remove_drift.check_data(new_cube, data_cube, inargs.data_file)
new_cube.metadata = data_cube.metadata
metadata_dict = {inargs.data_files[0]: data_history[0],
inargs.coefficient_file: coefficient_a_cube.attributes['history']}
new_cube.attributes['history'] = cmdprov.new_log(infile_history=metadata_dict, git_repo=repo_dir)
iris.save(new_cube, inargs.outfile)
def main(inargs):
"""Run the program."""
cube, history = gio.combine_files(inargs.infiles, inargs.var, checks=True)
if inargs.annual:
cube = timeseries.convert_to_annual(cube, chunk=False)
if inargs.aggregation == 'sum':
cube = cube.collapsed('depth', iris.analysis.SUM)
else:
dim_coord_names = [coord.name() for coord in cube.dim_coords]
depth_coord = cube.coord('depth')
assert depth_coord.units in ['m', 'dbar'], "Unrecognised depth axis units"
if depth_coord.units == 'm':
vert_extents = spatial_weights.calc_vertical_weights_1D(depth_coord, dim_coord_names, cube.shape)
elif depth_coord.units == 'dbar':
vert_extents = spatial_weights.calc_vertical_weights_2D(depth_coord, cube.coord('latitude'), dim_coord_names, cube.shape)
cube = cube.collapsed('depth', iris.analysis.MEAN, weights=vert_extents)
cube.remove_coord('depth')
metadata_dict = {}
metadata_dict[inargs.infiles[0]] = history
cube.attributes['history'] = cmdprov.new_log(infile_history=metadata_dict, git_repo=repo_dir)
iris.save(cube, inargs.outfile)
def main(inargs):
"""Run the program."""
assert len(inargs.infiles) > 1
if inargs.variables:
variable_list = inargs.variables
else:
cube_list = iris.load(inargs.infiles[0])
variable_list = [cube.long_name for cube in cube_list]
cube_list = iris.cube.CubeList([])
for var in variable_list:
cube, history = gio.combine_files(inargs.infiles, var)
cube_list.append(cube)
metadata_dict = {inargs.infiles[-1]: history[-1]}
log_entry = cmdprov.new_log(infile_history=metadata_dict,
git_repo=repo_dir)
if len(cube_list) > 1:
iris.util.equalise_attributes(cube_list)
for cube in cube_list:
cube.attributes['history'] = log_entry
else:
cube_list = cube_list[0]
cube_list.attributes['history'] = log_entry
iris.save(cube_list, inargs.outfile)
def main(inargs):
"""Run the program."""
pr_cube, pr_history = gio.combine_files(inargs.pr_files, 'precipitation_flux', checks=True)
evap_cube, evap_history = gio.combine_files(inargs.evap_files, 'water_evapotranspiration_flux', checks=True)
assert pr_cube.shape == evap_cube.shape
pe_cube = pr_cube.copy()
pe_cube.data = pr_cube.data - evap_cube.data
iris.std_names.STD_NAMES['precipitation_minus_evaporation_flux'] = {'canonical_units': pe_cube.units}
pe_cube.standard_name = 'precipitation_minus_evaporation_flux'
pe_cube.long_name = 'precipitation minus evaporation flux'
pe_cube.var_name = 'pe'
pe_cube.attributes['history'] = cmdprov.new_log(git_repo=repo_dir)
iris.save(pe_cube, inargs.outfile)
def main(inargs):
"""Run the program."""
if inargs.var == 'ocean_volume':
data_cube = gio.get_ocean_weights(inargs.infile, sanity_check=False)
else:
data_cube, data_history = gio.combine_files(inargs.infile, inargs.var)
mask_cube, mask_history = gio.combine_files(inargs.mask_file, inargs.mask_var)
if inargs.mask_method == 'copy':
mask = copy_mask(mask_cube, data_cube.shape)
else:
mask = create_mask(mask_cube, data_cube.shape)
data_cube.data = numpy.ma.asarray(data_cube.data)
data_cube.data.mask = mask
data_cube.attributes['history'] = cmdprov.new_log(git_repo=repo_dir)
iris.save(data_cube, inargs.outfile)
def read_data(file_list, var, grid_point, convert_to_annual=False):
"""Read input data."""
cube, history = gio.combine_files(file_list, var)
if grid_point:
cube = select_point(cube, grid_point, timeseries=True)
if convert_to_annual:
cube = timeseries.convert_to_annual(cube)
return cube, history[0]
def main(inargs):
"""Run the program."""
cube, history = gio.combine_files(inargs.infiles, inargs.var, checks=True)
if inargs.surface:
coord_names = [coord.name() for coord in cube.dim_coords]
if 'depth' in coord_names:
cube = cube.extract(iris.Constraint(depth=0))
else:
print('no depth axis for surface extraction')
if inargs.annual:
cube = timeseries.convert_to_annual(cube, chunk=inargs.chunk)
log = cmdprov.new_log(infile_history={inargs.infiles[0]: history[0]},
git_repo=repo_dir)
dim_vals = {}
dim_vals['latitude'] = get_dim_vals(inargs.lats)
dim_vals['longitude'] = get_dim_vals(inargs.lons)
if inargs.levs:
dim_vals['depth'] = get_dim_vals(inargs.levs)
else:
dim_vals['depth'] = get_dim_vals(inargs.depth_bnds, bounds=True)
# Regrid from curvilinear to rectilinear if necessary
regrid_status = False
if inargs.lats:
horizontal_grid = grids.make_grid(dim_vals['latitude'],
dim_vals['longitude'])
cube, coord_names, regrid_status = grids.curvilinear_to_rectilinear(
cube, target_grid_cube=horizontal_grid)
# Regrid to new grid
if dim_vals['depth'] or not regrid_status:
sample_points = get_sample_points(cube, dim_vals)
cube = cube.interpolate(sample_points, iris.analysis.Linear())
coord_names = [coord.name() for coord in cube.dim_coords]
if 'latitude' in coord_names:
cube.coord('latitude').guess_bounds()
if 'longitude' in coord_names:
cube.coord('longitude').guess_bounds()
if inargs.levs:
cube = spatial_weights.guess_depth_bounds(cube)
else:
cube.coord('depth').bounds = get_depth_bounds(inargs.depth_bnds)
if numpy.isnan(numpy.min(cube.data)):
cube = remove_nans(cube)
# Reinstate time dim_coord if necessary
aux_coord_names = [coord.name() for coord in cube.aux_coords]
if 'time' in aux_coord_names:
cube = iris.util.new_axis(cube, 'time')
cube.attributes['history'] = log
iris.save(cube, inargs.outfile, fill_value=1e20)
def read_infiles(infiles, var, time_constraint, ensnum):
"""Combine multiple input files into one cube"""
cube, history = gio.combine_files(infiles, var)
cube = gio.check_time_units(cube)
cube = cube.extract(time_constraint)
new_aux_coord = iris.coords.AuxCoord(ensnum, long_name='ensemble_member', units='no_unit')
cube.add_aux_coord(new_aux_coord)
return cube, history[0]
def get_data(infile, variable, time_constraint):
"""Get the data for a particular model"""
cube, history = gio.combine_files(infile, variable, new_calendar='365_day')
#cube = iris.load_cube(infile, 'precipitation minus evaporation flux' & time_constraint)
cube = cube.extract(time_constraint)
iris.coord_categorisation.add_year(cube, 'time')
anomaly_data = cube.data - cube.data[0, :]
start_data = cube.data[0, :]
return cube, anomaly_data, start_data
def get_weights_data(file_list, var):
"""Read the weights data file/s"""
w_var = mom_vars[var] if var in mom_vars else var
if ('vol' in w_var) or ('area' in w_var):
assert len(file_list) == 1
w_cube = gio.get_ocean_weights(file_list[0])
history = w_cube.attributes['history']
else:
w_cube, history = gio.combine_files(file_list, var, checks=True)
return w_cube, history
def combine_infiles(inargs, time_constraint, depth_constraint):
"""Combine multiple input files into one cube"""
cube, history = gio.combine_files(inargs.infiles, inargs.var)
atts = cube[0].attributes
cube = cube.extract(time_constraint & depth_constraint)
cube = iris.util.squeeze(cube)
log = cmdprov.new_log(infile_history={inargs.infiles[0]: history[0]}, git_repo=repo_dir)
cube.attributes['history'] = log
return cube
def get_data(infile, var, time_constraint):
"""Get the data for a particular component"""
cube, history = gio.combine_files(infile, var, new_calendar='365_day')
cube = cube[:, -1, :]
cube.remove_coord('precipitation minus evaporation region')
if time_constraint:
cube = cube.extract(time_constraint)
cube.data = cube.data - cube.data[0, ::]
cube = cube[-1, ::]
cube.remove_coord('time')
return cube, history
def main(inargs):
"""Run the program."""
tas_cube, history = gio.combine_files(inargs.tas_files, inargs.var)
if inargs.annual:
tas_cube = timeseries.convert_to_annual(tas_cube)
area_data = spatial_weights.area_array(tas_cube)
coord_names = [coord.name() for coord in tas_cube.dim_coords]
tasga_cube = tas_cube.collapsed(coord_names[1:], iris.analysis.MEAN, weights=area_data)
tasga_cube.remove_coord(coord_names[1])
tasga_cube.remove_coord(coord_names[2])
tasga_cube.attributes['history'] = cmdprov.new_log(git_repo=repo_dir)
iris.save(tasga_cube, inargs.outfile)
def main(inargs):
"""Run the program."""
cube, history = gio.combine_files(inargs.infiles, inargs.var)
if inargs.annual:
cube = timeseries.convert_to_annual(cube, aggregation='mean', days_in_month=True)
if inargs.flux_to_mag:
cube = uconv.flux_to_magnitude(cube)
dim_coord_names = [coord.name() for coord in cube.dim_coords]
assert dim_coord_names[0] in ['time', 'year']
cube.data = numpy.cumsum(cube.data, axis=0)
cube.attributes['history'] = cmdprov.new_log(git_repo=repo_dir, infile_history={inargs.infiles[0]: history[0]})
iris.save(cube, inargs.outfile)
def main(inargs):
"""Run the program."""
cube, history = gio.combine_files(inargs.infiles, inargs.var)
start_cube = period_mean(cube.copy(), inargs.start_period)
end_cube = period_mean(cube.copy(), inargs.end_period)
outcube = cube[0, ::].copy()
outcube.remove_coord('time')
outcube.data = end_cube.data - start_cube.data
log = cmdprov.new_log(infile_history={inargs.infiles[0]: history[0]}, git_repo=repo_dir)
outcube.attributes['history'] = log
iris.save(outcube, inargs.outfile)
def get_data(infiles, var, data_type, time_constraint, agg_method, pct=False):
"""Get the data for a particular model"""
cube_list = iris.cube.CubeList([])
for ensnum, infile in enumerate(infiles):
cube, history = gio.combine_files(infile, var, new_calendar='365_day')
assert cube.units == 'kg', f'{infile} units not kg'
if time_constraint:
cube = cube.extract(time_constraint)
if data_type == 'cumulative_anomaly':
cube.data = cube.data - cube.data[0, ::]
cube = cube[-1, ::]
elif data_type == 'climatology':
cube = cube.collapsed('time', iris.analysis.MEAN)
cube.remove_coord('time')
new_aux_coord = iris.coords.AuxCoord(ensnum,
long_name='ensemble_member',
units='no_unit')
cube.add_aux_coord(new_aux_coord)
cube_list.append(cube)
if len(cube_list) > 1:
ens_cube = ensagg.calc_ensagg(cube_list, operator=agg_method)
else:
ens_cube = cube_list[0]
basins = [
'Atlantic', 'Pacific', 'Indian', 'Arctic', 'Marginal Seas', 'Land',
'Ocean', 'Globe'
]
pe_regions = ['SH-P', 'SH-E', 'T-P', 'NH-E', 'NH-P', 'Globe']
df = pd.DataFrame(ens_cube.data, columns=basins, index=pe_regions)
#df.loc['Globe', 'Globe'] = np.nan
df = df.iloc[::-1]
if pct:
if var in [
'precipitation_minus_evaporation_flux',
'water_flux_into_sea_water'
]:
total_pos = df['Globe']['NH Precip'] + df['Globe'][
'Tropical Precip'] + df['Globe']['SH Precip']
total_neg = abs(df['Globe']['NH Evap'] + df['Globe']['SH Evap'])
df = df / max(total_pos, total_neg)
else:
df = df / df['Globe']['Globe']
return df, history
def get_data(infile, var, time_constraint):
"""Get the data for a particular model"""
cube, history = gio.combine_files(infile, var, new_calendar='365_day')
cube = cube.extract(time_constraint)
iris.coord_categorisation.add_year(cube, 'time')
assert cube.shape[-1] == 8
data_array = np.zeros([cube.shape[0], 4])
data_array[:, 0] = cube.data[:, -1, 0] #atlantic
data_array[:, 1] = cube.data[:, -1, 2] #indian
data_array[:, 2] = cube.data[:, -1, 1] #pacific
data_array[:, 3] = cube.data[:, -1, 5] #land
start_data = data_array[0, :]
anomaly_data = data_array - start_data
return cube, anomaly_data, start_data
def main(inargs):
"""Run the program."""
temperature_cube, history = gio.combine_files(inargs.temperature_files,
inargs.var)
temperature_atts = temperature_cube.attributes
metadata_dict = {inargs.temperature_files[0]: history[0]}
level_subset = gio.iris_vertical_constraint(inargs.min_depth,
inargs.max_depth)
temperature_cube = temperature_cube.extract(level_subset)
if inargs.annual:
temperature_cube = timeseries.convert_to_annual(temperature_cube,
chunk=inargs.chunk)
if inargs.regrid:
area_cube = read_area_file(inargs.regrid)
temperature_cube, coord_names, regrid_status = grids.curvilinear_to_rectilinear(
temperature_cube, weights=area_cube.data)
volume_data = spatial_weights.volume_array(temperature_cube)
grid = 'y72x144'
else:
assert inargs.volume_file, "Must provide volume file if not regridding data"
volume_data, metadata_dict = get_volume(inargs.volume_file,
temperature_cube, level_subset,
metadata_dict)
coord_names = [coord.name() for coord in temperature_cube.dim_coords]
grid = None
ohc_cube = ohc(temperature_cube,
volume_data,
inargs.density,
inargs.specific_heat,
coord_names,
vertical_integral=inargs.vertical_integral,
chunk=inargs.chunk)
ohc_cube = add_metadata(temperature_cube, temperature_atts, ohc_cube,
inargs)
log = cmdprov.new_log(infile_history=metadata_dict, git_repo=repo_dir)
ohc_cube.attributes['history'] = log
iris.save(ohc_cube, inargs.outfile)
def read_global_variable(model, variable, ensemble, manual_files):
"""Read data for a global variable"""
manual = file_match(manual_files, model, variable, ensemble)
if manual or variable == 'massa':
file_list = manual
else:
file_list = clef_search(model, variable, ensemble)
if file_list:
cube, history = gio.combine_files(file_list, names[variable])
cube = timeseries.convert_to_annual(cube)
cube = time_check(cube)
extra_log.append(file_list)
else:
cube = None
return cube
def main(inargs):
"""Run the program."""
prw_cube, history = gio.combine_files(inargs.prw_files, 'atmosphere_mass_content_of_water_vapor')
if inargs.area_file:
area_cube = iris.load_cube(inargs.area_file, 'cell_area')
weights = uconv.broadcast_array(area_cube.data, [1, 2], prw_cube.shape)
else:
weights = spatial_weights.area_array(prw_cube)
coord_names = [coord.name() for coord in prw_cube.dim_coords]
massa_cube = prw_cube.collapsed(coord_names[1:], iris.analysis.SUM, weights=weights)
units = str(massa_cube.units)
massa_cube.units = units.replace('m-2', '')
massa_cube.remove_coord(coord_names[1])
massa_cube.remove_coord(coord_names[2])
massa_cube.attributes['history'] = cmdprov.new_log(git_repo=repo_dir)
iris.save(massa_cube, inargs.outfile)
def get_cube_list(infiles, agg, time_bounds=None, quick=False):
"""Read and process data."""
assert agg in ['clim', 'anom']
ensemble_cube_list = iris.cube.CubeList([])
for ensnum, ensemble_member in enumerate(infiles):
print(ensemble_member)
cube, history = gio.combine_files(
ensemble_member,
'precipitation_minus_evaporation_flux',
new_calendar='365_day')
cube = gio.check_time_units(cube)
if time_bounds:
time_constraint = gio.get_time_constraint(time_bounds)
cube = cube.extract(time_constraint)
elif quick:
cube = cube[0:120, ::]
if agg == 'clim':
cube = timeseries.convert_to_annual(cube,
aggregation='mean',
days_in_month=True)
cube = cube.collapsed('time', iris.analysis.MEAN)
elif agg == 'anom':
start_data = cube.data[0, ::]
cube = cube[-1, ::]
cube.data = cube.data - start_data
cube.remove_coord('time')
cube = regrid(cube)
new_aux_coord = iris.coords.AuxCoord(ensnum,
long_name='ensemble_member',
units='no_unit')
cube.add_aux_coord(new_aux_coord)
cube.cell_methods = ()
ensemble_cube_list.append(cube)
print("Total number of models:", len(ensemble_cube_list))
return ensemble_cube_list, history
def main(inargs):
"""Run the program."""
agg_functions = {'mean': iris.analysis.MEAN, 'sum': iris.analysis.SUM}
metadata_dict = {}
basin_cube = iris.load_cube(inargs.basin_file, 'region')
assert basin_cube.data.min() == 11
assert basin_cube.data.max() == 18
basin_numbers = numpy.array([1, 2, 3])
metadata_dict[inargs.basin_file] = basin_cube.attributes['history']
flag_values = '0 1 2'
flag_meanings = 'atlantic indo-pacific globe'
basin_coord = iris.coords.DimCoord(basin_numbers,
standard_name=basin_cube.standard_name,
long_name=basin_cube.long_name,
var_name=basin_cube.var_name,
units=basin_cube.units,
attributes={'flag_values': flag_values,
'flag_meanings': flag_meanings})
if inargs.weights:
weights_cube = gio.get_ocean_weights(inargs.weights)
metadata_dict[inargs.weights] = weights_cube.attributes['history']
output_cubelist = iris.cube.CubeList([])
for infile in inargs.infiles:
print(infile)
if inargs.var == 'ocean_volume':
cube = gio.get_ocean_weights(infile)
history = [cube.attributes['history']]
else:
cube, history = gio.combine_files(infile, inargs.var, checks=True)
assert cube.ndim in [3, 4]
coord_names = [coord.name() for coord in cube.dim_coords]
if inargs.annual:
cube = timeseries.convert_to_annual(cube, chunk=inargs.chunk)
assert basin_cube.shape == cube.shape[-2:]
basin_array = uconv.broadcast_array(basin_cube.data, [cube.ndim - 2, cube.ndim - 1], cube.shape)
basin_masks = {'atlantic': basin_array > 12,
'indo-pacific': (basin_array < 13) | (basin_array > 15),
'globe': basin_array > 16}
if inargs.weights:
assert weights_cube.data.shape == cube.shape[-3:]
if cube.ndim == 4:
weights_array = uconv.broadcast_array(weights_cube.data, [1, 3], cube.shape)
else:
weights_array = weights_cube.data
else:
weights_array = None
if cube.ndim == 3:
outdata = numpy.ma.zeros([cube.shape[0], len(basin_numbers)])
else:
outdata = numpy.ma.zeros([cube.shape[0], cube.shape[1], len(basin_numbers)])
for basin_index, basin_name in enumerate(['atlantic', 'indo-pacific', 'globe']):
temp_cube = cube.copy()
mask = basin_masks[basin_name]
temp_cube.data = numpy.ma.masked_where(mask, temp_cube.data)
if len(coord_names) == cube.ndim:
horiz_agg = temp_cube.collapsed(coord_names[-2:], agg_functions[inargs.agg], weights=weights_array).data
elif inargs.agg == 'mean':
horiz_agg = numpy.ma.average(temp_cube.data, axis=(-2,-1), weights=weights_array)
elif inargs.agg == 'sum':
horiz_agg = numpy.ma.sum(temp_cube.data, axis=(-2,-1))
if outdata.ndim == 2:
outdata[:, basin_index] = horiz_agg
else:
outdata[:, :, basin_index] = horiz_agg
coord_list = [(cube.dim_coords[0], 0)]
if cube.ndim == 4:
coord_list.append((cube.dim_coords[1], 1))
coord_list.append((basin_coord, 2))
else:
coord_list.append((basin_coord, 1))
outcube = iris.cube.Cube(outdata,
standard_name=cube.standard_name,
long_name=cube.long_name,
var_name=cube.var_name,
units=cube.units,
attributes=cube.attributes,
dim_coords_and_dims=coord_list)
output_cubelist.append(outcube)
iris.util.equalise_attributes(output_cubelist)
iris.util.unify_time_units(output_cubelist)
outcube = output_cubelist.concatenate_cube()
if history:
metadata_dict[inargs.infiles[-1]] = history[0]
outcube.attributes['history'] = cmdprov.new_log(infile_history=metadata_dict, git_repo=repo_dir)
iris.save(outcube, inargs.outfile)
def main(args):
"""Run the program."""
infiles = sorted(args.infiles)
var_names = {
'precipitation_flux': 'precipitation',
'water_evapotranspiration_flux': 'evaporation',
'precipitation_minus_evaporation_flux': 'P-E'
}
var_name = var_names[args.var]
ylabel = f"annual mean {var_name} (kg)"
fig, axes = plt.subplots(1, 2, figsize=(12, 6))
region_data = []
basin_data = []
nfiles = len(infiles)
print(f"Number models: {nfiles}")
for modelnum in range(nfiles):
cube, history = gio.combine_files(infiles[modelnum], args.var)
cube = cube.collapsed('time', iris.analysis.MEAN)
cube.remove_coord('time')
try:
model = cube.attributes['model_id']
except KeyError:
model = cube.attributes['source_id']
# cube(time, pereg, basin)
# pereg: SH_precip SH_evap tropical_precip NH_evap NH_precip globe
# basin: atlantic pacific indian arctic marginal_seas land ocean globe
region_data.append([model, 'SH-P', cube.data[0, -1]])
region_data.append([model, 'SH-E', cube.data[1, -1]])
region_data.append([model, 'T-P', cube.data[2, -1]])
region_data.append([model, 'NH-E', cube.data[3, -1]])
region_data.append([model, 'NH-P', cube.data[4, -1]])
basin_data.append([model, 'Atlantic', cube.data[-1, 0]])
basin_data.append([model, 'Indian', cube.data[-1, 2]])
basin_data.append([model, 'Pacific', cube.data[-1, 1]])
basin_data.append([model, 'Land', cube.data[-1, 5]])
region_df = pd.DataFrame(region_data,
columns=['model', 'P-E region', ylabel])
basin_df = pd.DataFrame(basin_data, columns=['model', 'basin', ylabel])
plot_data(axes[0],
region_df,
'P-E region',
ylabel,
'(a) meridional climatology',
model_dots=args.dots)
plot_data(axes[1],
basin_df,
'basin',
ylabel,
'(b) zonal climatology',
model_dots=args.dots)
plt.savefig(args.outfile, bbox_inches='tight')
metadata_dict = {infiles[-1]: history[0]}
log_text = cmdprov.new_log(infile_history=metadata_dict, git_repo=repo_dir)
log_file = re.sub('.png', '.met', args.outfile)
cmdprov.write_log(log_file, log_text)
