typing

Author: akoumjian

src/adam_impact_study/analysis/collect.py

@@ -302,7 +302,10 @@ def create_missing_window_results_worker(
 
 
 def create_missing_window_results(
-    observations: Observations, window_results: WindowResult, max_processes: int = 1, chunk_size: int = 100
+    observations: Observations,
+    window_results: WindowResult,
+    max_processes: int = 1,
+    chunk_size: int = 100,
 ) -> WindowResult:
     """Create missing window results for observations that do not have correspondingwindow results.
 
@@ -470,5 +473,7 @@ def collect_all_results(
     window_results = collect_all_window_results_new(
         run_dir, max_processes=max_processes
     )
-    window_results = create_missing_window_results(observations, window_results, max_processes=max_processes)
+    window_results = create_missing_window_results(
+        observations, window_results, max_processes=max_processes
+    )
     return impactor_orbits, observations, timings, window_results

src/adam_impact_study/analysis/main.py

@@ -455,7 +455,7 @@ def compute_observation_cadence(
 
 def _select_ip_at_discovery_time(
     orbit_id: str,
-    all_window_results: WindowResult, # Note! Pass in all windows
+    all_window_results: WindowResult,  # Note! Pass in all windows
     discovery_dates: DiscoveryDates,
 ) -> float:
     """
@@ -466,13 +466,17 @@ def _select_ip_at_discovery_time(
     if len(orbit_results) == 0 or len(orbit_discovery_date) == 0:
         return None
 
-    assert len(orbit_discovery_date) == 1, f"{orbit_id} had {len(orbit_discovery_date)} discovery dates, expected 1"
-    
+    assert (
+        len(orbit_discovery_date) == 1
+    ), f"{orbit_id} had {len(orbit_discovery_date)} discovery dates, expected 1"
+
     if pc.all(pc.is_null(orbit_discovery_date.discovery_date.days)).as_py():
         return None
 
     discovery_window = orbit_results.apply_mask(
-        orbit_results.observation_end.equals(orbit_discovery_date.discovery_date, precision="ms")
+        orbit_results.observation_end.equals(
+            orbit_discovery_date.discovery_date, precision="ms"
+        )
     )
 
     return discovery_window.impact_probability[0].as_py()
@@ -489,7 +493,9 @@ def summarize_impact_study_object_results(
     Summarize the impact study results for a single object.
     """
     impactor_orbit = impactor_orbits.select("orbit_id", orbit_id)
-    assert len(impactor_orbit) == 1, f"{orbit_id} had {len(impactor_orbit)} impactor orbits, expected 1"
+    assert (
+        len(impactor_orbit) == 1
+    ), f"{orbit_id} had {len(impactor_orbit)} impactor orbits, expected 1"
 
     orbit_observations = observations.select("orbit_id", orbit_id)
     orbit_results_timing = results_timing.select("orbit_id", orbit_id)
@@ -504,7 +510,6 @@ def summarize_impact_study_object_results(
         pc.equal(orbit_window_results.status, "complete")
     ).as_py()
 
-
     first_observation = orbit_observations.coordinates.time.min()
     last_observation = orbit_observations.coordinates.time.max()
 
@@ -691,14 +696,13 @@ def summarize_impact_study_results(
     futures = []
     results = ImpactorResultSummary.empty()
 
-    # Place our objects in the object store if we are using ray 
+    # Place our objects in the object store if we are using ray
     if use_ray:
         impactor_orbits_ref = ray.put(impactor_orbits)
         observations_ref = ray.put(observations)
         results_timing_ref = ray.put(results_timing)
         window_results_ref = ray.put(window_results)
 
-
     for orbit_id in unique_orbit_ids:
         if use_ray:
             futures.append(

src/adam_impact_study/analysis/plots.py

@@ -27,13 +27,11 @@ def plot_warning_time_histogram(
     summary = summary.apply_mask(summary.complete())
 
     fig, ax = plt.subplots(1, 1, dpi=200)
-    warning_time_max = (
-        pc.ceil(
-            pc.max(
-                pc.subtract(
-                    summary.orbit.impact_time.mjd(),
-                    summary.ip_threshold_1_percent.mjd(),
-                )
+    warning_time_max = pc.ceil(
+        pc.max(
+            pc.subtract(
+                summary.orbit.impact_time.mjd(),
+                summary.ip_threshold_1_percent.mjd(),
             )
         )
     )
@@ -195,7 +193,6 @@ def plot_arclength_by_diameter(
 
     unique_diameters = summary.orbit.diameter.unique().sort().to_pylist()
     colors = plt.cm.coolwarm(np.linspace(0, 1, len(unique_diameters)))
-    
 
     # Get the max arc length
     max_arc_length = pc.max(summary.arc_length()).as_py()
@@ -213,14 +210,12 @@ def plot_arclength_by_diameter(
             alpha=0.7,
             label=f"{diameter:.3f} km",
         )
-    
+
     ax.set_xlabel("Arc Length [days]")
     ax.set_ylabel("Count")
     ax.legend(frameon=False, bbox_to_anchor=(1.01, 0.75))
 
     return fig, ax
-    
-
 
 
 def plot_runtime_by_diameter(
@@ -922,7 +917,9 @@ def plot_arc_length_by_diameter_decade(
             orbits_at_diameter_and_decade = orbits_at_decade.select(
                 "orbit.diameter", diameter
             )
-            arc_length = orbits_at_diameter_and_decade.arc_length().to_numpy(zero_copy_only=False)
+            arc_length = orbits_at_diameter_and_decade.arc_length().to_numpy(
+                zero_copy_only=False
+            )
             mean_arc_length = np.mean(arc_length) if arc_length.size > 0 else 0
             arc_length_by_diameter_decade = qv.concatenate(
                 [

src/adam_impact_study/types.py

@@ -48,7 +48,7 @@ class ImpactorOrbits(qv.Table):
     impact_time = Timestamp.as_column()
     dynamical_class = qv.LargeStringColumn()
     ast_class = qv.LargeStringColumn()
-    diameter = qv.Float64Column() # km
+    diameter = qv.Float64Column()  # km
     albedo = qv.Float64Column()
     H_r = qv.Float64Column()
     u_r = qv.Float64Column()
@@ -293,7 +293,9 @@ def days_100_percent_to_impact(self) -> pa.FloatArray:
             self.orbit.impact_time.mjd(), self.ip_threshold_100_percent.mjd()
         )
 
-    def get_diameter_decade_data(self) -> Tuple[npt.NDArray[np.int64], npt.NDArray[np.int64], list]:
+    def get_diameter_decade_data(
+        self,
+    ) -> Tuple[npt.NDArray[np.int64], npt.NDArray[np.int64], list]:
         """
         Extract impact decades and common data needed for diameter-decade analysis.
         Returns