diff --git a/sasdata/quantities/math_operations_test.py b/sasdata/quantities/math_operations_test.py new file mode 100644 index 0000000..5bda5a2 --- /dev/null +++ b/sasdata/quantities/math_operations_test.py @@ -0,0 +1,152 @@ +""" Tests for math operations """ + +import pytest + +import numpy as np +from sasdata.quantities.quantity import NamedQuantity, tensordot, transpose +from sasdata.quantities import units + +order_list = [ + [0, 1, 2, 3], + [0, 2, 1], + [1, 0], + [0, 1], + [2, 0, 1], + [3, 1, 2, 0] +] + +@pytest.mark.parametrize("order", order_list) +def test_transpose_raw(order: list[int]): + """ Check that the transpose operation changes the order of indices correctly - uses sizes as way of tracking""" + + input_shape = tuple([i+1 for i in range(len(order))]) + expected_shape = tuple([i+1 for i in order]) + + input_mat = np.zeros(input_shape) + + measured_mat = transpose(input_mat, axes=tuple(order)) + + assert measured_mat.shape == expected_shape + + +@pytest.mark.parametrize("order", order_list) +def test_transpose_raw(order: list[int]): + """ Check that the transpose operation changes the order of indices correctly - uses sizes as way of tracking""" + input_shape = tuple([i + 1 for i in range(len(order))]) + expected_shape = tuple([i + 1 for i in order]) + + input_mat = NamedQuantity("testmat", np.zeros(input_shape), units=units.none) + + measured_mat = transpose(input_mat, axes=tuple(order)) + + assert measured_mat.value.shape == expected_shape + + +rng_seed = 1979 +tensor_product_with_identity_sizes = (4,6,5) + +@pytest.mark.parametrize("index, size", [tup for tup in enumerate(tensor_product_with_identity_sizes)]) +def test_tensor_product_with_identity_quantities(index, size): + """ Check the correctness of the tensor product by multiplying by the identity (quantity, quantity)""" + np.random.seed(rng_seed) + + x = NamedQuantity("x", np.random.rand(*tensor_product_with_identity_sizes), units=units.meters) + y = NamedQuantity("y", np.eye(size), units.seconds) + + z = tensordot(x, y, index, 0) + + # Check units + assert z.units == units.meters * units.seconds + + # Expected sizes - last index gets moved to end + output_order = [i for i in (0, 1, 2) if i != index] + [index] + output_sizes = [tensor_product_with_identity_sizes[i] for i in output_order] + + assert z.value.shape == tuple(output_sizes) + + # Restore original order and check + reverse_order = [-1, -1, -1] + for to_index, from_index in enumerate(output_order): + reverse_order[from_index] = to_index + + z_reordered = transpose(z, axes = tuple(reverse_order)) + + assert z_reordered.value.shape == tensor_product_with_identity_sizes + + # Check values + + mat_in = x.in_si() + mat_out = transpose(z, axes=tuple(reverse_order)).in_si() + + assert np.all(np.abs(mat_in - mat_out) < 1e-10) + + +@pytest.mark.parametrize("index, size", [tup for tup in enumerate(tensor_product_with_identity_sizes)]) +def test_tensor_product_with_identity_quantity_matrix(index, size): + """ Check the correctness of the tensor product by multiplying by the identity (quantity, matrix)""" + np.random.seed(rng_seed) + + x = NamedQuantity("x", np.random.rand(*tensor_product_with_identity_sizes), units.meters) + y = np.eye(size) + + z = tensordot(x, y, index, 0) + + assert z.units == units.meters + + # Expected sizes - last index gets moved to end + output_order = [i for i in (0, 1, 2) if i != index] + [index] + output_sizes = [tensor_product_with_identity_sizes[i] for i in output_order] + + assert z.value.shape == tuple(output_sizes) + + # Restore original order and check + reverse_order = [-1, -1, -1] + for to_index, from_index in enumerate(output_order): + reverse_order[from_index] = to_index + + z_reordered = transpose(z, axes = tuple(reverse_order)) + + assert z_reordered.value.shape == tensor_product_with_identity_sizes + + # Check values + + mat_in = x.in_si() + mat_out = transpose(z, axes=tuple(reverse_order)).in_si() + + assert np.all(np.abs(mat_in - mat_out) < 1e-10) + + +@pytest.mark.parametrize("index, size", [tup for tup in enumerate(tensor_product_with_identity_sizes)]) +def test_tensor_product_with_identity_matrix_quantity(index, size): + """ Check the correctness of the tensor product by multiplying by the identity (matrix, quantity)""" + np.random.seed(rng_seed) + + x = np.random.rand(*tensor_product_with_identity_sizes) + y = NamedQuantity("y", np.eye(size), units.seconds) + + z = tensordot(x, y, index, 0) + + assert z.units == units.seconds + + + # Expected sizes - last index gets moved to end + output_order = [i for i in (0, 1, 2) if i != index] + [index] + output_sizes = [tensor_product_with_identity_sizes[i] for i in output_order] + + assert z.value.shape == tuple(output_sizes) + + # Restore original order and check + reverse_order = [-1, -1, -1] + for to_index, from_index in enumerate(output_order): + reverse_order[from_index] = to_index + + z_reordered = transpose(z, axes = tuple(reverse_order)) + + assert z_reordered.value.shape == tensor_product_with_identity_sizes + + # Check values + + mat_in = x + mat_out = transpose(z, axes=tuple(reverse_order)).in_si() + + assert np.all(np.abs(mat_in - mat_out) < 1e-10) diff --git a/sasdata/quantities/quantity.py b/sasdata/quantities/quantity.py index 852469d..c4ba7da 100644 --- a/sasdata/quantities/quantity.py +++ b/sasdata/quantities/quantity.py @@ -23,15 +23,23 @@ ################### Quantity based operations, need to be here to avoid cyclic dependencies ##################### - -def transpose(a: Union["Quantity[ArrayLike]", ArrayLike]): - """ Transpose an array or an array based quantity """ +def transpose(a: Union["Quantity[ArrayLike]", ArrayLike], axes: tuple | None = None): + """ Transpose an array or an array based quantity, can also do reordering of axes""" if isinstance(a, Quantity): - return DerivedQuantity(value=np.transpose(a.value), - units=a.units, - history=QuantityHistory.apply_operation(Transpose, a.history)) + + if axes is None: + return DerivedQuantity(value=np.transpose(a.value, axes=axes), + units=a.units, + history=QuantityHistory.apply_operation(Transpose, a.history)) + + else: + return DerivedQuantity(value=np.transpose(a.value, axes=axes), + units=a.units, + history=QuantityHistory.apply_operation(Transpose, a.history, axes=axes)) + else: - return np.transpose(a) + return np.transpose(a, axes=axes) + def dot(a: Union["Quantity[ArrayLike]", ArrayLike], b: Union["Quantity[ArrayLike]", ArrayLike]): """ Dot product of two arrays or two array based quantities """ @@ -43,10 +51,10 @@ def dot(a: Union["Quantity[ArrayLike]", ArrayLike], b: Union["Quantity[ArrayLike # If its only one of them that is a quantity, convert the other one if not a_is_quantity: - a = Quantity(a, units.dimensionless) + a = Quantity(a, units.none) if not b_is_quantity: - b = Quantity(b, units.dimensionless) + b = Quantity(b, units.none) return DerivedQuantity( value=np.dot(a.value, b.value), @@ -57,6 +65,18 @@ def dot(a: Union["Quantity[ArrayLike]", ArrayLike], b: Union["Quantity[ArrayLike return np.dot(a, b) def tensordot(a: Union["Quantity[ArrayLike]", ArrayLike] | ArrayLike, b: Union["Quantity[ArrayLike]", ArrayLike], a_index: int, b_index: int): + """ Tensor dot product - equivalent to contracting two tensors, such as + + A_{i0, i1, i2, i3...} and B_{j0, j1, j2...} + + e.g. if a_index is 1 and b_index is zero, it will be the sum + + C_{i0, i2, i3 ..., j1, j2 ...} = sum_k A_{i0, k, i2, i3 ...} B_{k, j1, j2 ...} + + (I think, have to check what happens with indices TODO!) + + """ + a_is_quantity = isinstance(a, Quantity) b_is_quantity = isinstance(b, Quantity) @@ -65,10 +85,10 @@ def tensordot(a: Union["Quantity[ArrayLike]", ArrayLike] | ArrayLike, b: Union[" # If its only one of them that is a quantity, convert the other one if not a_is_quantity: - a = Quantity(a, units.dimensionless) + a = Quantity(a, units.none) if not b_is_quantity: - b = Quantity(b, units.dimensionless) + b = Quantity(b, units.none) return DerivedQuantity( value=np.tensordot(a.value, b.value, axes=(a_index, b_index)), @@ -791,11 +811,15 @@ def __eq__(self, other): # Matrix operations # -class Transpose(UnaryOperation): +class Transpose(Operation): """ Transpose operation - as per numpy""" serialisation_name = "transpose" + def __init__(self, a: Operation, axes: tuple[int] | None = None): + self.a = a + self.axes = axes + def evaluate(self, variables: dict[int, T]) -> T: return np.transpose(self.a.evaluate(variables)) @@ -806,9 +830,27 @@ def _clean(self): clean_a = self.a._clean() return Transpose(clean_a) + + def _serialise_parameters(self) -> dict[str, Any]: + if self.axes is None: + return { "a": self.a._serialise_json() } + else: + return { + "a": self.a._serialise_json(), + "axes": list(self.axes) + } + + @staticmethod def _deserialise(parameters: dict) -> "Operation": - return Transpose(Operation.deserialise_json(parameters["a"])) + if "axes" in parameters: + return Transpose( + a=Operation.deserialise_json(parameters["a"]), + axes=tuple(parameters["axes"])) + else: + return Transpose( + a=Operation.deserialise_json(parameters["a"])) + def _summary_open(self): return "Transpose" @@ -974,6 +1016,10 @@ def jacobian(self) -> list[Operation]: # Use the hash value to specify the variable of differentiation return [self.operation_tree.derivative(key) for key in self.reference_key_list] + def _recalculate(self): + """ Recalculate the value of this object - primary use case is for testing """ + return self.operation_tree.evaluate(self.references) + def variance_propagate(self, quantity_units: Unit, covariances: dict[tuple[int, int]: "Quantity"] = {}): """ Do standard error propagation to calculate the uncertainties associated with this quantity @@ -985,14 +1031,6 @@ def variance_propagate(self, quantity_units: Unit, covariances: dict[tuple[int, raise NotImplementedError("User specified covariances not currently implemented") jacobian = self.jacobian() - # jacobian_units = [quantity_units / self.references[key].units for key in self.reference_key_list] - # - # # Evaluate the jacobian - # # TODO: should we use quantities here, does that work automatically? - # evaluated_jacobian = [Quantity( - # value=entry.evaluate(self.si_reference_values), - # units=unit.si_equivalent()) - # for entry, unit in zip(jacobian, jacobian_units)] evaluated_jacobian = [entry.evaluate(self.references) for entry in jacobian]