Reformat to 110 line length

B612-Asteroid-Institute · Oct 11, 2023 · f0df522 · f0df522
1 parent 5c16206
commit f0df522
Show file tree

Hide file tree

Showing 62 changed files with 240 additions and 753 deletions.
diff --git a/adam_core/__init__.py b/adam_core/__init__.py
@@ -1,2 +1,5 @@
 from .__version__ import __version__
 
+__all__ = [
+    "__version__",
+]
diff --git a/adam_core/coordinates/cartesian.py b/adam_core/coordinates/cartesian.py
@@ -30,7 +30,6 @@
 
 
 class CartesianCoordinates(qv.Table):
-
     x = qv.Float64Column(nullable=True)
     y = qv.Float64Column(nullable=True)
     z = qv.Float64Column(nullable=True)
@@ -158,9 +157,7 @@ def sigma_v_mag(self) -> np.ndarray:
         """
         return np.sqrt(np.sum(self.covariance.sigmas[:, 3:6] ** 2, axis=1))
 
-    def rotate(
-        self, rotation_matrix: np.ndarray, frame_out: str
-    ) -> "CartesianCoordinates":
+    def rotate(self, rotation_matrix: np.ndarray, frame_out: str) -> "CartesianCoordinates":
         """
         Rotate Cartesian coordinates and their covariances by the given rotation matrix.
 
@@ -188,24 +185,16 @@ def rotate(
         coords_rotated = np.ma.dot(masked_coords, rotation_matrix.T, strict=False)
 
         # Extract covariances
-        masked_covariances = np.ma.masked_array(
-            self.covariance.to_matrix(), fill_value=0.0
-        )
+        masked_covariances = np.ma.masked_array(self.covariance.to_matrix(), fill_value=0.0)
         masked_covariances.mask = np.isnan(masked_covariances.data)
 
         # Rotate covariances
-        covariances_rotated = (
-            rotation_matrix @ masked_covariances.filled() @ rotation_matrix.T
-        )
+        covariances_rotated = rotation_matrix @ masked_covariances.filled() @ rotation_matrix.T
         # Reset the mask to the original mask
         covariances_rotated[masked_covariances.mask] = np.NaN
 
         # Check if any covariance elements are near zero, if so set them to zero
-        near_zero = len(
-            covariances_rotated[
-                np.abs(covariances_rotated) < COVARIANCE_ROTATION_TOLERANCE
-            ]
-        )
+        near_zero = len(covariances_rotated[np.abs(covariances_rotated) < COVARIANCE_ROTATION_TOLERANCE])
         if near_zero > 0:
             logger.debug(
                 f"{near_zero} covariance elements are within {COVARIANCE_ROTATION_TOLERANCE:.0e}"
@@ -292,9 +281,7 @@ def to_keplerian(self) -> "KeplerianCoordinates":
         return KeplerianCoordinates.from_cartesian(self)
 
     @classmethod
-    def from_keplerian(
-        cls, keplerian: "KeplerianCoordinates"
-    ) -> "CartesianCoordinates":
+    def from_keplerian(cls, keplerian: "KeplerianCoordinates") -> "CartesianCoordinates":
         return keplerian.to_cartesian()
 
     def to_spherical(self) -> "SphericalCoordinates":
@@ -303,7 +290,5 @@ def to_spherical(self) -> "SphericalCoordinates":
         return SphericalCoordinates.from_cartesian(self)
 
     @classmethod
-    def from_spherical(
-        cls, spherical: "SphericalCoordinates"
-    ) -> "CartesianCoordinates":
+    def from_spherical(cls, spherical: "SphericalCoordinates") -> "CartesianCoordinates":
         return spherical.to_cartesian()
diff --git a/adam_core/coordinates/cometary.py b/adam_core/coordinates/cometary.py
@@ -224,9 +224,7 @@ def to_cartesian(self) -> CartesianCoordinates:
                 tol=1e-15,
             )
         else:
-            covariances_cartesian = np.empty(
-                (len(coords_cartesian), 6, 6), dtype=np.float64
-            )
+            covariances_cartesian = np.empty((len(coords_cartesian), 6, 6), dtype=np.float64)
             covariances_cartesian.fill(np.nan)
 
         covariances_cartesian = CoordinateCovariances.from_matrix(covariances_cartesian)
@@ -279,9 +277,7 @@ def from_cartesian(cls, cartesian: CartesianCoordinates) -> "CometaryCoordinates
                 mu=mu,
             )
         else:
-            covariances_cometary = np.empty(
-                (len(coords_cometary), 6, 6), dtype=np.float64
-            )
+            covariances_cometary = np.empty((len(coords_cometary), 6, 6), dtype=np.float64)
             covariances_cometary.fill(np.nan)
 
         covariances_cometary = CoordinateCovariances.from_matrix(covariances_cometary)
