{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting and transferring discrepancy\n",
    "\n",
    "A discrepancy describes the part of a measured response that a physical model misses. This tutorial fits a lossy coax with a small conductor error, predicts the correction with its uncertainty, and carries it to a new line length. It then shows how a port discrepancy can carry a joint parameter and correction prediction into a terminated circuit.\n",
    "\n",
    "The [discrepancy prediction example](../examples/discrepancy_prediction.rst) is the short introduction. The code here uses small synthetic data sets so every result can be checked against a known line."
   ],
   "id": "discrepancy-01"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A reference line and a shorter target\n",
    "\n",
    "The reference line is 30 cm long; the target is 15 cm. The nominal line has 10% excess conductor loss. A port correction belongs to the reference response, while an internal correction changes characteristic impedance, attenuation, and phase per unit length. That distinction matters when the line length changes."
   ],
   "id": "discrepancy-02"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import sys\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.optimize import minimize as scipy_minimize\n",
    "import pmrf as prf\n",
    "from pmrf.materials import BulkConductor, ConstantDielectric\n",
    "from pmrf.models import (\n",
    "    CoaxialLine, GridLineDiscrepancy, GridPortDiscrepancy, LineCorrected,\n",
    "    Load, PortCorrected, line_internal_features, project_line_basis_joint,\n",
    "    reference_line_features,\n",
    ")\n",
    "from pmrf.objectives import MarginalLogLikelihood\n",
    "from pmrf.optimize import ScipyMinimize, minimize\n",
    "from pmrf.stats import (\n",
    "    GaussianLikelihood, GaussianProcess, Matern52Kernel, RBFKernel, posterior_covariance,\n",
    ")\n",
    "\n",
    "tutorial_dir = Path.cwd()\n",
    "if not (tutorial_dir / \"line_internal_discrepancy.py\").exists():\n",
    "    tutorial_dir = Path(\"docs/tutorials\")\n",
    "sys.path.insert(0, str(tutorial_dir.resolve()))\n",
    "from line_internal_discrepancy import (\n",
    "    MODEL_SIGMA, TRUTH_SIGMA, _basis, compare_transfer,\n",
    "    exact_values, frequency, line,\n",
    ")\n",
    "from port_discrepancy import ReciprocalEventTransform, reciprocal_predictor\n",
    "\n",
    "f = frequency()\n",
    "comparison = compare_transfer()\n",
    "for name, errors in comparison.items():\n",
    "    print(f\"{name}: RMS S21 error = {errors['rms_s21']:.4g}, eta bias = {errors['eta_bias']:.4g}\")"
   ],
   "id": "discrepancy-03"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The internal correction follows the shorter line's electrical length. The port correction uses the reference line's S-parameter error as if it belonged to the shorter line."
   ],
   "id": "discrepancy-04"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learn an internal correction\n",
    "\n",
    "Here the reference measurement supplies characteristic impedance $Z_c$ and propagation constant $\\gamma$. `reference_line_features` converts them to the four log-ratio channels used by `GridLineDiscrepancy`: real and imaginary $\\log Z_c$, log attenuation, and log phase constant. The likelihood variance is specified in this transformed event space.\n",
    "\n",
    "The nominal model has a free conductor parameter. `predict_joint` returns one Gaussian over that parameter and the correction grid, preserving their covariance."
   ],
   "id": "discrepancy-05"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reference = line(0.30, TRUTH_SIGMA)\n",
    "zc, gamma_length = reference.zc_and_gammaL(f)\n",
    "observed = reference_line_features(zc, gamma_length / 0.30)\n",
    "nominal = line(0.30, MODEL_SIGMA, free_sigma=True)\n",
    "reference_likelihood = MarginalLogLikelihood(\n",
    "    predictor=line_internal_features,\n",
    "    observed=observed,\n",
    "    likelihood=GaussianLikelihood(1e-10),\n",
    "    discrepancy=GaussianProcess(RBFKernel(0.5) * 0.1**2, jitter=1e-12),\n",
    ")\n",
    "names, joint = reference_likelihood.predict_joint(nominal, f, f)\n",
    "print(\"Joint dimensions:\", joint.mean().shape[0], \"parameter names:\", names)"
   ],
   "id": "discrepancy-06"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same prediction can be attached to a shorter line. The joint prior keeps the conductor parameter correlated with the discrepancy values. `truncate='unnormalised'` respects the conductor's validity range for MAP work; it is unsuitable for sampling or evidence calculations."
   ],
   "id": "discrepancy-07"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "carrier = GridLineDiscrepancy.from_prediction(f, joint, joint=True)\n",
    "short_line = LineCorrected(line(0.15, MODEL_SIGMA, free_sigma=True), carrier)\n",
    "short_line = prf.prior(\n",
    "    short_line, [*names, \"discrepancy.values\"], joint,\n",
    "    truncate=\"unnormalised\",\n",
    ")\n",
    "truth_short = line(0.15, TRUTH_SIGMA)\n",
    "rms_s21 = np.sqrt(np.mean(np.abs(np.asarray(short_line.s(f) - truth_short.s(f))[:, 1, 0])**2))\n",
    "print(\"Transferred RMS S21 error:\", float(rms_s21))"
   ],
   "id": "discrepancy-08"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The joint mean contains the conductor parameter first, then four correction channels. Its diagonal gives marginal uncertainty in each channel; the off-diagonal entries still matter when the correction is used with the conductor parameter."
   ],
   "id": "discrepancy-09"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean = np.asarray(joint.mean())[1:].reshape(2, 2, f.npoints)\n",
    "std = np.sqrt(np.diag(np.asarray(joint.covariance()))[1:]).reshape(2, 2, f.npoints)\n",
    "truth = np.asarray(exact_values(f))\n",
    "labels = (\"log Zc real\", \"log Zc imaginary\", \"log attenuation\", \"log phase\")\n",
    "fig, axes = plt.subplots(2, 2, figsize=(10, 6), sharex=True)\n",
    "for axis, label, predicted, sigma, actual in zip(axes.flat, labels, mean.reshape(4, -1), std.reshape(4, -1), truth.reshape(4, -1)):\n",
    "    axis.plot(f.f_scaled, actual, label=\"injected\")\n",
    "    axis.plot(f.f_scaled, predicted, label=\"predicted\")\n",
    "    axis.fill_between(f.f_scaled, predicted - 2*sigma, predicted + 2*sigma, alpha=0.2)\n",
    "    axis.set_title(label)\n",
    "axes[0, 0].legend()\n",
    "fig.tight_layout()"
   ],
   "id": "discrepancy-10"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## When only S-parameters are measured\n",
    "\n",
    "The preceding route uses $Z_c$ and $\\gamma$. If the reference only provides complex S-parameters, represent the internal correction with basis coefficients and fit them through the line response. The [supporting script](line_internal_discrepancy.py) defines the basis from an RBF spectral density; the fit and projection are shown below."
   ],
   "id": "discrepancy-11"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "basis = _basis(f)\n",
    "nominal_s = LineCorrected(line(0.30, MODEL_SIGMA), basis)\n",
    "observed_s = line(0.30, TRUTH_SIGMA).s(f)\n",
    "s_likelihood = MarginalLogLikelihood(\n",
    "    predictor=lambda model, grid: model.s(grid),\n",
    "    observed=observed_s,\n",
    "    likelihood=GaussianLikelihood(1e-8),\n",
    ")\n",
    "objective = lambda model: -s_likelihood(model, f) - prf.log_prior(model)\n",
    "fitted_s = minimize(objective, nominal_s, solver=ScipyMinimize(method=\"BFGS\"), max_iter=400).model\n",
    "fitted_basis = prf.update(\n",
    "    basis, \"coefficients\", value=prf.values(fitted_s)[\"discrepancy.coefficients\"]\n",
    ")\n",
    "linearization = s_likelihood.linearize(fitted_s, f)\n",
    "covariance = posterior_covariance([linearization], fitted_s)\n",
    "names_s, joint_s = project_line_basis_joint(\n",
    "    fitted_s, fitted_basis, linearization, covariance, f\n",
    ")\n",
    "carrier_s = GridLineDiscrepancy.from_prediction(f, joint_s, joint=True)\n",
    "short_s = LineCorrected(line(0.15, MODEL_SIGMA), carrier_s)\n",
    "rms_s = np.sqrt(np.mean(np.abs(np.asarray(short_s.s(f) - truth_short.s(f))[:, 1, 0])**2))\n",
    "mean_s = np.asarray(joint_s.mean()).reshape(2, 2, f.npoints)\n",
    "std_s = np.sqrt(np.diag(np.asarray(joint_s.covariance()))).reshape(2, 2, f.npoints)\n",
    "coverage = np.mean(np.abs(np.asarray(exact_values(f)) - mean_s) <= 2 * std_s)\n",
    "print(\"Transferred RMS S21 error:\", float(rms_s))\n",
    "print(\"Injected values within mean ± 2σ:\", float(coverage))"
   ],
   "id": "discrepancy-12"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Port discrepancy in a later circuit\n",
    "\n",
    "A port correction applies when the fitted component stays the same and its surrounding circuit changes. Here the reference fit estimates a bounded coax length and a reciprocal port discrepancy. The [supporting script](port_discrepancy.py) defines the event transform used for symmetric transmission and contains independent covariance checks. We use that transform below, then carry the joint prediction into a terminated circuit."
   ],
   "id": "discrepancy-13"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cable(length):\n",
    "    return CoaxialLine(\n",
    "        length=length,\n",
    "        d_in=prf.Fixed(1.12e-3), d_out=prf.Fixed(3.2e-3),\n",
    "        dielectric=ConstantDielectric(ep_r=prf.Fixed(1.384), tand=prf.Fixed(0.001)),\n",
    "        conductor=BulkConductor(sigma=prf.Fixed(1 / 1.6e-8)),\n",
    "    )\n",
    "\n",
    "reference_f = prf.Frequency(100, 500, 12, \"MHz\")\n",
    "transfer_f = prf.Frequency(120, 480, 6, \"MHz\")\n",
    "true_cable = cable(prf.Fixed(0.1))\n",
    "initial_cable = cable(prf.Bounded(0.05, 0.15, value=0.095))\n",
    "reference_port = MarginalLogLikelihood(\n",
    "    predictor=reciprocal_predictor,\n",
    "    observed=reciprocal_predictor(true_cable, reference_f),\n",
    "    likelihood=GaussianLikelihood(jnp.full((3, 2, reference_f.npoints), 1e-8)),\n",
    "    discrepancy=GaussianProcess(Matern52Kernel(80.0) * 1e-6, jitter=1e-12),\n",
    "    event_transform=ReciprocalEventTransform,\n",
    ")\n",
    "\n",
    "def port_loss(length):\n",
    "    return -reference_port(prf.update(initial_cable, {\"length\": length}), reference_f)\n",
    "\n",
    "value_gradient = jax.jit(jax.value_and_grad(port_loss))\n",
    "fit_result = scipy_minimize(\n",
    "    lambda x: tuple(np.asarray(v) for v in value_gradient(jnp.asarray(x))),\n",
    "    np.array([0.095]), jac=True, method=\"L-BFGS-B\", bounds=[(0.05, 0.15)],\n",
    "    options={\"maxiter\": 1000, \"gtol\": 1e-10, \"ftol\": 1e-14},\n",
    ")\n",
    "assert fit_result.success\n",
    "fitted_cable = prf.update(initial_cable, {\"length\": jnp.asarray(fit_result.x[0])})\n",
    "port_names, port_joint = reference_port.predict_joint(fitted_cable, reference_f, transfer_f)\n",
    "print(\"Fitted length:\", float(prf.values(fitted_cable)[\"length\"]))"
   ],
   "id": "discrepancy-14"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "port_mean = port_joint.mean()[1:].reshape(3, 2, transfer_f.npoints)\n",
    "corrected_cable = PortCorrected(\n",
    "    fitted_cable,\n",
    "    GridPortDiscrepancy(prf.Unconstrained(port_mean), transfer_f, (\"11\", \"22\", \"s21\")),\n",
    ")\n",
    "corrected_cable = prf.prior(\n",
    "    corrected_cable, [*port_names, \"discrepancy.values\"], port_joint,\n",
    "    truncate=\"unnormalised\",\n",
    ")\n",
    "load = Load(gamma=(75.0 - 50.0) / (75.0 + 50.0), z0=50.0)\n",
    "transfer_circuit = corrected_cable.terminated(load)\n",
    "observed_reflection = true_cable.terminated(load).s(transfer_f, z0=50.0)[:, 0, 0]\n",
    "transfer_likelihood = MarginalLogLikelihood(\n",
    "    predictor=lambda model, grid: model.s(grid, z0=50.0)[:, 0, 0],\n",
    "    observed=observed_reflection,\n",
    "    likelihood=GaussianLikelihood(1e-6),\n",
    ")\n",
    "transfer_linearization = transfer_likelihood.linearize(\n",
    "    transfer_circuit, transfer_f, space=\"declared\"\n",
    ")\n",
    "transfer_covariance = posterior_covariance(\n",
    "    [transfer_linearization], transfer_circuit, space=\"declared\"\n",
    ")\n",
    "print(\"Joint dimension:\", port_joint.mean().shape[0])\n",
    "print(\"Transferred log prior:\", float(prf.log_prior(transfer_circuit, space=\"declared\")))\n",
    "print(\"Posterior variance, first coordinate:\", float(transfer_covariance[0, 0]))"
   ],
   "id": "discrepancy-15"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The carrier follows what stays physically common between fits: use an internal line correction across lengths, or a port correction when moving the same component into another circuit. The joint distribution retains parameter–discrepancy covariance in each later calculation. The [port script](port_discrepancy.py) compares the final covariance with a NumPy Kalman update and checks the transferred prior's precision."
   ],
   "id": "discrepancy-16"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}