diff --git a/src/copapy/_tensors.py b/src/copapy/_tensors.py index 522b7e4..4eb3f43 100644 --- a/src/copapy/_tensors.py +++ b/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() diff --git a/src/copapy/_vectors.py b/src/copapy/_vectors.py index 2ba51ed..96af715 100644 --- a/src/copapy/_vectors.py +++ b/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: