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_init_source_name(self):
dim = Dimension('foo', source_name='actual_foo')
self.assertEqual(dim.source_name, 'actual_foo')
dim = Dimension('foo')
dim.set_name('new_foo')
self.assertEqual(dim.source_name, 'foo')
def test_get_unioned_spatial_average_differing_dimensions(self):
pa = self.get_geometryvariable()
to_weight = Variable(name='to_weight',
dimensions=pa.dimensions,
dtype=float)
to_weight.get_value()[0] = 5.0
to_weight.get_value()[1] = 10.0
pa.parent.add_variable(to_weight)
to_weight2 = Variable(name='to_weight2',
dimensions=[
Dimension('time', 10),
Dimension('level', 3), pa.dimensions[0]
],
dtype=float)
for time_idx in range(to_weight2.shape[0]):
for level_idx in range(to_weight2.shape[1]):
to_weight2.get_value()[time_idx, level_idx] = (
time_idx + 2) + (level_idx + 2)**(level_idx + 1)
pa.parent.add_variable(to_weight2)
unioned = pa.get_unioned(spatial_average=['to_weight', 'to_weight2'])
actual = unioned.parent[to_weight2.name]
self.assertEqual(actual.shape, (10, 3, 1))
self.assertEqual(to_weight2.shape, (10, 3, 2))
self.assertNumpyAll(actual.get_value(),
to_weight2.get_value()[:, :, 0].reshape(10, 3, 1))
self.assertEqual(actual.dimension_names,
('time', 'level', 'ocgis_geom_union'))
self.assertEqual(unioned.parent[to_weight.name].get_value()[0], 7.5)
def test_len(self):
dim = Dimension('foo')
self.assertEqual(len(dim), 0)
# Test size current is used for length.
dim = Dimension('unlimited', size=None, size_current=4)
self.assertEqual(len(dim), 4)
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 fixture_regular_ugrid_file(self, ufile, resolution, crs=None):
"""
Create a UGRID convention file from a structured, rectilinear grid. This will create element centers in addition
to element node connectivity.
:param str ufile: Path to the output UGRID NetCDF file.
:param float resolution: The resolution fo the structured grid to convert to UGRID.
:param crs: The coordinate system for the UGRID data.
:type crs: :class:`~ocgis.CRS`
"""
src_grid = self.get_gridxy_global(resolution=resolution, crs=crs)
polygc = src_grid.get_abstraction_geometry()
polygc.reshape([Dimension(name='element_count', size=polygc.size)])
polygc = polygc.convert_to(max_element_coords=4)
pointgc = src_grid.get_point()
pointgc.reshape([Dimension(name='element_count', size=pointgc.size)])
pointgc = pointgc.convert_to(xname='face_center_x', yname='face_center_y')
polygc.parent.add_variable(pointgc.parent.first(), force=True)
polygc.parent.dimension_map.update(pointgc.dimension_map)
pointgc.parent = polygc.parent
gu = GridUnstruct(geoms=[polygc, pointgc])
gu.parent.write(ufile)
def test_convert_to_empty(self):
dim = Dimension('three', 3, src_idx=np.array([10, 11, 12]), dist=True)
dim.convert_to_empty()
self.assertFalse(dim.is_unlimited)
self.assertTrue(dim.is_empty)
self.assertEqual(len(dim), 0)
self.assertEqual(dim.size, 0)
self.assertEqual(dim.size_current, 0)
def test_convert_to_empty(self):
dim = Dimension('three', 3, src_idx=np.array([10, 11, 12]), dist=True)
dim.convert_to_empty()
self.assertFalse(dim.is_unlimited)
self.assertTrue(dim.is_empty)
self.assertEqual(len(dim), 0)
self.assertEqual(dim.size, 0)
self.assertEqual(dim.size_current, 0)
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_renamed_dimensions(self):
d = [Dimension('a', 5), Dimension('b', 6)]
desired_after = deepcopy(d)
name_mapping = {'time': ['b']}
desired = [Dimension('a', 5), Dimension('time', 6)]
with renamed_dimensions(d, name_mapping):
for idx in range(len(d)):
self.assertEqual(d[idx], desired[idx])
self.assertEqual(d, desired)
self.assertEqual(desired_after, d)
def test_write_variable_collection(self):
# Attempt to write without a geometry variable.
v = Variable('a', value=[1, 2], dimensions='bb')
field = Field(variables=v)
path = self.get_temporary_file_path('out.shp')
with self.assertRaises(ValueError):
field.write(path, driver=DriverVector)
# Test writing a field with two-dimensional geometry storage.
value = [Point(1, 2), Point(3, 4), Point(5, 6), Point(6, 7), Point(8, 9), Point(10, 11)]
gvar = GeometryVariable(value=value, name='points', dimensions='ngeoms')
gvar.reshape([Dimension('lat', 2), Dimension('lon', 3)])
var1 = Variable(name='dummy', value=[6, 7, 8], dimensions=['a'])
var2 = Variable(name='some_lats', value=[41, 41], dimensions=['lat'])
var3 = Variable(name='some_lons', value=[0, 90, 280], dimensions=['lon'])
var4 = Variable(name='data', value=np.random.rand(4, 3, 2), dimensions=['time', 'lon', 'lat'])
field = Field(variables=[var1, var2, var3, var4], geom=gvar, is_data=['data'])
path = self.get_temporary_file_path('2d.shp')
field.write(path, iter_kwargs={'followers': ['some_lats', 'some_lons']}, driver=DriverVector)
read = RequestDataset(uri=path).get()
self.assertTrue(len(read) > 2)
self.assertEqual(list(read.keys()),
['data', 'some_lats', 'some_lons', constants.DimensionName.GEOMETRY_DIMENSION])
# Test writing a subset of the variables.
path = self.get_temporary_file_path('limited.shp')
value = [Point(1, 2), Point(3, 4), Point(5, 6)]
gvar = GeometryVariable(value=value, name='points', dimensions='points')
var1 = Variable('keep', value=[1, 2, 3], dimensions='points')
var2 = Variable('remove', value=[4, 5, 6], dimensions='points')
field = Field(variables=[var1, var2], geom=gvar, is_data=[var1])
field.write(path, variable_names=['keep'], driver=DriverVector)
read = RequestDataset(uri=path).get()
self.assertNotIn('remove', read)
# Test using append.
path = self.get_temporary_file_path('limited.shp')
value = [Point(1, 2), Point(3, 4), Point(5, 6)]
gvar = GeometryVariable(value=value, name='points', dimensions='points')
var1 = Variable('keep', value=[1, 2, 3], dimensions='points')
var2 = Variable('remove', value=[4, 5, 6], dimensions='points')
field = Field(variables=[var1, var2], geom=gvar, is_data=[var1, var2])
for idx in range(3):
sub = field[{'points': idx}]
if idx == 0:
write_mode = MPIWriteMode.WRITE
else:
write_mode = MPIWriteMode.APPEND
sub.write(path, write_mode=write_mode, driver=DriverVector)
self.assertOGRFileLength(path, idx + 1)
def test_getitem(self):
dim = Dimension('foo', size=50)
sub = dim[30:40]
self.assertEqual(len(sub), 10)
dim = Dimension('foo', size=None)
with self.assertRaises(IndexError):
dim[400:500]
dim = self.get_dimension(src_idx=np.arange(10))
sub = dim[4]
self.assertEqual(sub.size, 1)
sub = dim[4:5]
self.assertEqual(sub.size, 1)
sub = dim[4:6]
self.assertEqual(sub.size, 2)
sub = dim[[4, 5, 6]]
self.assertEqual(sub.size, 3)
sub = dim[[2, 4, 6]]
self.assertEqual(sub.size, 3)
self.assertNumpyAll(sub._src_idx, dim._src_idx[[2, 4, 6]])
sub = dim[:]
self.assertEqual(len(sub), len(dim))
dim = self.get_dimension()
sub = dim[2:]
self.assertEqual(len(sub), 8)
dim = self.get_dimension(src_idx=np.arange(10))
sub = dim[3:-1]
np.testing.assert_equal(sub._src_idx, [3, 4, 5, 6, 7, 8])
dim = self.get_dimension(src_idx=np.arange(10))
sub = dim[-3:]
self.assertEqual(sub._src_idx.shape[0], sub.size)
dim = self.get_dimension(src_idx=np.arange(10))
sub = dim[-7:-3]
self.assertEqual(sub._src_idx.shape[0], sub.size)
dim = self.get_dimension(src_idx=np.arange(10))
sub = dim[:-3]
self.assertEqual(sub._src_idx.shape[0], sub.size)
def test_grid(self):
# Test mask variable information is propagated through property.
grid = self.get_gridxy(with_xy_bounds=True)
self.assertTrue(grid.is_vectorized)
self.assertTrue(grid.has_bounds)
np.random.seed(1)
value = np.random.rand(*grid.shape)
select = value > 0.4
mask_var = create_spatial_mask_variable('nonstandard', select, grid.dimensions)
grid.set_mask(mask_var)
field = Field(grid=grid)
self.assertTrue(field.grid.has_bounds)
self.assertEqual(field.dimension_map.get_spatial_mask(), mask_var.name)
self.assertNumpyAll(field.grid.get_mask(), mask_var.get_mask())
# Test dimension map bounds are updated appropriately.
dim = Dimension('count', 2)
x = Variable(name='x', value=[1., 2.], dimensions=dim)
y = Variable(name='y', value=[1., 2.], dimensions=dim)
xb = Variable(name='xb', value=[[0., 1.5], [1.5, 2.5]], dimensions=[dim, 'bounds'])
yb = Variable(name='yb', value=[[0., 1.5], [1.5, 2.5]], dimensions=[dim, 'bounds'])
variables = [x, y, xb, yb]
dmap = DimensionMap()
dmap.set_variable(DMK.X, x, bounds=xb)
dmap.set_variable(DMK.Y, y, bounds=yb)
f = Field(dimension_map=dmap, variables=variables)
self.assertTrue(f.grid.has_bounds)
def _get_field_write_target_(cls, field):
dmap = field.dimension_map
driver = dmap.get_driver()
if driver == DriverKey.NETCDF_UGRID:
ret = cls._convert_from_ugrid_(field)
elif driver == DriverKey.NETCDF_ESMF_UNSTRUCT:
ret = field.copy()
# Coordinate variables were "sectioned" when loaded in. They need to be put back together before writing to
# disk. Only do this if they have been sectioned which only occurs when the coordinate variable is retrieved
# from the field.
for toponame, topo in dmap.iter_topologies():
yname = topo.get_variable(DMK.Y)
xname = topo.get_variable(DMK.X)
if xname != yname:
yvar = topo.get_variable(DMK.Y, parent=field)
new_coords = np.hstack(
(topo.get_variable(DMK.X,
parent=field).v().reshape(-1, 1),
yvar.v().reshape(-1, 1)))
assert new_coords.shape[1] == 2
new_dimensions = (field[yname].dimensions[0],
Dimension(name='coordDim', size=2))
new_name = yname[0:-2]
ret.add_variable(
Variable(name=new_name,
dimensions=new_dimensions,
value=new_coords,
attrs=yvar.attrs))
ret.remove_variable(yname)
ret.remove_variable(xname)
return ret
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_system_create_many_large_dimensions(self):
count = 10000
dims = [None] * count
for idx in range(count):
dims[idx] = Dimension(str(idx), size=500000000, src_idx='auto')
sliced_dims = [None] * len(dims)
for idx in range(count):
sliced_dims[idx] = dims[idx][50000:490000000]
def add_data_variable_to_grid(grid):
ydim, xdim = grid.dimensions
tdim = Dimension(name='time', size=None)
value = np.random.rand(31, ydim.size, xdim.size)
data = Variable(name='data', dimensions=[tdim, ydim, xdim], value=value)
tvar = TemporalVariable(name='time', value=list(range(31)), dimensions=tdim, attrs={'axis': 'T'})
grid.parent.add_variable(data)
grid.parent.add_variable(tvar)
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 test_set_extrapolated_bounds_empty(self):
"""Test bounds extrapolation works on empty objects."""
dimx = Dimension('x', 2, is_empty=True, dist=True)
x = Variable('x', dimensions=dimx)
y = Variable('3', dimensions=Dimension('y', 3))
grid = Grid(x, y)
self.assertTrue(dimx.is_empty)
self.assertTrue(x.is_empty)
self.assertTrue(grid.is_empty)
self.assertFalse(grid.has_bounds)
self.assertEqual(grid.abstraction, 'point')
grid.set_extrapolated_bounds('xbnds', 'ybnds', 'bounds')
self.assertEqual(grid.abstraction, 'polygon')
self.assertTrue(grid.has_bounds)
self.assertTrue(grid.is_empty)
def test_copy(self):
sd = self.get_dimension(src_idx=np.arange(10))
self.assertIsNotNone(sd._src_idx)
sd2 = sd.copy()
self.assertTrue(np.may_share_memory(sd._src_idx, sd2._src_idx))
sd3 = sd2[2:5]
self.assertEqual(sd, sd2)
self.assertNotEqual(sd2, sd3)
self.assertTrue(np.may_share_memory(sd2._src_idx, sd._src_idx))
self.assertTrue(np.may_share_memory(sd2._src_idx, sd3._src_idx))
self.assertTrue(np.may_share_memory(sd3._src_idx, sd._src_idx))
# Test setting new values on object.
dim = Dimension('five', 5)
cdim = dim.copy()
cdim._name = 'six'
self.assertEqual(dim.name, 'five')
self.assertEqual(cdim.name, 'six')
def test_copy(self):
sd = self.get_dimension(src_idx=np.arange(10))
self.assertIsNotNone(sd._src_idx)
sd2 = sd.copy()
self.assertTrue(np.may_share_memory(sd._src_idx, sd2._src_idx))
sd3 = sd2[2:5]
self.assertEqual(sd, sd2)
self.assertNotEqual(sd2, sd3)
self.assertTrue(np.may_share_memory(sd2._src_idx, sd._src_idx))
self.assertTrue(np.may_share_memory(sd2._src_idx, sd3._src_idx))
self.assertTrue(np.may_share_memory(sd3._src_idx, sd._src_idx))
# Test setting new values on object.
dim = Dimension('five', 5)
cdim = dim.copy()
cdim._name = 'six'
self.assertEqual(dim.name, 'five')
self.assertEqual(cdim.name, 'six')
def test_create_dimensions(self):
driver = self.get_driver()
actual = driver.create_dist().mapping[MPI_RANK]
desired = {None: {
'variables': {},
'dimensions': {
'ocgis_ngeom': Dimension(name='ocgis_ngeom', size=51, size_current=51, dist=True, src_idx='auto')},
'groups': {}}}
self.assertEqual(actual, desired)
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 test_get_nonempty_ranks(self):
from ocgis.variable.dimension import Dimension
comm, rank, size = get_standard_comm_state()
if size not in [1, 3]:
raise SkipTest('MPI_SIZE != 1 or 3')
target = Dimension('a')
live_ranks = get_nonempty_ranks(target, vm)
if MPI_RANK == 0:
self.assertEqual(live_ranks, tuple(range(size)))
if MPI_SIZE == 3:
targets = {0: Dimension('a', is_empty=True, dist=True),
1: Dimension('a'),
2: Dimension('a')}
with vm.scoped('ner', vm.ranks):
live_ranks = get_nonempty_ranks(targets[MPI_RANK], vm)
self.assertEqual(live_ranks, (1, 2))
def _get_dimensions_main_(self, group_metadata):
"""
:param dict group_metadata: Metadata dictionary for the target group.
:return: A sequence of dimension objects.
:rtype: sequence
"""
gmd = group_metadata['dimensions']
dims = [Dimension(dref['name'], size=dref['size'], src_idx='auto') for dref in list(gmd.values())]
return tuple(dims)
def test_init_source_name(self):
dim = Dimension('foo', source_name='actual_foo')
self.assertEqual(dim.source_name, 'actual_foo')
dim = Dimension('foo')
dim.set_name('new_foo')
self.assertEqual(dim.source_name, 'foo')
def test_get_dimensions(self):
driver = self.get_driver()
actual = driver.get_dist().mapping[MPI_RANK]
desired = {None: {
'variables': {'STATE_FIPS': {'dimensions': ('ocgis_geom',)},
'STATE_ABBR': {'dimensions': ('ocgis_geom',)},
'UGID': {'dimensions': ('ocgis_geom',)},
'ID': {'dimensions': ('ocgis_geom',)},
'STATE_NAME': {'dimensions': ('ocgis_geom',)}},
'dimensions': {
'ocgis_geom': Dimension(name='ocgis_geom', size=51, size_current=51, dist=True, src_idx='auto')},
'groups': {}}}
self.assertEqual(actual, desired)
def get_ocgfield_example(self):
dtime = Dimension(name='time')
t = TemporalVariable(value=[1, 2, 3, 4], name='the_time', dimensions=dtime, dtype=float)
t.set_extrapolated_bounds('the_time_bounds', 'bounds')
lon = Variable(value=[30., 40., 50., 60.], name='longitude', dimensions='lon')
lat = Variable(value=[-10., -20., -30., -40., -50.], name='latitude', dimensions='lat')
tas_shape = [t.shape[0], lat.shape[0], lon.shape[0]]
tas = Variable(value=np.arange(reduce_multiply(tas_shape)).reshape(*tas_shape),
dimensions=(dtime, 'lat', 'lon'), name='tas')
time_related = Variable(value=[7, 8, 9, 10], name='time_related', dimensions=dtime)
garbage1 = Variable(value=[66, 67, 68], dimensions='three', name='garbage1')
dmap = {'time': {'variable': t.name},
'x': {'variable': lon.name, DimensionMapKey.DIMENSION: [lon.dimensions[0].name]},
'y': {'variable': lat.name, DimensionMapKey.DIMENSION: [lat.dimensions[0].name]}}
field = Field(variables=[t, lon, lat, tas, garbage1, time_related], dimension_map=dmap, is_data=tas.name)
return field
def get_dimensions_from_netcdf_metadata(metadata, desired_dimensions):
new_dimensions = []
for dim_name in desired_dimensions:
dim = metadata['dimensions'][dim_name]
dim_length = dim['size']
if dim['isunlimited']:
length = None
length_current = dim_length
else:
length = dim_length
length_current = None
# tdk: identify best method to remove the need to set 'auto' when creating a source index
# new_dim = Dimension(dim_name, size=length, size_current=length_current)
new_dim = Dimension(dim_name,
size=length,
size_current=length_current,
src_idx='auto')
new_dimensions.append(new_dim)
return new_dimensions
def test_init(self):
dim = Dimension('foo')
self.assertEqual(dim.name, 'foo')
self.assertIsNone(dim.size)
self.assertIsNone(dim.size_current)
self.assertTrue(dim.is_unlimited)
self.assertEqual(len(dim), 0)
self.assertEqual(dim.bounds_local, (0, len(dim)))
self.assertEqual(dim.bounds_global, dim.bounds_local)
dim = Dimension('foo', size=23)
self.assertEqual(dim.size, 23)
src_idx = np.arange(0, 10, dtype=DataType.DIMENSION_SRC_INDEX)
dim = self.get_dimension(src_idx=src_idx)
self.assertNumpyAll(dim._src_idx, src_idx)
self.assertEqual(dim._src_idx.shape[0], 10)
# Test distributed dimensions require and size definition.
with self.assertRaises(ValueError):
Dimension('foo', dist=True)
# Test without size definition and a source index.
for src_idx in ['auto', 'ato', [0, 1, 2]]:
try:
dim = Dimension('foo', src_idx=src_idx)
except ValueError:
self.assertNotEqual(src_idx, 'auto')
else:
self.assertIsNone(dim._src_idx)
# Test a unique identifier.
dim = Dimension('foo', uid=10)
self.assertEqual(dim.uid, 10)
# Test using a tuple with conversion required.
desired = (0, 5)
dim = Dimension('test', 5, src_idx=list(desired))
self.assertEqual(dim._src_idx, desired)
# Test bounds are the default source index type.
dim = Dimension('test', 8, src_idx='auto')
self.assertEqual(dim._src_idx, (0, 8))
def test_convert_to_geometry_coordinates_polygons(self):
grid = create_gridxy_global(resolution=45.0)
geom = grid.get_abstraction_geometry()
geom.reshape(Dimension('n_elements', geom.size))
keywords = dict(pack=[True, False],
repeat_last_node=[False, True],
start_index=[0, 1])
for k in self.iter_product_keywords(keywords):
actual = geom.convert_to(pack=k.pack,
repeat_last_node=k.repeat_last_node,
start_index=k.start_index)
self.assertEqual(actual.cindex.attrs['start_index'], k.start_index)
self.assertEqual(actual.start_index, k.start_index)
self.assertEqual(actual.cindex.get_value()[0].min(), k.start_index)
self.assertEqual(actual.packed, k.pack)
self.assertIsNotNone(actual.cindex)
for actual_geom, desired_geom in zip(
actual.get_geometry_iterable(),
geom.get_value().flat):
self.assertEqual(actual_geom[1], desired_geom)
def create_dimensions(group_metadata):
"""
Create dimension objects. The key may differ from the dimension name. In which case, we can assume the dimension
is being renamed.
:param dict group_metadata: Metadata dictionary for the target group.
:rtype: list
"""
gmd = group_metadata.get('dimensions', {})
dims = {}
for k, v in gmd.items():
size = v['size']
if v.get('isunlimited', False):
size_current = size
size = None
else:
size_current = None
dims[k] = Dimension(v.get('name', k),
size=size,
size_current=size_current,
src_idx='auto',
source_name=k)
return dims
def test_is_matched_by_alias(self):
dim = Dimension('non_standard_time')
dim.append_alias('time')
self.assertTrue(dim.is_matched_by_alias('time'))
def get_unioned(self,
dimensions=None,
union_dimension=None,
spatial_average=None,
root=0):
"""
Unions _unmasked_ geometry objects. Collective across the current :class:`~ocgis.OcgVM`.
"""
# TODO: optimize!
# Get dimension names and lengths for the dimensions to union.
if dimensions is None:
dimensions = self.dimensions
dimension_names = get_dimension_names(dimensions)
dimension_lengths = [
len(self.parent.dimensions[dn]) for dn in dimension_names
]
# Get the variables to spatial average.
if spatial_average is not None:
variable_names_to_weight = get_variable_names(spatial_average)
else:
variable_names_to_weight = []
# Get the new dimensions for the geometry variable. The union dimension is always the last dimension.
if union_dimension is None:
from ocgis.variable.dimension import Dimension
union_dimension = Dimension(
constants.DimensionName.UNIONED_GEOMETRY, 1)
new_dimensions = []
for dim in self.dimensions:
if dim.name not in dimension_names:
new_dimensions.append(dim)
new_dimensions.append(union_dimension)
# Configure the return variable.
ret = self.copy()
if spatial_average is None:
ret = ret.extract()
ret.set_mask(None)
ret.set_value(None)
ret.set_dimensions(new_dimensions)
ret.allocate_value()
# Destination indices in the return variable are filled with non-masked, unioned geometries.
for dst_indices in product(
*
[list(range(dl)) for dl in get_dimension_lengths(new_dimensions)]):
dst_slc = {
new_dimensions[ii].name: dst_indices[ii]
for ii in range(len(new_dimensions))
}
# Select the geometries to union skipping any masked geometries.
to_union = deque()
for indices in product(
*[list(range(dl)) for dl in dimension_lengths]):
dslc = {
dimension_names[ii]: indices[ii]
for ii in range(len(dimension_names))
}
sub = self[dslc]
sub_mask = sub.get_mask()
if sub_mask is None:
to_union.append(sub.get_value().flatten()[0])
else:
if not sub_mask.flatten()[0]:
to_union.append(sub.get_value().flatten()[0])
# Execute the union operation.
processed_to_union = deque()
for geom in to_union:
if isinstance(geom, MultiPolygon) or isinstance(
geom, MultiPoint):
for element in geom:
processed_to_union.append(element)
else:
processed_to_union.append(geom)
unioned = cascaded_union(processed_to_union)
# Pull unioned geometries and union again for the final unioned geometry.
if vm.size > 1:
unioned_gathered = vm.gather(unioned)
if vm.rank == root:
unioned = cascaded_union(unioned_gathered)
# Fill the return geometry variable value with the unioned geometry.
to_fill = ret[dst_slc].get_value()
to_fill[0] = unioned
# Spatial average shared dimensions.
if spatial_average is not None:
# Get source data to weight.
for var_to_weight in filter(
lambda ii: ii.name in variable_names_to_weight,
list(self.parent.values())):
# Holds sizes of dimensions to iterate. These dimension are not squeezed by the weighted averaging.
range_to_itr = []
# Holds the names of dimensions to squeeze.
names_to_itr = []
# Dimension names that are squeezed. Also the dimensions for the weight matrix.
names_to_slice_all = []
for dn in var_to_weight.dimensions:
if dn.name in self.dimension_names:
names_to_slice_all.append(dn.name)
else:
range_to_itr.append(len(dn))
names_to_itr.append(dn.name)
# Reference the weights on the source geometry variable.
weights = self[{
nsa: slice(None)
for nsa in names_to_slice_all
}].weights
# Path if there are iteration dimensions. Checks for axes ordering in addition.
if len(range_to_itr) > 0:
# New dimensions for the spatially averaged variable. Unioned dimension is always last. Remove the
# dimensions aggregated by the weighted average.
new_dimensions = [
dim for dim in var_to_weight.dimensions
if dim.name not in dimension_names
]
new_dimensions.append(union_dimension)
# Prepare the spatially averaged variable.
target = ret.parent[var_to_weight.name]
target.set_mask(None)
target.set_value(None)
target.set_dimensions(new_dimensions)
target.allocate_value()
# Swap weight axes to make sure they align with the target variable.
swap_chain = get_swap_chain(dimension_names,
names_to_slice_all)
if len(swap_chain) > 0:
weights = weights.copy()
for sc in swap_chain:
weights = weights.swapaxes(*sc)
# The main weighting loop. Can get quite intensive with many, large iteration dimensions.
len_names_to_itr = len(names_to_itr)
slice_none = slice(None)
squeeze_out = [
ii for ii, dim in enumerate(var_to_weight.dimensions)
if dim.name in names_to_itr
]
should_squeeze = True if len(squeeze_out) > 0 else False
np_squeeze = np.squeeze
np_atleast_1d = np.atleast_1d
np_ma_average = np.ma.average
for nonweighted_indices in product(
*[list(range(ri)) for ri in range_to_itr]):
w_slc = {
names_to_itr[ii]: nonweighted_indices[ii]
for ii in range(len_names_to_itr)
}
for nsa in names_to_slice_all:
w_slc[nsa] = slice_none
data_to_weight = var_to_weight[w_slc].get_masked_value(
)
if should_squeeze:
data_to_weight = np_squeeze(
data_to_weight, axis=tuple(squeeze_out))
weighted_value = np_atleast_1d(
np_ma_average(data_to_weight, weights=weights))
target[w_slc].get_value()[:] = weighted_value
else:
target_to_weight = var_to_weight.get_masked_value()
# Sort to minimize floating point sum errors.
target_to_weight = target_to_weight.flatten()
weights = weights.flatten()
sindices = np.argsort(target_to_weight)
target_to_weight = target_to_weight[sindices]
weights = weights[sindices]
weighted_value = np.atleast_1d(
np.ma.average(target_to_weight, weights=weights))
target = ret.parent[var_to_weight.name]
target.set_mask(None)
target.set_value(None)
target.set_dimensions(new_dimensions)
target.set_value(weighted_value)
# Collect areas of live ranks and convert to weights.
if vm.size > 1:
# If there is no area information (points for example, we need to use counts).
if ret.area.data[0].max() == 0:
weight_or_proxy = float(self.size)
else:
weight_or_proxy = ret.area.data[0]
if vm.rank != root:
vm.comm.send(weight_or_proxy, dest=root)
else:
live_rank_areas = [weight_or_proxy]
for tner in vm.ranks:
if tner != vm.rank:
recv_area = vm.comm.recv(source=tner)
live_rank_areas.append(recv_area)
live_rank_areas = np.array(live_rank_areas)
rank_weights = live_rank_areas / np.max(live_rank_areas)
for var_to_weight in filter(
lambda ii: ii.name in variable_names_to_weight,
list(ret.parent.values())):
dimensions_to_itr = [
dim.name for dim in var_to_weight.dimensions
if dim.name != union_dimension.name
]
slc = {union_dimension.name: 0}
for idx_slc in var_to_weight.iter_dict_slices(
dimensions=dimensions_to_itr):
idx_slc.update(slc)
to_weight = var_to_weight[idx_slc].get_value().flatten(
)[0]
if vm.rank == root:
collected_to_weight = [to_weight]
if not vm.rank == root:
vm.comm.send(to_weight, dest=root)
else:
for tner in vm.ranks:
if not tner == root:
recv_to_weight = vm.comm.recv(source=tner)
collected_to_weight.append(recv_to_weight)
# Sort to minimize floating point sum errors.
collected_to_weight = np.array(collected_to_weight)
sindices = np.argsort(collected_to_weight)
collected_to_weight = collected_to_weight[sindices]
rank_weights = rank_weights[sindices]
weighted = np.atleast_1d(
np.ma.average(collected_to_weight,
weights=rank_weights))
var_to_weight[idx_slc].get_value()[:] = weighted
if vm.rank == root:
return ret
else:
return
def test_bounds_local(self):
dim = Dimension('a', 5)
dim.bounds_local = [0, 2]
self.assertEqual(dim.bounds_local, (0, 2))