def test_aggregate_input_da(dataarray_2d_example):
blocks = [("i", 3), ("j", 3)]
with pytest.raises(RuntimeError):
aggregate(dataarray_2d_example.compute(), blocks, func=np.nanmean)
def test_aggregate_input_blocks(dataarray_2d_example, blocks_fail):
with pytest.raises(RuntimeError):
aggregate(dataarray_2d_example, blocks_fail, func=np.nanmean)
def test_aggregate_regular_func(dataarray_2d_example, func, expected_result):
blocks = [("i", 3), ("j", 3)]
a = aggregate(dataarray_2d_example, blocks, func=func)
assert_allclose(a.data.compute(), expected_result)
def test_aggregate_regular_blocks(
dataarray_2d_example, blocks, expected_result
):
func = np.nanmean
a = aggregate(dataarray_2d_example, blocks, func=func)
assert_allclose(a.data, expected_result)
def KOC_Full(snap,
mean,
validfile,
tr_num,
bins,
kappa=63,
method='LT',
reset_frq=None,
reset_pha=None,
dt_model=None,
dt_tracer=None,
cut_time=7776000,
ref_date='No date',
debug=False):
# !!! totally hacky...this needs to be replaced.
axis_bins = [('X', bins[0][1]), ('Y', bins[0][1])]
area = mean.rA
area_sum = aggregate(area, bins, func=np.sum)
landmask_w = mean.hFacW.data != 0
landmask_s = mean.hFacS.data != 0
landmask_c = mean.hFacC != 0
landmask = np.logical_and(np.logical_and(landmask_w, landmask_s),
landmask_c)
validmask = xr.DataArray(readbin(validfile, area.shape),
dims=area.dims,
coords=area.coords)
mask = np.logical_and(validmask, landmask)
check_KOC_input(mean, snap, tr_num, method)
# snapshots averaged in space
if method == 'L':
data = snap
grid = xgcm.Grid(data)
grid_coarse = xgcm.Grid(grid_aggregate(grid._ds, axis_bins))
# Numerator
q = data['TRAC' + tr_num]
q_grad_sq = gradient_sq_amplitude(grid, q)
q_grad_sq_coarse = custom_coarse(q_grad_sq, area, bins, mask)
n = q_grad_sq_coarse
# Denominator
q_coarse = custom_coarse(q, area, bins, mask)
q_coarse_grad_sq = gradient_sq_amplitude(grid_coarse, q_coarse)
d = q_coarse_grad_sq
elif method == 'T':
data = mean
grid = xgcm.Grid(data)
grid_coarse = xgcm.Grid(grid_aggregate(grid._ds, axis_bins))
# Numerator
q_gradx_sq_mean = data['DXSqTr' + tr_num]
q_grady_sq_mean = data['DYSqTr' + tr_num]
q_grad_sq_mean = grid.interp(q_gradx_sq_mean, 'X') + \
grid.interp(q_grady_sq_mean, 'Y')
n = q_grad_sq_mean
# !!! this is not the right way to do it but its the same way ryan did
n = custom_coarse(n, area, bins, mask)
# Denominator
q_mean = data['TRAC' + tr_num]
q_mean_grad_sq = gradient_sq_amplitude(grid, q_mean)
d = q_mean_grad_sq
# !!! this is not the right way to do it but its the same way ryan did
d = custom_coarse(d, area, bins, mask)
elif method == 'LT':
data = mean
grid = xgcm.Grid(data)
grid_coarse = xgcm.Grid(grid_aggregate(grid._ds, axis_bins))
# Numerator
q_gradx_sq_mean = data['DXSqTr' + tr_num]
q_grady_sq_mean = data['DYSqTr' + tr_num]
q_grad_sq_mean = grid.interp(q_gradx_sq_mean, 'X') + \
grid.interp(q_grady_sq_mean, 'Y')
n = custom_coarse(q_grad_sq_mean, area, bins, mask)
# Denominator
q_mean = data['TRAC' + tr_num]
q_mean_coarse = custom_coarse(q_mean, area, bins, mask)
q_mean_grad_sq = gradient_sq_amplitude(grid_coarse, q_mean_coarse)
d = q_mean_grad_sq
# Calculate the gradient criterion
crit_q_mean = custom_coarse(data['TRAC' + tr_num], area, bins, mask)
crit_q_sq_mean = custom_coarse(data['TRACSQ' + tr_num], area, bins, mask)
crit_dict = gradient_criterion(grid_coarse, crit_q_mean, crit_q_sq_mean)
# Export the 'raw tracer fields' ###
raw = data['TRAC' + tr_num]
raw_coarse = custom_coarse(raw, area, bins, mask)
raw.coords['landmask'] = landmask
raw.coords['validmask'] = validmask
raw.coords['mask'] = mask
raw.validmask.attrs['long_name'] = 'Mask for valid Aviso data points'
raw.landmask.attrs['long_name'] = 'Land mask'
raw.mask.attrs['long_name'] = 'combination of land and validmask'
# Final edits for output
koc = n / d * kappa
# Count of aggregated valid cells per output pixel.
mask.data = da.from_array(mask.data, mask.data.shape)
mask_count = aggregate(mask, bins, func=np.sum)
# replace with coarse grid coords
co = matching_coords(grid_coarse._ds, koc.dims)
# !!! this is not necessarily enough (this needs to be automated to chose
# only the corrds with all dims matching i,g,time)
d = xr.DataArray(d.data, coords=co, dims=d.dims)
n = xr.DataArray(n.data, coords=co, dims=n.dims)
koc = xr.DataArray(koc.data, coords=co, dims=koc.dims)
raw_coarse = xr.DataArray(raw_coarse.data, coords=co, dims=raw_coarse.dims)
ds = xr.Dataset({
'KOC': koc,
'Numerator': n,
'Denominator': d,
'AveTracer': raw_coarse,
'RMSTracer': crit_dict['q_rms'],
'GradTracer': crit_dict['q_grad_abs'],
'gradient_criterion': crit_dict['crit']
})
ds.coords['mask_count'] = mask_count
ds.coords['area'] = area_sum
# Add attributes to ds
ds.KOC.attrs['long_name'] = 'Osborn-Cox Diffusivity'
ds.AveTracer.attrs['long_name'] = 'Coarsened Tracer'
ds.RMSTracer.attrs['long_name'] = 'Coarsened Tracer Variance'
ds.GradTracer.attrs['long_name'] = 'Coarsened Tracer Gradient (absolute)'
ds.gradient_criterion.attrs['long_name'] = \
'Gradient Criterion'
ds.Numerator.attrs['long_name'] = 'Mixing Enhancement'
ds.Denominator.attrs['long_name'] = 'Background Mixing'
ds.mask_count.attrs['long_name'] = 'number of ocean data points before \
coarsening'
# Determine reset properties
val_idx, _, _ = reset_cut(reset_frq, reset_pha, dt_model, dt_tracer,
mean.iter.data, cut_time)
ds.coords['valid_index'] = (['time'], val_idx)
ds['valid_index'].attrs = {'Description': 'Time mask eliminating spin up'}
return ds, raw
