105 lines
3.7 KiB
Python
105 lines
3.7 KiB
Python
from typing import Optional
|
|
|
|
import numpy as np
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class Coordinate(BaseModel):
|
|
x: float = 0.0
|
|
y: float = 0.0
|
|
z: float = 0.0
|
|
|
|
# Overload +
|
|
def __add__(self, other: "Coordinate") -> "Coordinate":
|
|
if isinstance(other, Coordinate):
|
|
return Coordinate(
|
|
x=self.x + other.x, y=self.y + other.y, z=self.z + other.z
|
|
)
|
|
return NotImplemented
|
|
|
|
# Overload -
|
|
def __sub__(self, other: "Coordinate") -> "Coordinate":
|
|
if isinstance(other, Coordinate):
|
|
return Coordinate(
|
|
x=self.x - other.x, y=self.y - other.y, z=self.z - other.z
|
|
)
|
|
return NotImplemented
|
|
|
|
# Overload *
|
|
def __mul__(self, other):
|
|
if isinstance(other, (int, float)): # Scalar multiplication
|
|
return Coordinate(x=self.x * other, y=self.y * other, z=self.z * other)
|
|
elif isinstance(other, Coordinate): # Dot product
|
|
return self.x * other.x + self.y * other.y + self.z * other.z
|
|
return NotImplemented
|
|
|
|
# Overload /
|
|
def __truediv__(self, scalar: float) -> "Coordinate":
|
|
if isinstance(scalar, (int, float)) and scalar != 0: # Avoid division by zero
|
|
return Coordinate(x=self.x / scalar, y=self.y / scalar, z=self.z / scalar)
|
|
elif scalar == 0:
|
|
raise ValueError("Cannot divide by zero")
|
|
return NotImplemented
|
|
|
|
def normalize(self) -> "Coordinate":
|
|
magnitude = np.sqrt(self.x**2 + self.y**2 + self.z**2)
|
|
if magnitude == 0:
|
|
raise ValueError("Cannot normalize a zero-magnitude vector.")
|
|
return self / magnitude
|
|
|
|
def rotate(self, angle_deg: float, axis: str) -> "Coordinate":
|
|
"""
|
|
Rotate the coordinate around a specified axis by a given angle in degrees.
|
|
|
|
:param angle_deg: The angle of rotation in degrees.
|
|
:param axis: The axis to rotate around ('x', 'y', or 'z').
|
|
:return: A new Coordinate after rotation.
|
|
"""
|
|
angle_rad = np.radians(angle_deg) # Convert angle to radians
|
|
c = np.cos(angle_rad)
|
|
s = np.sin(angle_rad)
|
|
|
|
if axis == "x":
|
|
# Rotate around x-axis (affects y, z)
|
|
y_new = self.y * c - self.z * s
|
|
z_new = self.y * s + self.z * c
|
|
return Coordinate(x=self.x, y=y_new, z=z_new)
|
|
elif axis == "y":
|
|
# Rotate around y-axis (affects x, z)
|
|
x_new = self.x * c + self.z * s
|
|
z_new = -self.x * s + self.z * c
|
|
return Coordinate(x=x_new, y=self.y, z=z_new)
|
|
elif axis == "z":
|
|
# Rotate around z-axis (affects x, y)
|
|
x_new = self.x * c - self.y * s
|
|
y_new = self.x * s + self.y * c
|
|
return Coordinate(x=x_new, y=y_new, z=self.z)
|
|
else:
|
|
raise ValueError("Invalid axis. Choose 'x', 'y', or 'z'.")
|
|
|
|
|
|
class SmargonCoordinate(BaseModel):
|
|
sh_mm: Optional[Coordinate] = None # SH coordinate in Smargon
|
|
phi_deg: Optional[float] = None
|
|
chi_deg: Optional[float] = None
|
|
|
|
def eq(self, other: "SmargonCoordinate", tol: float) -> bool:
|
|
return (
|
|
abs(self.sh_mm.x - other.sh_mm.x) < tol
|
|
and abs(self.sh_mm.y - other.sh_mm.y) < tol
|
|
and abs(self.sh_mm.z - other.sh_mm.z) < tol
|
|
and abs(self.phi_deg - other.phi_deg) < tol
|
|
and abs(self.chi_deg - other.chi_deg) < tol
|
|
)
|
|
|
|
def __eq__(self, other: object) -> bool:
|
|
if isinstance(other, SmargonCoordinate):
|
|
return self.eq(other, tol=0.1)
|
|
return NotImplemented
|
|
|
|
|
|
def positive_coords(value: Coordinate) -> Coordinate:
|
|
if value.x <= 0 or value.y <= 0:
|
|
raise ValueError("Coordinates must be positive")
|
|
return value
|