def test_raises_on_raw_array_missing(self): input_state = { 'air_temperature': DataArray( np.zeros([2, 2, 4]), dims=['x', 'y', 'z'], attrs={'units': 'degK'}, ) } input_properties = { 'air_temperature': { 'dims': ['x', 'y', 'z'], 'units': 'degK', } } raw_arrays = { 'foo': np.zeros([2, 2, 4]) } output_properties = { 'foo': { 'dims': ['x', 'y', 'z'], 'units': 'm', }, 'bar': { 'dims': ['x', 'y', 'z'], 'units': 'm', }, } try: restore_data_arrays_with_properties( raw_arrays, output_properties, input_state, input_properties ) except KeyError: pass else: raise AssertionError('should have raised KeyError')
def test_assumes_dims_like_own_name(self): input_state = { 'air_temperature': DataArray( np.zeros([2, 2, 4]), dims=['x', 'y', 'z'], attrs={'units': 'degK'}, ) } input_properties = { 'air_temperature': { 'dims': ['x', 'y', 'z'], 'units': 'degK', } } raw_arrays = get_numpy_arrays_with_properties(input_state, input_properties) output_properties = { 'air_temperature': { 'units': 'degK/s', } } return_value = restore_data_arrays_with_properties( raw_arrays, output_properties, input_state, input_properties ) assert isinstance(return_value, dict) assert len(return_value.keys()) == 1 assert isinstance(return_value['air_temperature'], DataArray) assert return_value['air_temperature'].attrs['units'] is 'degK/s' assert np.byte_bounds( return_value['air_temperature'].values) == np.byte_bounds( input_state['air_temperature'].values) assert (return_value['air_temperature'].values.base is input_state['air_temperature'].values) assert return_value['air_temperature'].shape == (2, 2, 4)
def test_returns_simple_value(self): input_state = { 'air_temperature': DataArray( np.zeros([2, 2, 4]), dims=['x', 'y', 'z'], attrs={'units': 'degK'}, ) } input_properties = { 'air_temperature': { 'dims': ['x', 'y', 'z'], 'units': 'degK', } } raw_arrays = get_numpy_arrays_with_properties(input_state, input_properties) raw_arrays = {key + '_tendency': value for key, value in raw_arrays.items()} output_properties = { 'air_temperature_tendency': { 'dims': ['x', 'y', 'z'], 'units': 'degK/s', } } return_value = restore_data_arrays_with_properties( raw_arrays, output_properties, input_state, input_properties ) assert isinstance(return_value, dict) assert len(return_value.keys()) == 1 assert isinstance(return_value['air_temperature_tendency'], DataArray) assert return_value['air_temperature_tendency'].attrs['units'] is 'degK/s' assert np.byte_bounds( return_value['air_temperature_tendency'].values) == np.byte_bounds( input_state['air_temperature'].values) assert (return_value['air_temperature_tendency'].values.base is input_state['air_temperature'].values) assert return_value['air_temperature_tendency'].shape == (2, 2, 4)
def test_restores_new_dims_with_wildcard(self): input_state = { 'air_pressure': DataArray( np.zeros([2, 2, 4]), dims=['x', 'y', 'z'], attrs={'units': 'degK'}, ), } input_properties = { 'air_pressure': { 'dims': ['*'], 'units': 'degK', 'alias': 'p' }, } raw_arrays = { 'q': np.zeros([16, 2]) } output_properties = { 'q': { 'dims': ['*', 'new_dim'], 'units': 'm', }, } data_arrays = restore_data_arrays_with_properties( raw_arrays, output_properties, input_state, input_properties ) assert len(data_arrays.keys()) == 1 assert 'q' in data_arrays.keys() assert np.all(data_arrays['q'].values.flatten() == raw_arrays['q'].flatten()) assert np.byte_bounds( data_arrays['q'].values) == np.byte_bounds( raw_arrays['q']) assert data_arrays['q'].dims == ('x', 'y', 'z', 'new_dim') assert data_arrays['q'].shape == (2, 2, 4, 2)
def test_restores_aliased_name(self): input_state = { 'air_temperature': DataArray( np.zeros([2, 2, 4]), dims=['x', 'y', 'z'], attrs={'units': 'degK'}, ) } input_properties = { 'air_temperature': { 'dims': ['x', 'y', 'z'], 'units': 'degK', } } raw_arrays = { 'p': np.zeros([2, 2, 4]) } output_properties = { 'air_pressure': { 'dims': ['x', 'y', 'z'], 'units': 'm', 'alias': 'p', }, } data_arrays = restore_data_arrays_with_properties( raw_arrays, output_properties, input_state, input_properties ) assert len(data_arrays.keys()) == 1 assert 'air_pressure' in data_arrays.keys() assert np.all(data_arrays['air_pressure'].values == raw_arrays['p']) assert np.byte_bounds(data_arrays['air_pressure'].values) == np.byte_bounds(raw_arrays['p'])
def test_restores_scalar_array(self): T_array = np.array(0.) input_properties = { 'surface_temperature': { 'units': 'degK', 'dims': ['*'], }, } input_state = { 'surface_temperature': DataArray( T_array, dims=[], attrs={'units': 'degK'}, ), } raw_arrays = get_numpy_arrays_with_properties(input_state, input_properties) output_properties = { 'surface_temperature': { 'units': 'degK', } } return_value = restore_data_arrays_with_properties( raw_arrays, output_properties, input_state, input_properties ) assert len(return_value.keys()) == 1 assert 'surface_temperature' in return_value.keys() assert len(return_value['surface_temperature'].values.shape) == 0 assert return_value['surface_temperature'].attrs['units'] == 'degK'
def __call__(self, state): """ Gets tendencies and diagnostics from the passed model state. Copied from sympl develop branch (to-be v0.3.3), ignoring checks. Args ---- state : dict A model state dictionary. Returns ------- tendencies : dict A dictionary whose keys are strings indicating state quantities and values are the time derivative of those quantities in units/second at the time of the input state. diagnostics : dict A dictionary whose keys are strings indicating state quantities and values are the value of those quantities at the time of the input state. Raises ------ KeyError If a required quantity is missing from the state. InvalidStateError If state is not a valid input for the Prognostic instance. """ raw_state = get_numpy_arrays_with_properties(state, self.input_properties) raw_state['time'] = state['time'] raw_tendencies, raw_diagnostics = self.array_call(raw_state) tendencies = restore_data_arrays_with_properties( raw_tendencies, self.tendency_properties, state, self.input_properties) diagnostics = restore_data_arrays_with_properties( raw_diagnostics, self.diagnostic_properties, state, self.input_properties) return tendencies, diagnostics
def test_restores_with_dims(self): raw_arrays = { 'output1': np.ones([10]), } output_properties = { 'output1': { 'dims': ['dim1'], 'units': 'm' } } output = restore_data_arrays_with_properties( raw_arrays, output_properties, {}, {}) assert len(output) == 1 assert 'output1' in output.keys() assert isinstance(output['output1'], DataArray) assert len(output['output1'].dims) == 1 assert 'dim1' in output['output1'].dims assert 'units' in output['output1'].attrs.keys() assert output['output1'].attrs['units'] == 'm'
def test_restores_new_dims(self): input_state = {} input_properties = {} raw_arrays = { 'air_pressure': np.zeros([2, 2, 4]) } output_properties = { 'air_pressure': { 'dims': ['x', 'y', 'z'], 'units': 'm', }, } data_arrays = restore_data_arrays_with_properties( raw_arrays, output_properties, input_state, input_properties ) assert len(data_arrays.keys()) == 1 assert 'air_pressure' in data_arrays.keys() assert np.all(data_arrays['air_pressure'].values == raw_arrays['air_pressure']) assert np.byte_bounds( data_arrays['air_pressure'].values) == np.byte_bounds( raw_arrays['air_pressure'])
def test_restores_collected_horizontal_dimensions(self): random = np.random.RandomState(0) T_array = random.randn(3, 2, 4) input_state = { 'air_temperature': DataArray( T_array, dims=['x', 'y', 'z'], attrs={'units': 'degK'}, ) } input_properties = { 'air_temperature': { 'dims': ['z', '*'], 'units': 'degK', } } raw_arrays = get_numpy_arrays_with_properties(input_state, input_properties) raw_arrays = {key + '_tendency': value for key, value in raw_arrays.items()} output_properties = { 'air_temperature_tendency': { 'dims': ['z', '*'], 'units': 'degK/s', } } return_value = restore_data_arrays_with_properties( raw_arrays, output_properties, input_state, input_properties ) assert isinstance(return_value, dict) assert len(return_value.keys()) == 1 assert isinstance(return_value['air_temperature_tendency'], DataArray) assert return_value['air_temperature_tendency'].attrs['units'] is 'degK/s' assert np.byte_bounds( return_value['air_temperature_tendency'].values) == np.byte_bounds( input_state['air_temperature'].values) assert (return_value['air_temperature_tendency'].values.base is input_state['air_temperature'].values) assert return_value['air_temperature_tendency'].dims == ('z', 'x', 'y') assert return_value['air_temperature_tendency'].shape == (4, 3, 2) for i in range(4): assert np.all(return_value['air_temperature_tendency'][i, :, :] == T_array[:, :, i])
def __call__(self, state, timestep): """ Gets diagnostics from the current model state and steps the state forward in time according to the timestep. Args ---- state : dict A model state dictionary satisfying the input_properties of this object. timestep : timedelta The amount of time to step forward. Returns ------- diagnostics : dict Diagnostics from the timestep of the input state. new_state : dict A dictionary whose keys are strings indicating state quantities and values are the value of those quantities at the timestep after input state. Raises ------ KeyError If a required quantity is missing from the state. InvalidStateError If state is not a value input for the Stepper instance for other reasons. """ self._check_self_is_initialized() self._input_checker.check_inputs(state) raw_state = get_numpy_arrays_with_properties(state, self.input_properties) if self.uses_tracers: raw_state['tracers'] = self._tracer_packer.pack(state) raw_state['time'] = state['time'] tendencies, diagnostics = self._tendency_component(state, timestep) for name, value in tendencies.items(): if name in self.input_properties.keys(): tendencies[name] = value.to_units( self.input_properties[name]['units'] + ' s^-1') raw_diagnostics, raw_new_state = self.array_call( raw_state, timestep, prognostic_tendencies=tendencies) if self.uses_tracers: new_state = self._tracer_packer.unpack( raw_new_state.pop('tracers'), state) else: new_state = {} if self.tendencies_in_diagnostics: self._insert_tendencies_to_diagnostics(raw_state, raw_new_state, timestep, raw_diagnostics) diagnostics.update( restore_data_arrays_with_properties(raw_diagnostics, self.diagnostic_properties, state, self.input_properties, ignore_missing=True)) new_state.update( restore_data_arrays_with_properties(raw_new_state, self.output_properties, state, self.input_properties, ignore_missing=True)) gfs_output_quantities = list(self._gfs_output_properties.keys()) for tracer in self.prepend_tracers: gfs_output_quantities.append(tracer[0]) remaining = set(tendencies.keys()).difference(gfs_output_quantities) for name in remaining: new_state[name] = state[ name] + tendencies[name] * timestep.total_seconds() for key in state.keys(): if key not in new_state: new_state[key] = state[key] check_new_state = { name: quantity for (name, quantity) in new_state.items() if name in self.output_properties.keys() } self._diagnostic_checker.check_diagnostics(diagnostics) self._output_checker.check_outputs(check_new_state) return diagnostics, new_state