robosim.utils

29def get_project_root() -> Path:
30    """The project root"""
31    return Path(__file__).parents[2]

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

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

61@dataclass
62class RlModelWithTexture:
63    """basic raylib model with an optional texture applied to it"""
64    model: RlModel
65    texture: RlTexture | None
66    _bbox: tuple[np.ndarray, np.ndarray] | None = None
67
68    def __post_init__(self):
69        if self.texture is not None:
70            self.model.materials.maps[rl.MATERIAL_MAP_ALBEDO].texture = self.texture
71
72    @staticmethod
73    def from_path(model_path: Path, texture_path: Path | None = None) -> RlModelWithTexture:
74        """builds a RLModelWithTexture from a path"""
75        assert model_path.exists(), f"Model path: '{model_path}' doesn't exist."
76        logger.info(f"Loading model from '{model_path}'")
77        model = rl.LoadModel(str(model_path).encode())
78        texture = rl.LoadTexture(str(texture_path).encode()) if texture_path is not None else None
79        return RlModelWithTexture(model=model, texture=texture)
80
81    @property
82    def bbox(self) -> tuple[np.ndarray, np.ndarray]:
83        """gets the top left and bottom right 3D points of this mesh. Used for collisions. Cahed once after 1st call"""
84        if self._bbox is None:
85            mins, maxs = np.float32([1<<31, 1<<31, 1<<31]), np.float32([-1<<31, -1<<31, -1<<31])
86            for i in range(self.model.meshCount):
87                mesh: RlMesh = self.model.meshes[i] # of type struct Mesh &
88                buf = rl.ffi.buffer(mesh.vertices, mesh.vertexCount * 3 * 4)
89                vertices = np.frombuffer(buf, "float32").reshape(-1, 3)
90                mins = np.minimum(mins, np.min(vertices, axis=0))
91                maxs = np.maximum(maxs, np.max(vertices, axis=0))
92            self._bbox = mins, maxs
93        return self._bbox
94
95    def __del__(self):
96        if self.model is not None:
97            rl.UnloadModel(self.model)
98        if self.texture is not None:
99            rl.UnloadTexture(self.texture)

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

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

117class FixedSizeDict(OrderedDict):
118    """An dict with a fixed size. Useful for caching purposes."""
119    def __init__(self, *args, maxlen: int = 0, **kwargs):
120        self._maxlen = maxlen
121        super().__init__(*args, **kwargs)
122
123    def __setitem__(self, key, value):
124        existed_before_add = key in self
125        super().__setitem__(key, value)
126        if not existed_before_add and len(self) > self._maxlen:
127            self.popitem(False)

101def rl_get_device_and_renderer() -> tuple[str, str]:
102    """returns the device and renderer used by raylib for logging"""
103    gl_renderer = 0x1F01
104    gl_get_string = rl.ffi.cast("unsigned char *(*)(unsigned int)", rl.rlGetProcAddress(b"glGetString"))
105    gstr = lambda e: rl.ffi.string(rl.ffi.cast("char*", gl_get_string(e))).decode()
106    renderer = gstr(gl_renderer)
107    software = ("llvmpipe", "softpipe", "swrast", "lavapipe")
108    device = "CPU (software rasterizer)" if any(s in renderer.lower() for s in software) else "GPU"
109    return device, renderer

 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)