"""
NumPy-based 3D vector mathematics for molecular analysis.
This module provides efficient vector operations for high-performance
molecular geometry calculations, supporting both single vectors and
batch operations on multiple vectors simultaneously.
"""
from typing import List, Optional, Tuple, Union
import numpy as np
from ..constants import VectorDefaults
[docs]
class NPVec3D:
"""NumPy-based 3D vector class for molecular calculations.
This class provides comprehensive 3D vector operations using NumPy,
enabling efficient batch processing of multiple vectors simultaneously.
:param coords: Either x,y,z coordinates or numpy array of shape (3,) or (N,3)
:type coords: Union[float, np.ndarray]
:param y: Y coordinate (only if coords is a float)
:type y: Optional[float]
:param z: Z coordinate (only if coords is a float)
:type z: Optional[float]
"""
[docs]
def __init__(
self,
coords: Union[
float, np.ndarray, List[float], Tuple[float, float, float]
] = VectorDefaults.DEFAULT_X,
y: Optional[float] = None,
z: Optional[float] = None,
):
"""Initialize a 3D vector or batch of vectors."""
if isinstance(coords, (list, tuple)):
self._data = np.array(coords, dtype=np.float64)
elif isinstance(coords, np.ndarray):
self._data = coords.astype(np.float64)
elif y is not None and z is not None:
self._data = np.array([coords, y, z], dtype=np.float64)
else:
self._data = np.array(
[
VectorDefaults.DEFAULT_X,
VectorDefaults.DEFAULT_Y,
VectorDefaults.DEFAULT_Z,
],
dtype=np.float64,
)
# Ensure shape is correct
if self._data.ndim == 1 and len(self._data) == 3:
pass # Single vector
elif self._data.ndim == 2 and self._data.shape[1] == 3:
pass # Batch of vectors
else:
raise ValueError(
f"Invalid shape {self._data.shape}, expected (3,) or (N,3)"
)
@property
def x(self) -> Union[float, np.ndarray]:
"""X coordinate(s)."""
if self._data.ndim == 1:
return float(self._data[0])
return self._data[:, 0] # type: ignore[no-any-return]
@property
def y(self) -> Union[float, np.ndarray]:
"""Y coordinate(s)."""
if self._data.ndim == 1:
return float(self._data[1])
return self._data[:, 1] # type: ignore[no-any-return]
@property
def z(self) -> Union[float, np.ndarray]:
"""Z coordinate(s)."""
if self._data.ndim == 1:
return float(self._data[2])
return self._data[:, 2] # type: ignore[no-any-return]
@property
def shape(self) -> Tuple[int, ...]:
"""Shape of the underlying array."""
return self._data.shape # type: ignore[no-any-return]
@property
def is_batch(self) -> bool:
"""Whether this represents multiple vectors."""
return self._data.ndim == 2 # type: ignore[no-any-return]
def __repr__(self) -> str:
if self._data.ndim == 1:
return f"NPVec3D({self.x}, {self.y}, {self.z})"
return f"NPVec3D(batch_size={self._data.shape[0]})"
def __str__(self) -> str:
if self._data.ndim == 1:
return f"({self.x:.3f}, {self.y:.3f}, {self.z:.3f})"
return f"Batch of {self._data.shape[0]} vectors"
[docs]
def __add__(self, other: Union["NPVec3D", np.ndarray]) -> "NPVec3D":
"""Vector addition."""
if isinstance(other, NPVec3D):
return NPVec3D(self._data + other._data)
return NPVec3D(self._data + other)
[docs]
def __sub__(self, other: Union["NPVec3D", np.ndarray]) -> "NPVec3D":
"""Vector subtraction."""
if isinstance(other, NPVec3D):
return NPVec3D(self._data - other._data)
return NPVec3D(self._data - other)
[docs]
def __mul__(self, scalar: Union[float, np.ndarray]) -> "NPVec3D":
"""Scalar multiplication."""
return NPVec3D(self._data * scalar)
[docs]
def __rmul__(self, scalar: Union[float, np.ndarray]) -> "NPVec3D":
"""Reverse scalar multiplication."""
return self.__mul__(scalar)
[docs]
def __truediv__(self, scalar: Union[float, np.ndarray]) -> "NPVec3D":
"""Scalar division."""
return NPVec3D(self._data / scalar)
[docs]
def __eq__(self, other: object) -> bool:
"""Vector equality comparison."""
if not isinstance(other, NPVec3D):
return False
return bool(np.allclose(self._data, other._data, atol=1e-10))
[docs]
def __getitem__(
self, index: Union[int, slice, np.ndarray]
) -> Union[float, "NPVec3D"]:
"""Get component or subset."""
if isinstance(index, int) and self._data.ndim == 1:
return float(self._data[index])
result = self._data[index]
if isinstance(result, np.ndarray):
return NPVec3D(result)
return float(result)
[docs]
def dot(self, other: "NPVec3D") -> Union[float, np.ndarray]:
"""Dot product with another vector.
:param other: The other vector(s)
:type other: NPVec3D
:returns: Dot product result(s)
:rtype: Union[float, np.ndarray]
"""
if self._data.ndim == 1 and other._data.ndim == 1:
return float(np.dot(self._data, other._data))
elif self._data.ndim == 1:
# Single vector dot batch
return np.dot(other._data, self._data) # type: ignore[no-any-return]
elif other._data.ndim == 1:
# Batch dot single vector
return np.dot(self._data, other._data) # type: ignore[no-any-return]
else:
# Batch dot batch
return np.sum(self._data * other._data, axis=1) # type: ignore[no-any-return]
[docs]
def cross(self, other: "NPVec3D") -> "NPVec3D":
"""Cross product with another vector.
:param other: The other vector(s)
:type other: NPVec3D
:returns: Cross product vector(s)
:rtype: NPVec3D
"""
return NPVec3D(np.cross(self._data, other._data))
[docs]
def length(self) -> Union[float, np.ndarray]:
"""Calculate vector length/magnitude.
:returns: Euclidean length of the vector(s)
:rtype: Union[float, np.ndarray]
"""
if self._data.ndim == 1:
return float(np.linalg.norm(self._data))
return np.linalg.norm(self._data, axis=1) # type: ignore[no-any-return]
[docs]
def magnitude(self) -> Union[float, np.ndarray]:
"""Alias for length()."""
return self.length()
[docs]
def normalize(self) -> "NPVec3D":
"""Return normalized unit vector(s).
:returns: Unit vector(s) in the same direction
:rtype: NPVec3D
"""
mag = self.length()
if self._data.ndim == 1:
if mag == 0:
return NPVec3D(np.zeros(3))
return NPVec3D(self._data / mag)
else:
# Handle zero-length vectors
mag = np.where(mag == 0, 1, mag)
return NPVec3D(self._data / mag[:, np.newaxis])
[docs]
def unit_vector(self) -> "NPVec3D":
"""Alias for normalize()."""
return self.normalize()
[docs]
def distance_to(self, other: "NPVec3D") -> Union[float, np.ndarray]:
"""Calculate distance to another vector.
:param other: The target vector(s)
:type other: NPVec3D
:returns: Euclidean distance(s) between vectors
:rtype: Union[float, np.ndarray]
"""
return (self - other).length()
[docs]
def angle_to(self, other: "NPVec3D") -> Union[float, np.ndarray]:
"""Calculate angle to another vector in radians.
:param other: The target vector(s)
:type other: NPVec3D
:returns: Angle(s) between vectors in radians
:rtype: Union[float, np.ndarray]
"""
dot_product = self.dot(other)
mag_product = self.length() * other.length()
if isinstance(mag_product, np.ndarray):
# Avoid division by zero
mag_product = np.where(mag_product == 0, 1, mag_product)
cos_angle = dot_product / mag_product
# Clamp to avoid numerical errors
cos_angle = np.clip(cos_angle, -1.0, 1.0)
return np.arccos(cos_angle) # type: ignore[no-any-return]
else:
if mag_product == 0:
return 0.0
cos_angle = dot_product / mag_product
# Clamp to avoid numerical errors
cos_angle = max(-1.0, min(1.0, cos_angle))
return float(np.arccos(cos_angle))
[docs]
def to_array(self) -> np.ndarray:
"""Convert to numpy array.
:returns: Numpy array of coordinates
:rtype: np.ndarray
"""
return self._data.copy() # type: ignore[no-any-return]
[docs]
def to_list(self) -> List[float]:
"""Convert to list [x, y, z] (single vector only).
:returns: Vector components as a list
:rtype: List[float]
"""
if self._data.ndim != 1:
raise ValueError("to_list() only works for single vectors")
return self._data.tolist() # type: ignore[no-any-return]
[docs]
def to_tuple(self) -> Tuple[float, float, float]:
"""Convert to tuple (x, y, z) (single vector only).
:returns: Vector components as a tuple
:rtype: Tuple[float, float, float]
"""
if self._data.ndim != 1:
raise ValueError("to_tuple() only works for single vectors")
return tuple(self._data.tolist())
[docs]
@classmethod
def from_list(cls, coords: List[float]) -> "NPVec3D":
"""Create vector from list [x, y, z].
:param coords: List of coordinates
:type coords: List[float]
:returns: New NPVec3D instance
:rtype: NPVec3D
"""
return cls(coords)
[docs]
@classmethod
def from_tuple(cls, coords: Tuple[float, ...]) -> "NPVec3D":
"""Create vector from tuple (x, y, z).
:param coords: Tuple of coordinates
:type coords: Tuple[float, ...]
:returns: New NPVec3D instance
:rtype: NPVec3D
"""
return cls(coords) # type: ignore[arg-type]
[docs]
@classmethod
def from_atoms(cls, atoms: List) -> "NPVec3D":
"""Create batch vector from list of atoms.
:param atoms: List of atoms with x, y, z attributes
:type atoms: List
:returns: Batch NPVec3D instance
:rtype: NPVec3D
"""
coords = np.array([[atom.x, atom.y, atom.z] for atom in atoms])
return cls(coords)
[docs]
def compute_distance_matrix(
coords1: np.ndarray, coords2: Optional[np.ndarray] = None
) -> np.ndarray:
"""Compute pairwise distance matrix between two sets of coordinates.
:param coords1: First set of coordinates, shape (N, 3)
:type coords1: np.ndarray
:param coords2: Second set of coordinates, shape (M, 3). If None, computes self-distances
:type coords2: Optional[np.ndarray]
:returns: Distance matrix of shape (N, M) or (N, N)
:rtype: np.ndarray
"""
if coords2 is None:
coords2 = coords1
# Efficient pairwise distance computation
diff = coords1[:, np.newaxis, :] - coords2[np.newaxis, :, :]
distances = np.linalg.norm(diff, axis=2)
return distances # type: ignore[no-any-return]
[docs]
def batch_angle_between(
a: NPVec3D, b: NPVec3D, c: Optional[NPVec3D] = None
) -> Union[float, np.ndarray]:
"""Calculate angles between vectors (optimized for batches).
If c is provided: Calculate angle ABC where B is the vertex.
If c is None: Calculate angle between vectors a and b.
:param a: First vector(s) or point(s) A
:type a: NPVec3D
:param b: Second vector(s) or vertex point(s) B
:type b: NPVec3D
:param c: Optional third point(s) C for angle ABC
:type c: Optional[NPVec3D]
:returns: Angle(s) in radians
:rtype: Union[float, np.ndarray]
"""
if c is None:
return a.angle_to(b)
ba = a - b
bc = c - b
return ba.angle_to(bc)
[docs]
def batch_dihedral_angle(
a: NPVec3D, b: NPVec3D, c: NPVec3D, d: NPVec3D
) -> Union[float, np.ndarray]:
"""Calculate dihedral angles between planes ABC and BCD (optimized for batches).
:param a: First point(s) defining plane ABC
:type a: NPVec3D
:param b: Second point(s) defining both planes
:type b: NPVec3D
:param c: Third point(s) defining both planes
:type c: NPVec3D
:param d: Fourth point(s) defining plane BCD
:type d: NPVec3D
:returns: Dihedral angle(s) in radians
:rtype: Union[float, np.ndarray]
"""
# Vectors along the bonds
ba = a - b
bc = c - b
cd = d - c
# Normal vectors to the planes
n1 = ba.cross(bc)
n2 = bc.cross(cd)
# Calculate angle between normal vectors
angle = n1.angle_to(n2)
# Determine sign of angle
cross_dot = n1.cross(n2).dot(bc)
if isinstance(angle, np.ndarray):
# Batch operation
angle = np.where(cross_dot < 0, -angle, angle)
else:
# Single operation
if cross_dot < 0:
angle = -angle
return angle
