Source code for apsg.math._vector

import math
from abc import ABC, abstractmethod

import numpy as np

from apsg.config import apsg_conf
from apsg.helpers._helper import is_jsonable
from apsg.helpers._math import acosd, atan2d, cosd, sind
from apsg.helpers._notation import (
    geo2vec_linear,
    vec2geo_linear_signed,
)


class Vector(ABC):
    """Abstract base class for Vector2 and Vector3."""

    __slots__ = ("_coords", "_attrs")
    __shape__ = None

    @abstractmethod
    def __init__(self, *args, **kwargs):
        self._coords = (None, None, None)
        self._attrs = {}

    def __copy__(self):
        return type(self)(self._coords)

    copy = __copy__

    def __hash__(self):
        return hash((type(self).__name__,) + self._coords)

    def __array__(self, dtype=None, copy=None):
        return np.array(self._coords, dtype=dtype)

    def to_json(self):
        return {
            "datatype": type(self).__name__,
            "args": (self._coords,),
            "kwargs": self._attrs,
        }

    def __eq__(self, other):
        cls = type(self)
        other_arr = np.asarray(other)
        if other_arr.shape == cls.__shape__:
            other = cls(other_arr)
        else:
            return NotImplemented
        return np.allclose(self, other)

    def __ne__(self, other):
        return not self.__eq__(other)

    def __nonzero__(self):
        return any(self._coords)

    def __getitem__(self, key):
        return self._coords[key]

    def __iter__(self):
        return iter(self._coords)

    def __add__(self, other):
        return type(self)(np.add(self, other))

    __radd__ = __add__

    def __sub__(self, other):
        return type(self)(np.subtract(self, other))

    def __rsub__(self, other):
        return type(self)(np.subtract(other, self))

    def __mul__(self, other):
        return type(self)(np.multiply(self, other))

    __rmul__ = __mul__

    def __floordiv__(self, other):
        return type(self)(np.floor_divide(self, other))

    def __rfloordiv__(self, other):
        return type(self)(np.floor_divide(other, self))

    def __truediv__(self, other):
        return type(self)(np.true_divide(self, other))

    def __rtruediv__(self, other):
        return type(self)(np.true_divide(other, self))

    __pos__ = __copy__

    def __abs__(self):
        return math.sqrt(sum(map(lambda x: x * x, self._coords)))

    magnitude = __abs__

    @abstractmethod
    def dot(self, other):
        pass

    @abstractmethod
    def normalized(self):
        pass

    def label(self):
        """Return label."""

        return str(self)

    def is_unit(self):
        """Return true if the magnitude is 1."""

        return math.isclose(self.magnitude(), 1)

    def _ensure_same(self, other):
        cls = type(self)
        if np.asarray(other).shape == cls.__shape__:
            return cls(other)
        raise TypeError(f"Unsupported argument. Expecting {cls.__name__}")

    def angle(self, other):
        """Return the angle to the vector other."""

        other = self._ensure_same(other)
        return acosd(self.normalized().dot(other.normalized()))

    def project(self, other):
        """Return vector projection on the vector other."""

        other = self._ensure_same(other)
        n = other.normalized()
        return type(self)(self.dot(n) * n)

    proj = project

    def reject(self, other):
        """Return vector rejection on the vector other."""

        other = self._ensure_same(other)
        return self - self.project(other)

    @property
    def x(self):
        """Return x-component of the vector."""

        return self._coords[0]

    @property
    def y(self):
        """Return y-component of the vector."""

        return self._coords[1]