@@ -306,9 +302,7 @@ def to_keplerian(self) -> "KeplerianCoordinates":
         return KeplerianCoordinates.from_cartesian(self.to_cartesian())
 
     @classmethod
-    def from_keplerian(
-        cls, keplerian_coordinates: "KeplerianCoordinates"
-    ) -> "CometaryCoordinates":
+    def from_keplerian(cls, keplerian_coordinates: "KeplerianCoordinates") -> "CometaryCoordinates":
         return cls.from_cartesian(keplerian_coordinates.to_cartesian())
 
     def to_spherical(self) -> "SphericalCoordinates":
@@ -317,7 +311,5 @@ def to_spherical(self) -> "SphericalCoordinates":
         return SphericalCoordinates.from_cartesian(self.to_cartesian())
 
     @classmethod
-    def from_spherical(
-        cls, spherical_coordinates: "SphericalCoordinates"
-    ) -> "CometaryCoordinates":
+    def from_spherical(cls, spherical_coordinates: "SphericalCoordinates") -> "CometaryCoordinates":
         return cls.from_cartesian(spherical_coordinates.to_cartesian())
diff --git a/adam_core/coordinates/covariances.py b/adam_core/coordinates/covariances.py
@@ -123,9 +123,7 @@ def from_matrix(cls, covariances: np.ndarray) -> "CoordinateCovariances":
         """
         # cov = pa.FixedShapeTensorArray.from_numpy_ndarray(covariances)
         if covariances.shape[1:] != (6, 6):
-            raise ValueError(
-                f"Covariance matrices should have shape (N, 6, 6) but got {covariances.shape}"
-            )
+            raise ValueError(f"Covariance matrices should have shape (N, 6, 6) but got {covariances.shape}")
         cov = covariances.flatten()
         offsets = np.arange(0, (len(covariances) + 1) * 36, 36, dtype=np.int64)
         return cls.from_kwargs(values=pa.ListArray.from_arrays(offsets, cov))
@@ -323,9 +321,7 @@ def weighted_mean(samples: np.ndarray, W: np.ndarray) -> np.ndarray:
     return np.dot(W, samples)
 
 
-def weighted_covariance(
-    mean: np.ndarray, samples: np.ndarray, W_cov: np.ndarray
-) -> np.ndarray:
+def weighted_covariance(mean: np.ndarray, samples: np.ndarray, W_cov: np.ndarray) -> np.ndarray:
     """
     Calculate a covariance matrix from samples and their corresponding weights.
 
@@ -581,10 +577,7 @@ def covariances_from_df(
             covariances[:, i, j] = df[f"cov_{coord_names[i]}_{coord_names[j]}"].values
             covariances[:, j, i] = covariances[:, i, j]
         except KeyError:
-            logger.debug(
-                "No covariance column found for dimensions"
-                f" {coord_names[i]},{coord_names[j]}."
-            )
+            logger.debug("No covariance column found for dimensions" f" {coord_names[i]},{coord_names[j]}.")
 
     return covariances
 
@@ -685,14 +678,9 @@ def covariances_from_table(
 
     for i, j in zip(ii, jj):
         try:
-            covariances[:, i, j] = table[
-                f"cov_{coord_names[i]}_{coord_names[j]}"
-            ].values
+            covariances[:, i, j] = table[f"cov_{coord_names[i]}_{coord_names[j]}"].values
             covariances[:, j, i] = covariances[:, i, j]
         except KeyError:
-            logger.debug(
-                "No covariance column found for dimensions"
-                f" {coord_names[i]},{coord_names[j]}."
-            )
+            logger.debug("No covariance column found for dimensions" f" {coord_names[i]},{coord_names[j]}.")
 
     return covariances
diff --git a/adam_core/coordinates/keplerian.py b/adam_core/coordinates/keplerian.py
@@ -33,7 +33,6 @@
 
 
 class KeplerianCoordinates(qv.Table):
-
     a = qv.Float64Column()
     e = qv.Float64Column()
     i = qv.Float64Column()
@@ -210,9 +209,7 @@ def to_cartesian(self) -> CartesianCoordinates:
                 tol=1e-15,
             )
         else:
-            covariances_cartesian = np.empty(
-                (len(coords_cartesian), 6, 6), dtype=np.float64
-            )
+            covariances_cartesian = np.empty((len(coords_cartesian), 6, 6), dtype=np.float64)
             covariances_cartesian.fill(np.nan)
 
         covariances_cartesian = CoordinateCovariances.from_matrix(covariances_cartesian)
@@ -257,9 +254,7 @@ def from_cartesian(cls, cartesian: CartesianCoordinates):
                 mu=mu,
             )
         else:
