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