Format to 80 lines

adam_core/coordinates/cartesian.py

@@ -157,7 +157,9 @@ 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.
 
@@ -182,19 +184,29 @@ def rotate(self, rotation_matrix: np.ndarray, frame_out: str) -> "CartesianCoord
         masked_coords.mask = np.isnan(masked_coords.data)
 
         # Rotate coordinates
-        coords_rotated = np.ma.dot(masked_coords, rotation_matrix.T, strict=False)
+        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}"
@@ -220,7 +232,9 @@ def rotate(self, rotation_matrix: np.ndarray, frame_out: str) -> "CartesianCoord
         )
         return coords
 
-    def translate(self, vector: np.ndarray, origin_out: str) -> "CartesianCoordinates":
+    def translate(
+        self, vector: np.ndarray, origin_out: str
+    ) -> "CartesianCoordinates":
         """
         Translate Cartesian coordinates by the given vector.
 
@@ -272,7 +286,9 @@ def to_cometary(self) -> "CometaryCoordinates":
         return CometaryCoordinates.from_cartesian(self)
 
     @classmethod
-    def from_cometary(cls, cometary: "CometaryCoordinates") -> "CartesianCoordinates":
+    def from_cometary(
+        cls, cometary: "CometaryCoordinates"
+    ) -> "CartesianCoordinates":
         return cometary.to_cartesian()
 
     def to_keplerian(self) -> "KeplerianCoordinates":
@@ -281,7 +297,9 @@ 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":
@@ -290,5 +308,7 @@ 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()

adam_core/coordinates/cometary.py

@@ -224,10 +224,14 @@ 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)
+        covariances_cartesian = CoordinateCovariances.from_matrix(
+            covariances_cartesian
+        )
         coords = CartesianCoordinates.from_kwargs(
             x=coords_cartesian[:, 0],
             y=coords_cartesian[:, 1],
@@ -244,7 +248,9 @@ def to_cartesian(self) -> CartesianCoordinates:
         return coords
 
     @classmethod
-    def from_cartesian(cls, cartesian: CartesianCoordinates) -> "CometaryCoordinates":
+    def from_cartesian(
+        cls, cartesian: CartesianCoordinates
+    ) -> "CometaryCoordinates":
         from .transform import _cartesian_to_cometary, cartesian_to_cometary
 
         if cartesian.time is None:
@@ -277,10 +283,14 @@ 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)
+        covariances_cometary = CoordinateCovariances.from_matrix(
+            covariances_cometary
+        )
         coords = cls.from_kwargs(
             q=coords_cometary[:, 0],
             e=coords_cometary[:, 1],
@@ -302,7 +312,9 @@ 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":
@@ -311,5 +323,7 @@ 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())

adam_core/coordinates/conversions.py

@@ -5,7 +5,9 @@
 
 
 def _convert_coordinates_units(
-    coords: Union[np.ndarray, np.ma.masked_array], units: List, desired_units: List
+    coords: Union[np.ndarray, np.ma.masked_array],
+    units: List,
+    desired_units: List,
 ) -> Union[np.ndarray, np.ma.masked_array]:
     """
     Convert coordinate units to desired units.
@@ -84,7 +86,9 @@ def _convert_covariances_units(
     conversion_matrix = np.zeros((D, D), dtype=np.float64)
     for i in range(D):
         for j in range(D):
-            conversion_matrix[i, j] = covariance_units[i, j].to(desired_units[i, j])
+            conversion_matrix[i, j] = covariance_units[i, j].to(
+                desired_units[i, j]
+            )
 
     covariances_converted = conversion_matrix * covariances
 

adam_core/coordinates/covariances.py

@@ -123,7 +123,9 @@ 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))
@@ -321,7 +323,9 @@ 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.
 
