def test_get_distributed_slice_simple(self):
ompi = OcgDist()
dim = ompi.create_dimension('five', 5, dist=True, src_idx='auto')
ompi.update_dimension_bounds(min_elements=1)
with vm.scoped_by_emptyable('simple slice test', dim):
if not vm.is_null:
sub = dim.get_distributed_slice(slice(2, 4))
else:
sub = None
if sub is not None and not sub.is_empty:
self.assertEqual(sub.bounds_global, (0, 2))
else:
if dim.is_empty:
self.assertIsNone(sub)
else:
self.assertEqual(sub.bounds_global, (0, 0))
self.assertEqual(sub.bounds_local, (0, 0))
# Test global bounds are updated.
ompi = OcgDist()
dim = ompi.create_dimension('tester', 768, dist=False)
ompi.update_dimension_bounds()
sub = dim.get_distributed_slice(slice(73, 157))
self.assertEqual(sub.size, 84)
self.assertEqual(sub.bounds_global, (0, 84))
self.assertEqual(sub.bounds_local, (0, 84))
def test_create_index_variable_global(self):
raise SkipTest('not implemented')
dsrc_size = 5
ddst_size = 7
dsrc = Dimension('dsrc', dsrc_size, dist=True)
src_dist = OcgDist()
src_dist.add_dimension(dsrc)
src_dist.update_dimension_bounds()
ddst = Dimension('ddst', ddst_size, dist=True)
dst_dist = OcgDist()
dst_dist.add_dimension(ddst)
dst_dist.update_dimension_bounds()
if vm.rank == 0:
np.random.seed(1)
dst = np.random.rand(ddst_size)
src = np.random.choice(dst, size=dsrc_size, replace=False)
src = Variable(name='src', value=src, dimensions=dsrc.name)
# TODO: move create_ugid_global to create_global_index on a standard variable object
dst = GeometryVariable(name='dst', value=dst, dimensions=ddst.name)
else:
src, dst = [None] * 2
src = variable_scatter(src, src_dist)
dst = variable_scatter(dst, dst_dist)
actual = create_index_variable_global('index_array', src, dst)
self.assertNumpyAll(dst.get_value()[actual.get_value()],
src.get_value())
def test(self):
ompi = OcgDist()
src_idx = np.array([2, 3, 4, 5, 6])
dim = ompi.create_dimension('foo', size=5, group='subroot', dist=True, src_idx=src_idx)
self.assertEqual(dim, ompi.get_dimension(dim.name, group='subroot'))
self.assertEqual(dim.bounds_local, (0, len(dim)))
self.assertFalse(dim.is_empty)
ompi.update_dimension_bounds()
with self.assertRaises(ValueError):
ompi.update_dimension_bounds()
if ompi.size == 2:
if MPI_RANK == 0:
self.assertEqual(dim.bounds_local, (0, 3))
self.assertEqual(dim._src_idx.tolist(), [2, 3, 4])
elif MPI_RANK == 1:
self.assertEqual(dim.bounds_local, (3, 5))
self.assertEqual(dim._src_idx.tolist(), [5, 6])
elif ompi.size == 8:
if MPI_RANK <= 1:
self.assertTrue(len(dim) >= 2)
else:
self.assertTrue(dim.is_empty)
# Test with multiple dimensions.
d1 = Dimension('d1', size=5, dist=True, src_idx=np.arange(5, dtype=np.int32))
d2 = Dimension('d2', size=10, dist=False)
d3 = Dimension('d3', size=3, dist=True)
dimensions = [d1, d2, d3]
ompi = OcgDist()
for dim in dimensions:
ompi.add_dimension(dim)
ompi.update_dimension_bounds()
bounds_local = ompi.get_bounds_local()
if ompi.size <= 2:
desired = {(1, 0): ((0, 5), (0, 10), (0, 3)),
(2, 0): ((0, 3), (0, 10), (0, 3)),
(2, 1): ((3, 5), (0, 10), (0, 3))}
self.assertAsSetEqual(bounds_local, desired[(ompi.size, MPI_RANK)])
else:
if MPI_RANK <= 1:
self.assertTrue(dimensions[0]._src_idx.shape[0] <= 3)
for dim in dimensions:
if MPI_RANK >= 2:
if dim.name == 'd1':
self.assertTrue(dim.is_empty)
else:
self.assertFalse(dim.is_empty)
else:
self.assertFalse(dim.is_empty)
# Test adding an existing dimension.
ompi = OcgDist()
ompi.create_dimension('one')
with self.assertRaises(ValueError):
ompi.create_dimension('one')
def test_get_distributed_slice(self):
self.add_barrier = False
for d in [True, False]:
dist = OcgDist()
dim = dist.create_dimension('five', 5, dist=d, src_idx='auto')
dist.update_dimension_bounds()
if not dim.is_empty:
self.assertEqual(dim.bounds_global, (0, 5))
if dim.dist:
if MPI_RANK > 1:
self.assertTrue(dim.is_empty)
else:
self.assertFalse(dim.is_empty)
with vm.scoped_by_emptyable('dim slice', dim):
if not vm.is_null:
sub = dim.get_distributed_slice(slice(1, 3))
else:
sub = None
if dim.dist:
if not dim.is_empty:
self.assertIsNotNone(dim._src_idx)
else:
self.assertIsNone(sub)
if MPI_SIZE == 2:
desired_emptiness = {0: False, 1: True}[MPI_RANK]
desired_bounds_local = {0: (0, 2), 1: (0, 0)}[MPI_RANK]
self.assertEqual(sub.is_empty, desired_emptiness)
self.assertEqual(sub.bounds_local, desired_bounds_local)
if MPI_SIZE >= 5 and 0 < MPI_RANK > 2:
self.assertTrue(sub is None or sub.is_empty)
else:
self.assertEqual(len(dim), 5)
self.assertEqual(dim.bounds_global, (0, 5))
self.assertEqual(dim.bounds_local, (0, 5))
dist = OcgDist()
dim = dist.create_dimension('five', 5, dist=True, src_idx='auto')
dist.update_dimension_bounds()
with vm.scoped_by_emptyable('five slice', dim):
if not vm.is_null:
sub2 = dim.get_distributed_slice(slice(2, 4))
else:
sub2 = None
if MPI_SIZE == 3 and MPI_RANK == 2:
self.assertIsNone(sub2)
def test_init(self):
ompi = OcgDist(size=2)
self.assertEqual(len(ompi.mapping), 2)
dim_x = Dimension('x', 5, dist=False)
dim_y = Dimension('y', 11, dist=True)
var_tas = Variable('tas', value=np.arange(0, 5 * 11).reshape(5, 11), dimensions=(dim_x, dim_y))
thing = Variable('thing', value=np.arange(11) * 10, dimensions=(dim_y,))
vc = VariableCollection(variables=[var_tas, thing])
child = VariableCollection(name='younger')
vc.add_child(child)
childer = VariableCollection(name='youngest')
child.add_child(childer)
dim_six = Dimension('six', 6)
hidden = Variable('hidden', value=[6, 7, 8, 9, 0, 10], dimensions=dim_six)
childer.add_variable(hidden)
ompi.add_dimensions([dim_x, dim_y])
ompi.add_dimension(dim_six, group=hidden.group)
ompi.add_variables([var_tas, thing])
ompi.add_variable(hidden)
var = ompi.get_variable(hidden)
self.assertIsInstance(var, dict)
def test_create_or_get_group(self):
ocmpi = OcgDist()
_ = ocmpi._create_or_get_group_(['moon', 'base'])
_ = ocmpi._create_or_get_group_(None)
_ = ocmpi._create_or_get_group_('flower')
ocmpi.create_dimension('foo', group=['moon', 'base'])
ocmpi.create_dimension('end_of_days')
ocmpi.create_dimension('start_of_days')
actual = ocmpi.get_dimension('foo', ['moon', 'base'])
desired = Dimension('foo')
self.assertEqual(actual, desired)
desired = {0: {None: {'variables': {},
'dimensions': {'end_of_days': Dimension(name='end_of_days', size=None, size_current=None),
'start_of_days': Dimension(name='start_of_days', size=None,
size_current=None)},
'groups': {'flower': {'variables': {}, 'dimensions': {}, 'groups': {}},
'moon': {'variables': {}, 'dimensions': {},
'groups': {'base': {'variables': {},
'dimensions': {
'foo': Dimension(
name='foo',
size=None,
size_current=None)},
'groups': {}}}}}}}}
self.assertDictEqual(ocmpi.mapping, desired)
def test_create_unique_global_array(self):
dist = OcgDist()
dist.create_dimension('dim', 9, dist=True)
dist.update_dimension_bounds()
values = [[4, 2, 1, 2, 1, 4, 1, 4, 2],
[44, 25, 16, 27, 18, 49, 10, 41, 22],
[44, 25, 16, 27, 44, 49, 10, 41, 44],
[1, 1, 1, 1, 1, 1, 1, 1, 1]]
for v in values:
if vm.rank == 0:
index = Variable(name='cindex', value=v, dimensions='dim')
desired = np.unique(index.get_value())
desired_length = len(desired)
else:
index = None
index = variable_scatter(index, dist)
with vm.scoped_by_emptyable('not empty', index):
if not vm.is_null:
uvar = create_unique_global_array(index.get_value())
uvar_gathered = vm.gather(uvar)
if vm.rank == 0:
uvar_gathered = hgather(uvar_gathered)
self.assertEqual(len(uvar_gathered), desired_length)
self.assertEqual(set(uvar_gathered), set(desired))
def test_variable_gather(self):
dist = OcgDist()
three = dist.create_dimension('three', 3, src_idx=np.arange(3) * 10)
four = dist.create_dimension('four', 4, src_idx=np.arange(4, dtype=np.int32), dist=True)
dist.create_variable('four', dimensions=[three, four])
dist.update_dimension_bounds()
if MPI_RANK == 0:
np.random.seed(1)
mask_value = np.random.random(12).reshape(3, 4)
mask = Variable('mask', value=mask_value, dimensions=['three', 'four'])
else:
mask = None
mask = variable_scatter(mask, dist)
with vm.scoped('mask gather', dist.get_empty_ranks(inverse=True)):
if not vm.is_null:
mask_gather = variable_gather(mask)
else:
mask_gather = None
if MPI_RANK == 0:
self.assertNumpyAll(mask_gather.get_value(), mask_value)
self.assertNumpyAll(mask_gather.dimensions[0]._src_idx, np.arange(3) * 10)
self.assertNumpyAll(mask_gather.dimensions[1]._src_idx, np.arange(4, dtype=DataType.DIMENSION_SRC_INDEX))
for dim in mask_gather.dimensions:
self.assertFalse(dim.dist)
else:
self.assertIsNone(mask_gather)
def test_write_variable_collection_object_arrays(self):
"""Test writing variable length arrays in parallel."""
with vm.scoped('write', [0]):
if not vm.is_null:
path_actual = self.get_temporary_file_path('in.nc')
path_desired = self.get_temporary_file_path('out.nc')
value = [[1, 3, 5], [7, 9], [11]]
v = Variable(name='objects',
value=value,
fill_value=4,
dtype=ObjectType(int),
dimensions='values')
v.write(path_desired)
else:
v, path_actual, path_desired = [None] * 3
path_actual = MPI_COMM.bcast(path_actual)
path_desired = MPI_COMM.bcast(path_desired)
dest_mpi = OcgDist()
dest_mpi.create_dimension('values', 3, dist=True)
dest_mpi.update_dimension_bounds()
scattered = variable_scatter(v, dest_mpi)
outvc = VariableCollection(variables=[scattered])
with vm.scoped_by_emptyable('write', outvc):
if not vm.is_null:
outvc.write(path_actual)
if MPI_RANK == 0:
self.assertNcEqual(path_actual, path_desired)
def test_variable_scatter(self):
var_value = np.arange(5, dtype=float) + 50
var_mask = np.array([True, True, False, True, False])
dest_dist = OcgDist()
five = dest_dist.create_dimension('five', 5, src_idx=np.arange(5), dist=True)
bounds = dest_dist.create_dimension('bounds', 2)
dest_dist.update_dimension_bounds()
if MPI_RANK == 0:
local_dim = Dimension('local', 5, src_idx=np.arange(5))
dim_src_idx = local_dim._src_idx.copy()
var = Variable('the_five', value=var_value, mask=var_mask, dimensions=five.name)
var.set_extrapolated_bounds('the_five_bounds', 'bounds')
var_bounds_value = var.bounds.get_value()
else:
var, var_bounds_value, dim_src_idx = [None] * 3
svar = variable_scatter(var, dest_dist)
var_bounds_value = MPI_COMM.bcast(var_bounds_value)
dim_src_idx = MPI_COMM.bcast(dim_src_idx)
if MPI_RANK > 1:
self.assertIsNone(svar.get_value())
self.assertTrue(svar.is_empty)
else:
dest_dim = dest_dist.get_dimension('five')
self.assertNumpyAll(var_value[slice(*dest_dim.bounds_local)], svar.get_value())
self.assertNumpyAll(var_mask[slice(*dest_dim.bounds_local)], svar.get_mask())
self.assertNumpyAll(var_bounds_value[slice(*dest_dim.bounds_local)], svar.bounds.get_value())
self.assertNumpyAll(dim_src_idx[slice(*dest_dim.bounds_local)], svar.dimensions[0]._src_idx)
self.assertNumpyAll(dim_src_idx[slice(*dest_dim.bounds_local)], svar.bounds.dimensions[0]._src_idx)
def test_variable_gather_bounded_source_index(self):
self.add_barrier = False
dist = OcgDist()
dist.create_dimension('long', 7, dist=True, src_idx='auto')
dist.create_dimension('short', 3, src_idx='auto')
dist.update_dimension_bounds()
if vm.rank == 0:
var = Variable(name='test',
value=np.arange(21).reshape(7, 3),
dimensions=['long', 'short'])
else:
var = None
svar = variable_scatter(var, dist)
with vm.scoped_by_emptyable('gather test', svar):
if vm.is_null:
return
gvar = variable_gather(svar)
if vm.rank == 0:
actual = gvar.dimensions[0]._src_idx
desired = (0, 7)
self.assertEqual(actual, desired)
else:
self.assertIsNone(gvar)
def test_variable_scatter_ndimensions(self):
r = Dimension('realization', 3)
t = Dimension('time', 365)
l = Dimension('level', 10)
y = Dimension('y', 90, dist=True)
x = Dimension('x', 360, dist=True)
dimensions = [r, t, l, y, x]
dest_mpi = OcgDist()
for d in dimensions:
dest_mpi.add_dimension(d)
dest_mpi.update_dimension_bounds()
if MPI_RANK == 0:
local_dimensions = deepcopy(dimensions)
for l in local_dimensions:
l.dist = False
var = Variable('tas', dimensions=local_dimensions)
else:
var = None
svar = variable_scatter(var, dest_mpi)
self.assertTrue(svar.dist)
self.assertIsNotNone(svar)
if MPI_SIZE == 2:
self.assertEqual(svar.shape, (3, 365, 10, 90, 180))
def get_wrap_field(crs=None, unwrapped=True):
ompi = OcgDist()
ompi.create_dimension('x', 5, dist=False)
ompi.create_dimension('y', 7, dist=True)
ompi.create_dimension('time', size_current=4, dist=False)
ompi.update_dimension_bounds()
if MPI_RANK == 0:
row = Variable(value=[-60, -40, -20, 0, 20, 40, 60], name='y', dimensions='y')
if unwrapped:
col_value = [1, 90, 180, 225, 270]
else:
col_value = [-170, -85, 0, 85, 170]
col = Variable(value=col_value, name='x', dimensions='x')
grid = Grid(col, row)
value = np.zeros((4, 7, 5))
for col_idx in range(value.shape[-1]):
value[:, :, col_idx] = col_idx
time = TemporalVariable(name='time', value=[1, 2, 3, 4], dimensions='time')
var = Variable(name='foo', value=value, dimensions=['time', 'y', 'x'])
field = Field(grid=grid, is_data=var, crs=crs, time=time)
else:
field = None
field = variable_collection_scatter(field, ompi)
return field
def test_system_parallel_write_ndvariable(self):
"""Test a parallel CSV write with a n-dimensional variable."""
ompi = OcgDist()
ompi.create_dimension('time', 3)
ompi.create_dimension('extra', 2)
ompi.create_dimension('x', 4)
ompi.create_dimension('y', 7, dist=True)
ompi.update_dimension_bounds()
if MPI_RANK == 0:
path = self.get_temporary_file_path('foo.csv')
t = TemporalVariable(name='time',
value=[1, 2, 3],
dtype=float,
dimensions='time')
t.set_extrapolated_bounds('the_time_bounds', 'bounds')
extra = Variable(name='extra', value=[7, 8], dimensions='extra')
x = Variable(name='x',
value=[9, 10, 11, 12],
dimensions='x',
dtype=float)
x.set_extrapolated_bounds('x_bounds', 'bounds')
# This will have the distributed dimension.
y = Variable(name='y',
value=[13, 14, 15, 16, 17, 18, 19],
dimensions='y',
dtype=float)
y.set_extrapolated_bounds('y_bounds', 'bounds')
data = Variable(name='data',
value=np.random.rand(3, 2, 7, 4),
dimensions=['time', 'extra', 'y', 'x'])
vc = VariableCollection(variables=[t, extra, x, y, data])
else:
path, vc = [None] * 2
path = MPI_COMM.bcast(path)
vc = variable_collection_scatter(vc, ompi)
with vm.scoped_by_emptyable('write', vc):
if not vm.is_null:
vc.write(path,
iter_kwargs={
'variable': 'data',
'followers': ['time', 'extra', 'y', 'x']
},
driver=DriverCSV)
if MPI_RANK == 0:
desired = 169
with open(path, 'r') as f:
lines = f.readlines()
self.assertEqual(len(lines), desired)
def test_redistribute_by_src_idx(self):
if vm.size != 4:
raise SkipTest('vm.size != 4')
dist = OcgDist()
dim1 = dist.create_dimension('dim1', 5 * vm.size, dist=True)
dim2 = dist.create_dimension('dim2', 2, dist=False)
dist.update_dimension_bounds()
rank_value = np.arange(5) + (10 * (vm.rank + 1))
var1 = Variable(name='dvar1', value=rank_value, dimensions=dim1)
var2 = Variable(name='dvar2', dimensions=[dim1, dim2])
var1.parent.add_variable(var2)
path = self.get_temporary_file_path('out.nc')
var1.parent.write(path)
desired_idx = np.array([1, 7, 9, 10, 14])
vdesired_value = variable_gather(var1)
if vm.rank == 0:
desired_value = vdesired_value.get_value()[desired_idx]
desired_idx_ranks = {0: slice(1, 2),
1: [2, 4],
2: [0, 4]}
rd = RequestDataset(path)
rd.metadata['dimensions'][dim1.name]['dist'] = True
field = rd.create_field()
indvar = field[var1.name]
field[var2.name].load()
try:
rank_slice = desired_idx_ranks[vm.rank]
except KeyError:
sub = Variable(is_empty=True)
else:
sub = indvar[rank_slice]
self.barrier_print(sub.is_empty)
redistribute_by_src_idx(indvar, dim1.name, sub.dimensions_dict.get(dim1.name))
with vm.scoped_by_emptyable('gather for test', indvar):
if vm.is_null:
self.assertIn(vm.rank_global, [2, 3])
else:
self.assertIn(vm.rank_global, [0, 1])
for v in [indvar, indvar.parent[var2.name]]:
self.assertIsNone(v._value)
self.assertIsNone(v._mask)
self.assertIsNone(v._is_empty)
self.assertFalse(v._has_initialized_value)
self.rank_print(indvar)
actual_value = variable_gather(indvar)
if vm.rank == 0:
actual_value = actual_value.get_value()
self.assertNumpyAll(actual_value, desired_value)
def test_update_dimension_bounds(self):
ompi = OcgDist()
dim1 = ompi.create_dimension('five', 5, dist=True)
ompi.update_dimension_bounds()
if dim1.is_empty:
desired = (0, 0)
else:
desired = (0, 5)
self.assertEqual(dim1.bounds_global, desired)
if MPI_SIZE > 1:
if MPI_SIZE == 2:
if MPI_RANK == 0:
self.assertEqual(dim1.bounds_local, (0, 3))
else:
self.assertEqual(dim1.bounds_local, (3, 5))
# Test updating on single processor.
if MPI_SIZE == 1:
ompi = OcgDist(size=2)
ompi.create_dimension('five', 5, dist=True)
ompi.update_dimension_bounds()
dim = ompi.get_dimension('five')
self.assertEqual(dim.bounds_global, (0, 5))
for rank in range(2):
actual = ompi.get_dimension('five', rank=rank)
self.assertEqual(actual.bounds_global, (0, 5))
if rank == 0:
self.assertEqual(actual.bounds_local, (0, 3))
else:
self.assertEqual(actual.bounds_local, (3, 5))
# Test two dimensions.
ompi = OcgDist(size=2)
ompi.create_dimension('lat', 64, dist=True)
ompi.create_dimension('lon', 128, dist=True)
ompi.update_dimension_bounds()
for rank in range(2):
lat = ompi.get_dimension('lat', rank=rank)
self.assertEqual(lat.bounds_local, (0, 64))
lon = ompi.get_dimension('lon', rank=rank)
if rank == 0:
self.assertEqual(lon.bounds_local, (0, 64))
else:
self.assertEqual(lon.bounds_local, (64, 128))
def test_get_mask_from_intersects(self):
poly = wkt.loads(
'POLYGON((-98.26574367088608142 40.19952531645570559,-98.71764240506330168 39.54825949367089066,-99.26257911392406186 39.16281645569620906,-99.43536392405064817 38.64446202531645724,-98.78409810126584034 38.33876582278481493,-98.23916139240508016 37.71408227848101546,-97.77397151898735217 37.67420886075949937,-97.62776898734178133 38.15268987341772799,-98.39865506329114453 38.52484177215190186,-98.23916139240508016 39.33560126582278826,-97.73409810126582897 39.58813291139241386,-97.52143987341773368 40.27927215189873777,-97.52143987341773368 40.27927215189873777,-98.26574367088608142 40.19952531645570559))'
)
desired_mask = np.array([[True, True, False, True],
[True, False, True, True],
[True, True, False, True]])
dist = OcgDist()
xdim = dist.create_dimension('x', 4, dist=True)
ydim = dist.create_dimension('y', 3)
dist.create_dimension('bounds', 2)
dist.update_dimension_bounds()
if MPI_RANK == 0:
x = self.get_variable_x()
y = self.get_variable_y()
grid = Grid(x=x, y=y, abstraction='point', crs=WGS84())
pa = get_geometry_variable(grid)
else:
pa = None
pa = variable_scatter(pa, dist)
vm.create_subcomm_by_emptyable('test_get_mask_from_intersects',
pa,
is_current=True)
if vm.is_null:
self.assertTrue(pa.is_empty)
return
usi = [False]
if env.USE_SPATIAL_INDEX:
usi.append(True)
keywords = dict(use_spatial_index=usi)
for k in self.iter_product_keywords(keywords):
ret = pa.get_mask_from_intersects(
poly, use_spatial_index=k.use_spatial_index)
desired_mask_local = desired_mask[slice(*ydim.bounds_local),
slice(*xdim.bounds_local)]
if MPI_RANK > 1:
self.assertIsNone(ret)
else:
self.assertNumpyAll(desired_mask_local, ret)
# This does not test a parallel operation.
if MPI_RANK == 0:
# Test pre-masked values in geometry are okay for intersects operation.
value = [Point(1, 1), Point(2, 2), Point(3, 3)]
value = np.ma.array(value,
mask=[False, True, False],
dtype=object)
pa2 = GeometryVariable(value=value, dimensions='ngeom')
b = box(0, 0, 5, 5)
res = pa2.get_mask_from_intersects(
b, use_spatial_index=k.use_spatial_index)
self.assertNumpyAll(res, value.mask)
def get_OcgDist_01(self):
s = Dimension('first_dist', size=5, dist=True, src_idx='auto')
ompi = OcgDist()
ompi.add_dimension(s)
ompi.create_dimension('not_dist', size=8, dist=False)
ompi.create_dimension('another_dist', size=6, dist=True)
ompi.create_dimension('another_not_dist', size=100, dist=False)
ompi.create_dimension('coordinate_reference_system', size=0)
self.assertIsNotNone(s._src_idx)
return ompi
def create_slices_for_dimension(size, splits):
ompi = OcgDist(size=splits)
dimname = 'foo'
ompi.create_dimension(dimname, size, dist=True)
ompi.update_dimension_bounds()
slices = []
for rank in range(splits):
dimension = ompi.get_dimension(dimname, rank=rank)
slices.append(dimension.bounds_local)
return slices
def test_system_grid_chunking(self):
if vm.size != 4:
raise SkipTest('vm.size != 4')
from ocgis.spatial.grid_chunker import GridChunker
path = self.path_esmf_unstruct
rd_dst = RequestDataset(uri=path,
driver=DriverESMFUnstruct,
crs=Spherical(),
grid_abstraction='point',
grid_is_isomorphic=True)
rd_src = deepcopy(rd_dst)
resolution = 0.28125
chunk_wd = os.path.join(self.current_dir_output, 'chunks')
if vm.rank == 0:
os.mkdir(chunk_wd)
vm.barrier()
paths = {'wd': chunk_wd}
gc = GridChunker(rd_src,
rd_dst,
nchunks_dst=[8],
src_grid_resolution=resolution,
dst_grid_resolution=resolution,
optimized_bbox_subset=True,
paths=paths,
genweights=True)
gc.write_chunks()
dist = OcgDist()
local_ctr = Dimension(name='ctr', size=8, dist=True)
dist.add_dimension(local_ctr)
dist.update_dimension_bounds()
for ctr in range(local_ctr.bounds_local[0], local_ctr.bounds_local[1]):
ctr += 1
s = os.path.join(chunk_wd, 'split_src_{}.nc'.format(ctr))
d = os.path.join(chunk_wd, 'split_dst_{}.nc'.format(ctr))
sf = Field.read(s, driver=DriverESMFUnstruct)
df = Field.read(d, driver=DriverESMFUnstruct)
self.assertGreater(sf.grid.shape[0], df.grid.shape[0])
wgt = os.path.join(chunk_wd, 'esmf_weights_{}.nc'.format(ctr))
f = Field.read(wgt)
S = f['S'].v()
self.assertAlmostEqual(S.min(), 1.0)
self.assertAlmostEqual(S.max(), 1.0)
with vm.scoped('merge weights', [0]):
if not vm.is_null:
merged_weights = self.get_temporary_file_path(
'merged_weights.nc')
gc.create_merged_weight_file(merged_weights, strict=False)
f = Field.read(merged_weights)
S = f['S'].v()
self.assertAlmostEqual(S.min(), 1.0)
self.assertAlmostEqual(S.max(), 1.0)
def test_arange_from_dimension(self):
dist = OcgDist()
dim = dist.create_dimension('dim', size=7, dist=True)
dist.update_dimension_bounds()
actual = arange_from_dimension(dim, start=2, dtype=np.int64)
actual = vm.gather(actual)
if vm.rank == 0:
actual = hgather(actual)
desired = np.arange(2, 9, dtype=np.int64)
self.assertNumpyAll(actual, desired)
def create_dist(self, metadata=None):
"""
Create a distribution from global metadata. In general, this should not be overloaded by subclasses.
:param dict metadata: Global metadata to use for creating a distribution.
:rtype: :class:`ocgis.OcgDist`
"""
ompi = OcgDist(size=vm.size, ranks=vm.ranks)
# Convert metadata into a grouping consistent with the MPI dimensions.
if metadata is None:
metadata = self.metadata_source
metadata = {None: metadata}
for group_index in iter_all_group_keys(metadata):
group_meta = get_group(metadata, group_index)
# Add the dimensions to the distribution object.
dimensions = self.create_dimensions(group_meta)
# Only build a distribution if the group has more than one dimension.
if len(dimensions) == 0:
_ = ompi.get_group(group=group_index)
else:
for dimension_name, dimension_meta in list(
group_meta['dimensions'].items()):
target_dimension = dimensions[dimension_name]
target_dimension.dist = group_meta['dimensions'][
dimension_name].get('dist', False)
ompi.add_dimension(target_dimension, group=group_index)
try:
dimension_map = self.rd.dimension_map.get_group(
group_index)
except DimensionMapError:
# Likely a user-provided dimension map.
continue
# dimension_map = get_group(self.rd.dimension_map, group_index, has_root=False)
distributed_dimension_name = self.get_distributed_dimension_name(
dimension_map,
group_meta['dimensions'],
decomp_type=self.rd.decomp_type)
# Allow no distributed dimensions to be returned.
if distributed_dimension_name is not None:
for target_rank in range(ompi.size):
distributed_dimension = ompi.get_dimension(
distributed_dimension_name,
group=group_index,
rank=target_rank)
distributed_dimension.dist = True
ompi.update_dimension_bounds()
return ompi
def test_get_wrapped_state(self):
if sys.version_info.major == 3 and sys.version_info.minor == 5:
raise SkipTest('undefined behavior with Python 3.5')
ompi = OcgDist()
ompi.create_dimension('x', 5, dist=True)
ompi.create_dimension('y', 1)
ompi.update_dimension_bounds()
values = [{
'value': [-179, -90, 0, 90, 180],
'desired': WrappedState.WRAPPED
}, {
'value': [0, 90, 180, 270, 360],
'desired': WrappedState.UNWRAPPED
}, {
'value': [1, 2, 3, 4, 5],
'desired': WrappedState.UNKNOWN
}]
kwds = {'values': values, 'crs': [Spherical(), None]}
for k in self.iter_product_keywords(kwds):
ompi = deepcopy(ompi)
if MPI_RANK == 0:
vx = Variable(name='x',
value=k.values['value'],
dimensions='x')
vy = Variable(name='y', value=[0], dimensions='y')
else:
vx, vy = [None] * 2
vx = variable_scatter(vx, ompi)
vy = variable_scatter(vy, ompi)
grid = Grid(vx, vy)
field = Field(grid=grid, crs=k.crs)
with vm.scoped_by_emptyable('wrap', field):
if not vm.is_null:
wrapped_state = field.wrapped_state
else:
wrapped_state = None
if not field.is_empty:
if k.crs is None:
self.assertIsNone(wrapped_state)
else:
self.assertIsNotNone(wrapped_state)
if k.crs is None or field.is_empty:
self.assertIsNone(wrapped_state)
else:
self.assertEqual(wrapped_state, k.values['desired'])
def test_reduce_global(self):
pt = self.fixture(cindex=self.fixture_cindex(1), start_index=1)
self.assertEqual(pt.start_index, 1)
dist = OcgDist()
for d in pt.parent.dimensions.values():
d = d.copy()
if d.name == self.fixture_element_dimension.name:
d.dist = True
dist.add_dimension(d)
dist.update_dimension_bounds()
new_parent = variable_collection_scatter(pt.parent, dist)
vm.create_subcomm_by_emptyable('coordinate reduction',
new_parent,
is_current=True)
if vm.is_null:
return
pt.parent = new_parent
sub = pt.get_distributed_slice(slice(2, 5))
vm.create_subcomm_by_emptyable('distributed slice',
sub,
is_current=True)
if vm.is_null:
return
actual = sub.reduce_global()
actual_cindex = actual.cindex.extract()
actual_cindex = variable_gather(actual_cindex)
if vm.rank == 0:
actual_cindex = actual_cindex.get_value().flatten().tolist()
self.assertEqual(actual_cindex, [1, 2, 3])
gathered = [
variable_gather(c.extract()) for c in actual.coordinate_variables
]
if vm.rank == 0:
actual_coords = []
for c in gathered:
actual_coords.append(c.get_value().tolist())
desired = [[2.0, 3.0, 4.0], [8.0, 9.0, 10.0], [14.0, 15.0, 16.0]]
self.assertEqual(actual_coords, desired)
path = self.get_temporary_file_path('foo.nc')
actual.parent.write(path)
actual = Field.read(path)
self.assertEqual(actual['cindex'].attrs['start_index'], 1)
def test_update_dimension_bounds_with_source_indexing(self):
dist_size = 5
dist = OcgDist(size=dist_size)
dist.create_dimension('dim', 11, dist=True, src_idx='auto')
dist.update_dimension_bounds()
actual = []
for rank in range(dist_size):
rank_dim = dist.get_dimension('dim', rank=rank)
actual.append(rank_dim._src_idx)
desired = [(0, 3), (3, 5), (5, 7), (7, 9), (9, 11)]
self.assertEqual(actual, desired)
def test_reduce_reindex_coordinate_variables(self):
self.add_barrier = False
dist = OcgDist()
dist.create_dimension('dim', 12, dist=True)
dist.update_dimension_bounds()
global_cindex_arr = np.array([4, 2, 1, 2, 1, 4, 1, 4, 2, 5, 6, 7])
if vm.rank == 0:
var_cindex = Variable('cindex',
value=global_cindex_arr,
dimensions='dim')
else:
var_cindex = None
var_cindex = variable_scatter(var_cindex, dist)
vm.create_subcomm_by_emptyable('test', var_cindex, is_current=True)
if vm.is_null:
return
raise_if_empty(var_cindex)
coords = np.array([
0, 11, 22, 33, 44, 55, 66, 77, 88, 99, 100, 110, 120, 130, 140, 150
])
coords = Variable(name='coords', value=coords, dimensions='coord_dim')
new_cindex, u_indices = reduce_reindex_coordinate_variables(var_cindex)
desired = coords[global_cindex_arr].get_value()
if len(u_indices) > 0:
new_coords = coords[u_indices].get_value()
else:
new_coords = np.array([])
gathered_new_coords = vm.gather(new_coords)
gathered_new_cindex = vm.gather(new_cindex)
if vm.rank == 0:
gathered_new_coords = hgather(gathered_new_coords)
gathered_new_cindex = hgather(gathered_new_cindex)
actual = gathered_new_coords[gathered_new_cindex]
self.assertAsSetEqual(gathered_new_cindex.tolist(),
[2, 1, 0, 3, 4, 5])
desired_new_coords = [11, 22, 44, 55, 66, 77]
self.assertAsSetEqual(gathered_new_coords.tolist(),
desired_new_coords)
self.assertEqual(len(gathered_new_coords), len(desired_new_coords))
self.assertNumpyAll(actual, desired)
def test_update_dimension_bounds_single_simple_dimension(self):
ompi = OcgDist(size=2)
ompi.create_dimension('d1', 2, dist=True)
ompi.update_dimension_bounds(min_elements=2)
d1 = ompi.get_dimension('d1')
for rank in range(2):
actual = ompi.get_dimension(d1.name, rank=rank)
if rank == 0:
self.assertFalse(actual.is_empty)
self.assertEqual(actual.bounds_local, (0, 2))
self.assertEqual(id(d1), id(actual))
else:
self.assertTrue(actual.is_empty)
self.assertEqual(actual.bounds_local, (0, 0))
self.assertNotEqual(id(d1), id(actual))
def test_system_parallel_write_ndvariable(self):
"""Test a parallel vector GIS write with a n-dimensional variable."""
ompi = OcgDist()
ompi.create_dimension('time', 3)
ompi.create_dimension('extra', 2)
ompi.create_dimension('x', 4)
ompi.create_dimension('y', 7, dist=True)
ompi.update_dimension_bounds()
if MPI_RANK == 0:
path = self.get_temporary_file_path('foo.shp')
t = TemporalVariable(name='time', value=[1, 2, 3], dtype=float, dimensions='time')
t.set_extrapolated_bounds('the_time_bounds', 'bounds')
extra = Variable(name='extra', value=[7, 8], dimensions='extra')
x = Variable(name='x', value=[9, 10, 11, 12], dimensions='x', dtype=float)
x.set_extrapolated_bounds('x_bounds', 'bounds')
# This will have the distributed dimension.
y = Variable(name='y', value=[13, 14, 15, 16, 17, 18, 19], dimensions='y', dtype=float)
y.set_extrapolated_bounds('y_bounds', 'bounds')
data = Variable(name='data', value=np.random.rand(3, 2, 7, 4), dimensions=['time', 'extra', 'y', 'x'])
dimension_map = {'x': {'variable': 'x', 'bounds': 'x_bounds'},
'y': {'variable': 'y', 'bounds': 'y_bounds'},
'time': {'variable': 'time', 'bounds': 'the_time_bounds'}}
vc = Field(variables=[t, extra, x, y, data], dimension_map=dimension_map, is_data='data')
vc.set_abstraction_geom()
else:
path, vc = [None] * 2
path = MPI_COMM.bcast(path)
vc = variable_collection_scatter(vc, ompi)
with vm.scoped_by_emptyable('write', vc):
if not vm.is_null:
vc.write(path, driver=DriverVector)
MPI_COMM.Barrier()
desired = 168
rd = RequestDataset(path, driver=DriverVector)
sizes = MPI_COMM.gather(rd.get().geom.shape[0])
if MPI_RANK == 0:
self.assertEqual(sum(sizes), desired)
def test_variable_collection_scatter(self):
dest_mpi = OcgDist()
five = dest_mpi.create_dimension('five', 5, dist=True)
ten = dest_mpi.create_dimension('ten', 10)
dest_mpi.create_variable(name='five', dimensions=five)
dest_mpi.create_variable(name='all_in', dimensions=ten)
dest_mpi.create_variable(name='i_could_be_a_coordinate_system')
dest_mpi.update_dimension_bounds()
if MPI_RANK == 0:
var = Variable('holds_five', np.arange(5), dimensions='five')
var_empty = Variable('i_could_be_a_coordinate_system',
attrs={'reality': 'im_not'})
var_not_dist = Variable('all_in',
value=np.arange(10) + 10,
dimensions='ten')
vc = VariableCollection(variables=[var, var_empty, var_not_dist])
else:
vc = None
svc = variable_collection_scatter(vc, dest_mpi)
self.assertEqual(
svc['i_could_be_a_coordinate_system'].attrs['reality'], 'im_not')
if MPI_RANK < 2:
self.assertFalse(svc['all_in'].is_empty)
self.assertNumpyAll(svc['all_in'].get_value(), np.arange(10) + 10)
self.assertFalse(svc.is_empty)
self.assertFalse(svc['i_could_be_a_coordinate_system'].is_empty)
else:
self.assertTrue(svc['all_in'].is_empty)
self.assertTrue(svc.is_empty)
self.assertTrue(svc['i_could_be_a_coordinate_system'].is_empty)
if MPI_RANK == 0:
self.assertNumpyAll(var.get_value(), vc[var.name].get_value())
actual = svc['holds_five'].get_value()
if MPI_SIZE == 2:
desired = {0: np.arange(3), 1: np.arange(3, 5)}
self.assertNumpyAll(actual, desired[MPI_RANK])
actual = svc['holds_five'].is_empty
if MPI_RANK > 1:
self.assertTrue(actual)
else:
self.assertFalse(actual)
def test_create_ugid_global(self):
ompi = OcgDist()
m = ompi.create_dimension('m', 4)
n = ompi.create_dimension('n', 70, dist=True)
ompi.update_dimension_bounds()
gvar = GeometryVariable(name='geom', dimensions=(m, n))
ugid = gvar.create_ugid_global('gid')
if not gvar.is_empty:
self.assertEqual(gvar.dist, ugid.dist)
gathered = variable_gather(ugid)
if MPI_RANK == 0:
actual = gathered.get_value()
self.assertEqual(actual.size, len(set(actual.flatten().tolist())))