Merge pull request #10 from Nonannet/dev

New vector function added

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 ._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
+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",
 ]

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)
 
 

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)
 
-
+@overload
-def abs(x: T) -> T:
+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))

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

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()