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_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_packs_no_tracers(self):
dims = ['tracer', '*']
packer = TracerPacker(self.component, dims)
packed = packer.pack({})
assert isinstance(packed, np.ndarray)
assert packed.shape == (0, 0)