wehub-resource-sync
55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
|
|
MAX_INT64 = 2**63 - 1
|
|
MAX_INT32 = 2**31 - 1
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Point3DInput:
|
|
positions: npt.NDArray[np.float32]
|
|
colors: npt.NDArray[np.uint32]
|
|
radii: npt.NDArray[np.float32]
|
|
label: str = "some label"
|
|
|
|
@classmethod
|
|
def prepare(cls, seed: int, num_points: int) -> Point3DInput:
|
|
rng = np.random.default_rng(seed=seed)
|
|
|
|
return cls(
|
|
positions=rng.integers(0, MAX_INT64, (num_points, 3)).astype(dtype=np.float32),
|
|
colors=rng.integers(0, MAX_INT32, num_points, dtype=np.uint32),
|
|
radii=rng.integers(0, MAX_INT64, num_points).astype(dtype=np.float32),
|
|
)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Transform3DInput:
|
|
"""Input data for Transform3D benchmark with translation and mat3x3."""
|
|
|
|
translations: npt.NDArray[np.float32] # Shape: (num_time_steps, num_entities, 3)
|
|
mat3x3s: npt.NDArray[np.float32] # Shape: (num_time_steps, num_entities, 3, 3)
|
|
num_entities: int
|
|
num_time_steps: int
|
|
|
|
@classmethod
|
|
def prepare(cls, seed: int, num_entities: int, num_time_steps: int) -> Transform3DInput:
|
|
rng = np.random.default_rng(seed=seed)
|
|
|
|
# Generate translations in range [0, 10)
|
|
translations = rng.random((num_time_steps, num_entities, 3), dtype=np.float32) * 10.0
|
|
|
|
# Generate mat3x3 values in range [-1, 1)
|
|
mat3x3s = rng.random((num_time_steps, num_entities, 3, 3), dtype=np.float32) * 2.0 - 1.0
|
|
|
|
return cls(
|
|
translations=translations,
|
|
mat3x3s=mat3x3s,
|
|
num_entities=num_entities,
|
|
num_time_steps=num_time_steps,
|
|
)
|