from typing import Self, overload
import numpy as np
from . import _vectors as vec
[docs]
class CoordinateSystem:
r"""
Class representing a single coordinate system.
.. tip::
If you want to store an array of coordinate systems, consider
:class:`.CoordinateSystemArray`.
Parameters
----------
vector_1 : VectorLike
First vector :math:`\vec{v}_1` defining the coordinate system.
vector_2 : VectorLike
Second vector :math:`\vec{v}_2` defining the coordinate system.
Should not be parallel to `vector_1`.
vector_3 : VectorLike or None, default None
Optional third vector. If three vectors are provided, use them as the new
base vectors (after normalization).
Notes
-----
The three basis vectors :math:`\hat{x}, \hat{y}, \hat{z}` will be unit vectors along
the directions given by
.. math::
\vec{z} &= \vec{v_1} \\
\vec{y} &= \vec{v_z} \times \vec{v_2} \\
\vec{x} &= \vec{v_y} \times \vec{z}
If three vectors are provided, the directions will be given by
.. math::
\vec{x} &= \vec{v_1} \\
\vec{y} &= \vec{v_2} \\
\vec{z} &= \vec{v_3} \\
Attributes
----------
x_axis, y_axis, z_axis : :class:`.Vector`
x, y, z basis vectors of the coordinate system.
"""
def __init__(
self,
vector_1: vec.VectorLike,
vector_2: vec.VectorLike,
vector_3: vec.VectorLike | None = None,
):
vec1_ = vec.asvector(vector_1)
vec2_ = vec.asvector(vector_2)
if vector_3 is None:
self._z_axis = vec1_.norm()
self._y_axis = vec1_.cross(vec2_).norm()
self._x_axis = self._y_axis.cross(vec1_).norm()
else:
vec3_ = vec.asvector(vector_3)
self._x_axis = vec1_.norm()
self._y_axis = vec2_.norm()
self._z_axis = vec3_.norm()
@property
def x_axis(self) -> vec.Vector:
"""Unit vector defining the x axis"""
return self._x_axis
@property
def y_axis(self) -> vec.Vector:
"""Unit vector defining the y axis"""
return self._y_axis
@property
def z_axis(self) -> vec.Vector:
"""Unit vector defining the z axis"""
return self._z_axis
[docs]
def project_vector(self, vector: vec.VectorLike) -> vec.Vector:
"""
Project `vector` on the coordinate system.
Parameters
----------
vector : VectorLike
Returns
-------
:class:`.Vector`
The projected vector.
Examples
--------
::
>>> v1 = ap.Vector(1, 1, 0)
>>> v2 = ap.Vector(1, 0, 1)
>>> c = ap.CoordinateSystem(v1, v2)
>>> c.project_vector(v1)
Vector(0.0, 0.0, 1.414213562373095)
"""
vec_ = vec.asvector(vector)
out = vec.Vector(
vec_.dot(self.x_axis), vec_.dot(self.y_axis), vec_.dot(self.z_axis)
)
return out
[docs]
class CoordinateSystemArray:
r"""
Class representing a collection of coordinate systems.
.. tip::
If you want to store a single coordinate systems, consider
:class:`.CoordinateSystem`.
Parameters
----------
vectors_1 : VectorArrayLike
First set of vectors :math:`\vec{v}_1` defining the coordinate system.
vectors_2 : VectorArrayLike
Second set of vectors :math:`\vec{v}_2` defining the coordinate system.
Should not be parallel to `vectors_1`.
vector_3 : VectorArrayLike or None, default None
Optional third set of vectors. If three sets are provided, use them as the
new base vectors (after normalization).
Notes
-----
The three basis vectors :math:`\hat{x}, \hat{y}, \hat{z}` will be unit vectors
along the directions given by, respectively
.. math::
\vec{z} &= \vec{v_1} \\
\vec{y} &= \vec{v_z} \times \vec{v_2} \\
\vec{x} &= \vec{v_y} \times \vec{z}
If three vectors are provided, the directions will be given by
.. math::
\vec{x} &= \vec{v_1} \\
\vec{y} &= \vec{v_2} \\
\vec{z} &= \vec{v_3} \\
Attributes
----------
x_axis, y_axis, z_axis : :class:`.VectorArray`
x, y, z basis vectors of the coordinate system.
"""
def __init__(
self,
vectors_1: vec.VectorArrayLike,
vectors_2: vec.VectorArrayLike,
vectors_3: vec.VectorArrayLike | None = None,
):
vec1_ = vec.asvectorarray(vectors_1)
vec2_ = vec.asvectorarray(vectors_2)
if vectors_3 is None:
self._z_axis = vec1_.norm(copy=False)
self._y_axis = vec1_.cross(vec2_).norm(copy=False)
self._x_axis = self._y_axis.cross(vec1_).norm(copy=False)
else:
vec3_ = vec.asvectorarray(vectors_3)
self._x_axis = vec1_.norm(copy=False)
self._y_axis = vec2_.norm(copy=False)
self._z_axis = vec3_.norm(copy=False)
@property
def x_axis(self) -> vec.VectorArray:
"""Unit vectors defining the x axis"""
return self._x_axis
@property
def y_axis(self) -> vec.VectorArray:
"""Unit vectors defining the y axis"""
return self._y_axis
@property
def z_axis(self) -> vec.VectorArray:
"""Unit vectors defining the z axis"""
return self._z_axis
[docs]
def project_vectors(self, vectors: vec.VectorArrayLike) -> vec.VectorArray:
"""
Project `vectors` on the coordinate system.
Parameters
----------
vectors : VectorArrayLike
If a single vector is passed, project it into each coordinate system.
Returns
-------
:class:`.Vector`
The projected vectors.
Examples
--------
::
>>> v1 = ap.Vector(1, 1, 0)
>>> v2 = ap.Vector(1, 0, 1)
>>> c = ap.CoordinateSystemArray((v1, v2), (v2, v1))
>>> c.project_vectors(v1)
VectorArray([[0. 0. 1.41421356]
[1.22474487 0. 0.70710678]])
>>> c.project_vectors(v2)
VectorArray([[1.22474487 0. 0.70710678]
[0. 0. 1.41421356]])
>>> c.project_vectors((v1, v2))
VectorArray([[0. 0. 1.41421356]
[0. 0. 1.41421356]])
"""
vec_ = vec.asvectorarray(vectors)
if len(vec_) == 1:
vec_ = vec.VectorArray(np.repeat(vec_.asarray(), len(self), axis=0))
out = vec.VectorArray(np.empty_like(vec_.asarray()))
out.x = vec_.dot(self.x_axis)
out.y = vec_.dot(self.y_axis)
out.z = vec_.dot(self.z_axis)
return out
def __len__(self) -> int:
return len(self.x_axis)
@overload
def __getitem__(self, i: int) -> CoordinateSystem: ...
@overload
def __getitem__(self, i: slice) -> Self: ...
def __getitem__(self, i: int | slice) -> CoordinateSystem | Self:
if not (isinstance(i, int) or isinstance(i, slice)):
raise TypeError(f"i must be int or slice, but is {type(i)}")
cls_ = CoordinateSystem if isinstance(i, int) else type(self)
return cls_(self.x_axis[i], self.y_axis[i], self.z_axis[i]) # type: ignore