orbx 1.0.0__py3-none-any.whl

orbx/Configs.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+
+@dataclass
+class OrbitalConstants:
+    EARTH_RADIUS_M = 6371000  # meters
+    GM_EARTH = 3.986004418e14  # m^3/s^2
+    SECONDS_IN_DAY = 86400

orbx/Models.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+import numpy as np
+import pandas as pd 
+
+@dataclass
+class Satellite:
+    """Represents a satellite with its orbital parameters"""
+    sat_no: str
+    line1: str
+    line2: str
+    inclination: float
+    apogee: float
+    raan: float
+    argument_of_perigee: float
+    eccentricity: float
+    mean_motion: float
+    
+    def __post_init__(self):
+        # Normalize satellite number
+        self.sat_no = str(self.sat_no).replace('.0', '').zfill(5)
+        
+        
+# @dataclass
+# class ClusterResult:
+#     labels: np.ndarray
+#     density_df: pd.DataFrame
+

orbx/__init__.py

@@ -0,0 +1,5 @@
+from orbx.clustering import cluster
+from orbx.density import density
+from orbx.synthetic_orbits import synthetic_orbit
+
+__all__ = ["cluster", "synthetic_orbit", "density"]

orbx/clustering/Cluster.py

@@ -0,0 +1,48 @@
+import pandas as pd
+import io
+import warnings
+from contextlib import redirect_stdout
+import numpy as np
+from .Core import Core
+from .Schema import Schema
+
+
+def cluster(
+    df: pd.DataFrame,
+    min_samples: int = 3,
+    min_cluster_size: int = 2,
+    verbose: bool = False,
+    ) -> np.ndarray:
+
+    """
+    Arguments
+    ----------
+    df: pandas.DataFrame
+        This should contain the TLEs you want to cluster (at a minimum). Your df should contain
+        a "line1" and "line2" in each row, corresponding to a satellite.
+    min_samples: int, default 3
+        HDBSCAN minimum samples parameter. Refer to HDBSCAN documentation for more details.
+    min_cluster_size: int, default 2
+        HDBSCAN minimum cluster size parameter. Refer to HDBSCAN documentation for more details.
+    verbose: bool, default False
+        If ``True``, print internal clustering progress and allow warnings.
+        If ``False`` (default), suppress stdout and ``FutureWarning`` noise.
+    """
+
+
+
+    # Handle input validation 
+    """ 
+    Core().cluster() is expecting a dataframe with the columns "line1" and "line2"
+    """
+    Schema().validate(df)
+
+    if verbose:
+        labels = Core().cluster(df, min_samples, min_cluster_size)
+    else:
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", category=FutureWarning)
+            with redirect_stdout(io.StringIO()):
+                labels = Core().cluster(df, min_samples, min_cluster_size)
+    
+    return labels

orbx/clustering/Core.py

@@ -0,0 +1,44 @@
+import numpy as np
+import pandas as pd
+from orbx.clustering.Schema import Schema
+from orbx.tools.distance_matrix import get_distance_matrix
+from orbx.clustering.algorithm_wrappers.cluster_wrapper import ClusterWrapper
+from orbx.clustering.data_handling.DataHandler import DataHandler
+
+class Core:
+    def __init__(self):
+
+        self.cluster_wrapper = ClusterWrapper()
+        self.data_handler = DataHandler()
+
+    """
+    The cluster function will take in a dataframe of TLEs and return a ClusterResult.
+    """
+    def cluster(self, df: pd.DataFrame, min_samples: int = 3, min_cluster_size: int = 2) -> np.ndarray:
+
+        """
+        Needs to take the input df
+        which supposedly has only the line1 line2 cols, and then 
+        extract the keplerian elements. These must then be concatenated to the
+        original df
+        """
+        
+        # Validate the input dataframe
+        # Schema().validate(df)
+        
+        df = self.data_handler.tle_to_keplerian(df)
+        
+        distance_matrix, key = get_distance_matrix(df)
+        df = self._reorder_dataframe(df, key)
+        X = self.data_handler.get_points(df)
+        
+        labels = self.cluster_wrapper.run_hdbscan(distance_matrix, X, min_samples, min_cluster_size)
+        
+        return labels
+    
+    def _reorder_dataframe(self, df: pd.DataFrame, key: dict) -> pd.DataFrame:
+        """Reorder dataframe to match key order (this is just overly cautious)"""
+        idx_satNo = key["idx_satNo_dict"]
+
+        satNos_in_order = [idx_satNo[i] for i in range(len(idx_satNo))]
+        return df.set_index("satNo").loc[satNos_in_order].reset_index()

orbx/clustering/Schema.py

@@ -0,0 +1,41 @@
+import pandas as pd
+from sgp4.io import twoline2rv
+import sgp4.earth_gravity as earth_gravity
+
+REQUIRED_COLUMNS = {"line1", "line2"}
+
+class Schema:
+
+    def validate(self, df: pd.DataFrame) -> None:
+        """
+        Validates input DataFrame and computes all derived orbital elements.
+        
+        Required columns:
+            line1 (str): TLE line 1
+            line2 (str): TLE line 2
+        
+        Optional columns:
+            name (str): Satellite name. Defaults to sat_id if not provided.
+        
+        Returns a DataFrame with all orbital elements computed and ready for clustering.
+        """
+        missing = REQUIRED_COLUMNS - set(df.columns)
+        if missing:
+            raise ValueError(
+                f"Input DataFrame is missing required columns: {missing}\n"
+                f"DataFrame must contain at minimum 'line1' and 'line2'."
+            )
+        if df.empty: 
+            raise ValueError("Input DataFrame is empty")
+        if df[["line1", "line2"]].isnull().any().any():
+            raise ValueError("Input DataFrame contains null values in 'line1' or 'line2' columns")
+        
+        # Next we want to validate the two line elements.
+        
+        for _, row in df.iterrows():
+            try:
+                _ = twoline2rv(row['line1'], row['line2'], earth_gravity.wgs72)
+            except ValueError as e:
+                raise ValueError(f"Invalid TLE format for TLE: {row['line1']}, {row['line2']}: {e}")
+
+        return None

orbx/clustering/__init__.py

@@ -0,0 +1,3 @@
+from .Cluster import cluster
+
+__all__ = ["cluster"]

orbx/clustering/algorithm_wrappers/DBSCANWrapper.py

@@ -0,0 +1,216 @@
+import numpy as np
+from sklearn.cluster import DBSCAN
+from orbx.Models import ClusterResult
+from tqdm import tqdm
+import pandas as pd
+
+
+class DBSCANWrapper:
+
+    def __init__(self):
+        # Broad sweep, but selection will prefer small values
+        self.min_samples_range = [2]
+        self.min_samples = 2
+        self.target_cluster_range = (200, 500)
+        self.target_cluster_center = 350
+        self.max_noise_fraction = 0.6
+
+    def run(self, distance_matrix: np.ndarray, X: np.ndarray):
+        return self.fit(distance_matrix, X)
+
+    def _evaluate(self, X, distance_matrix, eps, min_samples):
+        model = DBSCAN(
+            eps=eps,
+            min_samples=min_samples,
+            metric="precomputed",
+            n_jobs=-1
+        )
+
+        labels = model.fit_predict(distance_matrix)
+        n_clusters = len(set(labels) - {-1})
+
+        if n_clusters < 2:
+            return -1.0, labels
+
+        try:
+            score = self.quality_metrics.dbcv_score_wrapper(X, labels)
+            return score, labels
+        except Exception:
+            return -1.0, labels
+
+    def _feasible(self, n_clusters, noise_count, n_total):
+        if not (self.target_cluster_range[0] <= n_clusters <= self.target_cluster_range[1]):
+            return False
+        if noise_count / n_total > self.max_noise_fraction:
+            return False
+        return True
+
+    def fit(self, distance_matrix: np.ndarray, X: np.ndarray) -> ClusterResult:
+        best_score = -np.inf
+        best_labels = None
+        best_params = None
+
+        n_total = distance_matrix.shape[0]
+
+        # =========================
+        # Stage 1: coarse eps sweep
+        # =========================
+        dist_vals = distance_matrix[distance_matrix > 0]
+
+        low, high = np.percentile(dist_vals, [0.1, 3.0])
+        low = max(low, 0.15)
+        high = min(high, 4.0)
+
+        coarse_eps = np.geomspace(low, high, 80)
+
+        coarse_results = []
+
+        for eps in tqdm(coarse_eps, desc="DBSCAN coarse eps sweep"):
+            score, labels = self._evaluate(X, distance_matrix, eps, self.min_samples)
+            n_clusters = len(set(labels) - {-1})
+            noise_count = (labels == -1).sum()
+
+            coarse_results.append(
+                (eps, score, n_clusters, noise_count)
+            )
+
+        coarse_df = pd.DataFrame(
+            coarse_results,
+            columns=["eps", "dbcv", "clusters", "noise"]
+        )
+
+        # =========================
+        # Find eps focus region
+        # =========================
+        focus = coarse_df[
+            coarse_df.clusters.between(
+                self.target_cluster_range[0],
+                self.target_cluster_range[1]
+            )
+        ]
+
+        if focus.empty:
+            raise RuntimeError("No eps values produced feasible cluster counts")
+
+        eps_center = focus.eps.median()
+
+        # =========================
+        # Stage 2: focused eps sweep
+        # =========================
+        eps_low = eps_center / 1.5
+        eps_high = eps_center * 1.5
+
+        eps_values = np.geomspace(eps_low, eps_high, 120)
+
+        sweep_results = []
+
+        for eps in tqdm(eps_values, desc="DBSCAN focused eps sweep"):
+            score, labels = self._evaluate(X, distance_matrix, eps, self.min_samples)
+            n_clusters = len(set(labels) - {-1})
+            noise_count = (labels == -1).sum()
+
+            sweep_results.append(
+                (2, eps, score, n_clusters, noise_count)
+            )
+
+            if score < 0:
+                continue
+
+            penalty = abs(n_clusters - self.target_cluster_center) / self.target_cluster_center
+            adjusted_score = score - 0.5 * penalty
+
+            if self._feasible(n_clusters, noise_count, n_total):
+                if adjusted_score > best_score:
+                    best_score = adjusted_score
+                    best_labels = labels
+                    best_params = (eps, 2)
+
+        if best_labels is None:
+            raise RuntimeError("DBSCAN failed to find a domain-feasible clustering")
+
+        # =========================
+        # Diagnostics plot (focused)
+        # =========================
+        # output_dir = Path("data")
+        # output_dir.mkdir(parents=True, exist_ok=True)
+
+        # df = pd.DataFrame(
+        #     sweep_results,
+        #     columns=["min_samples", "eps", "dbcv", "clusters", "noise"]
+        # )
+
+        # fig, ax1 = plt.subplots(figsize=(9, 5))
+
+        # # ---- Left axis: number of clusters ----
+        # ax1.plot(
+        #     df.eps,
+        #     df.clusters,
+        #     alpha=0.7,
+        #     label="Clusters",
+        #     color="tab:blue"
+        # )
+        # ax1.set_xscale("log")
+        # ax1.set_xlabel("eps (log scale)")
+        # ax1.set_ylabel("Number of clusters", color="tab:blue")
+        # ax1.tick_params(axis="y", labelcolor="tab:blue")
+
+        # # Vertical reference line
+        # ax1.axvline(
+        #     eps_center,
+        #     color="gray",
+        #     linestyle="--",
+        #     alpha=0.6,
+        #     label="eps center"
+        # )
+
+        # # ---- Right axis: DBCV score ----
+        # ax2 = ax1.twinx()
+        # ax2.plot(
+        #     df.eps,
+        #     df.dbcv,
+        #     alpha=0.7,
+        #     label="DBCV",
+        #     color="tab:orange"
+        # )
+        # ax2.set_ylabel("DBCV score", color="tab:orange")
+        # ax2.tick_params(axis="y", labelcolor="tab:orange")
+
+        # # ---- Combined legend ----
+        # lines_1, labels_1 = ax1.get_legend_handles_labels()
+        # lines_2, labels_2 = ax2.get_legend_handles_labels()
+        # ax1.legend(
+        #     lines_1 + lines_2,
+        #     labels_1 + labels_2,
+        #     loc="best",
+        #     fontsize=9
+        # )
+
+        # ax1.set_title("DBSCAN (focused): clusters & DBCV vs eps")
+
+        # output_path = output_dir / "dbscan_eps_clusters_dbcv_focused.png"
+        # plt.tight_layout()
+        # plt.savefig(output_path, dpi=200)
+        # plt.close(fig)
+
+
+        # =========================
+        # Final reporting
+        # =========================
+        # print(f"Saved plot to {output_path}")
+        print(
+            f"Best DBSCAN params → eps={best_params[0]:.6f}, min_samples=2"
+        )
+        print(
+            f"DBSCAN found {len(set(best_labels) - {-1})} clusters "
+            f"(noise points: {(best_labels == -1).sum()})"
+        )
+        print(f"Adjusted selection score: {best_score:.4f}")
+
+        return ClusterResult(
+            best_labels,
+            len(set(best_labels) - {-1}),
+            (best_labels == -1).sum(),
+            best_score,
+            self.cluster_wrapper.quality_metrics.s_dbw_score_wrapper(X, best_labels)
+        )
+