-            covariances_keplerian = np.empty(
-                (len(coords_keplerian), 6, 6), dtype=np.float64
-            )
+            covariances_keplerian = np.empty((len(coords_keplerian), 6, 6), dtype=np.float64)
             covariances_keplerian.fill(np.nan)
 
         covariances_keplerian = CoordinateCovariances.from_matrix(covariances_keplerian)
@@ -283,9 +278,7 @@ def to_cometary(self) -> "CometaryCoordinates":
         return CometaryCoordinates.from_cartesian(self.to_cartesian())
 
     @classmethod
-    def from_cometary(
-        cls, cometary_coordinates: "CometaryCoordinates"
-    ) -> "KeplerianCoordinates":
+    def from_cometary(cls, cometary_coordinates: "CometaryCoordinates") -> "KeplerianCoordinates":
         return cls.from_cartesian(cometary_coordinates.to_cartesian())
 
     def to_spherical(self) -> "SphericalCoordinates":
@@ -294,7 +287,5 @@ def to_spherical(self) -> "SphericalCoordinates":
         return SphericalCoordinates.from_cartesian(self.to_cartesian())
 
     @classmethod
-    def from_spherical(
-        cls, spherical_coordinates: "SphericalCoordinates"
-    ) -> "KeplerianCoordinates":
+    def from_spherical(cls, spherical_coordinates: "SphericalCoordinates") -> "KeplerianCoordinates":
         return cls.from_cartesian(spherical_coordinates.to_cartesian())
diff --git a/adam_core/coordinates/origin.py b/adam_core/coordinates/origin.py
@@ -59,7 +59,6 @@ class OriginGravitationalParameters(float, Enum):
 # TODO: Replace with DictionaryColumn or similar
 #       Investigate whether this class is even necessary
 class Origin(qv.Table):
-
     code = qv.StringColumn()
 
     def __init__(self, table: pa.Table, mu: Optional[float] = None):
@@ -89,19 +88,15 @@ def __ne__(self, other: object) -> np.ndarray:
     @property
     def mu(self):
         if self._mu is None:
-            logger.debug(
-                "Origin.mu called without mu set. Finding mu in OriginGravitationalParameters."
-            )
+            logger.debug("Origin.mu called without mu set. Finding mu in OriginGravitationalParameters.")
             codes = np.array(self.code)
             if len(np.unique(codes)) > 1:
                 raise ValueError("Origin.mu called on table with multiple origins.")
 
             try:
                 return OriginGravitationalParameters[codes[0]].value
             except KeyError:
-                raise ValueError(
-                    "Origin.mu called on table with unrecognized origin code."
-                )
+                raise ValueError("Origin.mu called on table with unrecognized origin code.")
         else:
             return self._mu
 

diff --git a/adam_core/coordinates/residuals.py b/adam_core/coordinates/residuals.py
@@ -31,16 +31,13 @@
 
 
 class Residuals(qv.Table):
-
     values = qv.ListColumn(pa.float64(), nullable=True)
     chi2 = qv.Float64Column(nullable=True)
     dof = qv.Int64Column(nullable=True)
     probability = qv.Float64Column(nullable=True)
 
     @classmethod
