Docstrings for vector and tensor type extended

src/copapy/_tensors.py

@@ -128,7 +128,14 @@ class tensor(ArrayType[TNum]):
         return self.shape[0]
 
     def get_scalar(self: 'tensor[TNum]', *key: int) -> TNum | value[TNum]:
-        """Get a single scalar value from the tensor given multi-dimensional indices."""
+        """Get a single scalar value from the tensor given multi-dimensional indices.
+
+        Arguments:
+            *key: Variable number of indices for each dimension.
+
+        Returns:
+            The scalar value at the specified indices.
+        """
         assert len(key) == self.ndim, f"Expected {self.ndim} indices, got {len(key)}"
         flat_idx = self._get_flat_index(key)
         return self.values[flat_idx]
@@ -567,7 +574,11 @@ class tensor(ArrayType[TNum]):
     @overload
     def trace(self: 'tensor[float]') -> float | value[float]: ...
     def trace(self) -> Any:
-        """Calculate the trace (sum of diagonal elements)."""
+        """Calculate the trace (sum of diagonal elements).
+
+        Returns:
+            The sum of diagonal elements.
+        """
         assert self.ndim == 2 and self.shape[0] == self.shape[1], "Trace is only defined for square matrices"
         return mixed_sum(self.get_scalar(i, i) for i in range(self.shape[0]))
 
@@ -615,7 +626,11 @@ class tensor(ArrayType[TNum]):
         return self.reshape(-1)
 
     def size(self) -> int:
-        """Return total number of elements."""
+        """Count number of elements over all dimensions.
+
+        Returns:
+            Total number of elements.
+        """
         size = 1
         for dim in self.shape:
             size *= dim
@@ -820,7 +835,11 @@ class tensor(ArrayType[TNum]):
         return tensor(result_vals, self.shape)
 
     def homogenize(self) -> 'tensor[TNum]':
-        """Convert all elements to copapy values if any element is a copapy value."""
+        """Convert all elements to Copapy values if any element is a Copapy value.
+
+        Returns:
+            Tensor with all elements converted to Copapy values, or the input tensor if no value is a Copapy value.
+        """
         if any(isinstance(val, value) for val in self.values):
             homogenized: tuple[value[Any], ...] = tuple(value(val) if not isinstance(val, value) else val for val in self.values)
             return tensor(homogenized, self.shape)
@@ -833,7 +852,13 @@ class tensor(ArrayType[TNum]):
 
 
 def zeros(shape: Sequence[int] | int) -> tensor[int]:
-    """Create a zero tensor of given shape."""
+    """Create a zero tensor of given shape.
+    
+    Arguments:
+        shape: shape of the tensor to create.
+
+    Returns:
+        New tensor of given shape, initialized with zeros."""
     if isinstance(shape, int):
         shape = (shape,)
 
@@ -845,7 +870,13 @@ def zeros(shape: Sequence[int] | int) -> tensor[int]:
 
 
 def ones(shape: Sequence[int] | int) -> tensor[int]:
-    """Create a tensor of ones with given shape."""
+    """Create a tensor of ones with given shape.
+    
+    Arguments:
+        shape: shape of the tensor to create.
+
+    Returns:
+        New tensor of given shape, initialized with ones."""
     if isinstance(shape, int):
         shape = (shape,)
 
@@ -930,7 +961,14 @@ def concat(tensors: Sequence[tensor[U] | vector[U]], axis: int = 0) -> tensor[U]
 
 
 def flatten(t: tensor[U]) -> tensor[U]:
-    """Flatten a tensor to a 1D tensor."""
+    """Flatten a tensor to a 1D tensor.
+
+    Arguments:
+        t: n-dimensional tensor.
+
+    Returns:
+        A 1D tensor containing all elements from the input tensor.
+    """
     return t.flatten()
 
 

src/copapy/_vectors.py

@@ -54,7 +54,14 @@ class vector(ArrayType[TNum]):
         return iter(self.values)
 
     def get_scalar(self, index: int) -> TNum | value[TNum]:
-        """Get a single scalar value from the vector."""
+        """Get a single scalar value from the vector.
+
+        Arguments:
+            index: The index of the element to retrieve.
+
+        Returns:
+            The scalar value at the specified index.
+        """
         return self.values[index]
 
     @overload
@@ -200,6 +207,14 @@ class vector(ArrayType[TNum]):
     @overload
     def dot(self, other: 'vector[int] | vector[float]') -> float | int | value[float] | value[int]: ...
     def dot(self, other: 'vector[int] | vector[float]') -> Any:
+        """Calculate the dot product of this vector with another.
+
+        Arguments:
+            other: Another vector of the same length.
+
+        Returns:
+            The dot product as a scalar value.
+        """
         assert len(self.values) == len(other.values), "Vectors must be of same length."
         return mixed_sum(a * b for a, b in zip(self.values, other.values))
 
@@ -216,7 +231,14 @@ class vector(ArrayType[TNum]):
         return self.dot(other)
 
     def cross(self: 'vector[float]', other: 'vector[float]') -> 'vector[float]':
-        """3D cross product"""
+        """Calculate the cross product of this vector with another 3D vector.
+
+        Arguments:
+            other: Another 3D vector.
+
+        Returns:
+            A new vector perpendicular to both vectors.
+        """
         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
@@ -286,20 +308,37 @@ class vector(ArrayType[TNum]):
     @overload
     def sum(self: 'vector[float]') -> float | value[float]: ...
     def sum(self) -> Any:
-        """Sum of all vector elements."""
+        """Sum of all vector elements.
+
+        Returns:
+            The sum of all vector elements as a scalar value.
+        """
         return mixed_sum(self.values)
 
     def magnitude(self) -> 'float | value[float]':
-        """Magnitude (length) of the vector."""
+        """Magnitude (length) of the vector.
+
+        Returns:
+            The magnitude of the vector as a scalar value.
+        """
         s = mixed_sum(a * a for a in self.values)
         return cp.sqrt(s)
 
     def normalize(self) -> 'vector[float]':
-        """Returns a normalized (unit length) version of the vector."""
+        """Calculate a normalized (unit length) version of the vector.
+
+        Returns:
+            A normalized (unit length) vector.
+        """
         mag = self.magnitude() + epsilon
         return self / mag
 
     def homogenize(self) -> 'vector[TNum]':
+        """Convert all elements to Copapy values if any element is a Copapy value.
+
+        Returns:
+            Vector with homogeneous element type.
+        """
         if any(isinstance(val, value) for val in self.values):
             return vector(mixed_homogenize(self))
         else: