Added notebook comparing original shift inversion to new one

10529f7a · Wolfgang Kastaun · e0211d4a · 10529f7a
Commit 10529f7a authored 1 month ago by Wolfgang Kastaun
--- a/docs/notebooks/shift_inversion_comparison.ipynb
+++ b/docs/notebooks/shift_inversion_comparison.ipynb
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "3efcb1ab-d4fb-4d84-aaa6-9222b0feb9a2",
+   "metadata": {},
+   "source": [
+    "# Compare original shift inversion algorithm to new one"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7a3b9115-11a1-4f19-aaba-36d246860aa9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib ipympl\n",
+    "from matplotlib import pyplot as plt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6cca7a2a-1853-4555-9f12-c27fbb30158e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c7043bd6-dcde-4840-bb1a-fa819d1c3936",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lisainstrument import dsp"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3d80234a-1eb8-4d8c-b76e-c4fb6fa0f3d8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lisainstrument.shift_inversion_numpy import make_shift_inverse_lagrange_numpy"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d8e70bed-5bb0-496c-a00a-730613971b2f",
+   "metadata": {},
+   "source": [
+    "Below is the original shift inversion method from the Instrument class. Only modifications are to make it self-contained, and removed logging."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "556511dc-17b4-42bd-867e-0fc9da431fd3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def legacy_invert_scet_wrt_tps(\n",
+    "    scet_wrt_tps: np.ndarray,\n",
+    "    clockinv_tolerance: float,\n",
+    "    physics_fs: float,\n",
+    "    interpolation_order: int,\n",
+    "    clockinv_maxiter: int = 5,\n",
+    "):\n",
+    "    edge = min(100, len(scet_wrt_tps) // 2 - 1)\n",
+    "    error = 0\n",
+    "\n",
+    "    niter = 0\n",
+    "    next_inverse = scet_wrt_tps\n",
+    "    while not niter or error > clockinv_tolerance:\n",
+    "        if niter >= clockinv_maxiter:\n",
+    "            msg = \"Legacy fixed point iter did not converge\"\n",
+    "            raise RuntimeError(msg)\n",
+    "        inverse = next_inverse\n",
+    "\n",
+    "        next_inverse = dsp.timeshift(\n",
+    "            scet_wrt_tps, -inverse * physics_fs, interpolation_order\n",
+    "        )\n",
+    "    \n",
+    "        error = np.max(np.abs((inverse - next_inverse)[edge:-edge]))\n",
+    "        niter += 1\n",
+    "\n",
+    "    return inverse"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "127215c1-5082-4994-8ef4-f6eb75dcbe11",
+   "metadata": {},
+   "source": [
+    "Setting up some example data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ed10795d-2d67-4715-9c18-a6e3993b2605",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "order = 31\n",
+    "nsamp = 3000\n",
+    "fsample = 16.0\n",
+    "dt = 1 / fsample\n",
+    "\n",
+    "f_mod = 0.007\n",
+    "a_mod = 1e-2 / (2 * np.pi * f_mod)\n",
+    "max_it = 5\n",
+    "tol = 1e-10"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d28a7126-215b-49d5-83fd-b6d6783160fb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "def dx_from_x(x):\n",
+    "    return np.cos(2 * np.pi * f_mod * x) * a_mod\n",
+    "\n",
+    "xi = np.arange(nsamp) * dt\n",
+    "dxi = dx_from_x(xi)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ccc111a8-3492-4569-a667-8731f60219ac",
+   "metadata": {},
+   "source": [
+    "Set up new shift inversion operator (unchunked pure numpy version)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "932daa39-b57d-4d97-b3c3-81471560c46a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "op_np = make_shift_inverse_lagrange_numpy(\n",
+    "        order=order,\n",
+    "        fsample=fsample,\n",
+    "        max_abs_shift=a_mod * 1.01,\n",
+    "        max_iter=max_it,\n",
+    "        tolerance=tol,\n",
+    "    )\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4b15c610-c7f2-44ea-b8e6-7e7fc62c5bf5",
+   "metadata": {},
+   "source": [
+    "Compute shift inversion with both"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5c3e3fea-5fdf-49ad-98d3-78062a04d801",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ai_np = op_np(dxi)\n",
+    "ai_leg = legacy_invert_scet_wrt_tps(dxi, tol/100, fsample, order, max_it+2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fba2473e-3a8c-48c5-9cb5-13175098b70f",
+   "metadata": {},
+   "source": [
+    "Expected result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eb55600f-0cd2-4702-bd95-217d4722e248",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ai_ex = dx_from_x(xi - ai_np)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fc7a324c-4532-4c66-8317-3620d55e6d3c",
+   "metadata": {},
+   "source": [
+    "The original version seems to have more boundary artifacts"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f8b72f4b-00bb-4f57-ae42-5a4293a03d65",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.close()\n",
+    "plt.plot(xi, ai_np, \"k+\")\n",
+    "plt.plot(xi, ai_leg, \"rx\")\n",
+    "plt.plot(xi, ai_ex, \"g-\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c665ea55-9d2b-4e8e-8ad9-e0113541633a",
+   "metadata": {},
+   "source": [
+    "The error within the valid points is within the tolerance for both."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "74d29668-57f2-4800-8e5b-bc82253a2185",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.close()\n",
+    "plt.semilogy(xi, np.abs(ai_np-ai_ex), \"k+\")\n",
+    "plt.semilogy(xi, np.abs(ai_leg-ai_ex), \"rx\")\n",
+    "plt.axhline(y=tol)\n",
+    "plt.axvline(x=xi[op_np.margin_left])\n",
+    "plt.axvline(x=xi[-op_np.margin_right])\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "912c544e-e926-41cb-ba93-d414c538c712",
+   "metadata": {},
+   "source": [
+    "To investigate the differences inside the valid region, fix a bug in the legacy method. It used to return the second-best guess instead of the latest iteration result. Also, we make the margin size for the error measure a parameter."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "45775517-6b5b-4b78-a345-f51e5bba6439",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def modified_invert_scet_wrt_tps(\n",
+    "    scet_wrt_tps: np.ndarray,\n",
+    "    clockinv_tolerance: float,\n",
+    "    physics_fs: float,\n",
+    "    edge : int,\n",
+    "    interpolation_order: int,\n",
+    "    clockinv_maxiter: int = 15,\n",
+    "):\n",
+    "    \n",
+    "    error = 0\n",
+    "\n",
+    "    niter = 0\n",
+    "    next_inverse = scet_wrt_tps\n",
+    "    while not niter or error > clockinv_tolerance:\n",
+    "        if niter >= clockinv_maxiter:\n",
+    "            msg = \"Legacy fixed point iter did not converge\"\n",
+    "            raise RuntimeError(msg)\n",
+    "        inverse = next_inverse\n",
+    "\n",
+    "        next_inverse = dsp.timeshift(\n",
+    "            scet_wrt_tps, -inverse * physics_fs, interpolation_order\n",
+    "        )\n",
+    "    \n",
+    "        error = np.max(np.abs((inverse - next_inverse)[edge:-edge]))\n",
+    "        niter += 1\n",
+    "    print( niter )\n",
+    "    return next_inverse"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f711b48f-fffa-4732-b609-7c2d500f59ff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#ai_mod = modified_invert_scet_wrt_tps(dxi, tol, fsample, op_np.margin_left, order)\n",
+    "ai_mod = modified_invert_scet_wrt_tps(dxi, tol, fsample, 100, order)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5566a405-1994-4f25-905c-6ed6d2ee9f6a",
+   "metadata": {},
+   "source": [
+    "Now the error inside the valid region agrees perfectly (the margin size was not the issue)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cba3a269-ae9b-4236-87a3-be7966af0dc8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.close()\n",
+    "plt.semilogy(xi, np.abs(ai_np-ai_ex), \"k+\")\n",
+    "plt.semilogy(xi, np.abs(ai_mod-ai_ex), \"rx\")\n",
+    "plt.axhline(y=tol)\n",
+    "plt.axvline(x=xi[op_np.margin_left])\n",
+    "plt.axvline(x=xi[-op_np.margin_right])\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9153fb52-c7de-4352-bf61-c1cda7cf7944",
+   "metadata": {},
+   "source": [
+    "The boundary artifacts are still larger for the legacy algorithm. Reason yet unkown."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "63689287-a4bf-4f23-957a-51f4046fd0ad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.close()\n",
+    "plt.plot(xi, ai_np, \"k+\")\n",
+    "plt.plot(xi, ai_mod, \"rx\")\n",
+    "plt.plot(xi, ai_ex, \"g-\")\n",
+    "plt.xlim(xi[-20], xi[-1]+dt)\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "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"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
+%% Cell type:markdown id:3efcb1ab-d4fb-4d84-aaa6-9222b0feb9a2 tags:
+# Compare original shift inversion algorithm to new one
+%% Cell type:code id:7a3b9115-11a1-4f19-aaba-36d246860aa9 tags:
+``` python
+%matplotlib ipympl
+from matplotlib import pyplot as plt
+```
+%% Cell type:code id:6cca7a2a-1853-4555-9f12-c27fbb30158e tags:
+``` python
+import numpy as np
+```
+%% Cell type:code id:c7043bd6-dcde-4840-bb1a-fa819d1c3936 tags:
+``` python
+from lisainstrument import dsp
+```
+%% Cell type:code id:3d80234a-1eb8-4d8c-b76e-c4fb6fa0f3d8 tags:
+``` python
+from lisainstrument.shift_inversion_numpy import make_shift_inverse_lagrange_numpy
+```
+%% Cell type:markdown id:d8e70bed-5bb0-496c-a00a-730613971b2f tags:
+Below is the original shift inversion method from the Instrument class. Only modifications are to make it self-contained, and removed logging.
+%% Cell type:code id:556511dc-17b4-42bd-867e-0fc9da431fd3 tags:
+``` python
+def legacy_invert_scet_wrt_tps(
+    scet_wrt_tps: np.ndarray,
+    clockinv_tolerance: float,
+    physics_fs: float,
+    interpolation_order: int,
+    clockinv_maxiter: int = 5,
+):
+    edge = min(100, len(scet_wrt_tps) // 2 - 1)
+    error = 0
+    niter = 0
+    next_inverse = scet_wrt_tps
+    while not niter or error > clockinv_tolerance:
+        if niter >= clockinv_maxiter:
+            msg = "Legacy fixed point iter did not converge"
+            raise RuntimeError(msg)
+        inverse = next_inverse
+        next_inverse = dsp.timeshift(
+            scet_wrt_tps, -inverse * physics_fs, interpolation_order
+        )
+        error = np.max(np.abs((inverse - next_inverse)[edge:-edge]))
+        niter += 1
+    return inverse
+```
+%% Cell type:markdown id:127215c1-5082-4994-8ef4-f6eb75dcbe11 tags:
+Setting up some example data
+%% Cell type:code id:ed10795d-2d67-4715-9c18-a6e3993b2605 tags:
+``` python
+order = 31
+nsamp = 3000
+fsample = 16.0
+dt = 1 / fsample
+f_mod = 0.007
+a_mod = 1e-2 / (2 * np.pi * f_mod)
+max_it = 5
+tol = 1e-10
+```
+%% Cell type:code id:d28a7126-215b-49d5-83fd-b6d6783160fb tags:
+``` python
+def dx_from_x(x):
+    return np.cos(2 * np.pi * f_mod * x) * a_mod
+xi = np.arange(nsamp) * dt
+dxi = dx_from_x(xi)
+```
+%% Cell type:markdown id:ccc111a8-3492-4569-a667-8731f60219ac tags:
+Set up new shift inversion operator (unchunked pure numpy version)
+%% Cell type:code id:932daa39-b57d-4d97-b3c3-81471560c46a tags:
+``` python
+op_np = make_shift_inverse_lagrange_numpy(
+        order=order,
+        fsample=fsample,
+        max_abs_shift=a_mod * 1.01,
+        max_iter=max_it,
+        tolerance=tol,
+    )
+```
+%% Cell type:markdown id:4b15c610-c7f2-44ea-b8e6-7e7fc62c5bf5 tags:
+Compute shift inversion with both
+%% Cell type:code id:5c3e3fea-5fdf-49ad-98d3-78062a04d801 tags:
+``` python
+ai_np = op_np(dxi)
+ai_leg = legacy_invert_scet_wrt_tps(dxi, tol/100, fsample, order, max_it+2)
+```
+%% Cell type:markdown id:fba2473e-3a8c-48c5-9cb5-13175098b70f tags:
+Expected result
+%% Cell type:code id:eb55600f-0cd2-4702-bd95-217d4722e248 tags:
+``` python
+ai_ex = dx_from_x(xi - ai_np)
+```
+%% Cell type:markdown id:fc7a324c-4532-4c66-8317-3620d55e6d3c tags:
+The original version seems to have more boundary artifacts
+%% Cell type:code id:f8b72f4b-00bb-4f57-ae42-5a4293a03d65 tags:
+``` python
+plt.close()
+plt.plot(xi, ai_np, "k+")
+plt.plot(xi, ai_leg, "rx")
+plt.plot(xi, ai_ex, "g-")
+```
+%% Cell type:markdown id:c665ea55-9d2b-4e8e-8ad9-e0113541633a tags:
+The error within the valid points is within the tolerance for both.
+%% Cell type:code id:74d29668-57f2-4800-8e5b-bc82253a2185 tags:
+``` python
+plt.close()
+plt.semilogy(xi, np.abs(ai_np-ai_ex), "k+")
+plt.semilogy(xi, np.abs(ai_leg-ai_ex), "rx")
+plt.axhline(y=tol)
+plt.axvline(x=xi[op_np.margin_left])
+plt.axvline(x=xi[-op_np.margin_right])
+```
+%% Cell type:markdown id:912c544e-e926-41cb-ba93-d414c538c712 tags:
+To investigate the differences inside the valid region, fix a bug in the legacy method. It used to return the second-best guess instead of the latest iteration result. Also, we make the margin size for the error measure a parameter.
+%% Cell type:code id:45775517-6b5b-4b78-a345-f51e5bba6439 tags:
+``` python
+def modified_invert_scet_wrt_tps(
+    scet_wrt_tps: np.ndarray,
+    clockinv_tolerance: float,
+    physics_fs: float,
+    edge : int,
+    interpolation_order: int,
+    clockinv_maxiter: int = 15,
+):
+    error = 0
+    niter = 0
+    next_inverse = scet_wrt_tps
+    while not niter or error > clockinv_tolerance:
+        if niter >= clockinv_maxiter:
+            msg = "Legacy fixed point iter did not converge"
+            raise RuntimeError(msg)
+        inverse = next_inverse
+        next_inverse = dsp.timeshift(
+            scet_wrt_tps, -inverse * physics_fs, interpolation_order
+        )
+        error = np.max(np.abs((inverse - next_inverse)[edge:-edge]))
+        niter += 1
+    print( niter )
+    return next_inverse
+```
+%% Cell type:code id:f711b48f-fffa-4732-b609-7c2d500f59ff tags:
+``` python
+#ai_mod = modified_invert_scet_wrt_tps(dxi, tol, fsample, op_np.margin_left, order)
+ai_mod = modified_invert_scet_wrt_tps(dxi, tol, fsample, 100, order)
+```
+%% Cell type:markdown id:5566a405-1994-4f25-905c-6ed6d2ee9f6a tags:
+Now the error inside the valid region agrees perfectly (the margin size was not the issue).
+%% Cell type:code id:cba3a269-ae9b-4236-87a3-be7966af0dc8 tags:
+``` python
+plt.close()
+plt.semilogy(xi, np.abs(ai_np-ai_ex), "k+")
+plt.semilogy(xi, np.abs(ai_mod-ai_ex), "rx")
+plt.axhline(y=tol)
+plt.axvline(x=xi[op_np.margin_left])
+plt.axvline(x=xi[-op_np.margin_right])
+```
+%% Cell type:markdown id:9153fb52-c7de-4352-bf61-c1cda7cf7944 tags:
+The boundary artifacts are still larger for the legacy algorithm. Reason yet unkown.
+%% Cell type:code id:63689287-a4bf-4f23-957a-51f4046fd0ad tags:
+``` python
+plt.close()
+plt.plot(xi, ai_np, "k+")
+plt.plot(xi, ai_mod, "rx")
+plt.plot(xi, ai_ex, "g-")
+plt.xlim(xi[-20], xi[-1]+dt)
+```