def test_ensure_no_shared_keys_with_shared_keys():
try:
ensure_no_shared_keys({'a': 1, 'b': 2}, {'e': 2, 'a': 1})
except SharedKeyError:
pass
except Exception as err:
raise err
else:
raise AssertionError(
'No exception raised but expected SharedKeyError.')
def test_ensure_no_shared_keys_with_no_shared_keys():
ensure_no_shared_keys({'a': 1, 'b': 2}, {'c': 2, 'd': 1})
ensure_no_shared_keys({'c': 2, 'd': 1}, {'a': 1, 'b': 2})
def test_ensure_no_shared_keys_empty_dicts():
ensure_no_shared_keys({}, {})
def test_ensure_no_shared_keys_one_empty_dict():
ensure_no_shared_keys({'a': 1, 'b': 2}, {})
ensure_no_shared_keys({}, {'a': 1, 'b': 2})
def get_default_state(component_list,
x={},
y={},
mid_levels={},
interface_levels={},
initial_state={}):
"""
Return a state dictionary required to run the model.
The model comprises of components in :code:`component_list`. If coordinate
values in :code:`x`, :code:`y`, :code:`mid_levels` and :code:`interface_levels`
are not provided, a single column
centered at 0 degrees south, 0 degrees east, with 30 vertical levels is returned.
Args:
component_list (iterable): The _components for which a default
state is required, in the order that they are called.
It is assumed diagnostic outputs are updated into
the state and passed into the next component, and that the state
from the previous component is passed into the next component. The
function will attempt to determine required inputs from the series
of _components.
x (dict,optional): A dictionary containing keys :code:`label`, :code:`values`,
and :code:`units`. :code:`label` refers to the name the coordinate axis will assume.
:code:`values` refers to the array of coordinate values. :code:`units` refers to the
units of the coordinate values.
If :code:`x` is an empty dictionary, a single default value of 0 degrees longitude
is used.
y (dict,optional): A dictionary containing keys :code:`label`, :code:`values`,
and :code:`units`. :code:`label` refers to the name the coordinate axis will assume.
:code:`values` refers to the array of coordinate values. :code:`units` refers to the
units of the coordinate values.
If :code:`y` is an empty dictionary, a single default value of 0 degrees latitude
is used.
mid_levels (dict,optional): A dictionary containing keys :code:`label`, :code:`values`,
and :code:`units`. :code:`label` refers to the name the coordinate axis will assume.
:code:`values` refers to the array of coordinate values. :code:`units` refers to the
units of the coordinate values.
If :code:`mid_levels` is an empty dictionary, 30 levels of arbitrary units are used.
interface_levels (dict,optional): A dictionary containing keys :code:`label`, :code:`values`,
and :code:`units`. :code:`label` refers to the name the coordinate axis will assume.
:code:`values` refers to the array of coordinate values. :code:`units` refers to the
units of the coordinate values.
If :code:`interface_levels` is an empty dictionary, 31 levels of arbitrary units are used.
initial_state (dict, optional): A dictionary containing some quantities that will
also be added to the final output state. Must not contain any quantities that
the output state will overwrite.
Returns:
default_state (dict):
A state dictionary containing the requested
quantities using the provided coordinate state.
Raises:
ValueError:
if any of the following conditions are satisfied:
* if :code:`component_list` is empty
* if the shape of :code:`x['values']` and :code:`y['values']` is not the same
* if only one of :code:`mid_levels` or :code:`interface_levels` is specified
* if vertical coordinates are not one dimensional
* if length of :code:`mid_levels['values']` is not one less than length
of :code:`interface['values']`
"""
if len(component_list) == 0:
raise ValueError('Component list must contain at least one component')
if (len(mid_levels.keys()) == 0) != (len(interface_levels.keys()) == 0):
raise ValueError('Both mid and interface levels must be specified')
output_state = {}
# Create 2D coordinate arrays
if len(x.keys()) == 0:
x_coordinate_values = np.zeros((1, ))
x_coordinate_label = 'longitude'
x_coordinate_units = 'degrees_east'
else:
x_coordinate_values = x['values']
x_coordinate_label = x['label']
x_coordinate_units = x['units']
if len(y.keys()) == 0:
y_coordinate_values = np.zeros((1, ))
y_coordinate_label = 'latitude'
y_coordinate_units = 'degrees_north'
else:
y_coordinate_values = y['values']
y_coordinate_label = y['label']
y_coordinate_units = y['units']
if len(mid_levels.keys()) == 0:
mid_levels_coordinate_values = np.arange(30)
mid_levels_coordinate_label = 'mid_levels'
mid_levels_coordinate_units = ''
else:
mid_levels_coordinate_values = mid_levels['values']
mid_levels_coordinate_label = mid_levels['label']
mid_levels_coordinate_units = mid_levels['units']
if len(interface_levels.keys()) == 0:
interface_levels_coordinate_values = np.arange(31)
interface_levels_coordinate_label = 'interface_levels'
interface_levels_coordinate_units = ''
else:
interface_levels_coordinate_values = interface_levels['values']
interface_levels_coordinate_label = interface_levels['label']
interface_levels_coordinate_units = interface_levels['units']
if not x_coordinate_values.ndim == y_coordinate_values.ndim:
raise ValueError('x and y coordinates must have the same shape')
if mid_levels_coordinate_values.ndim > 1:
raise ValueError(
'vertical coordinate mid_levels must be one dimensional.')
if interface_levels_coordinate_values.ndim > 1:
raise ValueError(
'vertical coordinate interface_levels must be one dimensional.')
if len(mid_levels_coordinate_values
) != len(interface_levels_coordinate_values) - 1:
raise ValueError(
'Interface levels must have one value more than mid levels')
use_2d_coordinate = False
two_dim_coord_dict = {}
if x_coordinate_values.ndim == 2:
if not x_coordinate_values.shape == y_coordinate_values.shape:
raise ValueError(
'If x and y are 2d coordinates, they must have the same shape')
two_dim_coord_dict[x_coordinate_label] = {}
two_dim_coord_dict[x_coordinate_label]['values'] = x_coordinate_values
two_dim_coord_dict[x_coordinate_label]['logical_dims'] = (
'logical_x_coordinate', 'logical_y_coordinate')
two_dim_coord_dict[x_coordinate_label]['units'] = x_coordinate_units
two_dim_coord_dict[y_coordinate_label] = {}
two_dim_coord_dict[y_coordinate_label]['values'] = y_coordinate_values
two_dim_coord_dict[y_coordinate_label]['logical_dims'] = (
'logical_x_coordinate', 'logical_y_coordinate')
two_dim_coord_dict[y_coordinate_label]['units'] = y_coordinate_units
output_state['x'] = DataArray(
x_coordinate_values,
dims=two_dim_coord_dict[x_coordinate_label]['logical_dims'],
attrs={
'units': two_dim_coord_dict[x_coordinate_label]['units'],
'label': x_coordinate_label
})
output_state['y'] = DataArray(
y_coordinate_values,
dims=two_dim_coord_dict[y_coordinate_label]['logical_dims'],
attrs={
'units': two_dim_coord_dict[y_coordinate_label]['units'],
'label': y_coordinate_label
})
x_coordinate_values = np.arange(
two_dim_coord_dict[x_coordinate_label]['values'].shape[0])
y_coordinate_values = np.arange(
two_dim_coord_dict[x_coordinate_label]['values'].shape[1])
x_coordinate_label = 'logical_x_coordinate'
y_coordinate_label = 'logical_y_coordinate'
x_coordinate_units = ''
y_coordinate_units = ''
use_2d_coordinate = True
x_coordinate = DataArray(x_coordinate_values,
dims=(x_coordinate_label),
attrs={
'units': x_coordinate_units,
'label': x_coordinate_label
})
output_state[x_coordinate_label] = x_coordinate
add_direction_names(x=x_coordinate_label)
y_coordinate = DataArray(y_coordinate_values,
dims=(y_coordinate_label),
attrs={
'units': y_coordinate_units,
'label': y_coordinate_label
})
output_state[y_coordinate_label] = y_coordinate
add_direction_names(y=y_coordinate_label)
mid_levels_coordinate = DataArray(mid_levels_coordinate_values,
dims=(mid_levels_coordinate_label, ),
attrs={
'units': mid_levels_coordinate_units,
'label': mid_levels_coordinate_label
})
output_state[mid_levels_coordinate_label] = mid_levels_coordinate
output_state['mid_levels'] = mid_levels_coordinate
add_direction_names(z=mid_levels_coordinate_label)
interface_levels_coordinate = DataArray(
interface_levels_coordinate_values,
dims=(interface_levels_coordinate_label, ),
attrs={
'units': interface_levels_coordinate_units,
'label': interface_levels_coordinate_label
})
output_state[
interface_levels_coordinate_label] = interface_levels_coordinate
output_state['interface_levels'] = interface_levels_coordinate
add_direction_names(z=interface_levels_coordinate_label)
if not use_2d_coordinate:
output_state['x'] = x_coordinate
output_state['y'] = y_coordinate
quantity_list = set()
temporary_description = copy.deepcopy(climt_quantity_descriptions)
additional_dimensions = {}
additional_descriptions = {}
for component in component_list:
quantity_list = quantity_list.union(set(component.inputs))
if hasattr(component, 'extra_dimensions'):
ensure_no_shared_keys(additional_dimensions,
component.extra_dimensions)
for dimension in component.extra_dimensions.keys():
if component.extra_dimensions[dimension].ndim > 1:
raise NotImplementedError(
'Two dimensional coordinates in extra_dimensions not yet supported'
)
output_state[dimension] = DataArray(
component.extra_dimensions[dimension], dims=(dimension, ))
additional_dimensions[dimension] = output_state[dimension]
if hasattr(component, 'quantity_descriptions'):
ensure_no_shared_keys(additional_descriptions,
component.quantity_descriptions)
additional_descriptions.update(component.quantity_descriptions)
temporary_description.update(additional_descriptions)
for name in quantity_list:
if name in output_state.keys():
continue
output_state[name] = get_default_values(
name, x_coordinate, y_coordinate, mid_levels_coordinate,
interface_levels_coordinate, initial_state, temporary_description,
additional_dimensions)
if use_2d_coordinate:
for physical_dimension in two_dim_coord_dict.keys():
output_state[name].coords[physical_dimension] = (
two_dim_coord_dict[physical_dimension]['logical_dims'],
two_dim_coord_dict[physical_dimension]['values'])
output_state[name][physical_dimension].attrs['units'] = \
two_dim_coord_dict[physical_dimension]['units']
ensure_no_shared_keys(initial_state, output_state)
output_state.update(initial_state)
if 'time' not in initial_state:
output_state['time'] = datetime(1, 1, 1)
return output_state
