def test_packer_does_not_change_output_properties_after_init(self):
np.random.seed(0)
dims = ['tracer', '*']
packer = TracerPacker(self.component, dims)
assert len(self.component.output_properties) == 0
register_tracer('tracer1', 'g/m^3')
assert len(self.component.output_properties) == 0
def test_packs_prepended_and_normal_tracers_register_after_init(self):
self.component = self.component.__class__(
prepend_tracers=[('tracer1', 'g/m^3')])
register_tracer('tracer2', 'J/m^3')
input_state = {
'time': timedelta(0),
'tracer1': DataArray(
np.random.randn(5),
dims=['dim1'],
attrs={'units': 'g/m^3'},
),
'tracer2': DataArray(
np.random.randn(5),
dims=['dim1'],
attrs={'units': 'J/m^3'},
),
}
unpacked = self.call_component(input_state)
packed = self.component.state_given['tracers']
assert isinstance(packed, np.ndarray)
assert packed.shape == (2, 5)
assert np.all(packed[0, :] == input_state['tracer1'].values)
assert np.all(packed[1, :] == input_state['tracer2'].values)
assert len(unpacked) == 2
assert np.all(
unpacked['tracer1'].values == input_state['tracer1'].values)
assert np.all(
unpacked['tracer2'].values == input_state['tracer2'].values)
def test_reregister_tracer(self):
register_tracer('tracer1', 'm')
register_tracer('tracer1', 'm')
d = get_tracer_unit_dict()
assert len(d) == 1
assert 'tracer1' in d
assert d['tracer1'] == 'm'
def test_packs_two_tracers(self):
np.random.seed(0)
register_tracer('tracer1', 'g/m^3')
register_tracer('tracer2', 'kg')
input_state = {
'time': timedelta(0),
'tracer1': DataArray(
np.random.randn(5),
dims=['dim1'],
attrs={'units': 'g/m^3'},
),
'tracer2': DataArray(
np.random.randn(5),
dims=['dim1'],
attrs={'units': 'kg'}
),
}
unpacked = self.call_component(input_state)
packed = self.component.state_given['tracers']
assert isinstance(packed, np.ndarray)
assert packed.shape == (2, 5)
assert np.all(packed[0, :] == input_state['tracer1'].values)
assert np.all(packed[1, :] == input_state['tracer2'].values)
assert len(unpacked) == 2
assert np.all(unpacked['tracer1'].values == input_state['tracer1'].values)
assert np.all(unpacked['tracer2'].values == input_state['tracer2'].values)
def test_packer_does_not_change_tendency_properties(self):
np.random.seed(0)
dims = ['tracer', '*']
register_tracer('tracer1', 'g/m^3')
packer = TracerPacker(self.component, dims)
assert 'tracer1' not in self.component.tendency_properties
assert len(self.component.tendency_properties) == 0
def test_register_two_tracers(self):
register_tracer('tracer1', 'm')
register_tracer('tracer2', 'degK')
d = get_tracer_unit_dict()
assert len(d) == 2
assert 'tracer1' in d
assert 'tracer2' in d
assert d['tracer1'] == 'm'
assert d['tracer2'] == 'degK'
def test_packs_one_tracer_registered_after_init(self):
np.random.seed(0)
dims = ['tracer', '*']
state = {'tracer1': DataArray(np.random.randn(5), dims=['dim1'], attrs={'units': 'g/m^3'})}
packer = TracerPacker(self.component, dims)
register_tracer('tracer1', 'g/m^3')
packed = packer.pack(state)
assert isinstance(packed, np.ndarray)
assert packed.shape == (1, 5)
assert np.all(packed[0, :] == state['tracer1'].values)
def test_packer_allows_overlap_input_registered_before_init(self):
self.component = self.component.__class__(
input_properties={
'name': {
'units': 'm',
'dims': ['*'],
}
}
)
register_tracer('name', 'm')
packer = TracerPacker(self.component, ['tracer', '*'])
def test_packer_wont_overwrite_output_registered_before_init(self):
self.component = MockStepper(
output_properties={
'name': {
'units': 'm',
'dims': ['*'],
}
}
)
register_tracer('name', 'm')
with self.assertRaises(InvalidPropertyDictError):
packer = TracerPacker(self.component, ['tracer', '*'])
def test_packer_wont_overwrite_tendency_registered_after_init(self):
self.component = MockTendencyComponent(
tendency_properties={
'name': {
'units': 'm',
'dims': ['*'],
}
}
)
packer = TracerPacker(self.component, ['tracer', '*'])
with self.assertRaises(InvalidPropertyDictError):
register_tracer('name', 'm')
def test_packs_two_tracers(self):
np.random.seed(0)
dims = ['tracer', '*']
register_tracer('tracer1', 'g/m^3')
register_tracer('tracer2', 'kg')
state = {
'tracer1': DataArray(np.random.randn(5), dims=['dim1'], attrs={'units': 'g/m^3'}),
'tracer2': DataArray(np.random.randn(5), dims=['dim1'], attrs={'units': 'kg'})
}
packer = TracerPacker(self.component, dims)
packed = packer.pack(state)
assert isinstance(packed, np.ndarray)
assert packed.shape == (2, 5)
def test_restores_one_3d_tracer(self):
np.random.seed(0)
register_tracer('tracer1', 'g/m^3')
input_state = {
'time': timedelta(0),
'tracer1': DataArray(
np.random.randn(1, 2, 3),
dims=['dim1', 'dim2', 'dim3'],
attrs={'units': 'g/m^3'},
)
}
unpacked = self.call_component(input_state)
assert len(unpacked) == 1
assert unpacked['tracer1'].shape == input_state['tracer1'].shape
assert np.all(unpacked['tracer1'].values == input_state['tracer1'].values)
def test_packs_one_3d_tracer(self):
np.random.seed(0)
register_tracer('tracer1', 'g/m^3')
input_state = {
'time': timedelta(0),
'tracer1': DataArray(
np.random.randn(1, 2, 3),
dims=['dim1', 'dim2', 'dim3'],
attrs={'units': 'g/m^3'},
)
}
self.call_component(input_state)
packed = self.component.state_given['tracers']
assert isinstance(packed, np.ndarray)
assert packed.shape == (1, 6)
assert np.all(packed[0, :] == input_state['tracer1'].values.flatten())
def test_packs_one_3d_tracer(self):
np.random.seed(0)
dims = ['tracer', 'latitude', 'longitude', 'mid_levels']
register_tracer('tracer1', 'g/m^3')
state = {
'tracer1': DataArray(
np.random.randn(2, 3, 4),
dims=['latitude', 'longitude', 'mid_levels'],
attrs={'units': 'g/m^3'}
)
}
packer = TracerPacker(self.component, dims)
packed = packer.pack(state)
assert isinstance(packed, np.ndarray)
assert packed.shape == (1, 2, 3, 4)
assert np.all(packed[0, :, :, :] == state['tracer1'].values)
def test_packing_differing_dims(self):
np.random.seed(0)
register_tracer('tracer1', 'g/m^3')
register_tracer('tracer2', 'kg')
input_state = {
'time': timedelta(0),
'tracer1': DataArray(
np.random.randn(2),
dims=['dim1'],
attrs={'units': 'g/m^3'},
),
'tracer2': DataArray(
np.random.randn(2),
dims=['dim2'],
attrs={'units': 'kg'}
),
}
unpacked = self.call_component(input_state)
packed = self.component.state_given['tracers']
assert isinstance(packed, np.ndarray)
assert packed.shape == (2, 4)
assert len(unpacked) == 2
def test_packs_three_tracers_in_order_registered(self):
np.random.seed(0)
dims = ['tracer', '*']
register_tracer('tracer1', 'g/m^3')
register_tracer('tracer2', 'kg'),
register_tracer('tracer3', 'kg/m^3')
state = {
'tracer1': DataArray(np.random.randn(5), dims=['dim1'], attrs={'units': 'g/m^3'}),
'tracer2': DataArray(np.random.randn(5), dims=['dim1'], attrs={'units': 'kg'}),
'tracer3': DataArray(np.random.randn(5), dims=['dim1'], attrs={'units': 'kg/m^3'}),
}
packer = TracerPacker(self.component, dims)
packed = packer.pack(state)
assert isinstance(packed, np.ndarray)
assert packed.shape == (3, 5)
assert np.all(packed[0, :] == state['tracer1'].values)
assert np.all(packed[1, :] == state['tracer2'].values)
assert np.all(packed[2, :] == state['tracer3'].values)
def test_unpacks_three_tracers_in_order_registered(self):
np.random.seed(0)
dims = ['tracer', '*']
register_tracer('tracer1', 'g/m^3')
register_tracer('tracer2', 'kg'),
register_tracer('tracer3', 'kg/m^3')
state = {
'tracer1': DataArray(np.random.randn(5), dims=['dim1'], attrs={'units': 'g/m^3'}),
'tracer2': DataArray(np.random.randn(5), dims=['dim1'], attrs={'units': 'kg'}),
'tracer3': DataArray(np.random.randn(5), dims=['dim1'], attrs={'units': 'kg/m^3'}),
}
packer = TracerPacker(self.component, dims)
packed = packer.pack(state)
unpacked = packer.unpack(packed, state)
assert isinstance(unpacked, dict)
assert len(unpacked) == 3
assert np.all(unpacked['tracer1'] == state['tracer1'])
assert np.all(unpacked['tracer2'] == state['tracer2'])
assert np.all(unpacked['tracer3'] == state['tracer3'])
def test_reregister_tracer_different_units(self):
register_tracer('tracer1', 'm')
with self.assertRaises(ValueError):
register_tracer('tracer1', 'degK')
def test_does_not_change_input_properties(self):
np.random.seed(0)
register_tracer('tracer1', 'g/m^3')
assert 'tracer1' not in self.component.input_properties
