robolib.utils

27def get_project_root() -> Path:
28    """The project root (this file lives at src/robolib/utils/utils.py, so root is 3 levels up)"""
29    return Path(__file__).parents[3]

23def fmt(arr: np.ndarray, precision: int=2) -> str:
24    """formats a numpy array for logging as tuple"""
25    return "(" + ", ".join(f"{float(x):.{precision}g}" for x in arr) + ")"

31@dataclass
32class FPVData:
33    """The first perso view camera data. Latest frame (as bytes), a lock (for concurrency) and a way to compress it"""
34    frame: bytes
35    frame_shape: tuple[int, int, int]
36    lock: threading.Lock = field(default_factory=threading.Lock)
37    frame_id: int = 0 # used for de-dup on client side
38    _frame_compressed: bytes | None = None
39
40    @property
41    def frame_compressed(self) -> bytes:
42        """compress the frame on demand when the client requests a get_state_with_frame. NOTE: use lock on call."""
43        if self._frame_compressed is None:
44            self._frame_compressed = zlib.compress(self.frame, level=1)
45        return self._frame_compressed

49def make_arr(*data):
50    """makes a numpy array from variable data, e.g. make_arr(1,2,3) <=> np.array([1,2,3], "float32")"""
51    return np.float32(data)

53class FixedSizeDict(OrderedDict):
54    """An dict with a fixed size. Useful for caching purposes."""
55    def __init__(self, *args, maxlen: int = 0, **kwargs):
56        self._maxlen = maxlen
57        super().__init__(*args, **kwargs)
58
59    def __setitem__(self, key, value):
60        existed_before_add = key in self
61        super().__setitem__(key, value)
62        if not existed_before_add and len(self) > self._maxlen:
63            self.popitem(False)

65def ecs_data_as_np(x: list | str | dict | np.ndarray | None, dtype: str) -> np.ndarray | None:
66    """converts a list to a f32/int32 np array for microecs compatibility"""
67    if x is None:
68        return np.array([None], "object")
69    res = np.array(x, dtype=dtype)
70    res = res[None] if res.dtype == "object" and len(res.shape) == 0 else res
71    return res

23@dataclass
24class RlModelWithTexture:
25    """basic raylib model with an optional texture applied to it"""
26    model: RlModel
27    texture: RlTexture | None
28    _bbox: tuple[np.ndarray, np.ndarray] | None = None
29
30    def __post_init__(self):
31        if self.texture is not None:
32            self.model.materials.maps[rl.MATERIAL_MAP_ALBEDO].texture = self.texture
33
34    @staticmethod
35    def from_path(model_path: Path | str, texture_path: Path | str | None = None) -> RlModelWithTexture:
36        """builds a RLModelWithTexture from a path"""
37        model_path = Path(model_path)
38        texture_path = texture_path and Path(texture_path)
39        assert not model_path.is_absolute(), model_path
40        assert texture_path is None or not texture_path.is_absolute(), texture_path
41        assert model_path.exists(), f"Model path: '{model_path}' doesn't exist."
42        logger.info(f"Loading model from '{model_path}'")
43        model = rl.LoadModel(str(model_path).encode())
44        texture = rl.LoadTexture(str(texture_path).encode()) if texture_path is not None else None
45        return RlModelWithTexture(model=model, texture=texture)
46
47    @property
48    def bbox(self) -> tuple[np.ndarray, np.ndarray]:
49        """gets the top left and bottom right 3D points of this mesh. Used for collisions. Cahed once after 1st call"""
50        if self._bbox is None:
51            mins, maxs = np.float32([1<<31, 1<<31, 1<<31]), np.float32([-1<<31, -1<<31, -1<<31])
52            for i in range(self.model.meshCount):
53                mesh: RlMesh = self.model.meshes[i] # of type struct Mesh &
54                buf = rl.ffi.buffer(mesh.vertices, mesh.vertexCount * 3 * 4)
55                vertices = np.frombuffer(buf, "float32").reshape(-1, 3)
56                mins = np.minimum(mins, np.min(vertices, axis=0))
57                maxs = np.maximum(maxs, np.max(vertices, axis=0))
58            self._bbox = mins, maxs
59        return self._bbox
60
61    def __del__(self):
62        if self.model is not None:
63            rl.UnloadModel(self.model)
64        if self.texture is not None:
65            rl.UnloadTexture(self.texture)

