python/src/rwsims/motions.py

34 lines
978 B
Python
Raw Normal View History

2024-06-19 17:10:49 +00:00
from __future__ import annotations
import numpy as np
from numpy.random import Generator
from numpy.typing import ArrayLike
def omega_q(delta: float, eta: float, theta: float, phi: float) -> ArrayLike:
cos_theta = np.cos(theta)
sin_theta = np.sin(theta)
return 2 * np.pi * delta * (3 * cos_theta * cos_theta - 1 + eta * sin_theta*sin_theta * np.cos(2*phi))
def draw_orientation(delta: float, eta: float, rng: Generator, size: int = 1) -> ArrayLike:
z_theta, z_phi = rng.random((2, size))
theta = np.arccos(1 - 2 * z_theta)
phi = 2 * np.pi * z_phi
return omega_q(delta, eta, theta, phi)
class RandomJump:
def __init__(self, delta: float, eta: float, rng: Generator | None = None):
self._delta = delta
self._eta = eta
self._rng = rng
def __repr__(self):
return 'Random Jump'
def jump(self, size: int = 1) -> ArrayLike:
return draw_orientation(self._delta, self._eta, self._rng, size=size)