Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def prepare_month_sarah(era5_func, xs, ys, year, month, template_sis, template_sid, resolution):
with xr.open_mfdataset(template_sis.format(year=year, month=month)) as ds_sis, \
xr.open_mfdataset(template_sid.format(year=year, month=month)) as ds_sid:
ds = xr.merge([ds_sis, ds_sid])
ds = _rename_and_clean_coords(ds, add_lon_lat=False)
ds = ds.sel(x=as_slice(xs), y=as_slice(ys))
def interpolate(ds, dim='time'):
def _interpolate1d(y):
nan = np.isnan(y)
if nan.all(): return y
x = lambda z: z.nonzero()[0]
y[nan] = np.interp(x(nan), x(~nan), y[~nan])
return y
def _interpolate(a):
return a.map_blocks(partial(np.apply_along_axis, _interpolate1d, -1), dtype=a.dtype)
data_vars = ds.data_vars.values() if isinstance(ds, xr.Dataset) else (ds,)
resolutions = ["6min", "2min"]
if resolution not in resolutions:
raise ValueError(
"Invalid seafloor age grid resolution '{}'. Must be one of {}.".format(
resolution, resolutions
)
)
fname_age = REGISTRY.fetch(
"age.3.{}.nc.bz2".format(resolution[0]), processor=Decompress()
)
fname_error = REGISTRY.fetch(
"ageerror.3.{}.nc.bz2".format(resolution[0]), processor=Decompress()
)
if not load:
return [fname_age, fname_error]
grid = xr.merge(
[
xr.open_dataset(fname_age, **kwargs).rename(z="age") / 100,
xr.open_dataset(fname_error, **kwargs).rename(z="uncertainty") / 100,
]
)
# Add more metadata and fix some names
grid = grid.rename(x="longitude", y="latitude")
grid.attrs["title"] = "Age of oceanic lithosphere"
grid.attrs["doi"] = "10.1029/2007GC001743"
grid.age.attrs["long_name"] = "Age of oceanic lithosphere"
grid.age.attrs["units"] = "million_years"
grid.uncertainty.attrs["long_name"] = "Age uncertainty"
grid.uncertainty.attrs["units"] = "million_years"
return grid
This get the data for a set of features from a module. All modules in
`atlite.datasets` are allowed.
"""
parameters = cutout.data.attrs
lock = SerializableLock()
datasets = []
get_data = datamodules[module].get_data
for feature in features:
feature_data = get_data(cutout, feature, tmpdir=tmpdir, lock=lock, **parameters)
datasets.append(feature_data)
if len(datasets) >= 1 and isinstance(datasets[0], Delayed):
datasets = dask.compute(*datasets)
ds = xr.merge(datasets, compat='equals')
for v in ds:
ds[v].attrs['module'] = module
return ds
{
var_name: "/".join([group_name, var_name])
for var_name in group.variables
},
)
# Add the group name also to the dimensions:
group = group.rename({
dim: "/".join([group_name, dim])
for dim in group.dims
if dim not in group.coords
})
datasets.append(group)
return xarray.merge(datasets)
print ('we had files but did not load any matching %s ' % \
var_to_load)
else:
# update some metadata for fun.
g = update_ecco_dataset_geospatial_metadata(g)
g = update_ecco_dataset_temporal_coverage_metadata(g)
# Add this DataArray to the DataSet g_all
# -- if this is the first DataArray, then make g_all, g
if len(g_all) == 0:
g_all = g
# otherwise merge g into g_all
else:
g_all = xr.merge((g_all,g))
# when you merge DataArrays or DataSets some metadata can
# be lost. This bit ensures that all the metadata is carried
# over to g_all
for g_attr_key in g.attrs.keys():
if g_attr_key not in g_all.attrs.keys():
g_all.attrs[g_attr_key] = g.attrs[g_attr_key]
# finished loading all variable DataArrays
# update metadata again
if len(g_all) > 0:
g_all = update_ecco_dataset_geospatial_metadata(g_all)
g_all = update_ecco_dataset_temporal_coverage_metadata(g_all)
g_all.attrs = OrderedDict(sorted(g_all.attrs.items()))
else:
print('we ended up with nothing loaded!')
# grain size and mean sediment density calculations
# FUTURE: make it possible to circumvent this if mean grain size
# has already been calculated (e.g. during 'zeroing' runs)
# Calculate mean values for density and grain size (weighted by volume).
sel_parcels = current_parcels.where(
current_parcels.element_id != self.OUT_OF_NETWORK
)
d_weighted = sel_parcels.D * sel_parcels.volume
rho_weighted = sel_parcels.density * sel_parcels.volume
d_weighted.name = "d_weighted"
rho_weighted.name = "rho_weighted"
grouped_by_element = xr.merge(
(sel_parcels.element_id, sel_parcels.volume, d_weighted, rho_weighted)
).groupby("element_id")
d_avg = grouped_by_element.sum().d_weighted / grouped_by_element.sum().volume
rho_avg = (
grouped_by_element.sum().rho_weighted / grouped_by_element.sum().volume
)
self._d_mean_active = np.zeros(self._grid.size("link"))
self._d_mean_active[d_avg.element_id.values.astype(int)] = d_avg.values
self._rhos_mean_active = np.zeros(self._grid.size("link"))
self._rhos_mean_active[rho_avg.element_id.values.astype(int)] = rho_avg.values
format(spc, lspc))
continue
if verbose:
print(" -> adding {} with scale {}".\
format(spc, lspc_dict[lspc][spc]))
darr.values = darr.values + \
ds_new[varname].values * lspc_dict[lspc][spc]
num_spc = num_spc + 1
# Replace values with NaN is no species found in dataset
if num_spc == 0:
print('No constituent species found in file. Setting to NaN.')
darr.values = np.full(darr.shape, np.nan)
# Merge new variable into dataset
ds_new = xr.merge([ds_new, darr])
return ds_new
#*******************************************************************************
print('Concatenating files')
def get_datespan(ds_all):
"""
:param ds_all:
:return: date_range a string of values to use in the filename of the model
"""
year_start = ds_all.time.data[0].astype('datetime64[Y]').astype(int) + 1970
year_end = ds_all.time.data[len(ds_all.time.data)-1].astype('datetime64[Y]').astype(int) + 1970
month_start = ds_all.time.data[0].astype('datetime64[M]').astype(int) % 12 + 1
month_end = ds_all.time.data[len(ds_all.time.data)-1].astype('datetime64[M]').astype(int) % 12 + 1
date_range = str(year_start)+str(month_start).zfill(2)+'_'+str(year_end)+str(month_end).zfill(2)
return date_range
ds_all = xr.merge([xr.open_dataset(f) for f in flist])
print('- Done')
#*******************************************************************************
#Creating output file
#*******************************************************************************
print('Creating output file')
ds_all.to_netcdf(fout)
print('- All monthly files from ' + MODEL + ' combined and saved to: ' + fout)
update_node_at_cell(self.ds, node_at_cell)
update_nodes_at_face(self.ds, nodes_at_face)
rename = {
"mesh": "dual",
"node": "corner",
"link": "face",
"patch": "cell",
"x_of_node": "x_of_corner",
"y_of_node": "y_of_corner",
"nodes_at_link": "corners_at_face",
"links_at_patch": "faces_at_cell",
"max_patch_links": "max_cell_faces",
}
self._ds = xr.merge([self._ds, self._dual.ds.rename(rename)])
self._origin = (0., 0.)
self._frozen = False
self.freeze()
if kwds.get("sort", True):
self.sort()
fnames = [
REGISTRY.fetch(
"{}_slab2_{}.grd".format(ZONES[zone]["fname_indicator"], dataset)
)
for dataset in DATASETS
]
if not load:
return fnames
arrays = [xr.open_dataarray(f).rename(x="longitude", y="latitude") for f in fnames]
for array, dataset in zip(arrays, DATASETS):
array.name = dataset
# Change long_name and add units of each array
array.attrs["long_name"] = DATASETS[dataset]["name"]
array.attrs["units"] = DATASETS[dataset]["units"]
# Merge arrays into a single xr.Dataset
grid = xr.merge(arrays)
# Change units of thickness, depth and depth_uncertainty to meters
# Also change units of the actual_range attribute to meters
for field in ("thickness", "depth", "depth_uncertainty"):
grid[field] *= 1000
grid[field].attrs["actual_range"] *= 1000
# Change long_name and units of longitude and latitude coords
grid.longitude.attrs["long_name"] = "Longitude"
grid.longitude.attrs["units"] = "degrees"
grid.latitude.attrs["long_name"] = "Latitude"
grid.latitude.attrs["units"] = "degrees"
# Add attributes to the xr.Dataset
grid.attrs.update(
{
"title": "Slab2 model - Zone: {}".format(ZONES[zone]["name"]),
"zone": zone,
"zone_full_name": ZONES[zone]["name"],