@@ -379,7 +383,9 @@ def transform_covariances_sampling(
     """
     covariances_out = []
     for coord, covariance in zip(coords, covariances):
-        samples, W, W_cov = sample_covariance_random(coord, covariance, num_samples)
+        samples, W, W_cov = sample_covariance_random(
+            coord, covariance, num_samples
+        )
         samples_converted = func(samples)
         covariances_out.append(np.cov(samples_converted.T))
 
@@ -414,7 +420,9 @@ def transform_covariances_jacobian(
         Transformed covariance matrices.
     """
     jacobian = calc_jacobian(coords, _func, **kwargs)
-    covariances = jacobian @ covariances @ np.transpose(jacobian, axes=(0, 2, 1))
+    covariances = (
+        jacobian @ covariances @ np.transpose(jacobian, axes=(0, 2, 1))
+    )
     return covariances
 
 
@@ -479,7 +487,9 @@ def sigmas_from_df(
             sigmas[:, i] = df[f"sigma_{coord_names[i]}"].values
 
         except KeyError:
-            logger.debug(f"No sigma column found for dimension {coord_names[i]}.")
+            logger.debug(
+                f"No sigma column found for dimension {coord_names[i]}."
+            )
 
     return sigmas
 
@@ -574,10 +584,15 @@ def covariances_from_df(
 
     for i, j in zip(ii, jj):
         try:
-            covariances[:, i, j] = df[f"cov_{coord_names[i]}_{coord_names[j]}"].values
+            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
 
@@ -678,9 +693,14 @@ 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

adam_core/coordinates/keplerian.py

@@ -209,10 +209,14 @@ 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)
+        covariances_cartesian = CoordinateCovariances.from_matrix(
+            covariances_cartesian
+        )
         coords = CartesianCoordinates.from_kwargs(
             x=coords_cartesian[:, 0],
             y=coords_cartesian[:, 1],
@@ -254,10 +258,14 @@ 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)
+        covariances_keplerian = CoordinateCovariances.from_matrix(
+            covariances_keplerian
+        )
         coords = cls.from_kwargs(
             a=coords_keplerian[:, 0],
             e=coords_keplerian[:, 4],
@@ -278,7 +286,9 @@ 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":
@@ -287,5 +297,7 @@ 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())

adam_core/coordinates/origin.py

@@ -88,15 +88,21 @@ 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.")
+                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
 

adam_core/coordinates/residuals.py

@@ -37,7 +37,9 @@ class Residuals(qv.Table):
     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.
@@ -92,7 +94,9 @@ def calculate(cls, observed: CoordinateType, predicted: CoordinateType) -> "Resi
             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
@@ -120,7 +124,9 @@ def calculate(cls, observed: CoordinateType, predicted: CoordinateType) -> "Resi
         )
 
 
-def calculate_chi2(residuals: np.ndarray, covariances: np.ndarray) -> np.ndarray:
+def calculate_chi2(
+    residuals: np.ndarray, covariances: np.ndarray
+) -> np.ndarray:
     """
     Calculate the vectorized normalized squared residual for each residual and covariance pair.
     This normalized residual is equivalent to the Mahalanobis distance squared.
@@ -145,13 +151,17 @@ def calculate_chi2(residuals: np.ndarray, covariances: np.ndarray) -> np.ndarray
         Icarus, 166, 248-270. https://doi.org/10.1016/S0019-1035(03)00051-4
     """
     W = np.linalg.inv(covariances)
-    chi2 = np.einsum("ij,ji->i", np.einsum("ij,ijk->ik", residuals, W), residuals.T)
+    chi2 = np.einsum(
+        "ij,ji->i", np.einsum("ij,ijk->ik", residuals, W), residuals.T
+    )
     return chi2
 
 
 def _batch_coords_and_covariances(
     coords: np.ndarray, covariances: np.ndarray
-) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray], List[np.ndarray]]:
+) -> Tuple[
+    List[np.ndarray], List[np.ndarray], List[np.ndarray], List[np.ndarray]
+]:
     """
     Batch coordinates and covariances into groups that have the same
     number of dimensions that have missing values (represented by NaNs).