orbx/clustering/algorithm_wrappers/HDBSCANWrapper.py

@@ -0,0 +1,37 @@
+import numpy as np
+from sklearn.cluster import HDBSCAN
+
+class HDBSCANWrapper:
+
+    def __init__(self):
+        pass
+
+    def run(self, distance_matrix: np.ndarray, X: np.ndarray, min_samples: int = 3, min_cluster_size: int = 2) -> np.ndarray:
+        # return labels, best_score
+        return self.fit(distance_matrix, X, min_samples, min_cluster_size)
+
+    def fit(self, distance_matrix: np.ndarray, X: np.ndarray, min_samples: int = 3, min_cluster_size: int = 2) -> np.ndarray:
+    
+        print(f"Running HDBSCAN: ms={min_samples}", flush=True)
+
+        clusterer = HDBSCAN(
+            min_cluster_size=min_cluster_size,
+            min_samples=min_samples,
+            metric="precomputed",
+            cluster_selection_method="eom",
+            n_jobs=-1,
+        )
+
+        labels = clusterer.fit_predict(distance_matrix)
+
+        unique_clusters = set(labels) - {-1}
+        n_clusters = len(unique_clusters)
+        if n_clusters < 2:
+            n_orbits = len(labels)
+            n_noise = int(np.sum(labels == -1))
+            raise ValueError(
+                f"HDBSCAN produced {n_clusters} cluster(s) from {n_orbits} orbit(s) "
+                f"({n_noise} marked as noise). Adjust your dataset or tune the HDBSCAN parameters."
+            )
+
+        return labels

orbx/clustering/algorithm_wrappers/OPTICSWrapper.py

