diff --git a/src/copapy/__init__.py b/src/copapy/__init__.py index dfe53fe..2670500 100644 --- a/src/copapy/__init__.py +++ b/src/copapy/__init__.py @@ -1,7 +1,7 @@ from ._target import Target from ._basic_types import NumLike, variable, generic_sdb, iif -from ._vectors import vector -from ._math import sqrt, abs, sin, cos, tan, asin, acos, atan, atan2, log, exp, pow, get_42 +from ._vectors import vector, distance, scalar_projection, angle_between, rotate_vector, vector_projection +from ._math import sqrt, abs, sin, cos, tan, asin, acos, atan, atan2, log, exp, pow, get_42, clamp, min, max __all__ = [ "Target", @@ -22,5 +22,13 @@ __all__ = [ "log", "exp", "pow", - "get_42" + "get_42", + "clamp", + "min", + "max", + "distance", + "scalar_projection", + "angle_between", + "rotate_vector", + "vector_projection", ] diff --git a/src/copapy/_basic_types.py b/src/copapy/_basic_types.py index 58fefe6..5782198 100644 --- a/src/copapy/_basic_types.py +++ b/src/copapy/_basic_types.py @@ -7,9 +7,11 @@ NumLike: TypeAlias = 'variable[int] | variable[float] | variable[bool] | int | f unifloat: TypeAlias = 'variable[float] | float' uniint: TypeAlias = 'variable[int] | int' unibool: TypeAlias = 'variable[bool] | bool' +uniboolint: TypeAlias = 'variable[bool] | bool | variable[int] | int' TCPNum = TypeVar("TCPNum", bound='variable[Any]') TNum = TypeVar("TNum", int, float, bool) +TVarNumb: TypeAlias = 'variable[Any] | int | float | bool' stencil_cache: dict[tuple[str, str], stencil_database] = {} @@ -113,6 +115,8 @@ class variable(Generic[TNum], Net): @overload def __add__(self, other: NumLike) -> 'variable[float] | variable[int]': ... def __add__(self, other: NumLike) -> Any: + if isinstance(other, int | float) and other == 0: + return self return add_op('add', [self, other], True) @overload @@ -120,6 +124,8 @@ class variable(Generic[TNum], Net): @overload def __radd__(self, other: float) -> 'variable[float]': ... def __radd__(self, other: NumLike) -> Any: + if isinstance(other, int | float) and other == 0: + return self return add_op('add', [self, other], True) @overload @@ -185,27 +191,27 @@ class variable(Generic[TNum], Net): def __neg__(self: TCPNum) -> TCPNum: return cast(TCPNum, add_op('sub', [variable(0), self])) - def __gt__(self, other: NumLike) -> 'variable[bool]': + def __gt__(self, other: TVarNumb) -> 'variable[bool]': ret = add_op('gt', [self, other]) return variable(ret.source, dtype='bool') - def __lt__(self, other: NumLike) -> 'variable[bool]': + def __lt__(self, other: TVarNumb) -> 'variable[bool]': ret = add_op('gt', [other, self]) return variable(ret.source, dtype='bool') - def __ge__(self, other: NumLike) -> 'variable[bool]': + def __ge__(self, other: TVarNumb) -> 'variable[bool]': ret = add_op('ge', [self, other]) return variable(ret.source, dtype='bool') - def __le__(self, other: NumLike) -> 'variable[bool]': + def __le__(self, other: TVarNumb) -> 'variable[bool]': ret = add_op('ge', [other, self]) return variable(ret.source, dtype='bool') - def __eq__(self, other: NumLike) -> 'variable[bool]': # type: ignore + def __eq__(self, other: TVarNumb) -> 'variable[bool]': # type: ignore ret = add_op('eq', [self, other], True) return variable(ret.source, dtype='bool') - def __ne__(self, other: NumLike) -> 'variable[bool]': # type: ignore + def __ne__(self, other: TVarNumb) -> 'variable[bool]': # type: ignore ret = add_op('ne', [self, other], True) return variable(ret.source, dtype='bool') @@ -249,34 +255,34 @@ class variable(Generic[TNum], Net): return super().__hash__() # Bitwise and shift operations for cp[int] - def __lshift__(self, other: uniint) -> 'variable[int]': + def __lshift__(self, other: uniboolint) -> 'variable[int]': return add_op('lshift', [self, other]) - def __rlshift__(self, other: uniint) -> 'variable[int]': + def __rlshift__(self, other: uniboolint) -> 'variable[int]': return add_op('lshift', [other, self]) - def __rshift__(self, other: uniint) -> 'variable[int]': + def __rshift__(self, other: uniboolint) -> 'variable[int]': return add_op('rshift', [self, other]) - def __rrshift__(self, other: uniint) -> 'variable[int]': + def __rrshift__(self, other: uniboolint) -> 'variable[int]': return add_op('rshift', [other, self]) - def __and__(self, other: uniint) -> 'variable[int]': + def __and__(self, other: uniboolint) -> 'variable[int]': return add_op('bwand', [self, other], True) - def __rand__(self, other: uniint) -> 'variable[int]': + def __rand__(self, other: uniboolint) -> 'variable[int]': return add_op('rwand', [other, self], True) - def __or__(self, other: uniint) -> 'variable[int]': + def __or__(self, other: uniboolint) -> 'variable[int]': return add_op('bwor', [self, other], True) - def __ror__(self, other: uniint) -> 'variable[int]': + def __ror__(self, other: uniboolint) -> 'variable[int]': return add_op('bwor', [other, self], True) - def __xor__(self, other: uniint) -> 'variable[int]': + def __xor__(self, other: uniboolint) -> 'variable[int]': return add_op('bwxor', [self, other], True) - def __rxor__(self, other: uniint) -> 'variable[int]': + def __rxor__(self, other: uniboolint) -> 'variable[int]': return add_op('bwxor', [other, self], True) diff --git a/src/copapy/_math.py b/src/copapy/_math.py index 0a15858..6e57f6a 100644 --- a/src/copapy/_math.py +++ b/src/copapy/_math.py @@ -70,7 +70,7 @@ def pow(x: VecNumLike, y: VecNumLike) -> Any: result of x**y """ if isinstance(x, vector) or isinstance(y, vector): - return map2(x, y, pow) + return _map2(x, y, pow) if isinstance(y, int) and 0 <= y < 8: if y == 0: return 1 @@ -218,7 +218,7 @@ def atan2(x: VecNumLike, y: VecNumLike) -> Any: Result in radian """ if isinstance(x, vector) or isinstance(y, vector): - return map2(x, y, atan2) + return _map2(x, y, atan2) if isinstance(x, variable) or isinstance(y, variable): return add_op('atan2', [x, y]) return math.atan2(x, y) @@ -278,8 +278,11 @@ def get_42(x: NumLike) -> variable[float] | float: return add_op('get_42', [x, x]) return float((int(x) * 3.0 + 42.0) * 5.0 + 21.0) - -def abs(x: T) -> T: +@overload +def abs(x: U) -> U: ... +@overload +def abs(x: variable[U]) -> variable[U]: ... +def abs(x: U | variable[U]) -> Any: """Absolute value function Arguments: @@ -292,7 +295,93 @@ def abs(x: T) -> T: return ret # pyright: ignore[reportReturnType] -def map2(self: VecNumLike, other: VecNumLike, func: Callable[[Any, Any], variable[U] | U]) -> vector[U]: +@overload +def clamp(x: variable[U], min_value: U | variable[U], max_value: U | variable[U]) -> variable[U]: ... +@overload +def clamp(x: U | variable[U], min_value: variable[U], max_value: U | variable[U]) -> variable[U]: ... +@overload +def clamp(x: U | variable[U], min_value: U | variable[U], max_value: variable[U]) -> variable[U]: ... +@overload +def clamp(x: U, min_value: U, max_value: U) -> U: ... +@overload +def clamp(x: vector[U], min_value: 'U | variable[U]', max_value: 'U | variable[U]') -> vector[U]: ... +def clamp(x: U | variable[U] | vector[U], min_value: U | variable[U], max_value: U | variable[U]) -> Any: + """Clamp function to limit a value between a minimum and maximum. + + Arguments: + x: Input value + min_value: Minimum limit + max_value: Maximum limit + + Returns: + Clamped value of x + """ + if isinstance(x, vector): + return vector(clamp(comp, min_value, max_value) for comp in x.values) + + return (x < min_value) * min_value + \ + (x > max_value) * max_value + \ + ((x >= min_value) & (x <= max_value)) * x + + +@overload +def min(x: variable[U], y: U | variable[U]) -> variable[U]: ... +@overload +def min(x: U | variable[U], y: variable[U]) -> variable[U]: ... +@overload +def min(x: U, y: U) -> U: ... +def min(x: U | variable[U], y: U | variable[U]) -> Any: + """Minimum function to get the smaller of two values. + + Arguments: + x: First value + y: Second value + + Returns: + Minimum of x and y + """ + return (x < y) * x + (x >= y) * y + + +@overload +def max(x: variable[U], y: U | variable[U]) -> variable[U]: ... +@overload +def max(x: U | variable[U], y: variable[U]) -> variable[U]: ... +@overload +def max(x: U, y: U) -> U: ... +def max(x: U | variable[U], y: U | variable[U]) -> Any: + """Maximum function to get the larger of two values. + + Arguments: + x: First value + y: Second value + + Returns: + Maximum of x and y + """ + return (x > y) * x + (x <= y) * y + + +@overload +def lerp(v1: variable[U], v2: U | variable[U], t: U | variable[U]) -> variable[U]: ... +@overload +def lerp(v1: U | variable[U], v2: variable[U], t: U | variable[U]) -> variable[U]: ... +@overload +def lerp(v1: U | variable[U], v2: U | variable[U], t: variable[U]) -> variable[U]: ... +@overload +def lerp(v1: U, v2: U, t: U) -> U: ... +@overload +def lerp(v1: vector[U], v2: vector[U], t: 'U | variable[U]') -> vector[U]: ... +def lerp(v1: U | variable[U] | vector[U], v2: U | variable[U] | vector[U], t: U | variable[U]) -> Any: + """Linearly interpolate between two values or vectors v1 and v2 by a factor t.""" + if isinstance(v1, vector) or isinstance(v2, vector): + assert isinstance(v1, vector) and isinstance(v2, vector), "None or both v1 and v2 must be vectors." + assert len(v1.values) == len(v2.values), "Vectors must be of the same length." + return vector(lerp(vv1, vv2, t) for vv1, vv2 in zip(v1.values, v2.values)) + return v1 * (1 - t) + v2 * t + + +def _map2(self: VecNumLike, other: VecNumLike, func: Callable[[Any, Any], variable[U] | U]) -> vector[U]: """Applies a function to each element of the vector and a second vector or scalar.""" if isinstance(self, vector) and isinstance(other, vector): return vector(func(x, y) for x, y in zip(self.values, other.values)) diff --git a/src/copapy/_vectors.py b/src/copapy/_vectors.py index 6f70149..4664912 100644 --- a/src/copapy/_vectors.py +++ b/src/copapy/_vectors.py @@ -81,7 +81,7 @@ class vector(Generic[T]): @overload def __mul__(self: 'vector[int]', other: VecIntLike) -> 'vector[int]': ... @overload - def __mul__(self: 'vector[float]', other: 'vector[int] | float | int | variable[int]') -> 'vector[float]': ... + def __mul__(self: 'vector[float]', other: VecNumLike) -> 'vector[float]': ... @overload def __mul__(self, other: VecNumLike) -> 'vector[int] | vector[float]': ... def __mul__(self, other: VecNumLike) -> Any: @@ -118,7 +118,7 @@ class vector(Generic[T]): @overload def dot(self, other: 'vector[int] | vector[float]') -> float | int | variable[float] | variable[int]: ... def dot(self, other: 'vector[int] | vector[float]') -> Any: - assert len(self.values) == len(other.values) + assert len(self.values) == len(other.values), "Vectors must be of same length." return sum(a * b for a, b in zip(self.values, other.values)) # @ operator @@ -135,7 +135,7 @@ class vector(Generic[T]): def cross(self: 'vector[float]', other: 'vector[float]') -> 'vector[float]': """3D cross product""" - assert len(self.values) == 3 and len(other.values) == 3 + assert len(self.values) == 3 and len(other.values) == 3, "Both vectors must be 3-dimensional." a1, a2, a3 = self.values b1, b2, b3 = other.values return vector([ @@ -162,6 +162,9 @@ class vector(Generic[T]): """Returns a normalized (unit length) version of the vector.""" mag = self.magnitude() + epsilon return self / mag + + def __neg__(self) -> 'vector[float] | vector[int]': + return vector(-a for a in self.values) def __iter__(self) -> Iterable[variable[T] | T]: return iter(self.values) @@ -169,3 +172,56 @@ class vector(Generic[T]): def map(self, func: Callable[[Any], variable[U] | U]) -> 'vector[U]': """Applies a function to each element of the vector and returns a new vector.""" return vector(func(x) for x in self.values) + + +# Utility functions for 3D vectors with two arguments + +def cross_product(v1: vector[float], v2: vector[float]) -> vector[float]: + """Calculate the cross product of two 3D vectors.""" + return v1.cross(v2) + + +def dot_product(v1: vector[float], v2: vector[float]) -> 'float | variable[float]': + """Calculate the dot product of two vectors.""" + return v1.dot(v2) + + +def distance(v1: vector[float], v2: vector[float]) -> 'float | variable[float]': + """Calculate the Euclidean distance between two vectors.""" + diff = v1 - v2 + return diff.magnitude() + + +def scalar_projection(v1: vector[float], v2: vector[float]) -> 'float | variable[float]': + """Calculate the scalar projection of v1 onto v2.""" + dot_prod = v1.dot(v2) + mag_v2 = v2.magnitude() + epsilon + return dot_prod / mag_v2 + + +def vector_projection(v1: vector[float], v2: vector[float]) -> vector[float]: + """Calculate the vector projection of v1 onto v2.""" + dot_prod = v1.dot(v2) + mag_v2_squared = v2.magnitude() ** 2 + epsilon + scalar_proj = dot_prod / mag_v2_squared + return v2 * scalar_proj + + +def angle_between(v1: vector[float], v2: vector[float]) -> 'float | variable[float]': + """Calculate the angle in radians between two vectors.""" + dot_prod = v1.dot(v2) + mag_v1 = v1.magnitude() + mag_v2 = v2.magnitude() + cos_angle = dot_prod / (mag_v1 * mag_v2 + epsilon) + return cp.acos(cos_angle) + + +def rotate_vector(v: vector[float], axis: vector[float], angle: 'float | variable[float]') -> vector[float]: + """Rotate vector v around a given axis by a specified angle using Rodrigues' rotation formula.""" + k = axis.normalize() + cos_angle = cp.cos(angle) + sin_angle = cp.sin(angle) + term1 = v * cos_angle + term2 = k.cross(v) * sin_angle + term3 = k * (k.dot(v)) * (1 - cos_angle) + return term1 + term2 + term3 diff --git a/tests/test_vector.py b/tests/test_vector.py index e59749d..db44649 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -1,3 +1,4 @@ +import math import copapy as cp import pytest @@ -20,8 +21,10 @@ def test_compiled_vectors(): t4 = ((t3 * t1) * 2).sum() t5 = ((t3 * t1) * 2).magnitude() + t6 = cp.angle_between(cp.vector([cp.variable(5.0), 0.0, 0.0]), cp.vector([5.0, 5.0, 0.0])) + tg = cp.Target() - tg.compile(t2, t4, t5) + tg.compile(t2, t4, t5, t6) tg.run() assert isinstance(t2, cp.variable) @@ -33,6 +36,70 @@ def test_compiled_vectors(): assert isinstance(t5, cp.variable) assert tg.read_value(t5) == pytest.approx(((10/1*2)**2 + (12/2*2)**2 + (14/3*2)**2) ** 0.5, 0.001) # pyright: ignore[reportUnknownMemberType] + assert isinstance(t6, cp.variable) + assert tg.read_value(t6) == pytest.approx(math.pi / 4, 0.001), tg.read_value(t6) # pyright: ignore[reportUnknownMemberType] + + +def test_non_compiled_vector_operations(): + v1 = cp.vector([1.0, 2.0, 3.0]) + v2 = cp.vector([4.0, 5.0, 6.0]) + + # dot product + assert v1.dot(v2) == pytest.approx(1*4 + 2*5 + 3*6) # pyright: ignore[reportUnknownMemberType] + + # cross product + cross = v1.cross(v2) + assert isinstance(cross, cp.vector) + assert cross.values == pytest.approx((-3.0, 6.0, -3.0)) # pyright: ignore[reportUnknownMemberType] + + # magnitude + mag = v1.magnitude() + assert mag == pytest.approx((1**2 + 2**2 + 3**2) ** 0.5) # pyright: ignore[reportUnknownMemberType] + + # normalize + norm = v1.normalize() + norm_mag = norm.magnitude() + assert norm_mag == pytest.approx(1.0) # pyright: ignore[reportUnknownMemberType] + + # distance + dist = cp.distance(v1, v2) + assert dist == pytest.approx(((1-4)**2 + (2-5)**2 + (3-6)**2) ** 0.5) # pyright: ignore[reportUnknownMemberType] + + # scalar projection + scalar_proj = cp.scalar_projection(v1, v2) + expected_scalar_proj = v1 @ v2 / (v2.magnitude() + 1e-20) + assert scalar_proj == pytest.approx(expected_scalar_proj) # pyright: ignore[reportUnknownMemberType] + + # vector projection + vector_proj = cp.vector_projection(v1, v2) + assert isinstance(vector_proj, cp.vector) + # Check direction and magnitude + proj_mag = vector_proj.magnitude() + tmp = cp.abs(scalar_proj) + assert proj_mag == pytest.approx(tmp) # pyright: ignore[reportUnknownMemberType] + + # angle between + angle = cp.angle_between(v1, v2) + import math + expected_angle = cp.acos(v1 @ v2 / ((v1.magnitude()) * (v2.magnitude()) + 1e-20)) + assert angle == pytest.approx(expected_angle) # pyright: ignore[reportUnknownMemberType] + + # rotate_vector + axis = cp.vector([0.0, 0.0, 1.0]) + angle_rad = math.pi / 2 + rotated = cp.rotate_vector(v1, axis, angle_rad) + assert isinstance(rotated, cp.vector) + # Rotating [1,2,3] 90deg around z should give [-2,1,3] (approx) + assert rotated.values[0] == pytest.approx(-2.0, abs=1e-6) # pyright: ignore[reportUnknownMemberType] + assert rotated.values[1] == pytest.approx(1.0, abs=1e-6) # pyright: ignore[reportUnknownMemberType] + assert rotated.values[2] == pytest.approx(3.0, abs=1e-6) # pyright: ignore[reportUnknownMemberType] + + +if __name__ == "__main__": + test_vectors_init() + test_compiled_vectors() + test_vector_operations() + print('Finished!') if __name__ == "__main__": test_compiled_vectors()