Clarify checkpoint data types in link_test_orbit

Author: moeyensj

thor/main.py

@@ -3,13 +3,11 @@
 import pathlib
 import time
 from dataclasses import dataclass
-from typing import Any, Iterable, Iterator, List, Literal, Optional, Union
+from typing import Any, Iterable, Iterator, List, Literal, Optional, Tuple
 
 import quivr as qv
 import ray
-from adam_core.coordinates import CartesianCoordinates
 from adam_core.propagator import PYOORB
-from adam_core.time import Timestamp
 
 from .clusters import ClusterMembers, Clusters, cluster_and_link
 from .config import Config, initialize_config
@@ -26,6 +24,16 @@
 
 logger = logging.getLogger("thor")
 
+VALID_STAGES = Literal[
+    "filter_observations",
+    "range_and_transform",
+    "cluster_and_link",
+    "initial_orbit_determination",
+    "differential_correction",
+    "recover_orbits",
+    "complete",
+]
+
 
 def initialize_use_ray(config: Config) -> bool:
     use_ray = False
@@ -43,25 +51,17 @@ def initialize_use_ray(config: Config) -> bool:
 
 @dataclass
 class CheckpointData:
-    stage: Literal[
-        "filter_observations",
-        "range_and_transform",
-        "cluster_and_link",
-        "initial_orbit_determination",
-        "differential_correction",
-        "recover_orbits",
-        "complete",
-    ]
-    filtered_observations: Optional[Observations] = None
-    transformed_detections: Optional[TransformedDetections] = None
-    clusters: Optional[Clusters] = None
-    cluster_members: Optional[ClusterMembers] = None
-    iod_orbits: Optional[FittedOrbits] = None
-    iod_orbit_members: Optional[FittedOrbitMembers] = None
-    od_orbits: Optional[FittedOrbits] = None
-    od_orbit_members: Optional[FittedOrbitMembers] = None
-    recovered_orbits: Optional[FittedOrbits] = None
-    recovered_orbit_members: Optional[FittedOrbitMembers] = None
+    stage: VALID_STAGES
+    filtered_observations: Observations = Observations.empty()
+    transformed_detections: TransformedDetections = TransformedDetections.empty()
+    clusters: Clusters = Clusters.empty()
+    cluster_members: ClusterMembers = ClusterMembers.empty()
+    iod_orbits: FittedOrbits = FittedOrbits.empty()
+    iod_orbit_members: FittedOrbitMembers = FittedOrbitMembers.empty()
+    od_orbits: FittedOrbits = FittedOrbits.empty()
+    od_orbit_members: FittedOrbitMembers = FittedOrbitMembers.empty()
+    recovered_orbits: FittedOrbits = FittedOrbits.empty()
+    recovered_orbit_members: FittedOrbitMembers = FittedOrbitMembers.empty()
 
 
 def initialize_test_orbit(
@@ -86,7 +86,7 @@ def load_initial_checkpoint_values(
 
     We want to avoid loading objects into memory that are not required.
     """
-    stage = "filter_observations"
+    stage: VALID_STAGES = "filter_observations"
     # Without a checkpoint directory, we always start at the beginning
     if test_orbit_directory is None:
         return CheckpointData(stage=stage)
@@ -245,8 +245,8 @@ class LinkTestOrbitStageResult:
         "differential_correction",
         "recover_orbits",
     ]
-    result: Iterable[Any]
-    path: Optional[Iterable[str]] = None
+    result: Iterable[qv.AnyTable]
+    path: Tuple[Optional[str], ...] = (None,)
 
 
 def link_test_orbit(
@@ -284,11 +284,11 @@ def link_test_orbit(
 
     test_orbit_directory = None
     if working_dir is not None:
-        working_dir = pathlib.Path(working_dir)
+        working_dir_path = pathlib.Path(working_dir)
         logger.info(f"Using working directory: {working_dir}")
-        test_orbit_directory = pathlib.Path(working_dir, test_orbit.orbit_id)
+        test_orbit_directory = pathlib.Path(working_dir_path, test_orbit.orbit_id)
         test_orbit_directory.mkdir(parents=True, exist_ok=True)
-        inputs_dir = pathlib.Path(working_dir, "inputs")
+        inputs_dir = pathlib.Path(working_dir_path, "inputs")
         inputs_dir.mkdir(parents=True, exist_ok=True)
 
     initialize_test_orbit(test_orbit, working_dir)
@@ -317,7 +317,7 @@ def link_test_orbit(
 
     if checkpoint.stage == "complete":
         logger.info("Found recovered orbits in checkpoint, exiting early...")
-        path = None
+        path: Tuple[Optional[str], ...] = (None,)
         if test_orbit_directory:
             path = (
                 os.path.join(test_orbit_directory, "recovered_orbits.parquet"),

thor/observations/filters.py

@@ -304,7 +304,11 @@ def filter_observations(
     if len(filtered_observations) > 0:
         filtered_observations = qv.defragment(filtered_observations)
         filtered_observations = filtered_observations.sort_by(
-            ["coordinates.time.days", "coordinates.time.nanos", "coordinates.origin.code"]
+            [
+                "coordinates.time.days",
+                "coordinates.time.nanos",
+                "coordinates.origin.code",
+            ]
         )
 
     return filtered_observations

thor/observations/states.py

@@ -28,8 +28,7 @@ def calculate_state_ids(
 
     # Append index column so we can maintain the original order
     table = table.append_column(
-        pa.field("index", pa.int64()),
-        pa.array(np.arange(0, len(table)))
+        pa.field("index", pa.int64()), pa.array(np.arange(0, len(table)))
     )
 
     # Select only the relevant columns
@@ -63,6 +62,5 @@ def calculate_state_ids(
         unique_time_origins, ["time.days", "time.nanos", "origin.code"]
     ).sort_by([("index", "ascending")])
 
-
     # Now return the state IDs
     return coordinates_with_states.column("state_id").combine_chunks()

thor/tests/test_main.py

@@ -238,12 +238,7 @@ def test_link_test_orbit(
     else:
         integration_config.max_processes = 1
 
-    (
-        test_orbit,
-        observations,
-        obs_ids_expected,
-        integration_config,
-    ) = setup_test_data(
+    (test_orbit, observations, obs_ids_expected, integration_config,) = setup_test_data(
         object_id, orbits, observations, integration_config, max_arc_length=14
     )
 
@@ -270,12 +265,7 @@ def test_benchmark_link_test_orbit(
     else:
         integration_config.max_processes = 1
 
-    (
-        test_orbit,
-        observations,
-        obs_ids_expected,
-        integration_config,
-    ) = setup_test_data(
+    (test_orbit, observations, obs_ids_expected, integration_config,) = setup_test_data(
         object_id, orbits, observations, integration_config, max_arc_length=14
     )