{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "id": "992386dc",
   "metadata": {},
   "source": [
    "# Controllers\n\nControllers map a reference $r_k$ and a state estimate $\\hat{x}_k$ to a control input $u_k$.\n`h` always returns $u_k$ so the output can be fed directly into a plant block.\n\n<a href=\"../f_h_functions.html\">\n    <img src=\"../../../_static/figures/block_controller.svg\" />\n</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eeff40e",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## PID Controller\nClassic **proportional\u2013integral\u2013derivative** controller.See the [pid](pid/pid.ipynb) notebook for derivation, API reference, and examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ad9dbff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jaxtyping import Float, jaxtyped\nfrom beartype import beartype\n\n\n@jaxtyped(typechecker=beartype)\ndef pid_f(\n    ck: tuple[float, float],\n    rk: float,\n    xhatk: float,\n    dt: float,\n) -> tuple[float, float]:\n    \"\"\"Advance the PID state; returns (current error, updated integral).\n\n    Args:\n        ck:    PID state \u2014 2-tuple of (e_{k-1}, e_int), previous error and accumulated integral.\n        rk:    Reference (setpoint) at step k.\n        xhatk: Current scalar state estimate.\n        KP:    Proportional gain.\n        KI:    Integral gain.\n        KD:    Derivative gain.\n        dt:    Timestep in seconds.\n\n    Returns:\n        tuple[float, float] \u2014 (e_k, e_int_k), the current error and updated integral.\n        Fed back as ck on the next step; not published directly to ROS.\n    \"\"\"\n    _, e_int = ck\n    ek = rk - xhatk\n    return (ek, e_int + ek * dt)\n\n\n@jaxtyped(typechecker=beartype)\ndef pid_h(\n    ck: tuple[float, float],\n    rk: float,\n    xhatk: float,\n    KP: float,\n    KI: float,\n    KD: float,\n    dt: float,\n) -> float:\n    \"\"\"Compute the PID control output u_k = KP*e + KI*e_int + KD*de/dt.\n\n    Args:\n        ck:    PID state \u2014 2-tuple of (e_{k-1}, e_int), previous error and accumulated integral.\n        rk:    Reference (setpoint) at step k.\n        xhatk: Current scalar state estimate.\n        KP:    Proportional gain.\n        KI:    Integral gain.\n        KD:    Derivative gain.\n        dt:    Timestep in seconds.\n\n    Returns:\n        float \u2014 control output u_k.\n        Published as Float64 via py2ros_float64.\n    \"\"\"\n    e_prev, e_int = ck\n    ek = rk - xhatk\n    return KP * ek + KI * e_int + KD * (ek - e_prev) / dt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09f9256e",
   "metadata": {},
   "source": [
    "## LQR Controller\n**Linear Quadratic Regulator** \u2014 optimal state-feedback for linear systems. See the [lqr](lqr/lqr.ipynb) notebook for derivation, API reference, and examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdcb7134",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\nfrom typing import Callable\nfrom jaxtyping import Float, jaxtyped\nfrom beartype import beartype\nfrom scipy.linalg import solve_discrete_are\n\n\n@jaxtyped(typechecker=beartype)\ndef lqr_gain(\n    A: Float[np.ndarray, \"n n\"],\n    B: Float[np.ndarray, \"n m\"],\n    Q: Float[np.ndarray, \"n n\"],\n    R: Float[np.ndarray, \"m m\"],\n) -> Float[np.ndarray, \"m n\"]:\n    \"\"\"Compute the discrete-time LQR gain K = (R + B'PB)^{-1} B'PA via DARE.\n\n    Args:\n        A: State transition matrix (n, n) of the linearised plant.\n        B: Control input matrix (n, m) \u2014 maps u_k into state space.\n        Q: State cost matrix (n, n), symmetric positive semi-definite.\n        R: Control cost matrix (m, m), symmetric positive definite.\n\n    Returns:\n        ndarray(m, n) \u2014 optimal gain matrix K. Compute once; pass as a static\n        parameter to lqr_h at every step.\n    \"\"\"\n    P = solve_discrete_are(A, B, Q, R)\n    return np.linalg.inv(R + B.T @ P @ B) @ B.T @ P @ A\n\n\n@jaxtyped(typechecker=beartype)\ndef lqr_h(\n    xhatk: Float[np.ndarray, \"n\"],\n    xrefk: Float[np.ndarray, \"n\"],\n    K: Float[np.ndarray, \"m n\"],\n) -> Float[np.ndarray, \"m\"]:\n    \"\"\"Compute the LQR control output u_k = -K(x\u0302_k - x_ref_k).\n\n    Args:\n        xhatk: Current state estimate (n,) from an estimator or sensor.\n        xrefk: Reference (target) state (n,).\n        K:     Precomputed LQR gain matrix (m, n) from lqr_gain.\n\n    Returns:\n        ndarray(m,) \u2014 control input u_k.\n        Published as Float64MultiArray via py2ros_float64_multiarray.\n    \"\"\"\n    return -K @ (xhatk - xrefk)"
   ]
  }
 ]
}