{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1a7b6b76",
   "metadata": {},
   "source": [
    "# Observers\n\nObservers estimate the internal state $\\hat{x}_k$ from plant observations $y_k$ and control inputs $u_k$.\n`h` always returns $\\hat{x}_k$ so the output can be fed directly into a controller block.\n\n<a href=\"../f_h_functions.html\">\n    <img src=\"../../../_static/figures/block_observer.svg\" />\n</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8cf74c2",
   "metadata": {},
   "source": [
    "## Kalman Filter (KF)\n\nOptimal estimator for **linear, Gaussian** systems. See the [kalman_filter](kf/kalman_filter.ipynb) notebook for derivation, API reference, and examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b1b7a5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\nfrom typing import Callable\nfrom jaxtyping import Float, jaxtyped\nfrom beartype import beartype\n\n\n@jaxtyped(typechecker=beartype)\ndef kf_f(\n    zk: tuple[Float[np.ndarray, \"n\"], Float[np.ndarray, \"n n\"]],\n    uk: float,\n    yk: Float[np.ndarray, \"m\"],\n    Fk: Float[np.ndarray, \"n n\"],\n    Bk: Float[np.ndarray, \"n\"],\n    Hk: Float[np.ndarray, \"m n\"],\n    Qk: Float[np.ndarray, \"n n\"],\n    Rk: Float[np.ndarray, \"m m\"],\n) -> tuple[Float[np.ndarray, \"n\"], Float[np.ndarray, \"n n\"]]:\n    \"\"\"Kalman filter predict+update; returns updated Kalman filter state.\n\n    Args:\n        zk: KF state \u2014 2-tuple of (x\u0302_k, P_k), state estimate and error covariance.\n        uk: Scalar control input.\n        yk: Measurement vector (m,).\n        Fk: State transition matrix (n, n) \u2014 maps x_k \u2192 x_{k+1}.\n        Bk: Control input matrix (n,) \u2014 maps scalar u_k into state space.\n        Hk: Observation matrix (m, n) \u2014 maps state to measurement space.\n        Qk: Process noise covariance (n, n) \u2014 encodes model uncertainty.\n        Rk: Measurement noise covariance (m, m) \u2014 encodes sensor noise.\n\n    Returns:\n        tuple \u2014 (x\u0302_{k+1}, P_{k+1}), the updated state estimate and error covariance.\n        Fed back as zk on the next step; not published directly to ROS.\n    \"\"\"\n    xk_hat, Pk = zk\n    x_pred = Fk @ xk_hat + Bk * uk\n    P_pred = Fk @ Pk @ Fk.T + Qk\n    S = Hk @ P_pred @ Hk.T + Rk\n    K = P_pred @ Hk.T @ np.linalg.inv(S)\n    yk = np.atleast_1d(yk)\n    x_upd = x_pred + K @ (yk - Hk @ x_pred)\n    P_upd = (np.eye(len(xk_hat)) - K @ Hk) @ P_pred\n    return (x_upd, P_upd)\n\n\n@jaxtyped(typechecker=beartype)\ndef kf_h(\n    zk: tuple[Float[np.ndarray, \"n\"], Float[np.ndarray, \"n n\"]],\n) -> Float[np.ndarray, \"n\"]:\n    \"\"\"Return the current state estimate from the KF state tuple.\n\n    Args:\n        zk: KF state \u2014 2-tuple of (x\u0302_k, P_k), state estimate and error covariance.\n\n    Returns:\n        ndarray(n,) \u2014 state estimate x\u0302_k.\n        Published as Float64MultiArray via py2ros_float64_multiarray.\n    \"\"\"\n    return zk[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fb31c7c",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## Extended Kalman Filter (EKF)\n\nExtends the KF to **nonlinear** systems by linearising around the current estimate\\nusing user-supplied Jacobians $F^J_k = \\\\partial f / \\\\partial x$ and $H^J_k = \\\\partial h / \\\\partial x$.\\nSee the [ekf](ekf/ekf.ipynb) notebook for derivation, API reference, and examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8e5e4ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\nfrom typing import Callable\nfrom jaxtyping import Float, jaxtyped\nfrom beartype import beartype\n\n\n@jaxtyped(typechecker=beartype)\ndef ekf_f(\n    zk: tuple[Float[np.ndarray, \"n\"], Float[np.ndarray, \"n n\"]],\n    uk: float,\n    yk: Float[np.ndarray, \"m\"],\n    f_nl: Callable,\n    h_nl: Callable,\n    Fjk: Float[np.ndarray, \"n n\"],\n    Hjk: Float[np.ndarray, \"m n\"],\n    Qk: Float[np.ndarray, \"n n\"],\n    Rk: Float[np.ndarray, \"m m\"],\n) -> tuple[Float[np.ndarray, \"n\"], Float[np.ndarray, \"n n\"]]:\n    \"\"\"EKF predict+update via Jacobian linearisation; returns updated state and covariance.\n\n    Args:\n        zk:   EKF state \u2014 2-tuple of (x\u0302_k, P_k), state estimate and error covariance.\n        uk:   Scalar control input.\n        yk:   Measurement vector (m,).\n        f_nl: Nonlinear state transition f(x, u) \u2192 x_next.\n        h_nl: Nonlinear observation function h(x) \u2192 y.\n        Fjk:  Jacobian \u2202f/\u2202x evaluated at x\u0302_{k-1} (n, n).\n        Hjk:  Jacobian \u2202h/\u2202x evaluated at x\u0302_{k|k-1} (m, n).\n        Qk:   Process noise covariance (n, n).\n        Rk:   Measurement noise covariance (m, m).\n\n    Returns:\n        tuple \u2014 (x\u0302_{k+1}, P_{k+1}), the updated state estimate and error covariance.\n        Fed back as zk on the next step; not published directly to ROS.\n    \"\"\"\n    x, P = zk\n    x_pred = f_nl(x, uk)\n    P_pred = Fjk @ P @ Fjk.T + Qk\n    S = Hjk @ P_pred @ Hjk.T + Rk\n    K = P_pred @ Hjk.T @ np.linalg.inv(S)\n    yk = np.atleast_1d(yk)\n    x_upd = x_pred + K @ (yk - h_nl(x_pred))\n    P_upd = (np.eye(len(x)) - K @ Hjk) @ P_pred\n    return (x_upd, P_upd)\n\n\n@jaxtyped(typechecker=beartype)\ndef ekf_h(\n    zk: tuple[Float[np.ndarray, \"n\"], Float[np.ndarray, \"n n\"]],\n) -> Float[np.ndarray, \"n\"]:\n    \"\"\"Return the current state estimate from the EKF state tuple.\n\n    Args:\n        zk: EKF state \u2014 2-tuple of (x\u0302_k, P_k), state estimate and error covariance.\n\n    Returns:\n        ndarray(n,) \u2014 state estimate x\u0302_k.\n        Published as Float64MultiArray via py2ros_float64_multiarray.\n    \"\"\"\n    return zk[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f3560eb",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## Unscented Kalman Filter (UKF)\nHandles **nonlinear** systems without Jacobians by propagating a set of deterministically chosen **sigma points** through the true nonlinear functions.See the [ukf](ukf/ukf.ipynb) notebook for derivation, API reference, and examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c5ab604",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jaxtyped(typechecker=beartype)\ndef ukf_f(\n    zk: tuple[Float[np.ndarray, \"n\"], Float[np.ndarray, \"n n\"]],\n    uk: float,\n    yk: Float[np.ndarray, \"m\"],\n    f_nl: Callable,\n    h_nl: Callable,\n    Qk: Float[np.ndarray, \"n n\"],\n    Rk: Float[np.ndarray, \"m m\"],\n    alpha: float = 1e-3,\n    beta: float = 2.0,\n    kappa: float = 0.0,\n) -> tuple[Float[np.ndarray, \"n\"], Float[np.ndarray, \"n n\"]]:\n    \"\"\"UKF predict+update via sigma-point propagation; returns updated state and covariance.\n\n    Args:\n        zk:    UKF state \u2014 2-tuple of (x\u0302_k, P_k), state estimate and error covariance.\n        uk:    Scalar control input.\n        yk:    Measurement vector (m,).\n        f_nl:  Nonlinear state transition f(x, u) \u2192 x_next.\n        h_nl:  Nonlinear observation function h(x) \u2192 y.\n        Qk:    Process noise covariance (n, n).\n        Rk:    Measurement noise covariance (m, m).\n        alpha: Sigma point spread around the mean (default 1e-3).\n        beta:  Distribution parameter; 2.0 optimal for Gaussian (default 2.0).\n        kappa: Secondary scaling parameter, typically 0 or 3-n (default 0.0).\n\n    Returns:\n        tuple \u2014 (x\u0302_{k+1}, P_{k+1}), the updated state estimate and error covariance.\n        Fed back as zk on the next step; not published directly to ROS.\n    \"\"\"\n    x, P = zk\n    n = len(x)\n    lam = alpha**2 * (n + kappa) - n\n\n    sqrtP = np.linalg.cholesky((n + lam) * P)\n    sigmas = np.column_stack(\n        [x] + [x + sqrtP[:, i] for i in range(n)] + [x - sqrtP[:, i] for i in range(n)]\n    )\n    Wm = np.full(2 * n + 1, 1.0 / (2 * (n + lam)))\n    Wm[0] = lam / (n + lam)\n    Wc = Wm.copy()\n    Wc[0] += 1 - alpha**2 + beta\n\n    # Predict\n    sf = np.column_stack([f_nl(sigmas[:, i], uk) for i in range(2 * n + 1)])\n    xp = sf @ Wm\n    Pp = (\n        sum(Wc[i] * np.outer(sf[:, i] - xp, sf[:, i] - xp) for i in range(2 * n + 1))\n        + Qk\n    )\n\n    # Update\n    yk = np.atleast_1d(yk)\n    sh = np.column_stack([np.atleast_1d(h_nl(sf[:, i])) for i in range(2 * n + 1)])\n    yp = sh @ Wm\n    S = (\n        sum(Wc[i] * np.outer(sh[:, i] - yp, sh[:, i] - yp) for i in range(2 * n + 1))\n        + Rk\n    )\n    Pxy = sum(Wc[i] * np.outer(sf[:, i] - xp, sh[:, i] - yp) for i in range(2 * n + 1))\n    K = Pxy @ np.linalg.inv(S)\n    return (xp + K @ (yk - yp), Pp - K @ S @ K.T)\n\n\n@jaxtyped(typechecker=beartype)\ndef ukf_h(\n    zk: tuple[Float[np.ndarray, \"n\"], Float[np.ndarray, \"n n\"]],\n) -> Float[np.ndarray, \"n\"]:\n    \"\"\"Return the current state estimate from the UKF state tuple.\n\n    Args:\n        zk: UKF state \u2014 2-tuple of (x\u0302_k, P_k), state estimate and error covariance.\n\n    Returns:\n        ndarray(n,) \u2014 state estimate x\u0302_k.\n        Published as Float64MultiArray via py2ros_float64_multiarray.\n    \"\"\"\n    return zk[0]"
   ]
  }
 ]
}