diff --git a/quantio/vector.py b/quantio/vector.py index ecd0ff6..5d43512 100644 --- a/quantio/vector.py +++ b/quantio/vector.py @@ -52,18 +52,11 @@ def __sub__(self, other: Vector[T] | np.ndarray) -> Vector[T]: def __mul__(self, other: Vector | np.ndarray | float) -> np.ndarray: """Multipy this vector with either another vector or a scalar.""" - if isinstance(other, (float, int)): - other_elements = np.array([other]) - elif isinstance(other, _QuantityBase): - other_elements = np.array([other._base_value]) - elif isinstance(other, Vector): - other_elements = other.to_numpy() - elif isinstance(other, np.ndarray): - other_elements = other - else: - raise TypeError - - return self.to_numpy() * other_elements + return self.to_numpy() * _other_to_numpy(other) + + def __truediv__(self, other: Vector | np.ndarray | float) -> np.ndarray: + """Multipy this vector with either another vector or a scalar.""" + return self.to_numpy() / _other_to_numpy(other) def __eq__(self, other: object) -> bool: """Assess if this object is the same as another.""" @@ -71,3 +64,19 @@ def __eq__(self, other: object) -> bool: return False return np.all(other._elements == self._elements) + + +def _other_to_numpy(other: Vector | np.ndarray | float) -> np.ndarray: + if isinstance(other, (float, int)): + return np.array([other]) + + if isinstance(other, _QuantityBase): + return np.array([other._base_value]) + + if isinstance(other, Vector): + return other.to_numpy() + + if isinstance(other, np.ndarray): + return other + + raise TypeError diff --git a/test/test_vector.py b/test/test_vector.py index 4bda0c6..7f776d2 100644 --- a/test/test_vector.py +++ b/test/test_vector.py @@ -107,5 +107,53 @@ def test_multiply__wrong_dimension(): vec1 * vec2 +def test_divide__with_vector_of_quantities(): + vec1: Vector[Length] = Vector([Length(meters=1), Length(meters=2)]) + vec2: Vector[Length] = Vector([Length(meters=3), Length(meters=4)]) + + actual = vec2 / vec1 + assert np.all(actual == np.array([3, 2])) + + +def test_divide__with_vector_of_float(): + vec1: Vector[float] = Vector([1, 2]) + vec2: Vector[float] = Vector([3, 4]) + + actual = vec2 / vec1 + assert np.all(actual == np.array([3, 2])) + + +def test_divide__with_array(): + vec: Vector[Length] = Vector([Length(meters=3), Length(meters=4)]) + array = np.array([1, 2]) + + actual = vec / array + assert np.all(actual == np.array([3, 2])) + + +def test_divide__with_scalar_float(): + vec: Vector[Length] = Vector([Length(meters=2), Length(meters=4)]) + scalar = 2 + + actual = vec / scalar + assert np.all(actual == np.array([1, 2])) + + +def test_divide__with_scalar_quantitiy(): + vec: Vector[Length] = Vector([Length(meters=2), Length(meters=4)]) + scalar = Length(meters=2) + + actual = vec / scalar + assert np.all(actual == np.array([1, 2])) + + +def test_divide__wrong_dimension(): + vec1: Vector[Length] = Vector([Length(meters=1), Length(meters=2), Length(meters=2)]) + vec2: Vector[Length] = Vector([Length(meters=3), Length(meters=4)]) + + with pytest.raises(ValueError): + vec1 / vec2 + + if __name__ == "__main__": pytest.main([__file__, "-v"])