67def rl_get_device_and_renderer() -> tuple[str, str]:
68    """returns the device and renderer used by raylib for logging"""
69    gl_renderer = 0x1F01
70    gl_get_string = rl.ffi.cast("unsigned char *(*)(unsigned int)", rl.rlGetProcAddress(b"glGetString"))
71    gstr = lambda e: rl.ffi.string(rl.ffi.cast("char*", gl_get_string(e))).decode()
72    renderer = gstr(gl_renderer)
73    software = ("llvmpipe", "softpipe", "swrast", "lavapipe")
74    device = "CPU (software rasterizer)" if any(s in renderer.lower() for s in software) else "GPU"
75    return device, renderer

19def vec3_from(arr: np.ndarray) -> RlVector3:
20    """convers the numpy array to a raylib vector"""
21    return rl.ffi.from_buffer("Vector3 *", np.ascontiguousarray(arr))[0]

 5class Clock:
 6    """clock used for physics with fixed DT in main loops"""
 7
 8    def __init__(self, dt: float, max_ticks: int):
 9        self.dt = dt
10        self.max_ticks = max_ticks
11        self.prev_time = rl.GetTime()
12        self.accumulator = 0
13
14    def tick(self):
15        """tick once by adding the delta between prev frame and now"""
16        now = rl.GetTime()
17        frame_time = now - self.prev_time
18        self.prev_time = now
19        self.accumulator += frame_time
20
21    def drain(self):
22        """drain the accumulator. in main loop: for _ in clock.drain(): ..."""
23        n_ticks = 0
24        while self.accumulator >= self.dt and n_ticks < self.max_ticks:
25            yield
26            self.accumulator -= self.dt
27            n_ticks += 1
28        self.accumulator = min(self.accumulator, self.dt) # Drop residual debt instead of it piling up across frames
29
30    def wait(self):
31        """waits the leftover time in case the previous tick ran too fast to maintain consistent FPS"""
32        rl.WaitTime(max(self.dt - (rl.GetTime() - self.prev_time), 0))
33
34    def wait_and_tick(self):
35        """calls wait() then tick(). Put this at the beginning of the main loop :)"""
36        self.wait()
37        self.tick()

11def norm(x: np.ndarray) -> np.ndarray:
12    """L2 normalize a vector"""
13    return (x / np.linalg.norm(x, ord=2)).astype("float32")

15def rot_from_forward_up(forward: Point3D, up: Point3D) -> Rot3x3:
16    """right hand rule: x x y = z, y x z = x, z x x = y; z = forward. (3, ) + (3, ) -> (3, 3)"""
17    z = norm(forward) # first, normalize forward, as we care about it's orientation only. z = ||z||
18    y_init = np.float32(up) # up = Y
19    x = norm(np.cross(y_init, z)) # second build x=||y_init x z|| as the cross of existing forward and up
20    y = np.cross(z, x) # third, get the new y as the cross of the already built x and z.
21    R = np.column_stack([x, y, z]) # NOTE: somehow it's left-handed and we make it right-handed here
22    assert is_rot(R), f"\n-{R=}\n-{R@R.T=}"
23    return R

46def pose_from_position_target_up(position: Point3D, target: Point3D, up: Point3D) -> Pose4x4:
47    """returns a 4x4 pose matrix from initial position+target+up (3, ) + (3, ) + (3, ) -> (4, 4)"""
48    pose = np.eye(4, dtype="float32")
49    pose[0:3, 0:3] = rot_from_forward_up(forward=np.float32(target) - position, up=up)
50    pose[0:3, 3] = position
51    return pose

25def pose_to_trans_euler(pose: Pose4x4) -> Point6D:
26    """converts a pose to a 6DoF translation + euler vector, mostly for printing reasons (4, 4) -> (6, )"""
27    return np.float32([*pose[0:3, 3], *tr2rpy(pose)])

29def pose_from_trans_euler(position_rpy: Point6D) -> Pose4x4:
30    """converts a position + a euler vector to a Pose. (6, ) -> (4, 4)"""
31    res = rpy2tr(position_rpy[-3:]).astype("float32")
32    res[0:3, 3] = position_rpy[0:3]
33    return res

