switch from pandas to polars

Author: MichaelSchuldes

omega_prime/cli.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from typing import Annotated
 
-import pandas as pd
+import polars as pl
 import typer
 
 import omega_prime
@@ -40,7 +40,7 @@ def from_csv(
     validate: bool = True,
     skip_odr_parse: bool = False,
 ):
-    df = pd.read_csv(input)
+    df = pl.read_csv(input)
     r = omega_prime.Recording(df, validate=validate)
     if odr is not None:
         r.map = omega_prime.asam_odr.MapOdr.from_file(odr, skip_parse=skip_odr_parse)
@@ -51,7 +51,7 @@ def from_csv(
 def validate(
     input: Annotated[Path, typer.Argument(help="Path to omega file to validate", exists=True, dir_okay=False)],
 ):
-    omega_prime.Recording.from_file(input)
+    omega_prime.Recording.from_file(input, validate=True)
     print(f"File {input} is valid.")
 
 

omega_prime/converters/lxd.py

@@ -9,7 +9,7 @@
 import multiprocessing as mp
 from ..asam_odr import MapOdr
 from ..recording import Recording
-
+import polars as pl
 
 __all__ = ["convert_lxd"]
 logger.configure(handlers=[{"sink": sys.stdout, "level": "WARNING"}])
@@ -29,10 +29,6 @@
 pedestrians = {"pedestrian": betterosi.MovingObjectType.TYPE_PEDESTRIAN}
 
 
-def wrap_angle(angle):
-    return (angle + np.pi) % (2 * np.pi) - np.pi
-
-
 class DatasetConverter:
     def __init__(self, dataset_dir: Path) -> None:
         self._dataset = Dataset(dataset_dir)
@@ -45,62 +41,70 @@ def get_recording_opendrive_path(self, recording_id: int) -> Path:
 
     def rec2df(self, recording_id):
         rec = self._dataset.get_recording(recording_id)
-        meta = rec._tracks_meta_data
         dt = 1 / rec.get_meta_data("frameRate")
-        meta["type"] = (
-            meta["class"]
-            .apply(
-                lambda x: betterosi.MovingObjectType.TYPE_VEHICLE
-                if x in vehicles
-                else betterosi.MovingObjectType.TYPE_PEDESTRIAN
+
+        meta = rec._tracks_meta_data
+        meta = meta.with_columns(
+            pl.col("class")
+            .map_elements(
+                (
+                    lambda x: betterosi.MovingObjectType.TYPE_VEHICLE
+                    if x in vehicles
+                    else betterosi.MovingObjectType.TYPE_PEDESTRIAN
+                ),
+                return_dtype=int,
             )
-            .values
-        )
-        meta.loc[:, "role"] = meta["class"].apply(
-            lambda x: betterosi.MovingObjectVehicleClassificationRole.ROLE_CIVIL if x in vehicles else -1
+            .alias("type"),
+            pl.col("class")
+            .map_elements(
+                (lambda x: betterosi.MovingObjectVehicleClassificationRole.ROLE_CIVIL if x in vehicles else -1),
+                return_dtype=int,
+            )
+            .alias("role"),
+            pl.col("class")
+            .map_elements((lambda x: vehicles[x] if x in vehicles else -1), return_dtype=int)
+            .alias("subtype"),
         )
-        meta["subtype"] = meta["class"].apply(lambda x: vehicles[x] if x in vehicles else -1)
-        meta = meta.rename(columns={"trackId": "idx"})
+        meta = meta.rename({"trackId": "idx"})
 
         tracks = rec._get_tracks_data()
         tracks = tracks.rename(
-            columns={
+            {
                 "xCenter": "x",
                 "yCenter": "y",
                 "xVelocity": "vel_x",
                 "yVelocity": "vel_y",
                 "xAcceleration": "acc_x",
                 "yAcceleration": "acc_y",
                 "trackId": "idx",
-                "width": "width",
-                "length": "length",
+                # "width": "width",
+                # "length": "length",
             }
         )
-        for k in ["acc_z", "z", "vel_z", "roll", "pitch"]:
-            tracks[k] = 0.0
-        tracks["height"] = 2.0
-        tracks["yaw"] = wrap_angle(np.deg2rad(tracks["heading"]))
-        tracks["total_nanos"] = np.array(tracks["frame"] * dt * NANOS_PER_SEC, dtype=int)
-        tracks = tracks.merge(meta[["role", "idx", "type", "subtype"]], on="idx")
-
-        tracks.loc[
-            (tracks["type"] == betterosi.MovingObjectType.TYPE_VEHICLE)
-            & (tracks["subtype"] == betterosi.MovingObjectVehicleClassificationType.TYPE_BICYCLE),
-            "width",
-        ] = 0.8
-        tracks.loc[
-            (tracks["type"] == betterosi.MovingObjectType.TYPE_VEHICLE)
-            & (tracks["subtype"] == betterosi.MovingObjectVehicleClassificationType.TYPE_BICYCLE),
-            "length",
-        ] = 2
-        tracks.loc[
-            (tracks["type"] == betterosi.MovingObjectType.TYPE_VEHICLE)
-            & (tracks["subtype"] == betterosi.MovingObjectVehicleClassificationType.TYPE_BICYCLE),
-            "height",
-        ] = 1.9
-        tracks.loc[(tracks["type"] == betterosi.MovingObjectType.TYPE_PEDESTRIAN), "width"] = 0.5
-        tracks.loc[(tracks["type"] == betterosi.MovingObjectType.TYPE_PEDESTRIAN), "length"] = 0.5
-        tracks.loc[(tracks["type"] == betterosi.MovingObjectType.TYPE_PEDESTRIAN), "height"] = 1.8
+        tracks = tracks.join(meta.select(["idx", "role", "type", "subtype"]), on="idx", how="left")
+        is_vehicle = pl.col("type") == betterosi.MovingObjectType.TYPE_VEHICLE
+        is_bicycle = pl.col("subtype") == betterosi.MovingObjectVehicleClassificationType.TYPE_BICYCLE
+        is_pedestrian = pl.col("type") == betterosi.MovingObjectType.TYPE_PEDESTRIAN
+        tracks = tracks.with_columns(
+            [pl.lit(0.0).alias(k) for k in ["acc_z", "z", "vel_z", "roll", "pitch"]]
+            + [
+                (((pl.col("heading") + np.pi) % (2 * np.pi)) - np.pi).alias("yaw"),
+                (pl.col("frame") * dt * NANOS_PER_SEC).cast(pl.Int64).alias("total_nanos"),
+                pl.when(is_vehicle & is_bicycle)
+                .then(0.8)
+                .when(is_pedestrian)
+                .then(0.5)
+                .otherwise(pl.col("width"))
+                .alias("width"),
+                pl.when(is_vehicle & is_bicycle)
+                .then(2.0)
+                .when(is_pedestrian)
+                .then(0.5)
+                .otherwise(pl.col("length"))
+                .alias("length"),
+                pl.when(is_vehicle & is_bicycle).then(1.9).when(is_pedestrian).then(1.8).otherwise(2.0).alias("height"),
+            ]
+        )
 
         return tracks
 
@@ -109,7 +113,7 @@ def convert_recording(args):
     converter, recording_id, out_filename = args
     tracks = converter.rec2df(recording_id)
     xodr_path = converter.get_recording_opendrive_path(recording_id)
-    rec = Recording(df=tracks, map=MapOdr.from_file(xodr_path))
+    rec = Recording(df=tracks, map=MapOdr.from_file(xodr_path), validate=False)
     rec.to_mcap(out_filename)
 
 

omega_prime/map.py

@@ -95,8 +95,6 @@ def create(cls, lane: betterosi.Lane):
             subtype=betterosi.LaneClassificationSubtype(lane.classification.subtype),
             successor_ids=[],
             predecessor_ids=[],
-            right_boundary_id=None,
-            left_boundary_id=None,
         )
 
     def plot(self, ax: plt.Axes):

omega_prime/recording.py

@@ -4,11 +4,12 @@
 
 import betterosi
 import numpy as np
-import pandas as pd
 import pyproj
 import shapely
 from matplotlib import pyplot as plt
 from matplotlib.patches import Polygon as PltPolygon
+import pandas as pd
+import polars as pl
 
 import pandera as pa
 import pandera.extensions as extensions
@@ -165,7 +166,7 @@ def __init__(self, recording, idx):
         self.idx = int(idx)
         self._recording = recording
 
-        self._df = self._recording._df.loc[self._recording._df["idx"] == self.idx]
+        self._df = self._recording._df.filter(pl.col("idx") == self.idx)
 
         for k in [
             "x",
@@ -182,23 +183,23 @@ def __init__(self, recording, idx):
             "pitch",
             "polygon",
         ]:
-            setattr(self, k, self._df.loc[:, k].values)
-        self.vel = np.linalg.norm([self.vel_x, self.vel_y], axis=0)
-        self.timestamps = self._df.loc[:, "total_nanos"].values / 1e9
+            setattr(self, k, self._df[:, k])
+        self.vel = np.linalg.norm(self._df["vel_x", "vel_y"], axis=0)
+        self.timestamps = self._df["total_nanos"] / 1e9
         for k in ["length", "width", "height"]:
-            setattr(self, f"{k}s", self._df.loc[:, k].values)
-            setattr(self, k, np.mean(self._df.loc[:, k].values))
+            setattr(self, f"{k}s", self._df[k])
+            setattr(self, k, self._df[k].mean())
 
-        self.type = betterosi.MovingObjectType(self._df.loc[:, "type"].iloc[0])
-        subtype_int = self._df.loc[:, "subtype"].iloc[0]
+        self.type = betterosi.MovingObjectType(self._df["type"][0])
+        subtype_int = self._df["subtype"][0]
         self.subtype = betterosi.MovingObjectVehicleClassificationType(subtype_int) if subtype_int != -1 else None
-        role_int = self._df.loc[:, "role"].iloc[0]
+        role_int = self._df["role"][0]
         self.role = betterosi.MovingObjectVehicleClassificationRole(role_int) if role_int != -1 else None
-        self.birth = int(self._df.loc[:, "frame"].iloc[0])
-        self.end = int(self._df.loc[:, "frame"].iloc[-1])
+        self.birth = int(self._df["frame"][0])
+        self.end = int(self._df["frame"][-1])
 
-    def set(self, k, val):
-        self._recording._df.loc[self._recording._df["idx"] == self.idx, k] = val
+    # def set(self, k, val):
+    #    self._recording._df.loc[self._recording._df["idx"] == self.idx, k] = val
 
     @property
     def nanos(self):
@@ -209,7 +210,7 @@ def plot(self, ax: plt.Axes):
         pass
 
     def plot_mv_frame(self, ax: plt.Axes, frame: int):
-        polys = self._df[self._df["frame"] == frame]["polygon"].values
+        polys = self._df.filter(pl.col("frame") == frame)["polygon"]
         for p in polys:
             ax.add_patch(PltPolygon(p.exterior.coords, fc="red", alpha=0.2))
 
@@ -241,10 +242,10 @@ def _get_polygons(df):
 
     @staticmethod
     def get_moving_object_ground_truth(
-        nanos: int, df: pd.DataFrame, host_vehicle=None, validate=True
+        nanos: int, df: pl.DataFrame, host_vehicle=None, validate=False
     ) -> betterosi.GroundTruth:
         if validate:
-            recording_moving_object_schema.validate(df, lazy=True)
+            recording_moving_object_schema.validate(df.to_pandas(), lazy=True)
 
         def get_object(row):
             return betterosi.MovingObject(
@@ -262,7 +263,7 @@ def get_object(row):
                 ),
             )
 
-        mvs = list(df.apply(get_object, axis=1).values)
+        mvs = [get_object(r) for r in df.iter_rows(named=True)]
         gt = betterosi.GroundTruth(
             version=betterosi.InterfaceVersion(version_major=3, version_minor=7, version_patch=9),
             timestamp=betterosi.Timestamp(seconds=int(nanos // 1_000_000_000), nanos=int(nanos % 1_000_000_000)),
@@ -273,14 +274,17 @@ def get_object(row):
         )
         return gt
 
-    def __init__(self, df, map=None, projections=None, host_vehicle=None, validate=True):
+    def __init__(self, df, map=None, projections=None, host_vehicle=None, validate=False):
         if validate:
-            recording_moving_object_schema.validate(df, lazy=True)
+            recording_moving_object_schema.validate(df.to_pandas(), lazy=True)
         super().__init__()
-        self.nanos2frame = {n: i for i, n in enumerate(df.total_nanos.unique())}
-        df["frame"] = df.total_nanos.map(self.nanos2frame)
-        if "polygon" not in df.columns:
-            df["polygon"] = self._get_polygons(df)
+        self.nanos2frame = {n: i for i, n in enumerate(df["total_nanos"].unique())}
+        mapping = pl.DataFrame({"total_nanos": list(self.nanos2frame.keys()), "frame": list(self.nanos2frame.values())})
+        if 'frame' in df.columns:
+            df = df.drop('frame')
+        df = df.join(mapping, on="total_nanos", how="left")
+        if 'polygon' not in df.columns:
+            df = df.with_columns([pl.Series(name="polygon", values=self._get_polygons(df))])
         self.projections = projections
         self._df = df
         self.map = map
@@ -290,7 +294,7 @@ def __init__(self, df, map=None, projections=None, host_vehicle=None, validate=T
     def to_osi_gts(self) -> list[betterosi.GroundTruth]:
         gts = [
             self.get_moving_object_ground_truth(nanos, group_df, host_vehicle=self.host_vehicle, validate=False)
-            for nanos, group_df in self._df.groupby("total_nanos")
+            for [nanos], group_df in self._df.group_by("total_nanos")
         ]
 
         if self.map is not None and isinstance(self.map, MapOsi | MapOsiCenterline):
@@ -351,11 +355,11 @@ def from_osi_gts(cls, gts: list[betterosi.GroundTruth], validate: bool = True):
                 )
                 for mv in gt.moving_object
             ]
-        df_mv = pd.DataFrame(mvs).sort_values(by=["total_nanos", "idx"]).reset_index(drop="index")
+        df_mv = pl.DataFrame(mvs).sort(["total_nanos", "idx"])
         return cls(df_mv, projections=projs, validate=validate)
 
     @classmethod
-    def from_file(cls, filepath, xodr_path: str | None = None, validate: bool = True, skip_odr_parse: bool = False):
+    def from_file(cls, filepath, xodr_path: str | None = None, validate: bool = False, skip_odr_parse: bool = False):
         gts = betterosi.read(
             filepath,
             return_ground_truth=True,
@@ -396,43 +400,47 @@ def to_mcap(self, filepath):
 
     def to_hdf(self, filename, key="moving_object"):
         #!pip install tables
-        self._df.drop(columns=["polygon", "frame"]).to_hdf(filename, key=key)
+        self._df.drop(columns=["polygon", "frame"]).to_pandas().to_hdf(filename, key=key)
 
     @classmethod
     def from_hdf(cls, filename, key="moving_object"):
-        df = pd.read_hdf(filename, key=key)
+        df = pl.DataFrame(pd.read_hdf(filename, key=key))
         return cls(df, map=None, host_vehicle=None)
 
     def interpolate(self, new_nanos: list[int] | None = None, hz: float | None = None):
         df = self._df
         if new_nanos is None and hz is None:
             new_nanos = np.linspace(
-                df.total_nanos.min(), df.total_nanos.max(), df.frame.max() - df.frame.min(), dtype=int
+                df["total_nanos"].min(), df["total_nanos"].max(), df["frame"].max() - df["frame"].min(), dtype=int
             )
         elif hz is not None:
-            step = 1_000_000_000 * hz
-            new_nanos = np.arange(start=df.total_nanos.min(), stop=df.total_nanos.max() + 1, step=step, dtype=int)
+            step = 1_000_000_000 / hz
+            new_nanos = np.arange(start=df["total_nanos"].min(), stop=df["total_nanos"].max() + 1, step=step, dtype=int)
         else:
             new_nanos = np.array(new_nanos)
         new_dfs = []
-        for idx, track_df in df.groupby("idx"):
+        for [idx], track_df in df.group_by("idx"):
             track_data = {}
             track_new_nanos = new_nanos[
-                track_df.frame.min() - df.frame.min() : track_df.frame.max() - df.frame.min() + 1
+                track_df["frame"].min() - df["frame"].min() : track_df["frame"].max() - df["frame"].min() + 1
             ]
             for c in ["x", "y", "z", "vel_x", "vel_y", "vel_z", "acc_x", "acc_y", "acc_z", "length", "width", "height"]:
                 track_data[c] = np.interp(track_new_nanos, track_df["total_nanos"], track_df[c])
             for c in ["type", "subtype", "role"]:
-                track_data[c] = nearest_interp(track_new_nanos, track_df["total_nanos"].values, track_df[c].values)
+                track_data[c] = nearest_interp(
+                    track_new_nanos, track_df["total_nanos"].to_numpy(), track_df[c].to_numpy()
+                )
             for c in ["roll", "pitch", "yaw"]:
                 track_data[c] = np.mod(
                     np.interp(track_new_nanos, track_df["total_nanos"], np.unwrap(track_df[c], period=np.pi)), np.pi
                 )
-            new_track_df = pd.DataFrame(track_data)
-            new_track_df["idx"] = idx
-            new_track_df["total_nanos"] = track_new_nanos
+            new_track_df = pl.DataFrame(track_data)
+            new_track_df = new_track_df.with_columns(
+                pl.Series(name="idx", values=np.ones_like(track_new_nanos) * idx),
+                pl.Series(name="total_nanos", values=track_new_nanos),
+            )
             new_dfs.append(new_track_df)
-        new_df = pd.concat(new_dfs)
+        new_df = pl.concat(new_dfs)
         return self.__init__(new_df, self.map, self.host_vehicle)
 
     def plot(self, ax=None, legend=False) -> plt.Axes:
@@ -452,6 +460,6 @@ def plot_frame(self, frame: int, ax=None):
         return ax
 
     def plot_mv_frame(self, ax: plt.Axes, frame: int):
-        polys = self._df[self._df["frame"] == frame]["polygon"].values
+        polys = self._df.filter(pl.col("frame") == frame)["polygon"]
         for p in polys:
             ax.add_patch(PltPolygon(p.exterior.coords, fc="red"))