diff --git a/ncore/impl/data/compat.py b/ncore/impl/data/compat.py index 2166685f..8505a78c 100644 --- a/ncore/impl/data/compat.py +++ b/ncore/impl/data/compat.py @@ -564,13 +564,16 @@ def model_parameters(self) -> Optional[ConcreteLidarModelParametersUnion]: """Returns parameters specific to the lidar's intrinsic model (optional as not mandatory)""" ... - def get_frame_ray_bundle_model_element(self, frame_index: int) -> Optional[npt.NDArray[np.uint16]]: + def get_frame_ray_bundle_model_element(self, frame_index: int) -> Optional[npt.NDArray[Any]]: """Returns the per-ray model elements for a ray bundle for a specific frame, if available. Args: frame_index: Index of the frame Returns: - Array of per-ray model elements [N,] or None if not available + Array of per-ray model element indices / sampling coordinates with a + leading per-ray dimension N and any trailing shape ((N,), (N, K), + (N, A, B, ...)); integer dtype or floating dtype up to float32. None + if not available. """ ... diff --git a/ncore/impl/data/v4/compat.py b/ncore/impl/data/v4/compat.py index 67ad492e..bebc876a 100644 --- a/ncore/impl/data/v4/compat.py +++ b/ncore/impl/data/v4/compat.py @@ -482,7 +482,7 @@ def model_parameters(self) -> Optional[ConcreteLidarModelParametersUnion]: return self._model_parameters @override - def get_frame_ray_bundle_model_element(self, frame_index: int) -> Optional[npt.NDArray[np.uint16]]: + def get_frame_ray_bundle_model_element(self, frame_index: int) -> Optional[npt.NDArray[Any]]: """Returns the per-ray model elements for a ray bundle for a specific frame, if available""" if self.lidar_reader.has_frame_ray_bundle_data( diff --git a/ncore/impl/data/v4/components.py b/ncore/impl/data/v4/components.py index bbc74e02..bb966b06 100644 --- a/ncore/impl/data/v4/components.py +++ b/ncore/impl/data/v4/components.py @@ -1497,8 +1497,8 @@ def store_frame( ], # per-ray unit-norm direction vectors in sensor frame at measure time (raw / not motion-compensated, needs to have unit norm) (float32, [N, 3]) timestamp_us: npt.NDArray[np.uint64], # per-ray timestamp in microseconds (uint64, [N]) model_element: Optional[ - npt.NDArray[np.uint16] - ], # per-ray model element indices, if applicable (uint16, [N, 2]) + npt.NDArray[Any] + ], # per-ray model element indices / sampling coordinates, if applicable. Shape is (N, ...): a leading per-ray dimension N followed by any number of trailing dimensions (none -> (N,), or (N, A, B, ...)). dtype may be any integer type (discrete model index) or any floating type up to float32 (continuous model sampling coordinate, e.g. float16 / float32); float64 (and wider) is disallowed as too storage-heavy. The shape and dtype are self-describing in storage and interpreted by the associated lidar model (e.g. (N, 2) (row, col) for the current structured spinning model; extra dimensions may encode zone, scan-direction / frame parity for solid-state models). # per-ray return data for R returns - non-existing values are indicated via NaNs distance_m: npt.NDArray[ np.float32 @@ -1537,8 +1537,17 @@ def store_frame( ray_data["timestamp_us"] = (timestamp_us, timestamp_us.shape) if model_element is not None: - assert model_element.shape == (n_rays, 2) - assert model_element.dtype == np.dtype("uint16") + assert model_element.ndim >= 1, "model_element must have a leading per-ray dimension" + assert model_element.shape[0] == n_rays, ( + "model_element must have one entry per ray (leading dimension == n_rays)" + ) + # Allow any integer dtype (discrete model index) or any floating + # type up to float32 (continuous model sampling coordinate, e.g. + # float16 / float32). float64 (and wider) is disallowed as too + # storage-heavy. + assert np.issubdtype(model_element.dtype, np.integer) or ( + np.issubdtype(model_element.dtype, np.floating) and model_element.dtype.itemsize <= 4 + ), f"model_element must be an integer type or a float type up to float32, got {model_element.dtype}" ray_data["model_element"] = (model_element, model_element.shape) ## Per return data diff --git a/ncore/impl/data/v4/components_test.py b/ncore/impl/data/v4/components_test.py index 31786170..534c1cfb 100644 --- a/ncore/impl/data/v4/components_test.py +++ b/ncore/impl/data/v4/components_test.py @@ -419,7 +419,7 @@ def normalize_points(vectors: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: lidar_writer.store_frame( ref_lidar_direction0, ref_lidar_timestamp_us0 := np.linspace(0 * 1e6, 0.5 * 1e6, num=5, dtype=np.uint64), - ref_lidar_model_element0 := np.arange(5 * 2, dtype=np.uint16).reshape((5, 2)), + ref_lidar_model_element0 := np.arange(5 * 3, dtype=np.uint16).reshape((5, 3)), ref_lidar_distance_m0 := lidar_distance_m0[np.newaxis, :], ref_lidar_intensity0 := np.random.default_rng().random((1, 5)).astype(np.float32), ref_lidar_timestamps_us0 := np.array([0 * 1e6, 0.5 * 1e6], dtype=np.uint64), @@ -3175,3 +3175,168 @@ def test_per_frame_generic_data_zero_dim(self) -> None: self.assertEqual(gd.shape, (0, 3)) tmpdir.cleanup() + + +@parameterized_class( + ("store_type"), + [ + ("itar",), + ("directory",), + ], +) +class TestLidarModelElementTupleDim(unittest.TestCase): + """Tests that per-ray model_element supports flexible shapes and dtypes. + + model_element keeps a leading per-ray dimension N but allows any trailing + shape: flat (N,), the legacy (N, 2) (row, col), or arbitrary (N, A, B, ...). + The dtype may be any integer type (discrete model index) or a floating type + up to float32 (continuous model sampling coordinate, e.g. float16 / float32); + float64 (and wider) is rejected as too storage-heavy. Shape and dtype are + self-describing in storage, so this is a backwards-compatible extension of + the previously fixed (N, 2) uint16 layout. + """ + + store_type: Literal["itar", "directory"] + + def setUp(self) -> None: + np.set_printoptions(floatmode="unique", linewidth=200, suppress=True) + + def _make_lidar_frame(self, n_rays: int) -> Dict[str, np.ndarray]: + rng = np.random.default_rng(7) + raw_pts = rng.random((n_rays, 3)).astype(np.float32) + 0.1 + norms = np.linalg.norm(raw_pts, axis=1) + return { + "direction": (raw_pts / norms[:, np.newaxis]).astype(np.float32), + "distance_m": norms[np.newaxis, :].astype(np.float32), + "intensity": rng.random((1, n_rays)).astype(np.float32), + "timestamp_us": np.linspace(0, 100_000, num=n_rays, dtype=np.uint64), + "frame_timestamps_us": np.array([0, 100_000], dtype=np.uint64), + } + + def _make_writer( + self, + ) -> Tuple[LidarSensorComponent.Writer, SequenceComponentGroupsWriter, tempfile.TemporaryDirectory]: + tmpdir = tempfile.TemporaryDirectory() + store_writer = SequenceComponentGroupsWriter( + output_dir_path=UPath(tmpdir.name), + store_base_name=(seq_id := "lidar-model-element-tuple-dim"), + sequence_id=seq_id, + sequence_timestamp_interval_us=HalfClosedInterval(0, 10_000_001), + store_type=self.store_type, + generic_meta_data={}, + ) + lidar_writer = store_writer.register_component_writer(LidarSensorComponent.Writer, "lidar_me", "lidars") + return lidar_writer, store_writer, tmpdir + + def _roundtrip(self, model_element: np.ndarray) -> np.ndarray: + """Store a single frame with the given model_element and read it back.""" + lidar_writer, store_writer, tmpdir = self._make_writer() + n_rays = model_element.shape[0] + frame = self._make_lidar_frame(n_rays) + lidar_writer.store_frame( + direction=frame["direction"], + timestamp_us=frame["timestamp_us"], + model_element=model_element, + distance_m=frame["distance_m"], + intensity=frame["intensity"], + frame_timestamps_us=frame["frame_timestamps_us"], + generic_data={}, + generic_meta_data={}, + ) + store_paths = store_writer.finalize() + reader = SequenceComponentGroupsReader(component_group_paths=store_paths) + lidar_reader = reader.open_component_readers(LidarSensorComponent.Reader)["lidar_me"] + frame_ts = int(lidar_reader.frames_timestamps_us[0, 1]) + read_back = lidar_reader.get_frame_ray_bundle_data(frame_ts, "model_element") + tmpdir.cleanup() + return read_back + + def _assert_store_rejected(self, model_element: np.ndarray) -> None: + lidar_writer, _store_writer, tmpdir = self._make_writer() + n_rays = model_element.shape[0] if model_element.ndim >= 1 else 0 + frame = self._make_lidar_frame(max(n_rays, 1)) + with self.assertRaises(AssertionError): + lidar_writer.store_frame( + direction=frame["direction"], + timestamp_us=frame["timestamp_us"], + model_element=model_element, + distance_m=frame["distance_m"], + intensity=frame["intensity"], + frame_timestamps_us=frame["frame_timestamps_us"], + generic_data={}, + generic_meta_data={}, + ) + tmpdir.cleanup() + + def test_legacy_n2_uint16_roundtrip(self) -> None: + """The original (N, 2) uint16 layout continues to work unchanged.""" + me = np.arange(4 * 2, dtype=np.uint16).reshape((4, 2)) + read_back = self._roundtrip(me) + self.assertEqual(read_back.shape, (4, 2)) + self.assertEqual(read_back.dtype, np.dtype("uint16")) + np.testing.assert_array_equal(read_back, me) + + def test_flat_n_roundtrip(self) -> None: + """A flat per-ray index (N,) round-trips.""" + me = np.arange(5, dtype=np.uint32) + read_back = self._roundtrip(me) + self.assertEqual(read_back.shape, (5,)) + self.assertEqual(read_back.dtype, np.dtype("uint32")) + np.testing.assert_array_equal(read_back, me) + + def test_higher_dim_roundtrip(self) -> None: + """An (N, A, B) model index round-trips with all trailing dims preserved.""" + me = np.arange(5 * 2 * 3, dtype=np.int32).reshape((5, 2, 3)) + read_back = self._roundtrip(me) + self.assertEqual(read_back.shape, (5, 2, 3)) + self.assertEqual(read_back.dtype, np.dtype("int32")) + np.testing.assert_array_equal(read_back, me) + + def test_float32_sampling_coordinate_roundtrip(self) -> None: + """A float32 continuous sampling coordinate (N, 2) round-trips.""" + me = (np.arange(5 * 2, dtype=np.float32).reshape((5, 2)) * 0.5).astype(np.float32) + read_back = self._roundtrip(me) + self.assertEqual(read_back.shape, (5, 2)) + self.assertEqual(read_back.dtype, np.dtype("float32")) + np.testing.assert_array_equal(read_back, me) + + def test_float16_sampling_coordinate_roundtrip(self) -> None: + """A float16 continuous sampling coordinate (N, 2) round-trips.""" + me = (np.arange(5 * 2, dtype=np.float16).reshape((5, 2)) * 0.5).astype(np.float16) + read_back = self._roundtrip(me) + self.assertEqual(read_back.shape, (5, 2)) + self.assertEqual(read_back.dtype, np.dtype("float16")) + np.testing.assert_array_equal(read_back, me) + + def test_rejects_float64(self) -> None: + """float64 model_element is rejected as too storage-heavy.""" + me = np.arange(3 * 2, dtype=np.float64).reshape((3, 2)) + self._assert_store_rejected(me) + + def test_rejects_zero_dim_scalar(self) -> None: + """A 0-d array (no per-ray dimension) is rejected.""" + me = np.array(3, dtype=np.uint16) + self._assert_store_rejected(me) + + def test_rejects_wrong_leading_dim(self) -> None: + """A model_element whose leading dim != n_rays is rejected.""" + # n_rays from the frame is 5, but provide 4 leading entries + lidar_writer, _store_writer, tmpdir = self._make_writer() + frame = self._make_lidar_frame(5) + me = np.arange(4 * 2, dtype=np.uint16).reshape((4, 2)) + with self.assertRaises(AssertionError): + lidar_writer.store_frame( + direction=frame["direction"], + timestamp_us=frame["timestamp_us"], + model_element=me, + distance_m=frame["distance_m"], + intensity=frame["intensity"], + frame_timestamps_us=frame["frame_timestamps_us"], + generic_data={}, + generic_meta_data={}, + ) + tmpdir.cleanup() + + +if __name__ == "__main__": + unittest.main()