def test_system_profile_iterator(self): # 24 seconds last profile raise SkipTest('profiling only') var1 = self.create_base_variable(shape=(1000, 1000), name='one') var2 = self.create_base_variable(shape=(1000, 1000), name='two') var3 = self.create_base_variable(shape=(1000, 1000), name='three') actual = list(Iterator(var1, followers=[var2, var3])) self.assertEqual(len(actual), 1000 * 1000)
def test_iter(self): var = self.create_base_variable() itr = Iterator(var) for _ in range(2): as_list = list(itr) self.assertEqual(len(as_list), var.shape[0] * var.shape[1]) self.assertIsInstance(as_list[0], OrderedDict) self.assertEqual(len(as_list[0]), 1)
def test_system_cf_data(self): raise SkipTest('profiling only') rd = self.test_data.get_rd('cancm4_tas') field = rd.get().get_field_slice({'time': slice(0, 100)}, strict=False) follower_names = ['time', 'lon', 'lat'] itr = Iterator(field['tas'], followers=[field[ii] for ii in follower_names]) print(field.shapes) list(itr)
def test_iter_primary_mask(self): var = Variable('a', [1, 2, 3], 'b', mask=np.ones(3, dtype=bool)) primary_mask = Variable('c', [4, 5, 6], 'b', mask=[True, False, True]) itr = Iterator(var, followers=[primary_mask], primary_mask=primary_mask, allow_masked=False) actual = list(itr) self.assertIsNone(actual[0]['a'])
def test_iter_followers(self): base_var = self.create_base_variable() follower_var = Variable(name='follower', value=np.random.rand(base_var.shape[1]), dimensions=base_var.dimensions[1].name) itr = Iterator(base_var, followers=[follower_var]) as_list = list(itr) self.assertEqual(len(as_list), base_var.shape[0] * base_var.shape[1]) actual = as_list[0] self.assertIn(base_var.name, actual) self.assertIn(follower_var.name, actual) # Test follower dimensions are not a subset of the base variable dimensions. follower_var = Variable(name='follower', value=np.random.rand(base_var.shape[1]), dimensions='unique') with self.assertRaises(ValueError): Iterator(base_var, followers=[follower_var])
def test_iter_melted(self): var = self.create_base_variable() with self.assertRaises(ValueError): Iterator(var, melted=var) var = Variable(name='lead', value=[1, 2], dimensions='dim') fvar = Variable(name='follower', value=[3, 4], dimensions='dim') emelted = Variable(name='just_melted', value=[5, 6], dimensions='dim') itr = Iterator(var, followers=[fvar, emelted], melted=[var, emelted]) desired = [ OrderedDict([('follower', 3), ('VARIABLE', 'lead'), ('VALUE', 1)]), OrderedDict([('follower', 3), ('VARIABLE', 'just_melted'), ('VALUE', 5)]), OrderedDict([('follower', 4), ('VARIABLE', 'lead'), ('VALUE', 2)]), OrderedDict([('follower', 4), ('VARIABLE', 'just_melted'), ('VALUE', 6)]) ] self.assertEqual(list(itr), desired)
def get_iter(self, *args, **kwargs): """ :rtype: :class:`~ocgis.variable.iterator.Iterator` """ should_add = kwargs.pop(KeywordArgument.ADD_GEOM_UID, False) if should_add and self.ugid is not None: followers = [self.ugid] else: followers = [] return Iterator(self, followers=followers, **kwargs)
def test_iter_mask(self): var = self.create_base_variable(shape=(2, 2)) mask = var.get_mask(create=True) mask[0, 1] = True var.set_mask(mask) for allow_masked in [True, False]: itr = Iterator(var, allow_masked=allow_masked) as_list = list(itr) if allow_masked: self.assertEqual(len(as_list), 4) self.assertIsNone(as_list[1][var.name]) else: self.assertEqual(len(as_list), 3)
def test_iter_follower_with_followers(self): leader = Variable(name='data', value=np.arange(6).reshape(2, 3), dimensions=['x', 'y']) x = Variable(name='x', value=[30., 40.], dimensions='x') x_follower = Variable(name='x_bounds', value=[25., 35.], dimensions='x') follower_iterator = Iterator(x, followers=[x_follower]) itr = Iterator(leader, followers=[follower_iterator]) desired = [ OrderedDict([('data', 0), ('x', 30.0), ('x_bounds', 25.0)]), OrderedDict([('data', 1), ('x', 30.0), ('x_bounds', 25.0)]), OrderedDict([('data', 2), ('x', 30.0), ('x_bounds', 25.0)]), OrderedDict([('data', 3), ('x', 40.0), ('x_bounds', 35.0)]), OrderedDict([('data', 4), ('x', 40.0), ('x_bounds', 35.0)]), OrderedDict([('data', 5), ('x', 40.0), ('x_bounds', 35.0)]) ] self.assertEqual(list(itr), desired)
def test_iter_repeater(self): var1 = Variable(name='var1', value=[1, 2, 3], dimensions='dim') var2 = Variable(name='var2', value=[1, 2, 3], dimensions='dim') var2.get_value()[:] *= 9 repeater = ('i_am', 'a_repeater') itr = Iterator(var1, followers=[var2], repeaters=[repeater]) desired = [ OrderedDict([('i_am', 'a_repeater'), ('var1', 1), ('var2', 9)]), OrderedDict([('i_am', 'a_repeater'), ('var1', 2), ('var2', 18)]), OrderedDict([('i_am', 'a_repeater'), ('var1', 3), ('var2', 27)]) ] actual = list(itr) self.assertEqual(actual, desired)
def test_iter_repeater_on_follower(self): var1 = Variable(name='f', value=[1, 2], dimensions='one', repeat_record=[('same', 'var1 owns this')]) var2 = Variable(name='g', value=[3, 4], dimensions='one', repeat_record=[('same', 'var2 owns this')]) itr = Iterator(var1, followers=[var2], melted=[var1, var2]) for row in itr: if row['VARIABLE'] == var2.name: self.assertEqual(row['same'], var2.repeat_record[0][1]) elif row['VARIABLE'] == var1.name: self.assertEqual(row['same'], var1.repeat_record[0][1])
def test_iter_with_bounds(self): var = Variable(name='bounded', value=[1, 2, 3, 4], dtype=float, dimensions='dim') var.set_extrapolated_bounds('the_bounds', 'bounds') lower = Variable(name='lower_bounds', value=var.bounds.get_value()[:, 0], dimensions=var.dimensions) upper = Variable(name='upper_bounds', value=var.bounds.get_value()[:, 1], dimensions=var.dimensions) itr = Iterator(var, followers=[lower, upper]) actual = list(itr) self.assertEqual(len(actual), var.shape[0]) self.assertEqual(len(actual[0]), 3)
def test_iter_formatter(self): def _formatter_(name, value, mask): if value is None: modified_value = None else: modified_value = value * 1000 value = str(value) ret = [(name, value), ('modified', modified_value)] return ret var = Variable(name='data', value=[1, 2, 3], mask=[False, True, False], dimensions='dim') itr = Iterator(var, formatter=_formatter_) as_list = list(itr) actual = as_list[1][var.name] self.assertIsNone(actual) self.assertEqual(as_list[2][var.name], str(var.get_value()[2])) self.assertEqual(as_list[0]['modified'], 1000)
def iter(self, **kwargs): """ :return: Yield record dictionaries for variables in the field applying standard names to dimensions by default. :rtype: dict """ if self.is_empty: return from ocgis.driver.registry import get_driver_class standardize = kwargs.pop(KeywordArgument.STANDARDIZE, KeywordArgument.Defaults.STANDARDIZE) tag = kwargs.pop(KeywordArgument.TAG, TagName.DATA_VARIABLES) driver = kwargs.get(KeywordArgument.DRIVER) primary_mask = kwargs.pop(KeywordArgument.PRIMARY_MASK, None) header_map = kwargs.pop(KeywordArgument.HEADER_MAP, None) melted = kwargs.pop(KeywordArgument.MELTED, False) variable = kwargs.pop(KeywordArgument.VARIABLE, None) followers = kwargs.pop(KeywordArgument.FOLLOWERS, None) allow_masked = kwargs.get(KeywordArgument.ALLOW_MASKED, False) if melted and not standardize: raise ValueError('"standardize" must be True when "melted" is True.') if KeywordArgument.ALLOW_MASKED not in kwargs: kwargs[KeywordArgument.ALLOW_MASKED] = False if driver is not None: driver = get_driver_class(driver) # Holds follower variables to pass to the generic iterator. if followers is None: followers = [] else: for ii, f in enumerate(followers): if not isinstance(f, Iterator): followers[ii] = get_variables(f, self)[0] if variable is None: # The primary variable(s) to iterate. tagged_variables = self.get_by_tag(tag, create=True) if len(tagged_variables) == 0: msg = 'Tag "{}" has no associated variables. Nothing to iterate.'.format(tag) raise ValueError(msg) variable = tagged_variables[0] if len(tagged_variables) > 1: followers += tagged_variables[1:] else: variable = get_variables(variable, self)[0] if self.geom is not None: if primary_mask is None: primary_mask = self.geom if standardize: add_geom_uid = True else: add_geom_uid = False followers.append(self.geom.get_iter(**{KeywordArgument.ADD_GEOM_UID: add_geom_uid, KeywordArgument.ALLOW_MASKED: allow_masked, KeywordArgument.PRIMARY_MASK: primary_mask})) geom = self.geom else: geom = None if self.realization is not None: followers.append(self.realization.get_iter(driver=driver, allow_masked=allow_masked, primary_mask=primary_mask)) if self.time is not None: followers.append(self.time.get_iter(add_bounds=True, driver=driver, allow_masked=allow_masked, primary_mask=primary_mask)) if self.level is not None: followers.append(self.level.get_iter(add_bounds=True, driver=driver, allow_masked=allow_masked, primary_mask=primary_mask)) # Collect repeaters from the target variable and followers. This initializes the iterator twice, but the # operation is not expensive. itr_for_repeaters = Iterator(variable, followers=followers) found = kwargs.get(KeywordArgument.REPEATERS) if found is not None: found = [ii[0] for ii in found] repeater_headers = itr_for_repeaters.get_repeaters(headers_only=True, found=found) if standardize: if header_map is None: header_map = OrderedDict() if len(repeater_headers) > 0: for k in repeater_headers: header_map[k] = k if self.geom is not None and self.geom.ugid is not None: header_map[self.geom.ugid.name] = self.geom.ugid.name if self.realization is not None: header_map[self.realization.name] = HeaderName.REALIZATION if self.time is not None: header_map[self.time.name] = HeaderName.TEMPORAL update_header_rename_bounds_names(HeaderName.TEMPORAL_BOUNDS, header_map, self.time) header_map['YEAR'] = 'YEAR' header_map['MONTH'] = 'MONTH' header_map['DAY'] = 'DAY' if self.level is not None: header_map[self.level.name] = HeaderName.LEVEL update_header_rename_bounds_names(HeaderName.LEVEL_BOUNDS, header_map, self.level) if melted: melted = tagged_variables else: melted = None kwargs[KeywordArgument.HEADER_MAP] = header_map kwargs[KeywordArgument.MELTED] = melted kwargs[KeywordArgument.VARIABLE] = variable kwargs[KeywordArgument.FOLLOWERS] = followers kwargs[KeywordArgument.GEOM] = geom for yld in super(Field, self).iter(**kwargs): yield yld