refactor: rename read functions and input positions instead of dataframe

Author: FelixAbrahamsson

kmm/positions/__init__.py

@@ -1,6 +1,5 @@
-from kmm.positions.kmm import kmm
-from kmm.positions.kmm2 import kmm2
-from kmm.positions.read_raw import read_raw
+from kmm.positions.read_kmm import read_kmm
+from kmm.positions.read_kmm2 import read_kmm2
 from kmm.positions.geodetic import geodetic
 from kmm.positions.wire_camera_positions import wire_camera_positions
 from kmm.positions.positions import Positions

kmm/positions/geodetic.py

@@ -1,27 +1,28 @@
-import pandas as pd
 from sweref99 import projections
 
+from kmm.positions.positions import Positions
+
 
 tm = projections.make_transverse_mercator("SWEREF_99_TM")
 
 
-def geodetic(positions: pd.DataFrame):
-    if len(positions) == 0:
-        positions = positions.assign(longitude=[], latitude=[])
+def geodetic(positions: Positions):
+    dataframe = positions.dataframe
+    if len(dataframe) == 0:
+        dataframe = dataframe.assign(longitude=[], latitude=[])
     else:
-        latitude, longitude = zip(*positions.apply(
+        latitude, longitude = zip(*dataframe.apply(
             lambda row: (tm.grid_to_geodetic(row["sweref99_tm_x"], row["sweref99_tm_y"])),
             axis="columns",
             result_type="reduce",
         ))
-        positions = positions.assign(longitude=longitude, latitude=latitude)
-    return positions
+        dataframe = dataframe.assign(longitude=longitude, latitude=latitude)
+    return positions.replace(dataframe=dataframe)
 
 
 def test_geodetic():
-    from kmm.positions.kmm2 import kmm2
 
-    df = kmm2("tests/ascending_B.kmm2")
-    df = geodetic(df)
+    positions = Positions.from_path("tests/ascending_B.kmm2")
+    df = geodetic(positions).dataframe
     assert ((df["latitude"] < 68) & (df["latitude"] > 55)).all()
     assert ((df["longitude"] < 25) & (df["longitude"] > 7)).all()

kmm/positions/positions.py

@@ -18,7 +18,14 @@ def from_path(path: Path):
         """
         Loads positions from .kmm or .kmm2 file.
         """
-        return Positions(dataframe=kmm.positions.read_raw(path))
+        if path.suffix == ".kmm":
+            dataframe = kmm.positions.read_kmm(path)
+        elif path.suffix == ".kmm2":
+            dataframe = kmm.positions.read_kmm2(path)
+        else:
+            raise ValueError(f"Unable to parse file type {path.suffix}")
+
+        return Positions(dataframe=dataframe)
 
     @staticmethod
     @validate_arguments
@@ -30,7 +37,6 @@ def read_sync_adjust(
         """
         Loads positions from .kmm or .kmm2 file + .hdr file, then performs
         frame index sync, position adjustment and geodetic coordinate transformation.
-        .
         """
         header = kmm.Header.from_path(header_path)
         return (
@@ -55,13 +61,12 @@ def sync_frame_index(self, header):
 
     def adjust(self, adjustment: kmm.PositionAdjustment, header):
         if adjustment == kmm.PositionAdjustment.WIRE_CAMERA:
-            dataframe = kmm.positions.wire_camera_positions(self.dataframe, header.car_direction)
+            return kmm.positions.wire_camera_positions(self, header.car_direction)
         else:
             raise ValueError(f"Unknown adjustment option {adjustment}")
-        return self.replace(dataframe=dataframe)
 
     def geodetic(self):
-        return self.replace(dataframe=kmm.positions.geodetic(self.dataframe))
+        return kmm.positions.geodetic(self)
 
 
 def test_read_kmm():

kmm/positions/kmm.py renamed to kmm/positions/read_kmm.py

@@ -5,7 +5,7 @@
 
 
 @validate_arguments
-def kmm(path: Path):
+def read_kmm(path: Path):
     return pd.read_csv(
         path,
         sep="\t",

kmm/positions/kmm2.py renamed to kmm/positions/read_kmm2.py

@@ -10,7 +10,7 @@
 
 
 @validate_arguments
-def kmm2(path: Path):
+def read_kmm2(path: Path):
 
     skiprows = [
         index

kmm/positions/read_raw.py

@@ -1,12 +1,12 @@
-import pandas as pd
 from pydantic import validate_arguments
 
 from kmm import CarDirection
+from kmm.positions.positions import Positions
 
 
 @validate_arguments(config=dict(arbitrary_types_allowed=True))
-def wire_camera_positions(positions: pd.DataFrame, direction: CarDirection):
-    ascending = kmm_ascending(positions)
+def wire_camera_positions(positions: Positions, direction: CarDirection):
+    ascending = kmm_ascending(positions.dataframe)
 
     if (
         (direction == CarDirection.A and ascending)
@@ -23,14 +23,16 @@ def wire_camera_positions(positions: pd.DataFrame, direction: CarDirection):
     else:
         raise ValueError
 
-    return (
-        positions
-        .assign(meter=lambda df: df["meter"] + correction)
+    return positions.replace(
+        dataframe=(
+            positions.dataframe
+            .assign(meter=lambda df: df["meter"] + correction)
+        )
     )
 
 
-def kmm_ascending(positions):
-    total_meter = positions["kilometer"] * 1000 + positions["meter"]
+def kmm_ascending(dataframe):
+    total_meter = dataframe["kilometer"] * 1000 + dataframe["meter"]
     diff = total_meter.values[:-1] - total_meter.values[1:]
     descending = (diff < 0).mean()
     ascending = (diff > 0).mean()
@@ -43,18 +45,24 @@ def kmm_ascending(positions):
 
 
 def test_camera_positions_kmm():
-    import kmm
-
-    df = kmm.positions.kmm("tests/ascending_B.kmm")
-    header = kmm.Header.from_path("tests/ascending_B.hdr")
-    df_calibrated = wire_camera_positions(df, header.car_direction)
-    assert df_calibrated["meter"].iloc[0] == df["meter"].iloc[0] - 8
+    from kmm import Header
+
+    positions = Positions.from_path("tests/ascending_B.kmm")
+    header = Header.from_path("tests/ascending_B.hdr")
+    positions_calibrated = wire_camera_positions(positions, header.car_direction)
+    assert (
+        positions_calibrated.dataframe["meter"].iloc[0]
+        == positions.dataframe["meter"].iloc[0] - 8
+    )
 
 
 def test_camera_positions_kmm2():
-    import kmm
-
-    df = kmm.positions.kmm2("tests/ascending_B.kmm2")
-    header = kmm.Header.from_path("tests/ascending_B.hdr")
-    df_calibrated = wire_camera_positions(df, header.car_direction)
-    assert df_calibrated["meter"].iloc[0] == df["meter"].iloc[0] - 8
+    from kmm import Header
+
+    positions = Positions.from_path("tests/ascending_B.kmm2")
+    header = Header.from_path("tests/ascending_B.hdr")
+    positions_calibrated = wire_camera_positions(positions, header.car_direction)
+    assert (
+        positions_calibrated.dataframe["meter"].iloc[0]
+        == positions.dataframe["meter"].iloc[0] - 8
+    )