jupyterlab-codex-sidebar 0.1.4 → 0.1.6

package/tests/e2e/fixtures/notebooks/tab2.ipynb

@@ -0,0 +1,252 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "tab2-cell",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "\n",
+    "def simulate_schrodinger_1d(\n",
+    "    n_grid: int = 1024,\n",
+    "    x_min: float = -20.0,\n",
+    "    x_max: float = 20.0,\n",
+    "    dt: float = 0.005,\n",
+    "    steps: int = 2500,\n",
+    "    save_every: int = 250,\n",
+    "    potential_type: str = \"harmonic\",\n",
+    "    omega: float = 0.2,\n",
+    "    barrier_height: float = 1.5,\n",
+    "    barrier_width: float = 1.5,\n",
+    "    x0: float = -7.0,\n",
+    "    sigma: float = 1.0,\n",
+    "    k0: float = 2.0,\n",
+    ") -> dict[str, object]:\n",
+    "    \"\"\"\n",
+    "    Simulate a 1D wave packet under the time-dependent Schrodinger equation\n",
+    "    using a split-operator FFT method in natural units (hbar = m = 1).\n",
+    "    \"\"\"\n",
+    "    if not isinstance(n_grid, int) or isinstance(n_grid, bool):\n",
+    "        raise TypeError(\"n_grid must be int\")\n",
+    "    if n_grid < 64:\n",
+    "        raise ValueError(\"n_grid must be >= 64\")\n",
+    "    if x_max <= x_min:\n",
+    "        raise ValueError(\"x_max must be greater than x_min\")\n",
+    "    if dt <= 0:\n",
+    "        raise ValueError(\"dt must be > 0\")\n",
+    "    if not isinstance(steps, int) or isinstance(steps, bool) or steps < 1:\n",
+    "        raise ValueError(\"steps must be an int >= 1\")\n",
+    "    if not isinstance(save_every, int) or isinstance(save_every, bool) or save_every < 1:\n",
+    "        raise ValueError(\"save_every must be an int >= 1\")\n",
+    "    if sigma <= 0:\n",
+    "        raise ValueError(\"sigma must be > 0\")\n",
+    "\n",
+    "    x = np.linspace(x_min, x_max, n_grid, endpoint=False)\n",
+    "    dx = x[1] - x[0]\n",
+    "\n",
+    "    if potential_type == \"harmonic\":\n",
+    "        if omega < 0:\n",
+    "            raise ValueError(\"omega must be >= 0\")\n",
+    "        potential = 0.5 * (omega**2) * (x**2)\n",
+    "    elif potential_type == \"barrier\":\n",
+    "        if barrier_width < 0:\n",
+    "            raise ValueError(\"barrier_width must be >= 0\")\n",
+    "        potential = np.where(np.abs(x) < barrier_width / 2.0, barrier_height, 0.0)\n",
+    "    else:\n",
+    "        raise ValueError(\"potential_type must be 'harmonic' or 'barrier'\")\n",
+    "\n",
+    "    psi = np.exp(-((x - x0) ** 2) / (2.0 * sigma**2)) * np.exp(1j * k0 * x)\n",
+    "\n",
+    "    def norm(wavefunc: np.ndarray) -> float:\n",
+    "        return float(np.sum(np.abs(wavefunc) ** 2) * dx)\n",
+    "\n",
+    "    def expected_x(wavefunc: np.ndarray) -> float:\n",
+    "        density = np.abs(wavefunc) ** 2\n",
+    "        return float(np.sum(x * density) * dx)\n",
+    "\n",
+    "    def total_energy(wavefunc: np.ndarray) -> float:\n",
+    "        grad = np.gradient(wavefunc, dx)\n",
+    "        kinetic = 0.5 * np.sum(np.abs(grad) ** 2) * dx\n",
+    "        potential_energy = np.sum(potential * (np.abs(wavefunc) ** 2)) * dx\n",
+    "        return float(np.real(kinetic + potential_energy))\n",
+    "\n",
+    "    psi /= np.sqrt(norm(psi))\n",
+    "    k = 2.0 * np.pi * np.fft.fftfreq(n_grid, d=dx)\n",
+    "    kinetic_phase = np.exp(-0.5j * (k**2) * dt)\n",
+    "    potential_half_phase = np.exp(-0.5j * potential * dt)\n",
+    "\n",
+    "    snapshots: list[tuple[float, np.ndarray]] = [(0.0, np.abs(psi) ** 2)]\n",
+    "    norm_history = [norm(psi)]\n",
+    "    x_expect_history = [expected_x(psi)]\n",
+    "    energy_history = [total_energy(psi)]\n",
+    "\n",
+    "    for step in range(1, steps + 1):\n",
+    "        psi = potential_half_phase * psi\n",
+    "        psi_k = np.fft.fft(psi)\n",
+    "        psi_k = kinetic_phase * psi_k\n",
+    "        psi = np.fft.ifft(psi_k)\n",
+    "        psi = potential_half_phase * psi\n",
+    "        psi /= np.sqrt(norm(psi))\n",
+    "\n",
+    "        current_norm = norm(psi)\n",
+    "        norm_history.append(current_norm)\n",
+    "        x_expect_history.append(expected_x(psi))\n",
+    "        energy_history.append(total_energy(psi))\n",
+    "\n",
+    "        if step % save_every == 0:\n",
+    "            snapshots.append((step * dt, np.abs(psi) ** 2))\n",
+    "\n",
+    "    if steps % save_every != 0:\n",
+    "        snapshots.append((steps * dt, np.abs(psi) ** 2))\n",
+    "\n",
+    "    return {\n",
+    "        \"x\": x,\n",
+    "        \"V\": potential,\n",
+    "        \"psi\": psi,\n",
+    "        \"dx\": dx,\n",
+    "        \"dt\": dt,\n",
+    "        \"steps\": steps,\n",
+    "        \"snapshots\": snapshots,\n",
+    "        \"final_norm\": norm(psi),\n",
+    "        \"norm_history\": np.array(norm_history),\n",
+    "        \"x_expect_history\": np.array(x_expect_history),\n",
+    "        \"energy_history\": np.array(energy_history),\n",
+    "    }\n",
+    "\n",
+    "\n",
+    "def validate_simulation(\n",
+    "    result: dict[str, object],\n",
+    "    norm_tolerance: float = 1e-6,\n",
+    "    energy_tolerance: float = 0.6,\n",
+    ") -> None:\n",
+    "    norms = np.asarray(result[\"norm_history\"], dtype=float)\n",
+    "    if not np.all(np.isclose(norms, 1.0, atol=norm_tolerance)):\n",
+    "        min_norm = float(np.min(norms))\n",
+    "        max_norm = float(np.max(norms))\n",
+    "        raise AssertionError(\n",
+    "            f\"norm preservation failed: min={min_norm:.8f}, max={max_norm:.8f}\"\n",
+    "        )\n",
+    "    energies = np.asarray(result[\"energy_history\"], dtype=float)\n",
+    "    if energies.size > 1:\n",
+    "        drift = float(np.max(energies) - np.min(energies))\n",
+    "        if drift > energy_tolerance:\n",
+    "            raise AssertionError(f\"energy drift too large: {drift:.6f}\")\n",
+    "\n",
+    "\n",
+    "def transmission_reflection_probabilities(\n",
+    "    result: dict[str, object],\n",
+    "    split_x: float = 0.0,\n",
+    ") -> tuple[float, float]:\n",
+    "    x = np.asarray(result[\"x\"], dtype=float)\n",
+    "    psi = np.asarray(result[\"psi\"], dtype=np.complex128)\n",
+    "    dx = float(result[\"dx\"])\n",
+    "    density = np.abs(psi) ** 2\n",
+    "    reflection = float(np.sum(density[x < split_x]) * dx)\n",
+    "    transmission = float(np.sum(density[x >= split_x]) * dx)\n",
+    "    return transmission, reflection\n",
+    "\n",
+    "\n",
+    "def plot_simulation(result: dict[str, object]) -> None:\n",
+    "    import matplotlib.pyplot as plt\n",
+    "\n",
+    "    x = np.asarray(result[\"x\"])\n",
+    "    potential = np.asarray(result[\"V\"])\n",
+    "    snapshots = result[\"snapshots\"]\n",
+    "\n",
+    "    max_density = max(float(np.max(density)) for _, density in snapshots)\n",
+    "    shifted = potential - float(np.min(potential))\n",
+    "    vmax = float(np.max(shifted))\n",
+    "    if vmax > 0:\n",
+    "        scaled_potential = (shifted / vmax) * (0.8 * max_density)\n",
+    "    else:\n",
+    "        scaled_potential = np.zeros_like(potential)\n",
+    "\n",
+    "    plt.figure(figsize=(10, 5))\n",
+    "    plt.plot(x, scaled_potential, \"k--\", linewidth=1.2, label=\"Scaled potential\")\n",
+    "    for time_point, density in snapshots:\n",
+    "        plt.plot(x, density, label=f\"t = {time_point:.2f}\")\n",
+    "    plt.title(\"1D Time-Dependent Schrodinger Simulation\")\n",
+    "    plt.xlabel(\"x\")\n",
+    "    plt.ylabel(r\"Probability density $|\\psi(x,t)|^2$\")\n",
+    "    plt.legend(ncol=2)\n",
+    "    plt.grid(alpha=0.25)\n",
+    "    plt.tight_layout()\n",
+    "    plt.show()\n",
+    "\n",
+    "\n",
+    "def run_demo(enable_plot: bool = False) -> None:\n",
+    "    harmonic_result = simulate_schrodinger_1d(\n",
+    "        potential_type=\"harmonic\",\n",
+    "        omega=0.2,\n",
+    "        x0=-7.0,\n",
+    "        k0=2.0,\n",
+    "    )\n",
+    "    validate_simulation(harmonic_result)\n",
+    "\n",
+    "    final_norm = float(harmonic_result[\"final_norm\"])\n",
+    "    x_expect = np.asarray(harmonic_result[\"x_expect_history\"], dtype=float)\n",
+    "    print(\"Harmonic oscillator\")\n",
+    "    print(f\"Final norm: {final_norm:.6f}\")\n",
+    "    print(\n",
+    "        \"Expectation x range:\",\n",
+    "        f\"[{float(np.min(x_expect)):.3f}, {float(np.max(x_expect)):.3f}]\",\n",
+    "    )\n",
+    "\n",
+    "    barrier_result = simulate_schrodinger_1d(\n",
+    "        potential_type=\"barrier\",\n",
+    "        barrier_height=1.5,\n",
+    "        barrier_width=1.5,\n",
+    "        x0=-10.0,\n",
+    "        k0=2.2,\n",
+    "        steps=3000,\n",
+    "    )\n",
+    "    validate_simulation(barrier_result)\n",
+    "    transmission, reflection = transmission_reflection_probabilities(\n",
+    "        barrier_result,\n",
+    "        split_x=0.0,\n",
+    "    )\n",
+    "    assert np.isclose(transmission + reflection, 1.0, atol=1e-5)\n",
+    "    print(\"Barrier tunneling\")\n",
+    "    print(\n",
+    "        \"Transmission / Reflection:\",\n",
+    "        f\"{transmission:.4f} / {reflection:.4f}\",\n",
+    "    )\n",
+    "    print(\"Validation passed.\")\n",
+    "\n",
+    "    if enable_plot:\n",
+    "        plot_simulation(harmonic_result)\n",
+    "        plot_simulation(barrier_result)\n",
+    "\n",
+    "\n",
+    "run_demo(enable_plot=False)"
+   ]
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "formats": "ipynb,py:percent"
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

package/tests/e2e/fixtures/notebooks/tab2.py

@@ -0,0 +1,231 @@
+# ---
+# jupyter:
+#   jupytext:
+#     formats: ipynb,py:percent
+#     text_representation:
+#       extension: .py
+#       format_name: percent
+#       format_version: '1.3'
+#       jupytext_version: 1.17.2
+#   kernelspec:
+#     display_name: Python 3 (ipykernel)
+#     language: python
+#     name: python3
+# ---
+
+# %%
+import numpy as np
+
+
+def simulate_schrodinger_1d(
+    n_grid: int = 1024,
+    x_min: float = -20.0,
+    x_max: float = 20.0,
+    dt: float = 0.005,
+    steps: int = 2500,
+    save_every: int = 250,
+    potential_type: str = "harmonic",
+    omega: float = 0.2,
+    barrier_height: float = 1.5,
+    barrier_width: float = 1.5,
+    x0: float = -7.0,
+    sigma: float = 1.0,
+    k0: float = 2.0,
+) -> dict[str, object]:
+    """
+    Simulate a 1D wave packet under the time-dependent Schrodinger equation
+    using a split-operator FFT method in natural units (hbar = m = 1).
+    """
+    if not isinstance(n_grid, int) or isinstance(n_grid, bool):
+        raise TypeError("n_grid must be int")
+    if n_grid < 64:
+        raise ValueError("n_grid must be >= 64")
+    if x_max <= x_min:
+        raise ValueError("x_max must be greater than x_min")
+    if dt <= 0:
+        raise ValueError("dt must be > 0")
+    if not isinstance(steps, int) or isinstance(steps, bool) or steps < 1:
+        raise ValueError("steps must be an int >= 1")
+    if not isinstance(save_every, int) or isinstance(save_every, bool) or save_every < 1:
+        raise ValueError("save_every must be an int >= 1")
+    if sigma <= 0:
+        raise ValueError("sigma must be > 0")
+
+    x = np.linspace(x_min, x_max, n_grid, endpoint=False)
+    dx = x[1] - x[0]
+
+    if potential_type == "harmonic":
+        if omega < 0:
+            raise ValueError("omega must be >= 0")
+        potential = 0.5 * (omega**2) * (x**2)
+    elif potential_type == "barrier":
+        if barrier_width < 0:
+            raise ValueError("barrier_width must be >= 0")
+        potential = np.where(np.abs(x) < barrier_width / 2.0, barrier_height, 0.0)
+    else:
+        raise ValueError("potential_type must be 'harmonic' or 'barrier'")
+
+    psi = np.exp(-((x - x0) ** 2) / (2.0 * sigma**2)) * np.exp(1j * k0 * x)
+
+    def norm(wavefunc: np.ndarray) -> float:
+        return float(np.sum(np.abs(wavefunc) ** 2) * dx)
+
+    def expected_x(wavefunc: np.ndarray) -> float:
+        density = np.abs(wavefunc) ** 2
+        return float(np.sum(x * density) * dx)
+
+    def total_energy(wavefunc: np.ndarray) -> float:
+        grad = np.gradient(wavefunc, dx)
+        kinetic = 0.5 * np.sum(np.abs(grad) ** 2) * dx
+        potential_energy = np.sum(potential * (np.abs(wavefunc) ** 2)) * dx
+        return float(np.real(kinetic + potential_energy))
+
+    psi /= np.sqrt(norm(psi))
+    k = 2.0 * np.pi * np.fft.fftfreq(n_grid, d=dx)
+    kinetic_phase = np.exp(-0.5j * (k**2) * dt)
+    potential_half_phase = np.exp(-0.5j * potential * dt)
+
+    snapshots: list[tuple[float, np.ndarray]] = [(0.0, np.abs(psi) ** 2)]
+    norm_history = [norm(psi)]
+    x_expect_history = [expected_x(psi)]
+    energy_history = [total_energy(psi)]
+
+    for step in range(1, steps + 1):
+        psi = potential_half_phase * psi
+        psi_k = np.fft.fft(psi)
+        psi_k = kinetic_phase * psi_k
+        psi = np.fft.ifft(psi_k)
+        psi = potential_half_phase * psi
+        psi /= np.sqrt(norm(psi))
+
+        current_norm = norm(psi)
+        norm_history.append(current_norm)
+        x_expect_history.append(expected_x(psi))
+        energy_history.append(total_energy(psi))
+
+        if step % save_every == 0:
+            snapshots.append((step * dt, np.abs(psi) ** 2))
+
+    if steps % save_every != 0:
+        snapshots.append((steps * dt, np.abs(psi) ** 2))
+
+    return {
+        "x": x,
+        "V": potential,
+        "psi": psi,
+        "dx": dx,
+        "dt": dt,
+        "steps": steps,
+        "snapshots": snapshots,
+        "final_norm": norm(psi),
+        "norm_history": np.array(norm_history),
+        "x_expect_history": np.array(x_expect_history),
+        "energy_history": np.array(energy_history),
+    }
+
+
+def validate_simulation(
+    result: dict[str, object],
+    norm_tolerance: float = 1e-6,
+    energy_tolerance: float = 0.6,
+) -> None:
+    norms = np.asarray(result["norm_history"], dtype=float)
+    if not np.all(np.isclose(norms, 1.0, atol=norm_tolerance)):
+        min_norm = float(np.min(norms))
+        max_norm = float(np.max(norms))
+        raise AssertionError(
+            f"norm preservation failed: min={min_norm:.8f}, max={max_norm:.8f}"
+        )
+    energies = np.asarray(result["energy_history"], dtype=float)
+    if energies.size > 1:
+        drift = float(np.max(energies) - np.min(energies))
+        if drift > energy_tolerance:
+            raise AssertionError(f"energy drift too large: {drift:.6f}")
+
+
+def transmission_reflection_probabilities(
+    result: dict[str, object],
+    split_x: float = 0.0,
+) -> tuple[float, float]:
+    x = np.asarray(result["x"], dtype=float)
+    psi = np.asarray(result["psi"], dtype=np.complex128)
+    dx = float(result["dx"])
+    density = np.abs(psi) ** 2
+    reflection = float(np.sum(density[x < split_x]) * dx)
+    transmission = float(np.sum(density[x >= split_x]) * dx)
+    return transmission, reflection
+
+
+def plot_simulation(result: dict[str, object]) -> None:
+    import matplotlib.pyplot as plt
+
+    x = np.asarray(result["x"])
+    potential = np.asarray(result["V"])
+    snapshots = result["snapshots"]
+
+    max_density = max(float(np.max(density)) for _, density in snapshots)
+    shifted = potential - float(np.min(potential))
+    vmax = float(np.max(shifted))
+    if vmax > 0:
+        scaled_potential = (shifted / vmax) * (0.8 * max_density)
+    else:
+        scaled_potential = np.zeros_like(potential)
+
+    plt.figure(figsize=(10, 5))
+    plt.plot(x, scaled_potential, "k--", linewidth=1.2, label="Scaled potential")
+    for time_point, density in snapshots:
+        plt.plot(x, density, label=f"t = {time_point:.2f}")
+    plt.title("1D Time-Dependent Schrodinger Simulation")
+    plt.xlabel("x")
+    plt.ylabel(r"Probability density $|\psi(x,t)|^2$")
+    plt.legend(ncol=2)
+    plt.grid(alpha=0.25)
+    plt.tight_layout()
+    plt.show()
+
+
+def run_demo(enable_plot: bool = False) -> None:
+    harmonic_result = simulate_schrodinger_1d(
+        potential_type="harmonic",
+        omega=0.2,
+        x0=-7.0,
+        k0=2.0,
+    )
+    validate_simulation(harmonic_result)
+
+    final_norm = float(harmonic_result["final_norm"])
+    x_expect = np.asarray(harmonic_result["x_expect_history"], dtype=float)
+    print("Harmonic oscillator")
+    print(f"Final norm: {final_norm:.6f}")
+    print(
+        "Expectation x range:",
+        f"[{float(np.min(x_expect)):.3f}, {float(np.max(x_expect)):.3f}]",
+    )
+
+    barrier_result = simulate_schrodinger_1d(
+        potential_type="barrier",
+        barrier_height=1.5,
+        barrier_width=1.5,
+        x0=-10.0,
+        k0=2.2,
+        steps=3000,
+    )
+    validate_simulation(barrier_result)
+    transmission, reflection = transmission_reflection_probabilities(
+        barrier_result,
+        split_x=0.0,
+    )
+    assert np.isclose(transmission + reflection, 1.0, atol=1e-5)
+    print("Barrier tunneling")
+    print(
+        "Transmission / Reflection:",
+        f"{transmission:.4f} / {reflection:.4f}",
+    )
+    print("Validation passed.")
+
+    if enable_plot:
+        plot_simulation(harmonic_result)
+        plot_simulation(barrier_result)
+
+
+run_demo(enable_plot=False)