def bw_sum(image, block_shape, keep_shape=False):
"""
Blockwise sum.
Modified from
https://github.com/dask/dask-image/pull/148#discussion_r444649473
"""
# add up and subtract shifted copy, like np.diff but shifted
integral_image = image.copy()
for ax in range(image.ndim):
integral_image = da.cumsum(integral_image, axis=ax)
window_sums = integral_image
for ax in range(image.ndim):
window_sums -= shift(window_sums, block_shape[ax], ax, fill_value=0)
# now sum is at the corner of each block, slice to get it
s = ['{}::{}'.format(d - 1, d) for d in block_shape]
window_sums = eval('window_sums[' + ', '.join(s) + ']')
# repeat to get same shape back
if keep_shape:
window_sums = repeat_block(window_sums, block_shape)
return window_sums
def count(a):
"""counts number of distinct "on" switches in a boolean 1d array
and assigns them cumulatively
"""
a = a.astype(int)
count = da.cumsum(da.insert(da.maximum(a[1:] - a[:-1], 0), a[0], 0,
axis=0),
axis=0)
return da.where(a == True, count, 0)
def dask_relabel_chunks(A):
"""
Relabel all the the chunks of an input array. It is assumed
that a `map_blocks` or `map_overlap`, has been previously applied
over `A`, and each of the chunks of `A` contains a local labelling.
The function labels all the chunks so that all the labels all
globally independent.
E.g. for an input dataset with 3 different chunks of local labels:
Chunk1: [1, 0, 0] # Maxlabel = 1
Chunk2: [0, 2, 1] # Maxlabel = 2
Chunk3: [2, 0, 1] # Maxlabel = 2
The relabelling of the chunks would look like:
newChunk1: [1, 0, 0]
newChunk2: [2, 4, 3] # Chunk2 + Maxlabel(newChunk1) + 1
newChunk3: [7, 5, 6] # Chunk3 + Maxlabel(newChunk2) + 1
Parameters
----------
A: dask.Array
An input array to be relabeled
Returns
-------
B: dask.Array
Dask array of the same shape, with chunks relabelled.
"""
inds = tuple(range(A.ndim))
max_per_block = da.blockwise(
np.max,
inds,
A,
inds,
axis=inds,
keepdims=True,
dtype=A.dtype,
adjust_chunks={i: 1
for i in inds},
)
block_index_global = da.cumsum(max_per_block.ravel()) * 3
def relabel(a, block_id=None):
bid = int(np.ravel_multi_index(block_id, A.numblocks))
if bid == 0:
return a
return a + block_index_global[bid - 1]
relabel = A.map_blocks(relabel, dtype=np.int64)
return relabel
def task_6():
# set up
N = np.logspace(2, 5, 6, dtype=np.int_)
E = np.zeros_like(N)
R = 100
# loop
for i, n in enumerate(N):
t, dt = np.linspace(0, 1, n+1, retstep=True)
dW = (da.sqrt(dt) * da.random.normal(0, 1, size=(R, n+1),
chunks=(R//2, n//2)))
W = da.cumsum(dW, axis=0) - dW[0, :]
# calculate X directly
X = EU_dask(t.shape[0], dt, R, dW)
X_exact = f_exact_dask(t, W)
# calculate error
E[i] = da.linalg.norm(X - X_exact).compute()
# plot
plt.loglog(1 / N, E, 'x-')
plt.xlabel(r"$\log 1/N$")
plt.ylabel(r"$\log E$")
plt.show()
def load_UM_cube(files,variable):
from .make_geopotential_height_coord import geopotential_height_coord,geopotential_height_coord_stag,geopotential_height_coord_short
filename_aux=os.path.join(os.path.dirname(files[0]),'LMCONSTANTS')
#Z=iris.load_cube(filename_aux,'Z')
lat=iris.load_cube(filename_aux,'lat')
lon=iris.load_cube(filename_aux,'lon')
#special treatment of accumulated precipitation if missing (in 1 min data)
print(variable)
if variable=='AccumPrecip':
cubes=iris.load(files,'pcp_accum',callback=callback)
else:
cubes=iris.load(files,variable,callback=callback)
for cube in cubes:
cube.attributes=[]
cube.coord('time').bounds=None
cube_out=cubes.concatenate_cube('time')
#special treatment of accumulated precipitation if missing (in 1 min data)
if variable=='AccumPrecip':
cube_out.data = cumsum(cube_out.core_data(), axis=0); # Make Cumulative Sum
# for 3D variables (x,y,z)
print(len(cube_out.shape))
if len(cube_out.shape) == 4:
model_level_coord=iris.coords.DimCoord(np.arange(cube_out.shape[1],0,-1),standard_name="model_level_number")
cube_out.add_dim_coord(model_level_coord,data_dim=1)
x_coord=iris.coords.DimCoord(np.arange(cube_out.shape[2]),long_name="y")
cube_out.add_dim_coord(x_coord,data_dim=2)
y_coord=iris.coords.DimCoord(np.arange(cube_out.shape[3]),long_name="x")
cube_out.add_dim_coord(y_coord,data_dim=3)
if (cube_out.shape[2]==lat.shape[1] and cube_out.shape[3]==lat.shape[2]):
lon_coord=iris.coords.AuxCoord(lon[0].data,standard_name="longitude")
cube_out.add_aux_coord(lon_coord,data_dims=(2,3))
lat_coord=iris.coords.AuxCoord(lat[0].data,standard_name="latitude")
cube_out.add_aux_coord(lat_coord,data_dims=(2,3))
#Z_coord=iris.coords.AuxCoord(Z[0].core_data(),standard_name="geopotential_height")
#cube.add_aux_coord(Z_coord,data_dims=(1,2,3))
if 'model_level_number' in [coord.name() for coord in cube_out.coords()]:
if cube_out.coord('model_level_number').shape[0]==95:
cube_out.add_aux_coord(deepcopy(geopotential_height_coord_stag),data_dims=cube_out.coord_dims('model_level_number'))
if cube_out.coord('model_level_number').shape[0]==94:
cube_out.add_aux_coord(deepcopy(geopotential_height_coord),data_dims=cube_out.coord_dims('model_level_number'))
if cube_out.coord('model_level_number').shape[0]==93:
cube_out.add_aux_coord(deepcopy(geopotential_height_coord_short),data_dims=cube_out.coord_dims('model_level_number'))
#the UM has top to bottom levels, flip around coordinates
coord=cube_out.coord('geopotential_height').points
cube_out.coord('geopotential_height').points=np.flip(coord,axis=0)
# for 2D variables (x,y)
elif len(np.shape(cube_out)) == 3:
x_coord=iris.coords.DimCoord(np.arange(cube_out.shape[1]),long_name="y")
cube_out.add_dim_coord(x_coord,data_dim=1)
y_coord=iris.coords.DimCoord(np.arange(cube_out.shape[2]),long_name="x")
cube_out.add_dim_coord(y_coord,data_dim=2)
if (cube_out.shape[1]==lat.shape[1] and cube_out.shape[2]==lat.shape[2]):
lon_coord=iris.coords.AuxCoord(lon[0].core_data(),standard_name="longitude")
cube_out.add_aux_coord(lon_coord,data_dims=(1,2))
lat_coord=iris.coords.AuxCoord(lat[0].core_data(),standard_name="latitude")
cube_out.add_aux_coord(lat_coord,data_dims=(1,2))
return cube_out