fisher_information.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Revise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import QuantumCollocation as QC\n",
    "import NamedTrajectories as NT\n",
    "import TrajectoryIndexingUtils as NTidx\n",
    "import Plots\n",
    "import LinearAlgebra as LA\n",
    "import SparseArrays as SA\n",
    "import ForwardDiff as FD\n",
    "using LaTeXStrings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "include(\"utils.jl\")\n",
    "include(\"system.jl\")\n",
    "include(\"objectives.jl\")\n",
    "include(\"constraints.jl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = 10.\n",
    "p_max = 5\n",
    "# E_R [kHz] found in Weidner thesis\n",
    "system = ShakenLatticeSystem1D(V, p_max; acc=0.0, include_acc_derivative=true, sparse=false)\n",
    "# middle index of statevector where p = 0\n",
    "mid = system.params[:mid]\n",
    "dim = system.params[:dim]\n",
    "#E_R = system.params[:E_R]\n",
    "#E_R = 1/0.05\n",
    "#print(\"time unit $(1/E_R) ms\\nE_R = $E_R kHz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system.H_drift_real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "duration = 2pi * 1.0 # in units of 1/E_R\n",
    "\n",
    "T = 301\n",
    "dt = duration / (T-1)\n",
    "dts = zeros(T) .+ dt\n",
    "dt_bound = (dt, dt)\n",
    "times = cumsum(dts) - dts;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_bound = fill(1., 2)\n",
    "dda_bound = fill(1000., 2)\n",
    "\n",
    "#a = vcat(ones(T)', zeros(T)') \n",
    "#a = vcat(cos.(2pi * times/duration)', sin.(2pi * times/duration)')\n",
    "#a = vcat(cos.(1. .+ 0.5*cos.(2π* 12. *times))', sin.(1. .+ 0.5*cos.(2π* 12. *times))')\n",
    "\n",
    "# phi_guess = 0.5 * (cos.(2pi * 4. * times) + cos.(2pi * 12. * times))\n",
    "# a = vcat(cos.(phi_guess)', sin.(phi_guess)')\n",
    "\n",
    "a = vcat(cos.(2. *sin.(11.5 *times))', sin.(2. *sin.(11.5 *times))')\n",
    "\n",
    "acc = collect(times')\n",
    "\n",
    "da = NT.derivative(a, dts)\n",
    "da[end, :] .= 1.\n",
    "# da[:, end] = da[:, end-1]\n",
    "dda = NT.derivative(da, dts)\n",
    "# dda[:, end] = dda[:, end-1] = dda[:, end-2]\n",
    "dda[end, :] .= 0.;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Z_split = NT.load_traj(\"interferometer/split_victor.jld2\")\n",
    "Z_mirror = NT.load_traj(\"interferometer/mirror_victor.jld2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z_split = NT.load_traj(\"interferometer/split_victor_opt2.jld2\")\n",
    "Z_mirror = NT.load_traj(\"interferometer/mirror_victor2_opt2.jld2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a, dts = get_interferometer(Z_split, Z_mirror, [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z = NT.load_traj(\"./interferometer/save.jld2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = Z.a\n",
    "dts = vec(Z.dts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z_split.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a[:,176:end] .= [1., 0.]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(dts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "duration = sum(dts) # in units of 1/E_R\n",
    "\n",
    "T = length(dts)\n",
    "dt = duration / (T-1)\n",
    "dts = zeros(T) .+ dt\n",
    "dt_bound = (dt, dt)\n",
    "times = cumsum(dts) - dts;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_bound = fill(1., 2)\n",
    "dda_bound = fill(1000., 2)\n",
    "\n",
    "acc = collect(times')\n",
    "\n",
    "da = NT.derivative(a, dts)\n",
    "# da[:, end] = da[:, end-1]\n",
    "dda = NT.derivative(da, dts)\n",
    "# dda[:, end] = dda[:, end-1] = dda[:, end-2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times, a')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times, dda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phi_guess = angle.(a[1,:] + im*a[2,:])\n",
    "phi_mod_clean!(phi_guess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times, phi_guess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi0 = get_bloch_state(system; lvl=0)\n",
    "if system.params[:accelerated]\n",
    "    append!(psi0, zeros(dim))\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.bar(-p_max:p_max, abs2.(psi0[1:dim]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi0_iso = QC.ket_to_iso(psi0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bloch_states = hcat([get_bloch_state(system; lvl=i) for i=0:dim-1]...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "time_flight = 2pi * 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z_mirror.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jumps = [(Z_split.T, time_flight), (Z_split.T+Z_mirror.T, time_flight)]\n",
    "#jumps = Tuple{Int, Float64}[]\n",
    "cuts = [jump[1] for jump in jumps]\n",
    "full_times = get_times(dts, jumps)\n",
    "G = get_shaken_lattice_propagator(system, times, jumps, 10000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "jumps = [(100, time_flight), (200, time_flight)]\n",
    "#jumps = Tuple{Int, Float64}[]\n",
    "cuts = [jump[1] for jump in jumps]\n",
    "full_times = get_times(dts, jumps)\n",
    "G = get_shaken_lattice_propagator(system, times, jumps, 10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "U = [QC.iso_vec_to_operator(QC.iso_operator_to_iso_vec(g)) for g in G]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.heatmap(sqrt.(abs.(U[1])), yflip=true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.heatmap(sqrt.(abs.(U[2])), yflip=true)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_iso = shaken_lattice_rollout(psi0_iso, a, dts, system, jumps, G)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "psi_iso = QC.rollout(psi0_iso, a, dts, system)#; integrator=exp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bloch_states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "B = blockdiagonal(bloch_states, bloch_states)'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi = hcat([QC.iso_to_ket(psi_iso[:,t]) for t=1:T]...)\n",
    "psi = B * psi\n",
    "pops = abs2.(psi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function format_plot(\n",
    "    p,\n",
    "    times=times,\n",
    "    jumps=jumps,\n",
    "    full_times=full_times,\n",
    ")\n",
    "    T = length(times)\n",
    "    Plots.xticks!(p, (times[1:div(T,10):end], string.(round.(full_times[1:div(T,10):end]; digits=1))))\n",
    "    Plots.xlims!(p, (times[1], times[end]))\n",
    "    rel_y = 0.98\n",
    "    for (cut, jump_time) in jumps\n",
    "        rel_x = times[cut] / times[end]\n",
    "        Plots.vline!(p, [times[cut]], color=:red, label=nothing)\n",
    "        Plots.annotate!(((rel_x, rel_y), (string(round(jump_time; digits=1)), 10, :red)))\n",
    "    end\n",
    "    Plots.xlabel!(p, L\"t \\, [1/\\omega_R]\")\n",
    "    return p\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Plots.plot(times, pops[1:dim,:]', labels=(-p_max:1:p_max)', legend=:outertopright, size=(500,300))\n",
    "Plots.plot!(\n",
    "    p,\n",
    "    title=L\"Bloch state population evolution $|B\\rangle$\"\n",
    ")\n",
    "format_plot(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times, pops[dim+1:end,:]', labels=(-p_max:1:p_max)', legend=:outertopright)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comps = (\n",
    "    psi_iso = psi_iso,\n",
    "    a = a,\n",
    "    acc = acc,\n",
    "    da = da,\n",
    "    dda = dda,\n",
    "    dts = dts\n",
    ")\n",
    "initial = (\n",
    "    psi_iso = psi0_iso,\n",
    "    a = [1.; 0.],\n",
    "    #da = da0\n",
    ")\n",
    "final = (;\n",
    "    a = [1.; 0.],\n",
    "    #da = zeros(2)\n",
    ")\n",
    "goal = (;)\n",
    "bounds = (\n",
    "    #a = a_bound,\n",
    "    dda = dda_bound,\n",
    "    dts = dt_bound,\n",
    ")\n",
    "\n",
    "Z_guess = nothing\n",
    "GC.gc()\n",
    "Z_guess = NT.NamedTrajectory(\n",
    "    comps;\n",
    "    controls=(:dda),\n",
    "    timestep=:dts,\n",
    "    bounds=bounds,\n",
    "    initial=initial,\n",
    "    final=final,\n",
    "    goal=goal\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z_guess = copy(Z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "pops_goal = abs2.(get_bloch_state(system; lvl=0))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "dim = system.params[:dim]\n",
    "J = NameComponentPathObjective(\n",
    "    :psi_iso,\n",
    "    1:dim,\n",
    "    [T],\n",
    "    [time -> pop for pop in pops_goal],\n",
    "    [x -> abs2.(x) for _=1:dim],\n",
    "    fill(1., dim, 1);\n",
    "    is_iso=true\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MZFI = (8pi*(full_times[end]/2)^2)^2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function Fisher(psi, dpsi)\n",
    "    eps = 0.0\n",
    "    P = abs2.(psi)\n",
    "    D = 2*real.(conj.(psi) .* dpsi)\n",
    "    F = (1 ./ (P .+ eps))' * D.^2\n",
    "    return F / MZFI\n",
    "end \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fisher_loss = psi_dpsi -> -Fisher(psi_dpsi[1:dim], psi_dpsi[dim+1:2*dim])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "-fisher_loss(QC.iso_to_ket(psi_iso[:,end]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_sensitivity_loss = psi_dpsi -> -0.5 * log10(Fisher(psi_dpsi[1:dim], psi_dpsi[dim+1:2*dim]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_sensitivity_loss(QC.iso_to_ket(psi_iso[:,end]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fisher_evol = [-fisher_loss(QC.iso_to_ket(psi_iso[:,t])) for t=1:T];\n",
    "p = Plots.plot(times, fisher_evol)\n",
    "Plots.vline!(p, times[cuts], color=:red, label=\"cut\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "flight_times = 2pi * collect(1:40)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#fishers = []\n",
    "for flight_time in flight_times\n",
    "    println(flight_time)\n",
    "    jumps = [(211, flight_time), (632, flight_time)]\n",
    "    cuts = [jump[1] for jump in jumps]\n",
    "    full_times = get_times(dts, jumps)\n",
    "    G = get_shaken_lattice_propagator(system, times, jumps, 10000)\n",
    "    psi_iso = shaken_lattice_rollout(psi0_iso, a, dts, system, jumps, G)\n",
    "    push!(fishers, fisher_loss(psi_iso[:,end]))\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(flight_times, hcat(-fishers, (4pi)^2 * flight_times.^4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#J = QC.QuantumObjective(name=:psi_iso, goals=QC.ket_to_iso(vcat(get_bloch_state(system; lvl=0), zeros(dim))), loss=:InfidelityLoss, Q=1e2)\n",
    "#J = QC.QuantumObjective(name=:psi_iso, goals=QC.ket_to_iso(get_bloch_state(system; lvl=3)), loss=:InfidelityLoss, Q=1e2)\n",
    "\n",
    "# J += NameComponentObjective(\n",
    "#     :psi_iso,\n",
    "#     [1:2*dim...],\n",
    "#     [1:T...],\n",
    "#     fisher_loss,\n",
    "#     ([1:T...] ./ T).^2;\n",
    "#     is_iso=true\n",
    "# )\n",
    "\n",
    "J = NameComponentObjective(\n",
    "    :psi_iso,\n",
    "    [1:2*dim...],\n",
    "    [T],\n",
    "    fisher_loss;\n",
    "    is_iso=true\n",
    ")\n",
    "\n",
    "J += QC.QuadraticRegularizer(:dda, Z_guess, 1e-8/T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "J.L(Z_guess.datavec, Z_guess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "integrators = nothing\n",
    "GC.gc()\n",
    "integrators = [\n",
    "    QC.QuantumStatePadeIntegrator(\n",
    "        system,\n",
    "        :psi_iso,\n",
    "        (:a, :acc),\n",
    "        :dts;\n",
    "        order=4\n",
    "    ),\n",
    "    QC.DerivativeIntegrator(\n",
    "        :a,\n",
    "        :da,\n",
    "        :dts,\n",
    "        Z_guess\n",
    "    ),\n",
    "    QC.DerivativeIntegrator(\n",
    "        :da,\n",
    "        :dda,\n",
    "        :dts,\n",
    "        Z_guess\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dynamics = QC.QuantumDynamics(\n",
    "    integrators,\n",
    "    Z_guess;\n",
    "    cuts=cuts\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "constraints = [\n",
    "    OmegaAbsConstraint(1.0, Z_guess, Z_guess.components[:a]),\n",
    "    vcat([get_link_constraints(\n",
    "        :psi_iso, \n",
    "        Z_guess, \n",
    "        c, \n",
    "        g, \n",
    "        (; a=[1.0,0.0]), \n",
    "        (; a=[1.0,0.0]); \n",
    "        hard_equality_constraint=true)\n",
    "        for (c, g) in zip(cuts, G)]...)...,\n",
    "    TimeAffineLinearControlConstraint(:acc, 1, Z_guess; jumps=jumps),\n",
    "    #custom_bounds_constraint(:a, Z_guess, vcat(cuts, cuts .+ 1), a_bound)\n",
    "    # NameComponentPathConstraint(\n",
    "    #     :a,\n",
    "    #     [3],\n",
    "    #     Z_guess,\n",
    "    #     [t -> t],\n",
    "    #     [x -> x]\n",
    "    # )\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ipopt options\n",
    "options = QC.Options(\n",
    "    max_iter=200,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# defining quantum control problem\n",
    "prob = nothing\n",
    "GC.gc()\n",
    "prob = QC.QuantumControlProblem(\n",
    "    system, \n",
    "    Z_guess, \n",
    "    J, \n",
    "    dynamics;\n",
    "    constraints=constraints,\n",
    "    options=options,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "QC.solve!(prob)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z = prob.trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z.psi_iso"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#psi_iso_rollout = QC.rollout(psi0_iso, vcat(Z.a, Z.acc'), dts, system; integrator=exp)\n",
    "#psi_iso_rollout = Z.psi_iso\n",
    "psi_iso_rollout = shaken_lattice_rollout(psi0_iso, Z.a, dts, system, jumps, G; integrator=exp)\n",
    "psi = hcat([QC.iso_to_ket(psi_iso_rollout[:,t]) for t=1:T]...)\n",
    "pops = abs2.(psi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_bloch = blockdiagonal(bloch_states', bloch_states') * psi\n",
    "pops_bloch = abs2.(psi_bloch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function format_plot(\n",
    "    p,\n",
    "    times=times,\n",
    "    jumps=jumps,\n",
    "    full_times=full_times,\n",
    ")\n",
    "    T = length(times)\n",
    "    Plots.xticks!(p, (times[1:div(T,10):end], string.(round.(full_times[1:div(T,10):end]; digits=1))))\n",
    "    Plots.xlims!(p, (times[1], times[end]))\n",
    "    rel_y = 0.98\n",
    "    for (cut, jump_time) in jumps\n",
    "        rel_x = times[cut] / times[end]\n",
    "        Plots.vline!(p, [times[cut]], color=:red, label=nothing)\n",
    "        Plots.annotate!(((rel_x, rel_y), (string(round(jump_time; digits=1)), 10, :red)))\n",
    "    end\n",
    "    Plots.xlabel!(p, L\"t \\, [1/\\omega_R]\")\n",
    "    return p\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Plots.plot(times, pops[1:dim,:]', labels=(-p_max:1:p_max)', legend=:outertopright)\n",
    "Plots.plot!(\n",
    "    p,\n",
    "    title=L\"Momentum state population evolution $|n\\rangle$\"\n",
    ")\n",
    "format_plot(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Plots.plot(times, pops[dim+1:end,:]', labels=(-p_max:1:p_max)', legend=:outertopright)\n",
    "format_plot(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Plots.plot(times, pops_bloch[1:dim,:]', labels=(0:dim-1)', legend=:outertopright)\n",
    "Plots.plot!(\n",
    "    p,\n",
    "    title=L\"Bloch state population evolution $|B\\rangle$\"\n",
    ")\n",
    "format_plot(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Plots.plot(times, pops_bloch[dim+1:end,:]', labels=(0:dim-1)', legend=:outertopright)\n",
    "format_plot(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times, Z.a')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times, Z.dda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times, Z.acc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phi = angle.(Z.a[1,:] + 1im * Z.a[2,:])\n",
    "phi_mod_clean!(phi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pi_ticks = LinRange(-pi, pi, 9)\n",
    "pi_lbls = [L\"-\\pi\", L\"-3\\pi/4\", L\"-\\pi/2\", L\"-\\pi/4\", L\"0\", L\"\\pi/4\", L\"\\pi/2\", L\"3\\pi/4\", L\"\\pi\"]\n",
    "p = Plots.plot(times, phi, ylim=(-pi, pi), yticks=(pi_ticks, pi_lbls), label=\"opt\")\n",
    "Plots.plot!(p, times, phi_guess, alpha=0.5, label=\"guess\")\n",
    "format_plot(p)\n",
    "Plots.plot!(\n",
    "    p,\n",
    "    title=L\"Phase protocol $\\varphi(t)$\",\n",
    "    ylabel=L\"\\varphi\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fisher_evol = [-fisher_loss(QC.iso_to_ket(psi_iso_rollout[:,t])) for t=1:T];\n",
    "p = Plots.plot(times, fisher_evol, label=nothing)\n",
    "format_plot(p)\n",
    "Plots.plot!(\n",
    "    p,\n",
    "    title=L\"Fisher information $F(a)$\",\n",
    "    xlabel=L\"t \\, [1/\\nu_R]\",\n",
    "    ylabel=L\"F(a)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_sensitivity_evol = [log_sensitivity_loss(QC.iso_to_ket(psi_iso_rollout[:,t])) for t=1:T]\n",
    "p = Plots.plot(times[2:end], log_sensitivity_evol[2:end], ylim=(-5., 5.), label=nothing)\n",
    "format_plot(p)\n",
    "Plots.plot!(\n",
    "    p,\n",
    "    title=L\"Log-Sensitivity $\\log_{10}(\\delta a) = -0.5 \\, \\log_{10}(F(a))$\",\n",
    "    ylabel=L\"\\log_{10}(\\delta a)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "F = fisher_evol[end]\n",
    "F, 1/sqrt(F)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#freqs = collect(12.464:0.001:12.468)\n",
    "freqs = collect(0.:0.1:40.0)\n",
    "phi_ft = fourier_time_freq(phi, times, freqs/2pi);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(freqs, abs2.(phi_ft))#, ylims=(0.000895, 0.0009))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import JLD2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jumps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "JLD2.save(\"./interferometer/176-5.0_352-5.0.jld2\", Z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z_guess = Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fisher range over a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc_range = LinRange(-0.001, 0.001, 81)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fisher_vals = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jumps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for acc_val in acc_range\n",
    "    println(acc_val)\n",
    "    system = ShakenLatticeSystem1D(V, p_max; acc=acc_val, include_acc_derivative=true)\n",
    "    psi_iso_final = shaken_lattice_rollout(psi0_iso, Z.a, dts, system, jumps, 10000; integrator=exp)[:,end]\n",
    "    push!(fisher_vals, -fisher_loss(psi_iso_final))\n",
    "end\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(acc_range, fisher_vals/MZFI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jumps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_full, dts_full = get_controls_dts(Z.a, vec(Z.dts), jumps, 10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "phi_full = angle.(a_full[1,:] + im*a_full[2,:])\n",
    "times_full = cumsum(dts_full) - dts_full"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times_full, phi_full)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_iso_full = shaken_lattice_rollout(psi0_iso, Z.a, vec(Z.dts), system, jumps, 10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_full = length(dts_full)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_full = hcat([QC.iso_to_ket(psi_iso_full[:,t]) for t=1:T_full]...)\n",
    "pops_full = abs2.(psi_full)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_bloch_full = bloch_states' * psi_full[1:dim,:]\n",
    "pops_bloch_full = abs2.(psi_bloch_full)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times_full, pops_full[1:dim,:]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times_full, pops_full[dim+1:end,:]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times_full, pops_bloch_full', label=(0:8)', legend=:outertopright)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Plots.plot(times, pops_bloch', labels=(0:dim-1)', legend=:outertopright)\n",
    "Plots.plot!(\n",
    "    p,\n",
    "    title=L\"Bloch state population evolution $|b\\rangle$\"\n",
    ")\n",
    "format_plot(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Z_save = copy(Z)\n",
    "Z_guess = Z_save"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## roll out repetition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function get_repeated_controls(a::AbstractMatrix, dts::AbstractVector, N::Int)\n",
    "    T = length(dts)\n",
    "    dts_long = repeat(dts, N)\n",
    "    a_long = repeat(a, 1, N)\n",
    "    times_long = cumsum(dts_long) - dts_long\n",
    "    for n=2:N\n",
    "        a_long[3,(n-1)*T+1:n*T] .+= (n-1)*times_long[T]\n",
    "    end \n",
    "    return (a_long, dts_long, times_long)\n",
    "end\n",
    "function get_repeated_controls(Z::NT.NamedTrajectory, N::Int)\n",
    "    return get_repeated_controls(Z.a, vec(Z.dts), N)\n",
    "end\n",
    "\n",
    "function get_repeated_controls_alternated(a::AbstractMatrix, dts::AbstractVector, N::Int)\n",
    "    T = length(dts)\n",
    "    dts_long = repeat(dts, N)\n",
    "    a_long = repeat(a, 1, N)\n",
    "    for n=2:2:N\n",
    "        a_long[1:2,(n-1)*T+1:n*T] = a[1:2,end:-1:1]\n",
    "    end\n",
    "    times_long = cumsum(dts_long) - dts_long\n",
    "    for n=2:N\n",
    "        a_long[3,(n-1)*T+1:n*T] .+= (n-1)*times_long[T]\n",
    "    end \n",
    "    return (a_long, dts_long, times_long)\n",
    "end\n",
    "function get_repeated_controls_alternated(Z::NT.NamedTrajectory, N::Int)\n",
    "    return get_repeated_controls_alternated(Z.a, vec(Z.dts), N)\n",
    "end\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_long, dts_long, times_long = get_repeated_controls(Z, N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times_long, a_long')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_iso_long = QC.rollout(psi0_iso, a_long, dts_long, system; integrator=exp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_long = hcat([QC.iso_to_ket(psi_iso_long[:,t]) for t=1:N*T]...)\n",
    "pops_long = hcat([abs2.(QC.iso_to_ket(psi_iso_long[:,t])) for t=1:N*T]...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times_long, pops_long[1:dim,:]', labels=(-p_max:1:p_max)', legend=:outertopright)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times_long, pops_long[dim+1:end,:]', labels=(-p_max:1:p_max)', legend=:outertopright)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P_expect = 2*pops_long[1:dim,:]' * collect(-p_max:p_max) + 1/4*system.params[:acc] * times_long\n",
    "Plots.plot(times_long, P_expect)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fisher_evol = [-fisher_loss(QC.iso_to_ket(psi_iso_long[:,t])) for t=1:N*T];\n",
    "Plots.plot(times_long, fisher_evol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sensitivity_evol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sensitivity_evol = 1 ./ sqrt.(fisher_evol);\n",
    "Plots.plot(times_long[1:end], sensitivity_evol[1:end], ylims=(0.0, 1.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_long, dts_long = get_repeated_controls(Z_guess, N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times_long, a_long')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_iso_long = QC.rollout(psi0_iso, a_long, dts_long, system)#; integrator=exp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_long = hcat([QC.iso_to_ket(psi_iso_long[:,t]) for t=1:N*T]...)\n",
    "pops_long = hcat([abs2.(QC.iso_to_ket(psi_iso_long[:,t])) for t=1:N*T]...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times_long, pops_long[1:dim,:]', labels=(-p_max:1:p_max)', legend=:outertopright)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plots.plot(times_long, pops_long[dim+1:end,:]', labels=(-p_max:1:p_max)', legend=:outertopright)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P_expect = 2*pops_long[1:dim,:]' * collect(-p_max:p_max) + 1/4*system.params[:acc] * times_long\n",
    "Plots.plot(times_long, P_expect)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fisher_evol = [-fisher_loss(QC.iso_to_ket(psi_iso_long[:,t])) for t=1:N*T];\n",
    "Plots.plot(times_long, fisher_evol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sensitivity_evol = 1 ./ sqrt.(fisher_evol);\n",
    "Plots.plot(times_long, sensitivity_evol, ylims=(0.0, 1.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sensitivity_evol"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.2",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}