@@ -0,0 +1,96 @@
+from sklearn.cluster import OPTICS
+import numpy as np
+from orbx.Models import ClusterResult
+import pickle
+
+class OPTICSWrapper:
+    def __init__(self):
+        self.min_samples_range = [3]
+
+        self.xi_values = np.geomspace(0.005, 0.2, 10)
+
+        self.max_eps = np.inf
+
+    """
+    Grid search over (min_samples, xi)
+    """
+    def run(self, distance_matrix, X) -> ClusterResult:
+        best_score = -np.inf
+        best_params = None
+        best_labels = None
+
+        n_total = distance_matrix.shape[0]
+
+        for min_samples in self.min_samples_range:
+            for xi in self.xi_values:
+                try:
+                    model = OPTICS(
+                        min_samples=min_samples,
+                        max_eps=self.max_eps,
+                        metric="precomputed",
+                        cluster_method="xi",
+                        xi=xi,
+                        n_jobs=-1,
+                    )
+
+                    labels = model.fit_predict(distance_matrix)
+
+                    n_clusters = len(set(labels) - {-1})
+                    noise_count = (labels == -1).sum()
+
+                    # acceptance = QualityMetrics.is_clustering_acceptable(labels.copy())
+                    # if not acceptance["acceptable"]:
+                    #     print(
+                    #         f"Rejected: min_samples={min_samples}, xi={xi:.4f} "
+                    #         f"({acceptance['fail_reasons']})"
+                    #     )
+                    #     continue
+
+                    score = self.quality_metrics.dbcv_score_wrapper(X, labels)
+
+                    print(
+                        f"min_samples={min_samples:2d}, "
+                        f"xi={xi:.4f}, "
+                        f"clusters={n_clusters:4d}, "
+                        f"noise={noise_count:4d}, "
+                        f"DBCV={score:.4f}"
+                    )
+
+                    if score > best_score:
+                        best_score = score
+                        best_params = (min_samples, xi)
+                        best_labels = labels
+
+                except Exception as e:
+                    print(
+                        f"Error: min_samples={min_samples}, xi={xi:.4f} → {e}"
+                    )
+                    continue
+
+        if best_labels is None:
+            raise RuntimeError("OPTICS failed to find a valid clustering")
+
+        best_min_samples, best_xi = best_params
+
+        print(
+            f"\nBest OPTICS params → min_samples={best_min_samples}, "
+            f"xi={best_xi:.4f}, "
+            f"DBCV={best_score:.4f}"
+        )
+
+        print(
+            f"OPTICS found {len(set(best_labels) - {-1})} clusters "
+            f"(noise points: {(best_labels == -1).sum()})"
+        )
+
+        # Save the labels as pkl  
+        with open("data/cluster_results/optics_labels.pkl", "wb") as f:
+            pickle.dump(best_labels, f)
+
+        return ClusterResult(
+            best_labels,
+            len(set(best_labels) - {-1}),
+            (best_labels == -1).sum(),
+            best_score,
+            self.cluster_wrapper.quality_metrics.s_dbw_score_wrapper(X, best_labels)
+        )

orbx/clustering/algorithm_wrappers/cluster_wrapper.py

@@ -0,0 +1,12 @@
+import numpy as np
+from .HDBSCANWrapper import HDBSCANWrapper
+
+class ClusterWrapper:
+    def __init__(self):
+        self.hdbscan = HDBSCANWrapper()
+   
+    def run_hdbscan(self, distance_matrix: np.ndarray, X: np.ndarray, min_samples: int = 3, min_cluster_size: int = 2) -> np.ndarray:
+        labels = self.hdbscan.run(distance_matrix.copy(), X.copy(), min_samples, min_cluster_size)
+        print(f"HDBSCAN found {len(set(labels) - {-1})} clusters")
+
+        return labels

orbx/clustering/data_handling/DataHandler.py

