Source code for nvtools.vector

r"""
Vector math utils.

.. code-block:: python

   import nvtools as nv

   # create new vectors
   a = nv.Vector(1, -1, 1)
   b = nv.Vector(2, 1, -3)

   # useful methods
   a.length
   a.angle(b)
   a.scalar_product(b)
   a.cross_product(b)
"""

import numpy as np


[docs] class NVVector: """NV unit vectors.""" @property def vectors(self): r""" NV unit vectors. .. math:: \hat{n}_1 = \frac{1}{\sqrt{3}} \begin{pmatrix} 1 \\ 1 \\ 1 \end{pmatrix} \enspace \hat{n}_2 = \frac{1}{\sqrt{3}} \begin{pmatrix} -1 \\ -1 \\ 1 \end{pmatrix} \enspace \hat{n}_3 = \frac{1}{\sqrt{3}} \begin{pmatrix} -1 \\ 1 \\ -1 \end{pmatrix} \enspace \hat{n}_4 = \frac{1}{\sqrt{3}} \begin{pmatrix} 1 \\ -1 \\ -1 \end{pmatrix} """ return ( Vector(1, 1, 1).normalize(), Vector(-1, -1, 1).normalize(), Vector(-1, 1, -1).normalize(), Vector(1, -1, -1).normalize(), )
[docs] class Vector: """N-dim Vector.""" def __init__(self, *args): """N-dim Vector.""" if isinstance(args[0], str): if args[0].lower() == "x": self.vals = np.array([1, 0, 0]) elif args[0].lower() == "y": self.vals = np.array([0, 1, 0]) elif args[0].lower() == "z": self.vals = np.array([0, 0, 1]) else: raise ValueError("unknown input") elif hasattr(args[0], "__iter__"): self.vals = np.array(args[0]) else: self.vals = np.array([*args]) def __len__(self) -> int: """Dimensionality.""" return len(self.vals) def __abs__(self) -> float: """Length.""" return self.length def __str__(self) -> str: """String representation.""" return f"Vector({self.vals})" def __repr__(self) -> str: """String representation.""" return self.__str__() def __iter__(self): """Iteration.""" yield from self.vals def __neg__(self): """Negation.""" return self * -1 def __add__(self, other): """Vector addition.""" return Vector(np.add(self.vals, other.vals)) def __sub__(self, other): """Vector subtraction.""" return Vector(np.subtract(self.vals, other.vals)) def __mul__(self, other): """LHS Multiplication.""" if isinstance(other, (int, float)): return Vector(self.vals * other) if isinstance(other, Vector): return self.scalar_product(other) return NotImplemented def __rmul__(self, other): """RHS Multiplication.""" if isinstance(other, (int, float)): return Vector(self.vals * other) return NotImplemented def __matmul__(self, other): """Cross product.""" if isinstance(other, Vector): return self.cross_product(other) return NotImplemented def __truediv__(self, scalar): """Scalar division.""" return Vector(self.vals / scalar) @property def dim(self) -> float: """Dimensionality.""" return len(self.vals) @property def x(self): """X component.""" return self.vals[0] @property def y(self): """Y component.""" return self.vals[1] @property def z(self): """Z component.""" return self.vals[2] def __eq__(self, other): """Check if nearly equal vector components.""" return np.allclose(self.vals, other.vals) @property def length(self) -> float: r""" Euclidian length. .. math:: \text{length} = \sqrt{\sum_i v_i} """ return np.sqrt(np.sum(np.power(self.vals, 2)))
[docs] def normalize(self) -> "Vector": """Normalize to unit length.""" return self / self.length
[docs] def rotate(self, u: "Vector", theta: float) -> "Vector": """ Rotate around arbitrary axis. :param Vector u: Rotation axis (Vector of arbitrary length) :param float theta: Rotation angle in radians """ u = u.normalize() ux = u.vals[0] uy = u.vals[1] uz = u.vals[2] sx = self.vals[0] sy = self.vals[1] sz = self.vals[2] cos = np.cos(theta) sin = np.sin(theta) x = ( (cos + np.power(ux, 2) * (1 - cos)) * sx + (ux * uy * (1 - cos) - uz * sin) * sy + (ux * uz * (1 - cos) + uy * sin) * sz ) y = ( (uy * ux * (1 - cos) + uz * sin) * sx + (cos + np.power(uy, 2) * (1 - cos)) * sy + (uy * uz * (1 - cos) - ux * sin) * sz ) z = ( (uz * ux * (1 - cos) - uy * sin) * sx + (uz * uy * (1 - cos) + ux * sin) * sy + (cos + np.power(uz, 2) * (1 - cos)) * sz ) return Vector((x, y, z))
[docs] def scalar_product(self, other) -> float: """Scalar Product between self and other.""" return np.dot(self.vals, other.vals)
[docs] def cross_product(self, other) -> "Vector": """Cross Product between self and other.""" return Vector(np.cross(self.vals, other.vals))
[docs] def angle(self, other) -> float: r""" Angle between self and other in radians. .. math:: \text{angle} = \arccos\left(\frac{\vec{v}_1 \cdot \vec{v}_2}{|\vec{v}_1| |\vec{v}_2|}\right) """ v1 = self.normalize() v2 = other.normalize() return np.arccos(np.clip(np.dot(v1.vals, v2.vals), -1.0, 1.0))
[docs] def angle_normal(self, other) -> float: r""" Angle between self and normal vector of a plane. .. math:: \text{angle normal} = \arcsin\left(\frac{\vec{v}_1 \cdot \vec{n}_2}{|\vec{v}_1| |\vec{n}_2|}\right) """ v1 = self.normalize() v2 = other.normalize() return np.arcsin(np.clip(np.dot(v1.vals, v2.vals), -1.0, 1.0))