"""Shared base class for new-generation Kalman attitude estimators.
All attitude filters compose the filter-neutral ``State``, ``Covariance``,
process-model, and ``MeasurementStack`` operations through this base class.
"""
from __future__ import annotations
from typing import Any, get_args
import warnings
import numpy as np
from ADCS.covariance import Covariance
from ADCS.estimators.process_model import propagate_state
from ADCS.estimators.process_noise import discretize_process_noise
from ADCS.orbits.universal_constants import TimeConstants
from ADCS.state import EstimatorState, QuaternionMode
__all__ = ["AttitudeEstimator"]
_QUATERNION_MODES = tuple(get_args(QuaternionMode))
[docs]
class AttitudeEstimator:
"""Shared prediction/correction lifecycle for the first simple filters.
The base implementation defaults to the physical attitude state. Filters
that explicitly support joint estimation of hardware biases or disturbance
parameters can opt in by setting ``supports_augmented_parameters``.
"""
supports_augmented_parameters = False
def __init__(
self,
satellite: Any,
state: EstimatorState,
*,
dt: float,
covariance_coordinates: str,
correction_mode: str,
measurement_quaternion_mode: str,
unmodeled_dynamics_psd: Any = 0.0,
) -> None:
if not isinstance(state, EstimatorState):
raise TypeError(f"state must be an EstimatorState, got {type(state).__name__}")
dt = float(dt)
if not np.isfinite(dt) or dt < 0.0:
raise ValueError("dt must be finite and non-negative")
if covariance_coordinates not in ("full", "tangent"):
raise ValueError("covariance_coordinates must be 'full' or 'tangent'")
self._validate_charts(
covariance_coordinates, correction_mode, measurement_quaternion_mode
)
self.satellite = satellite
self.dt = dt
self._covariance_coordinates = covariance_coordinates
self._correction_mode = correction_mode
self._measurement_quaternion_mode = measurement_quaternion_mode
self._validate_state(state)
self.unmodeled_dynamics_psd = np.array(
unmodeled_dynamics_psd, dtype=float, copy=True
)
self._state = self._normalize_initial_state(state)
# Keep the caller's per-step Q separate from the total Q exposed on a
# predicted state. Otherwise the model-generated part would be added
# again on the following prediction.
self._configured_process_noise = self._state.process_noise.copy()
self._diagnostics: dict[str, np.ndarray] = {}
self._previous_orbital_state: Any | None = None
@property
def covariance_coordinates(self) -> str:
"""Coordinates the covariance is kept in: ``"full"`` (EKF) or ``"tangent"`` (MEKF)."""
return self._covariance_coordinates
@property
def correction_mode(self) -> str:
"""Attitude chart used for the transition matrix, the correction, and the reset."""
return self._correction_mode
@property
def measurement_quaternion_mode(self) -> str:
"""Attitude chart used for quaternion residuals and their Jacobians."""
return self._measurement_quaternion_mode
@property
def state(self) -> EstimatorState:
"""Return an owned snapshot of the current estimate."""
return self._state.copy()
@property
def diagnostics(self) -> dict[str, np.ndarray]:
"""Return owned matrices and vectors from the latest operations."""
return {name: value.copy() for name, value in self._diagnostics.items()}
[docs]
def reset(self, state: EstimatorState) -> EstimatorState:
"""Replace the estimate after validating this filter's layout."""
self._validate_state(state)
self._state = self._normalize_initial_state(state)
self._configured_process_noise = self._state.process_noise.copy()
self._diagnostics = {}
self._previous_orbital_state = None
return self.state
@staticmethod
def _validate_charts(
covariance_coordinates: str,
correction_mode: str,
measurement_quaternion_mode: str,
) -> None:
"""Reject chart combinations that would silently mix conventions.
The measurement Jacobian is written in ``measurement_quaternion_mode``
and applied to a covariance kept in ``correction_mode``. With a tangent
covariance those are only interchangeable when they name the same chart;
with a full covariance the correction must be the four-element block.
"""
for name, mode in (
("correction_mode", correction_mode),
("measurement_quaternion_mode", measurement_quaternion_mode),
):
if mode not in _QUATERNION_MODES:
raise ValueError(f"{name} must be one of {_QUATERNION_MODES}, got {mode!r}")
if measurement_quaternion_mode == "full_quaternion":
raise ValueError(
"measurement_quaternion_mode must be a three-parameter attitude chart; "
"quaternion residuals are always minimal"
)
if covariance_coordinates == "full":
if correction_mode != "full_quaternion":
raise ValueError(
"a full covariance must be corrected with "
f"correction_mode='full_quaternion', got {correction_mode!r}"
)
elif correction_mode == "full_quaternion":
raise ValueError(
"a tangent covariance cannot be corrected with correction_mode='full_quaternion'"
)
elif correction_mode != measurement_quaternion_mode:
raise ValueError(
f"correction_mode {correction_mode!r} and measurement_quaternion_mode "
f"{measurement_quaternion_mode!r} must match when the covariance is kept "
"in tangent coordinates"
)
def _validate_state(self, state: EstimatorState) -> None:
"""Validate the state/hardware contract shared by reset and construction."""
if not isinstance(state, EstimatorState):
raise TypeError(f"state must be an EstimatorState, got {type(state).__name__}")
state.validate_layout(
wheel_momentum=len(self.satellite.rw_actuators),
actuator_bias=self.satellite.act_bias_len,
sensor_bias=self.satellite.att_sens_bias_len,
disturbance_parameter=self.satellite.dist_param_len,
)
if (
state.block_size("estimated_parameters")
and not self.supports_augmented_parameters
):
raise NotImplementedError(
f"{type(self).__name__} does not yet support estimated biases or "
"disturbance parameters"
)
expected_size = state.size(coordinates=self.covariance_coordinates)
if state.covariance.dimension != expected_size:
raise ValueError(
f"{type(self).__name__} requires a {expected_size}x{expected_size} "
f"{self.covariance_coordinates} covariance, got {state.covariance.shape}"
)
[docs]
def predict(
self,
control: Any,
orbital_state_start: Any,
orbital_state_end: Any,
*,
dt: float | None = None,
midpoint_orbital_state: Any | None = None,
) -> EstimatorState:
"""Propagate the nominal state and covariance one deterministic step.
The model-generated discrete process covariance and the prior state's
discrete ``process_noise`` (legacy ``int_cov``) are additive.
"""
step = self.dt if dt is None else float(dt)
if not np.isfinite(step) or step < 0.0:
raise ValueError("dt must be finite and non-negative")
self._check_step_against_orbit(step, orbital_state_start, orbital_state_end)
control = np.array(control, dtype=float, copy=True)
expected_control_shape = (self.satellite.control_len,)
if control.shape != expected_control_shape:
raise ValueError(
f"control must have shape {expected_control_shape}, got {control.shape}"
)
prior = self._state
# Augmented actuator/sensor/disturbance blocks are represented in the
# estimator state, while the historical spacecraft dynamics object
# stores their nominal values on hardware/disturbance instances. Keep
# the model synchronized before both nonlinear propagation and the
# Jacobian/process-noise construction. Legacy UKF estimators perform
# this operation for every sigma point; the EKF-family prediction uses
# the current nominal state once.
match_estimate = getattr(self.satellite, "match_estimate", None)
if match_estimate is not None:
match_estimate(prior, step)
predicted = propagate_state(
prior,
self.satellite,
control,
step,
orbital_state_start,
orbital_state_end,
midpoint_orbital_state=midpoint_orbital_state,
)
# Linearized prediction: the discrete process noise must include the
# actuator command-noise term (zero-order hold through B = dxdot/du);
# sigma-point filters carry it through control sigma points instead.
transition, process_noise = discretize_process_noise(
prior,
self.satellite,
control,
orbital_state_start,
step,
final_state=predicted,
include_control_noise=True,
unmodeled_dynamics_psd=self.unmodeled_dynamics_psd,
quaternion_mode=self.correction_mode,
quaternion_order="right",
)
process_noise = self._add_state_process_noise(process_noise)
predicted_covariance = prior.covariance.predicted_linear(
transition, process_noise
)
discrete_noise = Covariance(
process_noise,
form=prior.process_noise.form,
coordinates=prior.process_noise.coordinates,
psd_policy=prior.process_noise.psd_policy,
)
predicted.covariance = predicted_covariance
predicted.process_noise = discrete_noise
self._state = predicted
self._diagnostics.update(
transition=transition,
process_noise=process_noise,
predicted_covariance=predicted_covariance.as_matrix(),
)
return self.state
def _add_state_process_noise(self, process_noise: np.ndarray) -> np.ndarray:
"""Add the state's discrete per-step process covariance.
``discretize_process_noise`` returns the covariance generated by the
continuous model. The configured
``EstimatorState.process_noise`` (the legacy ``int_cov`` field) is an
independent, already-discrete covariance and therefore must be added
after discretization.
"""
model_noise = np.asarray(process_noise, dtype=float)
state_noise = self._configured_process_noise.as_matrix()
if model_noise.shape != state_noise.shape:
raise ValueError(
"discretized process noise and state process noise must have "
f"matching shapes, got {model_noise.shape} and {state_noise.shape}"
)
combined = model_noise + state_noise
return (combined + combined.T) / 2.0
[docs]
def correct(
self,
measurements: Any,
orbital_state: Any,
*,
enabled: Any | None = None,
time_s: float | None = None,
epoch_s: float = 0.0,
) -> EstimatorState:
"""Apply all selected finite measurements at the current state."""
stack = self.satellite.measurement_stack
candidate = stack.active_mask(
measurements,
enabled=enabled,
time_s=time_s,
epoch_s=epoch_s,
)
predicted_measurements = stack.predict(
self._state, orbital_state, active_mask=candidate
)
active = stack.active_mask(
measurements,
enabled=candidate,
time_s=time_s,
epoch_s=epoch_s,
predicted=predicted_measurements,
)
residual = stack.residual(
measurements,
predicted_measurements,
active,
quaternion_mode=self.measurement_quaternion_mode,
)
measurement_jacobian = stack.jacobian(
self._state,
orbital_state,
active,
quaternion_mode=self.measurement_quaternion_mode,
coordinates=self.covariance_coordinates,
)
measurement_noise = stack.covariance(
self._state,
active,
quaternion_mode=self.measurement_quaternion_mode,
)
prior = self._state
covariance_size = prior.covariance.dimension
if residual.size:
innovation_covariance = (
measurement_jacobian
[docs]
@ prior.covariance.as_matrix()
@ measurement_jacobian.T
+ measurement_noise.as_matrix()
)
gain, joseph_covariance = prior.covariance.updated_linear(
measurement_jacobian,
measurement_noise,
joseph=True,
)
correction = gain @ residual
reset_jacobian = prior.retraction_jacobian(
correction,
quaternion_mode=self.correction_mode,
quaternion_order="right",
)
corrected_covariance = prior.transport_covariance(
joseph_covariance,
correction,
quaternion_mode=self.correction_mode,
quaternion_order="right",
)
corrected = prior.plus(
correction,
quaternion_mode=self.correction_mode,
quaternion_order="right",
)
corrected.covariance = corrected_covariance
self._state = corrected
else:
innovation_covariance = np.zeros((0, 0))
gain = np.zeros((covariance_size, 0))
correction = np.zeros(covariance_size)
reset_jacobian = np.eye(covariance_size)
self._diagnostics.update(
active_mask=active,
predicted_measurement=predicted_measurements,
innovation=residual,
measurement_jacobian=measurement_jacobian,
measurement_noise=measurement_noise.as_matrix(),
innovation_covariance=innovation_covariance,
kalman_gain=gain,
correction=correction,
reset_jacobian=reset_jacobian,
corrected_covariance=self._state.covariance.as_matrix(),
)
return self.state
def step(
self,
measurements_or_control: Any,
orbital_state_or_measurements: Any,
orbital_state_start: Any | None = None,
orbital_state_end: Any | None = None,
*,
dt: float | None = None,
midpoint_orbital_state: Any | None = None,
enabled: Any | None = None,
time_s: float | None = None,
epoch_s: float = 0.0,
) -> EstimatorState:
"""Run the measurement stage, or preserve the old combined step call.
The staged contract calls ``step(measurements, orbital_state)`` after
:meth:`predict`. During migration, the historical
``step(control, measurements, orbital_state_start, orbital_state_end)``
form remains accepted and performs both stages.
"""
if (orbital_state_start is None) != (orbital_state_end is None):
raise TypeError(
"combined estimator.step requires both orbital_state_start and "
"orbital_state_end, or neither"
)
if orbital_state_start is None:
if not hasattr(orbital_state_or_measurements, "J2000"):
# A legacy step(control, measurements) call with matching lengths
# would otherwise be corrected as if the control were the
# measurement and the measurement the orbital state, silently.
raise TypeError(
"step(measurements, orbital_state) needs an Orbital_State as its "
"second argument; the combined form is "
"step(control, measurements, orbital_state_start, orbital_state_end)"
)
return self.correct(
measurements_or_control,
orbital_state_or_measurements,
enabled=enabled,
time_s=time_s,
epoch_s=epoch_s,
)
self.predict(
measurements_or_control,
orbital_state_start,
orbital_state_end,
dt=dt,
midpoint_orbital_state=midpoint_orbital_state,
)
return self.correct(
orbital_state_or_measurements,
orbital_state_end,
enabled=enabled,
time_s=time_s,
epoch_s=epoch_s,
)
[docs]
def update(
self,
u: Any | None = None,
sensors: Any | None = None,
os: Any | None = None,
) -> EstimatorState:
"""Return the posterior state, with a compatibility update adapter.
The staged simulation contract calls zero-argument ``update()`` after
``predict()`` and ``step()``. The argument-bearing form remains as a
temporary adapter for callers that still use the old one-call protocol.
"""
if u is None and sensors is None and os is None:
return self.state
if u is None or sensors is None or os is None:
raise TypeError("update requires u, sensors, and os when called with arguments")
if self._previous_orbital_state is None:
self._previous_orbital_state = os
self.correct(sensors, os)
else:
orbital_state_start = self._previous_orbital_state
self._previous_orbital_state = os
# No midpoint is passed on purpose: the propagation then averages the
# two orbital states, which is what the RK4 stages need. Passing the
# end state biased every stage toward the end of the interval.
self.predict(u, orbital_state_start, os)
self.correct(sensors, os)
return self.state
@staticmethod
def _check_step_against_orbit(step: float, start: Any, end: Any) -> None:
"""Warn when ``dt`` disagrees with the time between the two orbital states.
A filter built with one ``dt`` and driven by a simulation running at
another integrates the wrong interval and nothing else notices (a 10x
mismatch took a 16 degree mean error to 47). Two identical orbital
states (the static-orbit test pattern) carry no time information and
are left alone.
"""
start_time = getattr(start, "J2000", None)
end_time = getattr(end, "J2000", None)
if start_time is None or end_time is None:
return
gap = (float(end_time) - float(start_time)) * TimeConstants.cent2sec
if gap <= 0.0:
return
if step <= 0.0 or abs(gap - step) > 0.1 * step:
warnings.warn(
f"predict() integrates dt = {step:g} s but the orbital states are "
f"{gap:g} s apart; pass dt= or build the estimator with the "
"simulation step",
UserWarning,
stacklevel=3,
)
def _normalize_initial_state(self, state: EstimatorState) -> EstimatorState:
normalized = state.normalized()
if self.covariance_coordinates == "full":
normalizer = state.normalization_jacobian()
normalized.covariance = state.covariance.transformed(normalizer)
normalized.process_noise = state.process_noise.transformed(normalizer)
return normalized