@@ -0,0 +1,144 @@
+import pandas as pd
+import numpy as np
+from orbx.Configs import OrbitalConstants
+from orbx.Models import Satellite
+from math import pi
+
+class DataHandler:
+    
+    def __init__(self):
+        self.orbital_constants = OrbitalConstants()
+
+    def get_points(self, df: pd.DataFrame):
+            """Takes in a dataframe of Satellite objects. Converts each to a point in the 5D manifold embedded in 6D.
+            This is so the raw data can be passed into clustering algs, quality metrics, etc.
+
+            Args:
+                df (pd.DataFrame): _description_
+
+            Returns:
+                _type_: _description_
+            """
+
+            # Convert degrees -> radians
+            i = np.deg2rad(df["inclination"].values)
+            Omega = np.deg2rad(df["raan"].values)
+            omega = np.deg2rad(df["argument_of_perigee"].values)
+            e = df["eccentricity"].values
+            n = df["mean_motion"].values  # rev/day
+
+            # Constants
+            MU = self.orbital_constants.GM_EARTH  # m^3/s^2
+
+            # Semi-major axis from mean motion
+            n_rad = 2 * np.pi * n / 86400.0
+            a = (MU / n_rad**2) ** (1 / 3)
+
+            # Semi-latus rectum
+            p = a * (1 - e**2)
+            sqrt_p = np.sqrt(p)
+
+            # Angular momentum vector u
+            u = np.column_stack(
+                [
+                    sqrt_p * np.sin(i) * np.sin(Omega),
+                    -sqrt_p * np.sin(i) * np.cos(Omega),
+                    sqrt_p * np.cos(i),
+                ]
+            )
+
+            # LRL vector v
+            v = np.column_stack(
+                [
+                    e
+                    * sqrt_p
+                    * (
+                        np.cos(omega) * np.cos(Omega)
+                        - np.cos(i) * np.sin(omega) * np.sin(Omega)
+                    ),
+                    e
+                    * sqrt_p
+                    * (
+                        np.cos(omega) * np.sin(Omega)
+                        + np.cos(i) * np.sin(omega) * np.cos(Omega)
+                    ),
+                    e * sqrt_p * (np.sin(i) * np.sin(omega)),
+                ]
+            )
+
+            X = np.hstack([u, v])
+
+            return X
+
+    def tle_to_keplerian(self, input_df) -> pd.DataFrame:
+        
+        # Pre-allocate lists for keplerian elements
+        sat_nos = []
+        inclinations = []
+        apogees = []
+        raans = []
+        arguments_of_perigee = []
+        eccentricities = []
+        mean_motions = []
+        
+        # for each row, compute the keplerian elements and add to df
+        for index, row in input_df.iterrows():
+            try:
+                sat_obj = self._parse_tle_group(
+                    row['line1'],
+                    row['line2']
+                )
+            except ValueError as e:
+                raise ValueError(f"Error parsing TLE at row: {index}, TLE: {row['line1']}, {row['line2']}, Error: {e}")
+            
+            sat_nos.append(sat_obj.sat_no)
+            inclinations.append(sat_obj.inclination)
+            apogees.append(sat_obj.apogee)
+            raans.append(sat_obj.raan)
+            arguments_of_perigee.append(sat_obj.argument_of_perigee)
+            eccentricities.append(sat_obj.eccentricity)
+            mean_motions.append(sat_obj.mean_motion)
+        
+        # Add as new columns
+        input_df['satNo'] = sat_nos
+        input_df['inclination'] = inclinations
+        input_df['apogee'] = apogees
+        input_df['raan'] = raans
+        input_df['argument_of_perigee'] = arguments_of_perigee
+        input_df['eccentricity'] = eccentricities
+        input_df['mean_motion'] = mean_motions
+        
+        return input_df
+    
+    def _parse_tle_group(self, line1: str, line2: str) -> Satellite:
+
+        sat_no = line1[2:7].strip()
+
+        # should use sgp4 to extract keplerians
+        inclination = float(line2[8:16].strip())
+        mean_motion = float(line2[52:63].strip())
+        eccentricity = float("0." + line2[26:33].strip())
+        raan = float(line2[17:25].strip())
+        argument_of_perigee = float(line2[34:42].strip())
+        apogee = self._calculate_apogee(mean_motion, eccentricity)
+
+        return Satellite(
+            sat_no=sat_no,
+            line1=line1,
+            line2=line2,
+            inclination=inclination,
+            apogee=apogee,
+            raan=raan,
+            argument_of_perigee=argument_of_perigee,
+            eccentricity=eccentricity,
+            mean_motion=mean_motion
+        )
+    
+    def _calculate_apogee(self, mean_motion: float, eccentricity: float) -> float:
+        """Calculate apogee in kilometers from mean motion and eccentricity"""
+        n = mean_motion * 2 * pi / self.orbital_constants.SECONDS_IN_DAY
+        a = (self.orbital_constants.GM_EARTH / (n ** 2)) ** (1/3)
+        apogee_m = a * (1 + eccentricity) - self.orbital_constants.EARTH_RADIUS_M
+        return apogee_m / 1000  # Convert to km
+
+