add Extension mechanism and PET SMoS implementation (#34)

klit-la · web-flow · commit 8a1d7742104e · 2025-12-05T13:57:11.000+01:00
diff --git a/doc/source/reference/pandas/pose.rst b/doc/source/reference/pandas/pose.rst
@@ -7,7 +7,9 @@ Pose
 
 A ``Pose`` specialized for trajectory data representing the state of a traffic participant for a specific point in time.
 
-
+.. inheritance-diagram:: tasi.Pose
+    :top-classes: pandas.core.frame.DataFrame
+    :parts: 1
 
 Serialization / IO / conversion
 ************************************
diff --git a/doc/source/reference/pandas/trajectory.rst b/doc/source/reference/pandas/trajectory.rst
@@ -5,10 +5,13 @@ Trajectory
 
 .. currentmodule:: tasi
 
-A :py:class:`Trajectory` specialized for trajectory data representing the evolution of a traffic participant. A
-trajectory is a collection of :py:class:`Pose`s.
-
+A :py:class:`Trajectory` specialized for trajectory data representing the
+evolution of a traffic participant. A trajectory is a collection of
+:py:class:`Pose`s.
 
+.. inheritance-diagram:: tasi.Trajectory
+    :top-classes: pandas.core.frame.DataFrame
+    :parts: 1
 
 Serialization / IO / conversion
 ************************************
@@ -26,7 +29,7 @@ Indexing
    :toctree: api/
 
    Trajectory.att
-   TrajectoryDataset.trajectory
+
 
 Filtering
 ***********
@@ -54,7 +57,7 @@ Attributes
    Trajectory.id
    Trajectory.interval
    Trajectory.timestamps
-
+   Trajectory.smos
 
 .. note::
 
diff --git a/doc/source/reference/pandas/trajectorydataset.rst b/doc/source/reference/pandas/trajectorydataset.rst
@@ -7,6 +7,10 @@ TrajectoryDataset
 
 A ``Dataset`` specialized for trajectory data.
 
+.. inheritance-diagram:: tasi.TrajectoryDataset
+    :top-classes: pandas.core.frame.DataFrame, tasi.dataset.Dataset
+    :parts: 1
+
 
 Constructor
 ************
diff --git a/doc/source/reference/pydantic/geopose.rst b/doc/source/reference/pydantic/geopose.rst
@@ -20,7 +20,6 @@ Factory methods
    :toctree: api/
 
    GeoPosePublic.from_orm
-   GeoPosePublic.from_tasi
    GeoPosePublic.from_pose
    GeoPosePublic.model_validate
 
diff --git a/tasi/base.py b/tasi/base.py
@@ -63,15 +63,15 @@ class IndexMixin:
     ID_COLUMN = "id"
 
     @property
-    def ids(self) -> pd.Index:
+    def ids(self) -> List[int]:
         """Returns the unique ids in the dataset
 
         Returns:
             np.ndarray: A List of ids
         """
         idx: pd.MultiIndex = self.index  # type: ignore
 
-        return idx.get_level_values(self.ID_COLUMN).unique()
+        return idx.get_level_values(self.ID_COLUMN).unique()  # type: ignore
 
 
 class TASIBase(LocatableEntity):
@@ -231,6 +231,18 @@ def to_csv(
         finally:
             self.columns = old_columns
 
+    @property
+    def woid(self) -> pd.DataFrame:
+        """Without the object id on the index
+
+        This property will return a `DataFrame` without the `id` on the index
+        that can be used to compare trajectories with each other.
+
+        Returns:
+            pd.DataFrame: A DataFrame with the `id` as a column.
+        """
+        return self.reset_index("id")
+
 
 class PoseCollectionBase(CollectionBase):
 
diff --git a/tasi/calculus.py b/tasi/calculus.py
@@ -1,8 +1,7 @@
-from tasi.utils import has_extra
+from .utils import has_extra
 
 EXTRA = has_extra("performance")
 
-
 if EXTRA:
     from numba import jit, njit  # pyright: ignore[reportMissingImports]
 else:
diff --git a/tasi/dataset/base.py b/tasi/dataset/base.py
@@ -1,5 +1,5 @@
 from functools import wraps
-from typing import Iterable, List, Tuple, Union
+from typing import Iterable, List, Tuple, Union, overload
 
 import numpy as np
 import pandas as pd
@@ -54,7 +54,15 @@ def _pose_constructor(self):
 
         return Pose
 
-    def trajectory(self, index: Union[int, Iterable[int]], inverse: bool = False):
+    @overload
+    def trajectory(self, index: int, inverse: bool = False) -> Trajectory: ...
+
+    @overload
+    def trajectory(self, index: Iterable[int], inverse: bool = False) -> Self: ...
+
+    def trajectory(
+        self, index: Union[int, Iterable[int]], inverse: bool = False
+    ) -> Trajectory | Self:
         """
         Select trajectory data for specific indices, or exclude them if inverse is set to True.
 
@@ -67,7 +75,7 @@ def trajectory(self, index: Union[int, Iterable[int]], inverse: bool = False):
                 are excluded from the resulting dataset, and all other trajectories are included. Defaults to False.
 
         Returns:
-            TrajectoryDataset: A trajectory or multiple trajectories of the dataset.
+            tasi.Trajectory | TrajectoryDataset: A trajectory or multiple trajectories of the dataset.
         """
 
         if isinstance(index, (int, np.int_)):
diff --git a/tasi/dlr/dataset.py b/tasi/dlr/dataset.py
@@ -7,6 +7,7 @@
 from pathlib import Path
 from typing import List, Tuple, Union
 
+import numpy as np
 import pandas as pd
 import requests
 from tqdm import tqdm
@@ -458,7 +459,7 @@ def to_tasi(self) -> TrajectoryDataset:
             position=self.position,
             velocity=self.velocity,
             acceleration=self.acceleration,
-            heading=self.yaw,
+            heading=np.deg2rad(self.yaw),
             classifications=self.classifications,
             dimension=self.dimension,
         )
diff --git a/tasi/extensions/__init__.py b/tasi/extensions/__init__.py
diff --git a/tasi/extensions/base.py b/tasi/extensions/base.py
@@ -0,0 +1,20 @@
+from typing import Generic, TypeVar
+
+from tasi import Pose, Trajectory
+from tasi.pose.base import CollectionBase
+
+A = TypeVar("A", bound="CollectionBase", covariant=True)
+
+
+class ExtensionBase(Generic[A]):
+
+    def __init__(self, obj: A, *args, **kwargs) -> None:
+        self.obj: A = obj
+
+        super().__init__(*args, **kwargs)
+
+
+class PoseExtensionBase(ExtensionBase[Pose]): ...
+
+
+class TrajectoryExtensionBase(ExtensionBase[Trajectory]): ...
diff --git a/tasi/extensions/smos/__init__.py b/tasi/extensions/smos/__init__.py
@@ -0,0 +1,3 @@
+from .base import TrajectoryExtension
+
+__all__ = ["TrajectoryExtension"]
diff --git a/tasi/extensions/smos/base.py b/tasi/extensions/smos/base.py
@@ -0,0 +1,5 @@
+from .pet import PETExtension
+
+
+class TrajectoryExtension(PETExtension):
+    pass
diff --git a/tasi/extensions/smos/pet.py b/tasi/extensions/smos/pet.py
@@ -0,0 +1,49 @@
+import logging
+from typing import Tuple
+
+from tasi import Trajectory
+from tasi.smos.pet import PET
+
+from ..base import TrajectoryExtensionBase
+
+
+class PETExtension(TrajectoryExtensionBase):
+
+    def pet(
+        self,
+        other: Trajectory,
+        position: Tuple[Tuple[str, ...] | str, Tuple[str, ...] | str] = (
+            ("boundingbox", "center"),
+            ("boundingbox", "center"),
+        ),
+        return_first: bool = True,
+    ) -> PET | None:
+        """
+        Estimate the Post Encroachment Time (PET) between this trajectory and
+        the `other` trajectory according to the reference point(s).
+
+        Args:
+            other (ObjectTrajectory): The other trajectory
+            reference_point (Tuple[str, str]): Any of the pose 'position' points.
+                                               Defaults to 'center' position.
+
+        Returns:
+            PET: The PET between us and the other trajectory. None if
+            trying to estimate the PET with itself.
+
+        Notes:
+            The PET is a signed value. A :math:`PET < 0` indicates that the
+            current object crosses the intersection point after the other
+            object.
+
+        Raise:
+            RuntimeError: If there is no intersection point between both
+            trajectories.
+
+        """
+        if self.obj.equals(other):
+            logging.info("Cannot estimate PET with ourself")
+        else:
+            return PET.estimate(
+                self.obj, other, position=position, return_first=return_first
+            )
diff --git a/tasi/io/public/traffic_participant.py b/tasi/io/public/traffic_participant.py
@@ -11,8 +11,6 @@
 
 __all__ = ["TrafficParticipant"]
 
-__all__ = ["TrafficParticipant"]
-
 
 class TrafficParticipant(BaseModel, PublicEntityMixin):
 
diff --git a/tasi/smos/__init__.py b/tasi/smos/__init__.py
@@ -0,0 +1,3 @@
+from .pet import PET
+
+__all__ = ["PET"]
diff --git a/tasi/smos/base.py b/tasi/smos/base.py
@@ -0,0 +1,10 @@
+from datetime import datetime
+
+from tasi.io import Position
+from tasi.io.public.base import BaseModel
+
+
+class SMOS(BaseModel):
+
+    #: The SMOS value
+    value: float
diff --git a/tasi/smos/pet.py b/tasi/smos/pet.py
@@ -0,0 +1,90 @@
+from datetime import datetime
+from typing import Tuple, Union
+
+import numpy as np
+
+from tasi import Trajectory
+from tasi.io import Position
+from tasi.trajectory.geo import GeoTrajectory
+from tasi.utils.geo import BasicGeometry, MultiGeometry, geometry_to_coords
+
+from .base import SMOS
+
+
+class PET(SMOS):
+
+    #: The time the ego participant is at the intersection point
+    ego: datetime
+
+    #: The time the challenger participant is at the intersection point
+    challenger: datetime
+
+    #: The intersection point between the ego's and challenger's trajectory
+    point: Position
+
+    @classmethod
+    def estimate(
+        cls,
+        ego: Trajectory,
+        challenger: Trajectory,
+        position: Tuple[Tuple[str, ...] | str, Tuple[str, ...] | str] = (
+            "position",
+            "position",
+        ),
+        return_first: bool = True,
+    ):
+
+        # get the reference positions
+        ego_reference, challenger_reference = position
+
+        # we use their geometric representation, while representing the trajectory with a LineString
+        tj1: GeoTrajectory = ego.as_geo(aggregate=True, position=ego_reference)
+        tj2: GeoTrajectory = challenger.as_geo(
+            aggregate=True, position=challenger_reference
+        )
+
+        if not isinstance(ego_reference, str):
+            # we a sequence of str. The last element will be the column name
+            ego_ref = ego_reference[-1]
+        else:
+            ego_ref = ego_reference
+
+        if not isinstance(challenger_reference, str):
+            # we a sequence of str. The last element will be the column name
+            challenger_ref = challenger_reference[-1]
+        else:
+            challenger_ref = challenger_reference
+
+        # we use shapely to find the intersection points of both linestrings
+        intersections: Union[BasicGeometry, MultiGeometry] = (
+            tj1[ego_ref].intersection(tj2[challenger_ref], align=False).item()
+        )
+
+        # estimate the intersection point - note that the first intersection point
+        # is used
+        point = geometry_to_coords(intersections, return_first=return_first)
+
+        # find the closest point for both traffic participants
+        if point is None:
+            raise RuntimeError(
+                f"Failed to convert intersection {point} between trajectories into propery geometry"
+            )
+
+        # get the index of the position which is clostest to the current trajectory
+        ego_idx = np.nanargmin(
+            np.linalg.norm(point - ego[ego_reference].values, axis=1)
+        )
+
+        # get the index of the intersection point which is closest to other trajectory
+        challenger_idx = np.nanargmin(
+            np.linalg.norm(point - challenger[challenger_reference].values, axis=1)
+        )
+
+        return cls(
+            value=(
+                ego.timestamps[ego_idx] - challenger.timestamps[challenger_idx]
+            ).total_seconds(),
+            ego=ego.timestamps[ego_idx],
+            challenger=challenger.timestamps[challenger_idx],
+            point=Position(easting=point[0], northing=point[1]),
+        )
diff --git a/tasi/tests/test_geo.py b/tasi/tests/test_geo.py
@@ -27,6 +27,34 @@ def test_to_geo_trajectory(self):
         self.assertIsInstance(geotrajectory, gpd.GeoDataFrame)
         self.assertIsInstance(geotrajectory, GeoTrajectory)
 
+    def test_to_geo_trajectory_position(self):
+
+        geotrajectory = self.ds.trajectory(self.ds.ids[0]).as_geo(position="position")
+
+        self.assertIsInstance(geotrajectory, gpd.GeoDataFrame)
+        self.assertIsInstance(geotrajectory, GeoTrajectory)
+
+    def test_to_geo_trajectory_front(self):
+
+        geotrajectory = self.ds.trajectory(self.ds.ids[0]).as_geo(
+            position=("boundingbox", "front")
+        )
+
+        self.assertIsInstance(geotrajectory, gpd.GeoDataFrame)
+        self.assertIsInstance(geotrajectory, GeoTrajectory)
+
+    def test_to_geo_trajectory_front_and_position(self):
+
+        geotrajectory = self.ds.trajectory(self.ds.ids[0]).as_geo(
+            position=[("boundingbox", "front"), "position"]
+        )
+
+        self.assertIn("position", geotrajectory)
+        self.assertIn("front", geotrajectory)
+
+        self.assertIsInstance(geotrajectory, gpd.GeoDataFrame)
+        self.assertIsInstance(geotrajectory, GeoTrajectory)
+
     def test_to_geo_dataset(self):
 
         geodataset = self.ds.as_geo()
diff --git a/tasi/trajectory/base.py b/tasi/trajectory/base.py
diff --git a/tasi/trajectory/geo.py b/tasi/trajectory/geo.py
diff --git a/tasi/utils/__init__.py b/tasi/utils/__init__.py
diff --git a/tasi/utils/base.py b/tasi/utils/base.py
diff --git a/tasi/utils/geo.py b/tasi/utils/geo.py
