import numpy as np
from .exceptions import InvalidPropertyDictError
from .dataarray import DataArray
from .wildcard import (
get_wildcard_matches_and_dim_lengths, fill_dims_wildcard,
expand_array_wildcard_dims
)
def ensure_values_are_arrays(array_dict):
for name, value in array_dict.items():
if not isinstance(value, np.ndarray):
array_dict[name] = np.asarray(value)
def get_alias_or_name(name, output_properties, input_properties):
if 'alias' in output_properties[name].keys():
raw_name = output_properties[name]['alias']
elif name in input_properties.keys() and 'alias' in input_properties[name].keys():
raw_name = input_properties[name]['alias']
else:
raw_name = name
return raw_name
def check_array_shape(out_dims, raw_array, name, dim_lengths):
if len(out_dims) != len(raw_array.shape):
raise InvalidPropertyDictError(
'Returned array for {} has shape {} '
'which is incompatible with dims {} in properties'.format(
name, raw_array.shape, out_dims))
for dim, length in zip(out_dims, raw_array.shape):
if dim in dim_lengths.keys() and dim_lengths[dim] != length:
raise InvalidPropertyDictError(
'Dimension {} of quantity {} has length {}, but '
'another quantity has length {}'.format(
dim, name, length, dim_lengths[dim])
)
[docs]def restore_data_arrays_with_properties(
raw_arrays, output_properties, input_state, input_properties,
ignore_names=None, ignore_missing=False):
"""
Parameters
----------
raw_arrays : dict
A dictionary whose keys are quantity names and values are numpy arrays
containing the data for those quantities.
output_properties : dict
A dictionary whose keys are quantity names and values are dictionaries
with properties for those quantities. The property "dims" must be
present for each quantity not also present in input_properties. All
other properties are included as attributes on the output DataArray
for that quantity, including "units" which is required.
input_state : dict
A state dictionary that was used as input to a component for which
DataArrays are being restored.
input_properties : dict
A dictionary whose keys are quantity names and values are dictionaries
with input properties for those quantities. The property "dims" must be
present, indicating the dimensions that the quantity was transformed to
when taken as input to a component.
ignore_names : iterable of str, optional
Names to ignore when encountered in output_properties, will not be
included in the returned dictionary.
ignore_missing : bool, optional
If True, ignore any values in output_properties not present in
raw_arrays rather than raising an exception. Default is False.
Returns
-------
out_dict : dict
A dictionary whose keys are quantities and values are DataArrays
corresponding to those quantities, with data, shapes and attributes
determined from the inputs to this function.
Raises
------
InvalidPropertyDictError
When an output property is specified to have dims_like an input
property, but the arrays for the two properties have incompatible
shapes.
"""
raw_arrays = raw_arrays.copy()
if ignore_names is None:
ignore_names = []
if ignore_missing:
ignore_names = set(output_properties.keys()).difference(raw_arrays.keys()).union(ignore_names)
wildcard_names, dim_lengths = get_wildcard_matches_and_dim_lengths(
input_state, input_properties)
ensure_values_are_arrays(raw_arrays)
dims_from_out_properties = extract_output_dims_properties(
output_properties, input_properties, ignore_names)
out_dict = {}
for name, out_dims in dims_from_out_properties.items():
if name in ignore_names:
continue
raw_name = get_alias_or_name(name, output_properties, input_properties)
if '*' in out_dims:
for dim_name, length in zip(out_dims, raw_arrays[raw_name].shape):
if dim_name not in dim_lengths and dim_name != '*':
dim_lengths[dim_name] = length
out_dims_without_wildcard, target_shape = fill_dims_wildcard(
out_dims, dim_lengths, wildcard_names)
out_array = expand_array_wildcard_dims(
raw_arrays[raw_name], target_shape, name, out_dims)
else:
check_array_shape(out_dims, raw_arrays[raw_name], name, dim_lengths)
out_dims_without_wildcard = out_dims
out_array = raw_arrays[raw_name]
out_dict[name] = DataArray(
out_array,
dims=out_dims_without_wildcard,
attrs={'units': output_properties[name]['units']}
)
return out_dict
def extract_output_dims_properties(output_properties, input_properties, ignore_names):
return_array = {}
for name, properties in output_properties.items():
if name in ignore_names:
continue
elif 'dims' in properties.keys():
return_array[name] = properties['dims']
elif name not in input_properties.keys():
raise InvalidPropertyDictError(
'Output dims must be specified for {} in properties'.format(name))
elif 'dims' not in input_properties[name].keys():
raise InvalidPropertyDictError(
'Input dims must be specified for {} in properties'.format(name))
else:
return_array[name] = input_properties[name]['dims']
return return_array