35def renormalize_rotation_matrix(rot: Rot3x3, tol: float=1e10) -> Rot3x3:
36    """uses SVD to renormalize a rotation matrix if it drifted too much (i.e. during trajectory updates)"""
37    if is_rot(rot, tol=tol):
38        return rot
39    logger.debug(f"Renormalizing R as err is {np.linalg.norm((rot @ rot.T) - np.eye(3)):.7f}")
40    U, _, Vt = np.linalg.svd(rot)
41    D = np.diag([1, 1, np.linalg.det(U @ Vt)]) # force det=1 (not -1)
42    new_rot = U @ D @ Vt
43    assert is_rot(new_rot, tol=1e10), new_rot
44    return new_rot

53def get_closest_square(n: int) -> tuple[int, int]:
54    """
55    Given a stack of N images, find the closest square X>=N*N and return that.
56    Note: There are only 2 rows possible between x^2 and (x+1)^2 because (x+1)^2 = x^2 + 2*x + 1, thus we can add two
57    columns at most. If a 3rd column is needed, then closest lower bound is (x+1)^2 and we must use that.
58    Example: 9: 3*3; 12 -> 3*3 -> 3*4 (3 rows). 65 -> 8*8 -> 8*9. 73 -> 8*8 -> 8*9 -> 9*9
59    """
60    x = int(math.sqrt(n))
61    r, c = x, x
62    c = c + 1 if c * r < n else c
63    r = r + 1 if c * r < n else r
64    assert (c + 1) * r > n and c * (r + 1) > n
65    return r, c

67def make_radius(size: Point3D) -> float:
68    """creates a radius from a size and a scale"""
69    return math.sqrt(size[0]**2 + size[1]**2 + size[2]**2) / 2

193def btrexp(S: Mat4x4) -> Pose4x4:
194    """
195    Batched matrix exponential of an se(3) element -> SE(3) (a screw motion).
196    ``S`` is the batched (Nx)4x4 augmented skew-symmetric matrix produced by ``bskewa``.
197        >>> btrexp(bskewa(np.array([[1, 0, 0, 0, 0, 0]])))
198    """
199    n = len(S)
200    tw = bvexa(S) # (N, 6)
201    v = tw[:, 0:3] # (3, )
202    w = tw[:, 3:6] # (3, )
203
204    w_norm = np.linalg.norm(w, axis=1) # (N, )
205    v_norm = np.linalg.norm(v, axis=1) # (N, )
206    # theta is the screw magnitude: rotation dominates, else it's a pure translation
207    theta = np.where(w_norm >= 20 * _eps, w_norm, v_norm)[:, None] # (N, 1)
208
209    with np.errstate(divide="ignore", invalid="ignore"):   # 0/0 on zero-twist is expected; nan_to_num handles it
210        v = np.nan_to_num(v / theta, nan=0) # (N, 3)
211        w = np.nan_to_num(w / theta, nan=0) # (N, 3)
212
213    skw = bskew(w)                                                      # (N, 3, 3)
214    theta_n = theta[..., None]                                          # (N, 1, 1)
215    st, ct = np.sin(theta_n), np.cos(theta_n)                           # (N, 1, 1), (N, 1, 1)
216    eye = np.eye(3)[None].repeat(n, axis=0)                             # (N, 3, 3)
217    R = eye + st * skw + (1.0 - ct) * skw @ skw                         # (N, 3, 3) Rodrigues
218    V = eye * theta_n + (1.0 - ct) * skw + (theta_n - st) * skw @ skw   # (N, 3, 3) translation map
219
220    T = np.eye(4)[None].repeat(n, axis=0)                               # (N, 4, 4)
221    T[:, 0:3, 0:3] = R
222    T[:, 0:3, 3] = np.einsum("bij,bj->bi", V, v)
223    return T

182def bskewa(v: Point6D) -> Mat4x4:
183    """batched skewa"""
184    assert len(v.shape) == 2 and v.shape[1] == 6, v.shape
185    zero = np.zeros((len(v), ), "float32")
186    res = np.array([
187        [ zero,     -v[:, 5],    v[:, 4],  v[:, 0]],
188        [ v[:, 5],   zero,      -v[:, 3],  v[:, 1]],
189        [-v[:, 4],   v[:, 3],    zero,     v[:, 2]],
190        [ zero,      zero,       zero,     zero]], dtype=np.float64) # (4, 4, N)
191    return np.permute_dims(res, (2, 0, 1)) # (N, 4, 4)