65 lines
No EOL
2.3 KiB
Python
65 lines
No EOL
2.3 KiB
Python
import numpy as np
|
|
import zlib
|
|
import struct
|
|
from typing import Any, Callable
|
|
|
|
def crc32_array(a: np.ndarray) -> int:
|
|
if not a.size: return 0
|
|
ac = np.ascontiguousarray(a)
|
|
mv = memoryview(ac.data).cast("B")
|
|
return zlib.adler32(mv)
|
|
|
|
def crc32_list(l: list[np.ndarray]) -> int:
|
|
if not l: return 0
|
|
h = 0
|
|
for a in l:
|
|
h = zlib.adler32(struct.pack(">I", crc32_array(a)), h)
|
|
return h
|
|
|
|
def euler_rotation_matrix(rx: float, ry: float, rz: float) -> np.ndarray:
|
|
"""Computes a 3x3 rotation matrix (Rz @ Ry @ Rx) from Euler angles in radians."""
|
|
cx, sx = np.cos(rx), np.sin(rx)
|
|
cy, sy = np.cos(ry), np.sin(ry)
|
|
cz, sz = np.cos(rz), np.sin(rz)
|
|
Rx = np.array([[1.0, 0.0, 0.0], [0.0, cx, -sx], [0.0, sx, cx]], dtype=np.float64)
|
|
Ry = np.array([[cy, 0.0, sy], [0.0, 1.0, 0.0], [-sy, 0.0, cy]], dtype=np.float64)
|
|
Rz = np.array([[cz, -sz, 0.0], [sz, cz, 0.0], [0.0, 0.0, 1.0]], dtype=np.float64)
|
|
return Rz @ (Ry @ Rx)
|
|
|
|
from collections import OrderedDict
|
|
|
|
def _make_hashable(val: Any) -> Any:
|
|
if isinstance(val, np.ndarray):
|
|
return crc32_array(val)
|
|
elif isinstance(val, (list, tuple)):
|
|
if not val:
|
|
return ()
|
|
if isinstance(val[0], np.ndarray):
|
|
return crc32_list(list(val))
|
|
return tuple(_make_hashable(v) for v in val)
|
|
elif isinstance(val, dict):
|
|
return tuple(sorted((k, _make_hashable(v)) for k, v in val.items()))
|
|
elif isinstance(val, set):
|
|
return tuple(sorted(_make_hashable(v) for v in val))
|
|
return val
|
|
|
|
def ndarray_lru_cache(maxsize: int = 128) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
|
|
def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
|
|
cache: OrderedDict[tuple[Any, ...], Any] = OrderedDict()
|
|
def wrapper(*args, **kwargs):
|
|
key_parts = [_make_hashable(arg) for arg in args]
|
|
for k, v in sorted(kwargs.items()):
|
|
key_parts.append((k, _make_hashable(v)))
|
|
|
|
key = tuple(key_parts)
|
|
if key in cache:
|
|
cache.move_to_end(key)
|
|
return cache[key]
|
|
|
|
result = func(*args, **kwargs)
|
|
cache[key] = result
|
|
if len(cache) > maxsize:
|
|
cache.popitem(last=False)
|
|
return result
|
|
return wrapper
|
|
return decorator |