def common_test_setup(self, shape_3d=(0, 2), data_chunks=None):
# Construct a basic testcase with all-lazy mesh_cube and submesh_cubes
# full-mesh cube shape is 'shape_3d'
# data_chunks sets chunking of source cube, (else all-1-chunk)
n_outer, n_z = shape_3d
n_mesh = 20
mesh = sample_mesh(n_nodes=20, n_edges=0, n_faces=n_mesh)
mesh_cube = sample_mesh_cube(n_z=n_z, mesh=mesh)
# Fix index-coord name to the expected default for recombine_submeshes.
mesh_cube.coord("i_mesh_face").rename("i_mesh_index")
if n_outer:
# Crudely merge a set of copies to build an outer dimension.
mesh_cube.add_aux_coord(AuxCoord([0], long_name="outer"))
meshcubes_2d = []
for i_outer in range(n_outer):
cube = mesh_cube.copy()
cube.coord("outer").points = np.array([i_outer])
meshcubes_2d.append(cube)
mesh_cube = CubeList(meshcubes_2d).merge_cube()
if not data_chunks:
data_chunks = mesh_cube.shape[:-1] + (-1, )
mesh_cube.data = da.zeros(mesh_cube.shape, chunks=data_chunks)
n_regions = 4 # it doesn't divide neatly
region_len = n_mesh // n_regions
i_points = np.arange(n_mesh)
region_inds = [
np.where((i_points // region_len) == i_region)
for i_region in range(n_regions)
]
# Disturb slightly to ensure some gaps + some overlaps
region_inds = [list(indarr[0]) for indarr in region_inds]
region_inds[2] = region_inds[2][:-2] # missing points
region_inds[3] += region_inds[1][:2] # duplicates
self.mesh_cube = mesh_cube
self.region_inds = region_inds
self.region_cubes = [mesh_cube[..., inds] for inds in region_inds]
for i_cube, cube in enumerate(self.region_cubes):
for i_z in range(n_z):
# Set data='z' ; don't vary over other dimensions.
cube.data[..., i_z, :] = i_cube + 1000 * i_z + 1
cube.data = cube.lazy_data()
# Also construct an array to match the expected result (2d cases only).
# basic layer showing region allocation (large -ve values for missing)
expected = np.array([1.0, 1, 1, 1, 1] +
[4, 4] # points in #1 overlapped by #3
+ [2, 2, 2] + [3, 3, 3] +
[-99999, -99999] # missing points
+ [4, 4, 4, 4, 4])
# second layer should be same but +1000.
# NOTE: only correct if shape_3d=None; no current need to generalise this.
expected = np.stack([expected, expected + 1000])
# convert to masked array with missing points.
expected = np.ma.masked_less(expected, 0)
self.expected_result = expected
