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_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_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_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_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_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_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_raises_on_diagnostic_init(self):
diagnostic = MockDiagnosticComponent()
with self.assertRaises(TypeError):
TracerPacker(diagnostic, ['tracer', '*'])
def test_unpacks_no_tracers(self):
dims = ['tracer', '*']
packer = TracerPacker(self.component, dims)
unpacked = packer.unpack({}, {})
assert isinstance(unpacked, dict)
assert len(unpacked) == 0
def test_packs_no_tracers(self):
dims = ['tracer', '*']
packer = TracerPacker(self.component, dims)
packed = packer.pack({})
assert isinstance(packed, np.ndarray)
assert packed.shape == (0, 0)