-    def calculate(
-        cls, observed: CoordinateType, predicted: CoordinateType
-    ) -> "Residuals":
+    def calculate(cls, observed: CoordinateType, predicted: CoordinateType) -> "Residuals":
         """
         Calculate the residuals between the observed and predicted coordinates. The observed
         coordinate's covariance matrix is used to calculate the chi2 and degrees of freedom.
@@ -95,9 +92,7 @@ def calculate(
             batch_dimensions,
             batch_coords,
             batch_covariances,
-        ) = _batch_coords_and_covariances(
-            observed.values, observed.covariance.to_matrix()
-        )
+        ) = _batch_coords_and_covariances(observed.values, observed.covariance.to_matrix())
 
         for indices, dimensions, coords, covariances in zip(
             batch_indices, batch_dimensions, batch_coords, batch_covariances

diff --git a/adam_core/coordinates/spherical.py b/adam_core/coordinates/spherical.py
@@ -33,7 +33,6 @@
 
 
 class SphericalCoordinates(qv.Table):
-
     rho = qv.Float64Column(nullable=True)
     lon = qv.Float64Column(nullable=True)
     lat = qv.Float64Column(nullable=True)
@@ -47,9 +46,7 @@ class SphericalCoordinates(qv.Table):
 
     @property
     def values(self) -> np.ndarray:
-        return np.array(
-            self.table.select(["rho", "lon", "lat", "vrho", "vlon", "vlat"])
-        )
+        return np.array(self.table.select(["rho", "lon", "lat", "vrho", "vlon", "vlat"]))
 
     @property
     def sigma_rho(self):
@@ -164,9 +161,7 @@ def to_cartesian(self) -> CartesianCoordinates:
                 self.values, covariances_spherical, _spherical_to_cartesian
             )
         else:
-            covariances_cartesian = np.empty(
-                (len(coords_cartesian), 6, 6), dtype=np.float64
-            )
+            covariances_cartesian = np.empty((len(coords_cartesian), 6, 6), dtype=np.float64)
             covariances_cartesian.fill(np.nan)
 
         covariances_cartesian = CoordinateCovariances.from_matrix(covariances_cartesian)
@@ -197,9 +192,7 @@ def from_cartesian(cls, cartesian: CartesianCoordinates) -> "SphericalCoordinate
                 cartesian.values, cartesian_covariances, _cartesian_to_spherical
             )
         else:
-            covariances_spherical = np.empty(
-                (len(coords_spherical), 6, 6), dtype=np.float64
-            )
+            covariances_spherical = np.empty((len(coords_spherical), 6, 6), dtype=np.float64)
             covariances_spherical.fill(np.nan)
 
         covariances_spherical = CoordinateCovariances.from_matrix(covariances_spherical)
@@ -224,9 +217,7 @@ def to_cometary(self) -> "CometaryCoordinates":
         return CometaryCoordinates.from_cartesian(self.to_cartesian())
 
     @classmethod
-    def from_cometary(
-        cls, cometary_coordinates: "CometaryCoordinates"
-    ) -> "SphericalCoordinates":
+    def from_cometary(cls, cometary_coordinates: "CometaryCoordinates") -> "SphericalCoordinates":
         return cls.from_cartesian(cometary_coordinates.to_cartesian())
 
     def to_keplerian(self) -> "KeplerianCoordinates":
@@ -235,13 +226,9 @@ def to_keplerian(self) -> "KeplerianCoordinates":
         return KeplerianCoordinates.from_cartesian(self.to_cartesian())
 
     @classmethod
-    def from_keplerian(
-        cls, keplerian_coordinates: "KeplerianCoordinates"
-    ) -> "SphericalCoordinates":
+    def from_keplerian(cls, keplerian_coordinates: "KeplerianCoordinates") -> "SphericalCoordinates":
         return cls.from_cartesian(keplerian_coordinates.to_cartesian())
 
     @classmethod
-    def from_spherical(
-        cls, spherical_coordinates: "SphericalCoordinates"
-    ) -> "SphericalCoordinates":
+    def from_spherical(cls, spherical_coordinates: "SphericalCoordinates") -> "SphericalCoordinates":
         return spherical_coordinates
diff --git a/adam_core/coordinates/tests/conftest.py b/adam_core/coordinates/tests/conftest.py
@@ -6,43 +6,27 @@
 
 @pytest.fixture
 def orbital_elements():
-    orbital_elements_file = files("adam_core.utils.helpers.data").joinpath(
-        "elements_sun_ec.csv"
-    )
-    df = pd.read_csv(
-        orbital_elements_file, index_col=False, float_precision="round_trip"
-    )
+    orbital_elements_file = files("adam_core.utils.helpers.data").joinpath("elements_sun_ec.csv")
+    df = pd.read_csv(orbital_elements_file, index_col=False, float_precision="round_trip")
     return df
 
 
 @pytest.fixture
 def orbital_elements_equatorial():
-    orbital_elements_file = files("adam_core.utils.helpers.data").joinpath(
-        "elements_sun_eq.csv"
-    )
-    df = pd.read_csv(
-        orbital_elements_file, index_col=False, float_precision="round_trip"
-    )
+    orbital_elements_file = files("adam_core.utils.helpers.data").joinpath("elements_sun_eq.csv")
+    df = pd.read_csv(orbital_elements_file, index_col=False, float_precision="round_trip")
     return df
 
 
 @pytest.fixture
 def orbital_elements_barycentric():
-    orbital_elements_file = files("adam_core.utils.helpers.data").joinpath(
-        "elements_ssb_ec.csv"
-    )
-    df = pd.read_csv(
-        orbital_elements_file, index_col=False, float_precision="round_trip"
-    )
+    orbital_elements_file = files("adam_core.utils.helpers.data").joinpath("elements_ssb_ec.csv")
+    df = pd.read_csv(orbital_elements_file, index_col=False, float_precision="round_trip")
     return df
 
 
 @pytest.fixture
 def orbital_elements_barycentric_equatorial():
-    orbital_elements_file = files("adam_core.utils.helpers.data").joinpath(
-        "elements_ssb_eq.csv"
-    )
-    df = pd.read_csv(
-        orbital_elements_file, index_col=False, float_precision="round_trip"
-    )
+    orbital_elements_file = files("adam_core.utils.helpers.data").joinpath("elements_ssb_eq.csv")
+    df = pd.read_csv(orbital_elements_file, index_col=False, float_precision="round_trip")
     return df