def moments(quantity, moments_fn, mom_at_bottom=True):
"""
Create quantity with operation that evaluates moments_fn
:param quantity: Quantity
:param moments_fn: mlmc.moments.Moments child
:param mom_at_bottom: bool, if True moments_fn are underneath
:return: Quantity
"""
def eval_moments(x):
if mom_at_bottom:
mom = moments_fn.eval_all(x).transpose(
(0, 3, 1, 2)) # [M, R, N, 2]
else:
mom = moments_fn.eval_all(x).transpose(
(3, 0, 1, 2)) # [R, M, N, 2]
return mom.reshape((np.prod(mom.shape[:-2]), mom.shape[-2],
mom.shape[-1])) # [M, N, 2]
# Create quantity type which has moments_fn at the bottom
if mom_at_bottom:
moments_array_type = qt.ArrayType(shape=(moments_fn.size, ),
qtype=qt.ScalarType())
moments_qtype = quantity.qtype.replace_scalar(moments_array_type)
# Create quantity type that has moments_fn on the surface
else:
moments_qtype = qt.ArrayType(shape=(moments_fn.size, ),
qtype=quantity.qtype)
return mlmc.quantity.Quantity(quantity_type=moments_qtype,
input_quantities=[quantity],
operation=eval_moments)
def covariance(quantity, moments_fn, cov_at_bottom=True):
"""
Create quantity with operation that evaluates covariance matrix
:param quantity: Quantity
:param moments_fn: mlmc.moments.Moments child
:param cov_at_bottom: bool, if True cov matrices are underneath
:return: Quantity
"""
def eval_cov(x):
moments = moments_fn.eval_all(x)
mom_fine = moments[..., 0, :]
cov_fine = np.einsum('...i,...j', mom_fine, mom_fine)
if moments.shape[-2] == 1:
cov = np.array([cov_fine])
else:
mom_coarse = moments[..., 1, :]
cov_coarse = np.einsum('...i,...j', mom_coarse, mom_coarse)
cov = np.array([cov_fine, cov_coarse])
if cov_at_bottom:
cov = cov.transpose((1, 3, 4, 2, 0)) # [M, R, R, N, 2]
else:
cov = cov.transpose((3, 4, 1, 2, 0)) # [R, R, M, N, 2]
return cov.reshape(
(np.prod(cov.shape[:-2]), cov.shape[-2], cov.shape[-1]))
# Create quantity type which has covariance matrices at the bottom
if cov_at_bottom:
moments_array_type = qt.ArrayType(shape=(
moments_fn.size,
moments_fn.size,
),
qtype=qt.ScalarType())
moments_qtype = quantity.qtype.replace_scalar(moments_array_type)
# Create quantity type that has covariance matrices on the surface
else:
moments_qtype = qt.ArrayType(shape=(
moments_fn.size,
moments_fn.size,
),
qtype=quantity.qtype)
return mlmc.quantity.Quantity(quantity_type=moments_qtype,
input_quantities=[quantity],
operation=eval_cov)
def make_root_quantity(storage: SampleStorage, q_specs: List[QuantitySpec]):
"""
Create a root quantity that has QuantityStorage as the input quantity,
QuantityStorage is the only class that directly accesses the stored data.
Quantity type is created based on the q_spec parameter
:param storage: SampleStorage
:param q_specs: same as result format in simulation class
:return: QuantityStorage
"""
dict_types = []
for q_spec in q_specs:
scalar_type = qt.ScalarType(float)
array_type = qt.ArrayType(q_spec.shape, scalar_type)
field_type = qt.FieldType([(loc, array_type)
for loc in q_spec.locations])
ts_type = qt.TimeSeriesType(q_spec.times, field_type)
dict_types.append((q_spec.name, ts_type))
dict_type = qt.DictType(dict_types)
return QuantityStorage(storage, dict_type)
def _result_qtype(method, quantities):
"""
Determine QType from evaluation with given method and first few samples from storage
:param quantities: list of Quantities
:param method: ufunc function
:return: QType
"""
chunks_quantity_level = []
for q in quantities:
quantity_storage = q.get_quantity_storage()
# QuantityConst doesn't have QuantityStorage
if quantity_storage is None:
chunk_spec = ChunkSpec()
else:
chunk_spec = next(quantity_storage.chunks())
chunks_quantity_level.append(q.samples(chunk_spec))
result = method(*chunks_quantity_level) # numpy array of [M, <=10, 2]
qtype = qt.ArrayType(shape=result.shape[0],
qtype=Quantity._get_base_qtype(quantities))
return qtype
def wrap(value):
"""
Convert flat, bool or array (list) to Quantity
:param value: flat, bool, array (list) or Quantity
:return: Quantity
"""
if isinstance(value, Quantity):
return value
elif isinstance(value, (int, float)):
quantity = QuantityConst(quantity_type=qt.ScalarType(),
value=value)
elif isinstance(value, bool):
quantity = QuantityConst(quantity_type=qt.BoolType(), value=value)
elif isinstance(value, (list, np.ndarray)):
value = np.array(value)
qtype = qt.ArrayType(shape=value.shape, qtype=qt.ScalarType())
quantity = QuantityConst(quantity_type=qtype, value=value)
else:
raise ValueError(
"Values {} are not flat, bool or array (list)".format(value))
return quantity
def QArray(quantities):
flat_quantities = np.array(quantities).flatten()
qtype = Quantity._check_same_qtype(flat_quantities)
array_type = qt.ArrayType(np.array(quantities).shape, qtype)
return Quantity._concatenate(flat_quantities, qtype=array_type)