Source code for vectormath.vector

"""vector.py contains definitions for Vector and VectorArray classes"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import numpy as np


[docs]class BaseVector(np.ndarray): """Class to contain basic operations used by all Vector classes""" def __new__(cls, *args, **kwargs): """BaseVector cannot be created""" raise NotImplementedError('Please specify Vector2 or Vector3') def __array_finalize__(self, obj): #pylint: disable=no-self-use """This is called at the end of array creation obj depends on the context. Currently, this does not need to do anything regardless of context. See `subclassing docs <https://docs.scipy.org/numpy/user/basics.subclassing.html>`_. """ if obj is None: return @property def x(self): """x-component of vector""" return self[0] @x.setter def x(self, value): self[0] = value @property def y(self): """y-component of vector""" return self[1] @y.setter def y(self, value): self[1] = value @property def length(self): """Length of vector""" return float(np.sqrt(np.sum(self**2))) @length.setter def length(self, value): if not np.isscalar(value): raise ValueError('Length must be a scalar') value = float(value) if self.length != 0: new_length = value/self.length self *= new_length return if value != 0: raise ZeroDivisionError('Cannot resize vector of length 0 to ' 'nonzero length')
[docs] def as_length(self, value): """Return a new vector scaled to given length""" new_vec = self.copy() new_vec.length = value return new_vec
[docs] def as_percent(self, value): """Return a new vector scaled by given decimal percent""" new_vec = self.copy() new_vec.length = value * self.length return new_vec
[docs] def as_unit(self): """Return a new vector scaled to length 1""" new_vec = self.copy() new_vec.normalize() return new_vec
[docs] def normalize(self): """Scale the length of a vector to 1 in place""" self.length = 1 return self
[docs] def dot(self, vec): """Dot product with another vector""" if not isinstance(vec, self.__class__): raise TypeError('Dot product operand must be a vector') return np.dot(self, vec)
[docs] def cross(self, vec): """Cross product with another vector""" if not isinstance(vec, self.__class__): raise TypeError('Cross product operand must be a vector') return self.__class__(np.cross(self, vec))
def __mul__(self, multiplier): return self.__class__(self.view(np.ndarray) * multiplier)
[docs]class Vector3(BaseVector): """Primitive 3D vector defined from the origin New Vector3 can be created with: - another Vector3 - length-3 array - x, y, and y values - no input (returns [0., 0., 0.]) """ def __new__(cls, x=None, y=None, z=None): #pylint: disable=arguments-differ def read_array(X, Y, Z): """Build Vector3 from another Vector3, [x, y, z], or x/y/z""" if isinstance(X, cls) and Y is None and Z is None: return cls(X.x, X.y, X.z) if (isinstance(X, (list, tuple, np.ndarray)) and len(X) == 3 and Y is None and Z is None): return cls(X[0], X[1], X[2]) if X is None and Y is None and Z is None: return cls(0, 0, 0) if np.isscalar(X) and np.isscalar(Y) and np.isscalar(Z): xyz = np.r_[X, Y, Z] xyz = xyz.astype(float) return xyz.view(cls) raise ValueError('Invalid input for Vector3 - must be an instance ' 'of a Vector3, a length-3 array, 3 scalars, or ' 'nothing for [0., 0., 0.]') return read_array(x, y, z) @property def z(self): """z-component of vector""" return self[2] @z.setter def z(self, value): self[2] = value
[docs]class Vector2(BaseVector): """Primitive 2D vector defined from the origin New Vector2 can be created with: - another Vector2 - length-2 array - x and y values - no input (returns [0., 0.]) """ def __new__(cls, x=None, y=None): #pylint: disable=arguments-differ def read_array(X, Y): """Build Vector2 from another Vector2, [x, y], or x/y""" if isinstance(X, cls) and Y is None: return cls(X.x, X.y) if (isinstance(X, (list, tuple, np.ndarray)) and len(X) == 2 and Y is None): return cls(X[0], X[1]) if X is None and Y is None: return cls(0, 0) if np.isscalar(X) and np.isscalar(Y): xyz = np.r_[X, Y] xyz = xyz.astype(float) return xyz.view(cls) raise ValueError('Invalid input for Vector2 - must be an instance ' 'of a Vector2, a length-2 array, 2 scalars, or ' 'nothing for [0., 0.]') return read_array(x, y)
[docs]class BaseVectorArray(BaseVector): """Class to contain basic operations used by all VectorArray classes""" @property def x(self): """Array of x-component of vectors""" return self[:, 0] @x.setter def x(self, value): self[:, 0] = value @property def y(self): """Array of y-component of vectors""" return self[:, 1] @y.setter def y(self, value): self[:, 1] = value @property def nV(self): """Number of vectors""" return self.shape[0]
[docs] def normalize(self): """Scale the length of all vectors to 1 in place""" self.length = np.ones(self.nV) return self
@property def dims(self): """Tuple of different dimension names for Vector type""" raise NotImplementedError('Please use Vector2Array or Vector3Array') @property def length(self): """Array of vector lengths""" return np.sqrt(np.sum(self**2, axis=1)).view(np.ndarray) @length.setter def length(self, l): l = np.array(l) if self.nV != l.size: raise ValueError('Length vector must be the same number of ' 'elements as vector.') # This case resizes all vectors with nonzero length if np.all(self.length != 0): new_length = l/self.length for dim in self.dims: setattr(self, dim, new_length*getattr(self, dim)) return # This case only applies to single vectors if self.nV == 1 and l == 0: assert self.length == 0, \ 'Nonzero length should be resized in the first case' for dim in self.dims: setattr(self, dim, 0.) return # This case only applies if vectors with length == 0 # in an array are getting resized to 0 if self.nV > 1 and np.array_equal(self.length.nonzero(), l.nonzero()): #pylint: disable=no-member new_length = l/[x if x != 0 else 1 for x in self.length] for dim in self.dims: setattr(self, dim, new_length*getattr(self, dim)) return # Error if length zero array is resized to nonzero value raise ZeroDivisionError('Cannot resize vector of length 0 to ' 'nonzero length')
[docs] def dot(self, vec): """Dot product with another vector""" if not isinstance(vec, self.__class__): raise TypeError('Dot product operand must be a VectorArray') if self.nV != 1 and vec.nV != 1 and self.nV != vec.nV: raise ValueError('Dot product operands must have the same ' 'number of elements.') return np.sum((getattr(self, d)*getattr(vec, d) for d in self.dims), 1)
[docs]class Vector3Array(BaseVectorArray): """List of Vector3 A new Vector3Array can be created with: - another Vector3Array - x/y/z lists of equal length - n x 3 array - nothing (returns [[0., 0., 0.]]) """ def __new__(cls, x=None, y=None, z=None): #pylint: disable=arguments-differ def read_array(X, Y, Z): """Build Vector3Array from various inputs""" if isinstance(X, cls) and Y is None and Z is None: X = np.atleast_2d(X) return cls(X.x.copy(), X.y.copy(), X.z.copy()) if isinstance(X, (list, tuple)): X = np.array(X) if isinstance(Y, (list, tuple)): Y = np.array(Y) if isinstance(Z, (list, tuple)): Z = np.array(Z) if isinstance(X, np.ndarray) and Y is None and Z is None: X = np.squeeze(X) if X.size == 3: X = X.flatten() return cls(X[0], X[1], X[2]) elif len(X.shape) == 2 and X.shape[1] == 3: return cls( X[:, 0].copy(), X[:, 1].copy(), X[:, 2].copy() ) raise ValueError( 'Unexpected shape for vector init: {shp}'.format( shp=X.shape ) ) if np.isscalar(X) and np.isscalar(Y) and np.isscalar(Z): X, Y, Z = float(X), float(Y), float(Z) elif not (isinstance(X, type(Y)) and isinstance(X, type(Z))): raise TypeError('Must be the same types for x, y, and ' 'z for vector init') if isinstance(X, np.ndarray): if not (X.shape == Y.shape and X.shape == Z.shape): raise ValueError('Must be the same shapes for x, y, ' 'and z in vector init') vec_ndarray = np.c_[X, Y, Z] vec_ndarray = vec_ndarray.astype(float) return vec_ndarray.view(cls) if X is None: X, Y, Z = 0.0, 0.0, 0.0 vec_ndarray = np.r_[X, Y, Z].reshape((1, 3)) return np.asarray(vec_ndarray).view(cls) return read_array(x, y, z) def __getitem__(self, i): """Overriding _getitem__ allows coersion to Vector3 or ndarray""" item_out = super(Vector3Array, self).__getitem__(i) if np.isscalar(i): return item_out.view(Vector3) if isinstance(i, slice): return item_out return item_out.view(np.ndarray) @property def z(self): """Array of z-component of vectors""" return self[:, 2] @z.setter def z(self, value): self[:, 2] = value @property def dims(self): return ('x', 'y', 'z')
[docs] def cross(self, vec): """Cross product with another Vector3Array""" if not isinstance(vec, Vector3Array): raise TypeError('Cross product operand must be a Vector3Array') if self.nV != 1 and vec.nV != 1 and self.nV != vec.nV: raise ValueError('Cross product operands must have the same ' 'number of elements.') return Vector3Array(np.cross(self, vec))
[docs]class Vector2Array(BaseVectorArray): """List of Vector2 A new Vector2Array can be created with: - another Vector2Array - x/y lists of equal length - n x 2 array - nothing (returns [[0., 0.]]) """ def __new__(cls, x=None, y=None): #pylint: disable=arguments-differ def read_array(X, Y): """Build Vector2Array from various inputs""" if isinstance(X, cls) and Y is None: X = np.atleast_2d(X) return cls(X.x.copy(), X.y.copy()) if isinstance(X, (list, tuple)): X = np.array(X) if isinstance(Y, (list, tuple)): Y = np.array(Y) if isinstance(X, np.ndarray) and Y is None: X = np.squeeze(X) if X.size == 2: X = X.flatten() return cls(X[0], X[1]) elif len(X.shape) == 2 and X.shape[1] == 2: return cls( X[:, 0].copy(), X[:, 1].copy() ) raise ValueError( 'Unexpected shape for vector init: {shp}'.format( shp=X.shape ) ) if np.isscalar(X) and np.isscalar(Y): X, Y = float(X), float(Y) elif not isinstance(X, type(Y)): raise TypeError('Must be the same types for x and y ' 'for vector init') if isinstance(X, np.ndarray): if X.shape != Y.shape: raise ValueError('Must be the same shapes for x and y ' 'in vector init') vec_ndarray = np.c_[X, Y] vec_ndarray = vec_ndarray.astype(float) return vec_ndarray.view(cls) if X is None: X, Y = 0.0, 0.0 vec_ndarray = np.r_[X, Y].reshape((1, 2)) return np.asarray(vec_ndarray).view(cls) return read_array(x, y) def __getitem__(self, i): """Overriding _getitem__ allows coercion to Vector2 or ndarray""" item_out = super(Vector2Array, self).__getitem__(i) if np.isscalar(i): return item_out.view(Vector2) if isinstance(i, slice): return item_out return item_out.view(np.ndarray) @property def dims(self): return ('x', 'y')