[docs] class Vector2(Vector): """ A class to represent a 2D vector. There are different way to create ``Vector2`` object: - without arguments create default ``Vector2`` (0, 0, 1) - with single argument `v`, where - `v` could be Vector2-like object - `v` could be string 'x' or 'y' - principal axes of coordinate system - `v` could be tuple of (x, y) - vector components - `v` could be float - unit vector with given angle to 'x' axis - with 2 numerical arguments defining vector components Args: ang (float): angle between 'x' axis and vector in degrees Examples: >>> vec2() >>> vec2(1, -1) >>> vec2('y') >>> vec2(50) >>> v = vec2(1, -2) """ __slots__ = ("_coords", "_attrs") __shape__ = (2,) def __init__(self, *args, **kwargs): if len(args) == 0: coords = (1, 0) elif len(args) == 1: if np.asarray(args[0]).shape == Vector2.__shape__: coords = tuple(c.item() for c in np.asarray(args[0])) elif isinstance(args[0], str): if args[0].lower() == "x": coords = (1, 0) elif args[0].lower() == "y": coords = (0, 1) else: raise TypeError(f"Not valid arguments for {type(self).__name__}") else: coords = cosd(args[0]), sind(args[0]) elif len(args) == 2: coords = args else: raise TypeError(f"Not valid arguments for {type(self).__name__}") self._coords = tuple(coords) self._attrs = {} if kwargs: if not is_jsonable(kwargs): raise TypeError("Provided attributes are not serializable.") self._attrs = kwargs def __repr__(self): n = apsg_conf.ndigits return f"Vector2({round(self.x, n):g}, {round(self.y, n):g})" def __len__(self): return 2 def __neg__(self): return type(self)(-self.x, -self.y)
[docs] def normalized(self): """Returns normalized (unit length) vector.""" d = self.magnitude() if d: return type(self)(self.x / d, self.y / d) return self.copy()
uv = normalized shape = __shape__ @property def direction(self): """Returns direction of the vector in degrees.""" return atan2d(self.y, self.x) % 360
[docs] def dot(self, other): """ Calculate dot product with other vector. Args: other (Vector2): other vector Returns: float: Dot product of the two vectors. """ other = self._ensure_same(other) return self.x * other.x + self.y * other.y
def __matmul__(self, other): r = np.dot(self, other) if np.asarray(r).shape == Vector2.__shape__: return type(self)(r) else: return float(r) def __rmatmul__(self, other): r = np.dot(other, self) if np.asarray(r).shape == Vector2.__shape__: return type(self)(r) else: return float(r)
[docs] def cross(self, other): """Returns the scalar magnitude of the 2D cross product.""" other = self._ensure_same(other) return self.x * other.y - self.y * other.x
[docs] @classmethod def random(cls): """Create random 2D vector.""" return cls(360 * np.random.rand())
[docs] def rotate(self, theta): """Return the vector rotated counter-clockwise by angle theta in degrees.""" c, s = cosd(theta), sind(theta) return type(self)(c * self.x - s * self.y, s * self.x + c * self.y)
[docs] @classmethod def unit_x(cls): """Create unit length vector in x-direction.""" return cls(1, 0)
[docs] @classmethod def unit_y(cls): """Create unit length vector in y-direction.""" return cls(0, 1)
[docs] def transform(self, *args, **kwargs): """ Return affine transformation of vector `u` by matrix `F`. Args: F: transformation matrix Keyword Args: norm: normalize transformed vectors. [True or False] Default False Examples: # Reflexion of `y` axis. >>> F = [[1, 0], [0, -1]] >>> u = vec2([1, 1]) >>> u.transform(F) Vector2(1, -1) Returns: Vector2: Affine transformation of vector `u` by matrix `F`. """ r = Vector2(np.dot(args[0], self)) if kwargs.get("norm", False): r = r.normalized() return type(self)(r)
[docs] class Axial2(Vector2): # Do we need it? """ A class to represent a 2D axial vector. Note: the angle between axial data cannot be more than 90° """ def __eq__(self, other): cls = type(self) other_arr = np.asarray(other) if other_arr.shape == cls.__shape__: other = cls(other_arr) else: return NotImplemented return np.allclose(self, other) or np.allclose(self, -other) def __add__(self, other): if isinstance(other, Vector2): if super().dot(other) < 0: other = -other return type(self)(np.add(self, other)) __radd__ = __add__ def __sub__(self, other): if isinstance(other, Vector2): if super().dot(other) < 0: other = -other return type(self)(np.subtract(self, other)) def __rsub__(self, other): if isinstance(other, Vector2): if super().dot(other) < 0: other = -other return type(self)(np.subtract(other, self))
[docs] def dot(self, other): return abs(super().dot(other))
[docs] class Vector3(Vector): """ A class to represent a 3D vector. There are different way to create ``Vector3`` object: - without arguments create default ``Vector3`` (1, 0, 0) - with single argument `v`, where - `v` could be Vector3-like object - `v` could be string 'x', 'y' or 'z' - principal axes of coordinate system - `v` could be tuple of (x, y, z) - vector components - with 2 arguments plunge direction and plunge - with 3 numerical arguments defining vector components Args: azi (float): plunge direction of linear feature in degrees inc (float): plunge of linear feature in degrees Examples: >>> vec() >>> vec(1,2,-1) >>> vec('y') >>> vec(120, 30) >>> v = vec(1, -2, 1) """ __slots__ = ("_coords", "_attrs") __shape__ = (3,) def __init__(self, *args, **kwargs): if len(args) == 0: coords = (1, 0, 0) elif len(args) == 1: if np.asarray(args[0]).shape == Vector3.__shape__: coords = tuple(c.item() for c in np.asarray(args[0])) elif isinstance(args[0], str): if args[0].lower() == "x": coords = (1, 0, 0) elif args[0].lower() == "y": coords = (0, 1, 0) elif args[0].lower() == "z": coords = (0, 0, 1) else: raise TypeError(f"Not valid arguments for {type(self).__name__}") else: raise TypeError(f"Not valid arguments for {type(self).__name__}") elif len(args) == 2: coords = geo2vec_linear(*args) elif len(args) == 3: coords = args else: raise TypeError(f"Not valid arguments for {type(self).__name__}") self._coords = tuple(coords) self._attrs = {} if kwargs: if not is_jsonable(kwargs): raise TypeError("Provided attributes are not serializable.") self._attrs = kwargs @property def z(self): """Return z-component of the vector.""" return self._coords[2] def __repr__(self): if apsg_conf.vec2geo: azi, inc = self.geo return f"V:{azi:.0f}/{inc:.0f}" else: n = apsg_conf.ndigits return f"Vector3({round(self.x, n):g}, {round(self.y, n):g}, {round(self.z, n):g})" def __len__(self): return 3 def __neg__(self): return type(self)(-self.x, -self.y, -self.z) def __abs__(self): return math.sqrt(self.x**2 + self.y**2 + self.z**2)
[docs] def normalized(self): """Returns normalized (unit length) vector.""" d = self.magnitude() if d: return type(self)(self.x / d, self.y / d, self.z / d) return self.copy()
uv = normalized shape = __shape__
[docs] def dot(self, other): """ Calculate dot product with other vector. Args: other (Vector3): other vector Returns: float: Dot product of the two vectors. """ other = self._ensure_same(other) return self.x * other.x + self.y * other.y + self.z * other.z
def __matmul__(self, other): r = np.dot(self, other) if np.asarray(r).shape == Vector3.__shape__: return type(self)(r) else: return float(r) def __rmatmul__(self, other): r = np.dot(other, self) if np.asarray(r).shape == Vector3.__shape__: return type(self)(r) else: return float(r) def __pow__(self, other): if isinstance(other, Vector3): return self.cross(other) else: return type(self)(np.power(self, other))
[docs] def cross(self, other): """ Calculate cross product with other vector. Args: other (Vector3): other vector Returns: Vector3: Cross product of the two vectors. """ other = self._ensure_same(other) return type(self)( self.y * other.z - self.z * other.y, -self.x * other.z + self.z * other.x, self.x * other.y - self.y * other.x, )
[docs] def slerp(self, other, t): """Return a spherical linear interpolation between self and other vector.""" other = self._ensure_same(other) a, b = Vector3(self), Vector3(other) theta = a.angle(b) return type(self)(a * sind((1 - t) * theta) + b * sind(t * theta)) / sind(theta)
[docs] def lower(self): """Change vector direction to point towards positive Z direction.""" if self.z < 0: return -self else: return self
[docs] def is_upper(self): """Return True if vector points towards negative Z direction.""" return self.z < 0
@property def geo(self): """Return tuple of plunge direction and signed plunge.""" return vec2geo_linear_signed(self)
[docs] @classmethod def unit_x(cls): """Create unit length vector in x-direction.""" return cls(1, 0, 0)
[docs] @classmethod def unit_y(cls): """Create unit length vector in y-direction.""" return cls(0, 1, 0)
[docs] @classmethod def unit_z(cls): """Create unit length vector in z-direction.""" return cls(0, 0, 1)
[docs] @classmethod def random(cls): """Create random 3D vector.""" return cls(np.random.randn(3)).normalized()
[docs] def rotate(self, axis, theta): """Return the vector rotated around axis through angle theta. Right-hand rule.""" axis = self._ensure_same(axis) v = Vector3(self) # ensure vector k = Vector3(axis.uv()) return type(self)( cosd(theta) * v + sind(theta) * k.cross(v) + (1 - cosd(theta)) * k * (k.dot(v)) )
[docs] def angle(self, other): """Return the angle to the vector other.""" other = self._ensure_same(other) return acosd(np.clip(self.uv().dot(other.uv()), -1, 1))
[docs] def transform(self, F, **kwargs): """ Return affine transformation of vector `u` by matrix `F`. Args: F: transformation matrix Keyword Args: norm: normalize transformed vectors. [True or False] Default False Examples: # Reflexion of `y` axis. >>> F = [[1, 0, 0], [0, -1, 0], [0, 0, 1]] >>> u = Vector3([1, 1, 1]) >>> u.transform(F) Vector3(1, -1, 1) Returns: Vector3: Affine transformation of vector `u` by matrix `F`. """ r = Vector3(np.dot(F, self)) if kwargs.get("norm", False): r = r.normalized() return type(self)(r)
[docs] class Axial3(Vector3): """ A class to represent a 3D axial vector. Note: the angle between axial data cannot be more than 90° """ def __eq__(self, other): cls = type(self) other_arr = np.asarray(other) if other_arr.shape == cls.__shape__: other = cls(other_arr) else: return NotImplemented return np.allclose(self, other) or np.allclose(self, -other) def __add__(self, other): if isinstance(other, Vector3): if super().dot(other) < 0: other = -other return type(self)(np.add(self, other)) __radd__ = __add__ def __sub__(self, other): if isinstance(other, Vector3): if super().dot(other) < 0: other = -other return type(self)(np.subtract(self, other)) def __rsub__(self, other): if isinstance(other, Vector3): if super().dot(other) < 0: other = -other return type(self)(np.subtract(other, self))
[docs] def dot(self, other): return abs(super().dot(other))