170 lines
4.6 KiB
Python
170 lines
4.6 KiB
Python
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
|
|
import numpy as np
|
|
from numpy.random import Generator
|
|
|
|
|
|
class BaseMotion(ABC):
|
|
def __init__(self, delta: float, eta: float, rng: Generator):
|
|
self._delta = delta
|
|
self._eta = eta
|
|
|
|
self._rng = rng
|
|
|
|
@abstractmethod
|
|
def __repr__(self):
|
|
pass
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return self.__class__.__name__
|
|
|
|
def start(self):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def jump(self, size: int = 1) -> 'ArrayLike':
|
|
pass
|
|
|
|
|
|
class RandomJump(BaseMotion):
|
|
|
|
def __repr__(self):
|
|
return 'Random Jump'
|
|
|
|
def jump(self, size: int = 1) -> 'ArrayLike':
|
|
return omega_q(self._delta, self._eta, *draw_orientation(self._rng, size=size))
|
|
|
|
|
|
class TetrahedralJump(BaseMotion):
|
|
def __init__(self, delta: float, eta: float, rng: Generator):
|
|
super().__init__(delta, eta, rng)
|
|
|
|
self._orientation = None
|
|
self._start = None
|
|
|
|
def __repr__(self):
|
|
return 'Tetrahedral Jump'
|
|
|
|
def start(self):
|
|
self._orientation = self._make_tetrahedron()
|
|
self._start = self._rng.choice([0, 1, 2, 3])
|
|
|
|
def _make_tetrahedron(self) -> np.ndarray:
|
|
beta = np.arccos(-1/3) # tetrahedral angle 109.5 degrees
|
|
sin_beta = np.sin(beta)
|
|
cos_beta = np.cos(beta)
|
|
|
|
# corners of a tetrahedron
|
|
alpha = 2 * np.pi * self._rng.random()
|
|
corners = np.array([
|
|
[0, 0, 1],
|
|
[sin_beta * np.cos(alpha), sin_beta * np.sin(alpha), cos_beta],
|
|
[sin_beta * np.cos(alpha+2*np.pi/3), sin_beta * np.sin(alpha+2*np.pi/3), cos_beta],
|
|
[sin_beta * np.cos(alpha+4*np.pi/3), sin_beta * np.sin(alpha+4*np.pi/3), cos_beta]
|
|
])
|
|
|
|
# orientation in lab system
|
|
cos_theta0, phi0 = draw_orientation(self._rng)
|
|
|
|
rot = get_rotation_matrix(
|
|
corners[0],
|
|
np.array(spherical_to_xyz(1., np.arccos(cos_theta0), phi0)),
|
|
)
|
|
|
|
orientations = np.zeros(4)
|
|
for i in range(4):
|
|
corner_lab = rot @ corners[i]
|
|
_, theta_i, phi_i = xyz_to_spherical(*corner_lab)
|
|
orientations[i] = omega_q(self._delta, self._eta, theta_i, phi_i)
|
|
|
|
# print(orientations)
|
|
#
|
|
# theta0 = np.arccos(cos_theta0)
|
|
# v0 = np.array([np.sin(theta0) * np.cos(phi0), np.sin(theta0)*np.sin(theta0), cos_theta0])
|
|
# norm = np.linalg.norm(v0)
|
|
# print(norm)
|
|
#
|
|
#
|
|
# corners = np.zeros((4, 3))
|
|
# corners[0] = v0
|
|
|
|
|
|
return orientations
|
|
|
|
def jump(self, size: int = 1) -> 'ArrayLike':
|
|
jumps = self._rng.choice([1, 2, 3], size=size)
|
|
jumps = np.cumsum(jumps) + self._start
|
|
jumps %= 4
|
|
|
|
self._start = jumps[-1]
|
|
|
|
return self._orientation[jumps]
|
|
|
|
|
|
# Helper functions
|
|
|
|
def xyz_to_spherical(x_in: float, y_in: float, z_in: float) -> tuple[np.floating, float, float]:
|
|
r = np.linalg.norm([x_in, y_in, z_in])
|
|
theta: float = np.arccos(z_in)
|
|
phi: float = np.arctan2(y_in, x_in)
|
|
|
|
return r, theta, phi
|
|
|
|
|
|
def spherical_to_xyz(r: float, theta: float, phi: float) -> tuple[float, float, float]:
|
|
sin_theta = np.sin(theta)
|
|
return r*np.cos(phi)*sin_theta, r*np.sin(phi)*sin_theta, r*np.cos(theta)
|
|
|
|
|
|
def get_rotation_matrix(vec_in: np.ndarray, vec_out: np.ndarray):
|
|
rotation = np.eye(3)
|
|
|
|
# rotation by angle around given axis
|
|
cos_angle = np.dot(vec_in, vec_out)
|
|
|
|
# check for parallel / anti-parallel vectors
|
|
if cos_angle == 1:
|
|
return rotation
|
|
elif cos_angle == -1:
|
|
return -rotation
|
|
|
|
else:
|
|
axis = np.cross(vec_in, vec_out)
|
|
scale = np.linalg.norm(axis)
|
|
axis /= scale
|
|
sin_angle = scale / np.linalg.norm(vec_in) / np.linalg.norm(vec_out)
|
|
|
|
v_cross = np.array([
|
|
[0, -axis[2], axis[1]],
|
|
[axis[2], 0, -axis[0]],
|
|
[-axis[1], axis[0], 0],
|
|
])
|
|
|
|
rotation += sin_angle * v_cross
|
|
rotation += (1-cos_angle) * v_cross @ v_cross
|
|
|
|
return rotation
|
|
|
|
|
|
def omega_q(delta: float, eta: float, cos_theta: ArrayLike, phi: ArrayLike) -> ArrayLike:
|
|
# sin_theta = np.sin(cos_theta)
|
|
# cos_theta = np.cos(cos_theta)
|
|
sin_theta_sq = 1 - cos_theta**2
|
|
return np.pi * delta * (3 * cos_theta**2 - 1 + eta * sin_theta_sq * np.cos(2*phi))
|
|
|
|
|
|
def draw_orientation(rng: Generator, size: int | None = None) -> tuple[ArrayLike, ArrayLike]:
|
|
if size is not None:
|
|
z_theta, z_phi = rng.random((2, size))
|
|
else:
|
|
z_theta, z_phi = rng.random(2)
|
|
cos_theta = 1 - 2 * z_theta
|
|
phi = 2 * np.pi * z_phi
|
|
|
|
return cos_theta, phi
|
|
|
|
|