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
