{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "state-metrics-title",
   "metadata": {},
   "source": [
    "# State Metrics and Solver Diagnostics\n",
    "\n",
    "This tutorial shows how to compute common open-system quantities such as purity, entropy, fidelity, trace distance, populations, coherence, and Bloch-vector components. The same metric callbacks can be passed directly to solvers through `state_observables`; `mcsolve` aggregates the built-in linearizable metrics in the backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "state-metrics-imports",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import openquantumsim as oqs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "state-metrics-direct-note",
   "metadata": {},
   "source": [
    "The metric functions accept kets or density matrices whenever the quantity is well defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "state-metrics-direct",
   "metadata": {},
   "outputs": [],
   "source": [
    "up = np.array([1.0, 0.0], dtype=np.complex128)\n",
    "down = np.array([0.0, 1.0], dtype=np.complex128)\n",
    "plus = (up + down) / np.sqrt(2.0)\n",
    "mixed = 0.5 * np.eye(2, dtype=np.complex128)\n",
    "\n",
    "summary = {\n",
    "    \"purity(mixed)\": oqs.purity(mixed),\n",
    "    \"entropy(mixed)\": oqs.von_neumann_entropy(mixed),\n",
    "    \"fidelity(up, plus)\": oqs.fidelity(up, plus),\n",
    "    \"trace_distance(up, down)\": oqs.trace_distance(up, down),\n",
    "    \"bloch(plus)\": oqs.bloch_vector(plus),\n",
    "}\n",
    "summary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "state-metrics-mesolve-note",
   "metadata": {},
   "source": [
    "`state_metrics(...)` builds a dictionary of scalar callbacks. Deterministic solvers and single trajectories evaluate callbacks at each requested time without returning full states unless `Options(save_states=True)` is also set. `mcsolve` can aggregate built-in metrics such as purity, entropy, pure-state fidelity, populations, and Bloch components without saving every trajectory state.\n",
    "\n",
    "`mesolve` uses the Julia backend. Run `python setup_julia.py` once before executing this notebook locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "state-metrics-system",
   "metadata": {},
   "outputs": [],
   "source": [
    "gamma = 0.45\n",
    "atom = oqs.SpinSpace(0.5, label=\"atom\")\n",
    "\n",
    "H = 0.0 * oqs.sigmaz(atom)\n",
    "collapse = np.sqrt(gamma) * oqs.sigmam(atom)\n",
    "excited = oqs.basis(atom, \"up\")\n",
    "ground = oqs.basis(atom, \"down\")\n",
    "rho0 = oqs.ket2dm(excited)\n",
    "times = np.linspace(0.0, 8.0, 121)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "state-metrics-mesolve-run",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = oqs.state_metrics(\n",
    "    purity=True,\n",
    "    entropy=True,\n",
    "    linear_entropy=True,\n",
    "    participation_ratio=True,\n",
    "    fidelity_to=excited,\n",
    "    trace_distance_to=ground,\n",
    "    population_indices=[0, 1],\n",
    "    bloch_vector=True,\n",
    ")\n",
    "\n",
    "result = oqs.mesolve(\n",
    "    H,\n",
    "    rho0,\n",
    "    times,\n",
    "    c_ops=[collapse],\n",
    "    state_observables=metrics,\n",
    "    options=oqs.Options(rtol=1e-9, atol=1e-11, save_states=False),\n",
    ")\n",
    "\n",
    "sorted(result.state_observables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "state-metrics-mesolve-plot",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs = result.state_observables\n",
    "\n",
    "fig, axes = plt.subplots(2, 1, figsize=(7, 6), sharex=True)\n",
    "axes[0].plot(times, obs[\"population_0\"].real, label=\"excited population\")\n",
    "axes[0].plot(times, obs[\"population_1\"].real, label=\"ground population\")\n",
    "axes[0].plot(times, obs[\"fidelity\"].real, \"--\", label=\"fidelity to initial\")\n",
    "axes[0].set_ylabel(\"probability\")\n",
    "axes[0].legend()\n",
    "\n",
    "axes[1].plot(times, obs[\"purity\"].real, label=\"purity\")\n",
    "axes[1].plot(times, obs[\"entropy\"].real, label=\"entropy\")\n",
    "axes[1].plot(times, obs[\"linear_entropy\"].real, label=\"linear entropy\")\n",
    "axes[1].set_xlabel(\"time\")\n",
    "axes[1].legend()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "state-metrics-trajectory-note",
   "metadata": {},
   "source": [
    "The same callback dictionary works for a single Monte Carlo wave-function trajectory. A trajectory stays pure, but the populations and fidelity change discontinuously when a jump occurs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "state-metrics-trajectory-run",
   "metadata": {},
   "outputs": [],
   "source": [
    "trajectory_metrics = oqs.state_metrics(\n",
    "    purity=True,\n",
    "    fidelity_to=excited,\n",
    "    population_indices=[0, 1],\n",
    ")\n",
    "\n",
    "trajectory = oqs.single_trajectory(\n",
    "    H,\n",
    "    excited,\n",
    "    times,\n",
    "    c_ops=[collapse],\n",
    "    state_observables=trajectory_metrics,\n",
    "    options=oqs.Options(seed=2026, max_step=0.02, save_states=False),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "state-metrics-trajectory-plot",
   "metadata": {},
   "outputs": [],
   "source": [
    "traj_obs = trajectory.state_observables\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 3.5))\n",
    "ax.step(times, traj_obs[\"population_0\"].real, where=\"post\", label=\"excited\")\n",
    "ax.step(times, traj_obs[\"population_1\"].real, where=\"post\", label=\"ground\")\n",
    "ax.plot(times, traj_obs[\"purity\"].real, \"--\", label=\"purity\")\n",
    "ax.set_xlabel(\"time\")\n",
    "ax.set_ylabel(\"trajectory metric\")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "state-metrics-custom-note",
   "metadata": {},
   "source": [
    "You can mix built-in callbacks with custom scalar diagnostics. This is useful for problem-specific return probabilities, subsystem entropies, or distances to reference states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "state-metrics-custom",
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_metrics = oqs.state_metrics(purity=True, fidelity_to=excited)\n",
    "custom_metrics[\"ground_infidelity\"] = lambda rho: oqs.infidelity(rho, ground)\n",
    "\n",
    "custom_result = oqs.mesolve(\n",
    "    H,\n",
    "    rho0,\n",
    "    times,\n",
    "    c_ops=[collapse],\n",
    "    state_observables=custom_metrics,\n",
    "    options=oqs.Options(rtol=1e-9, atol=1e-11),\n",
    ")\n",
    "\n",
    "custom_result.state_observables[\"ground_infidelity\"].real[-1]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}