FIX: switch to scipy.interpolate.RegularGridInterpolator (#278)

* MAINT: list minimal requirement for `myst-parser` See description of f51313b.
ComPWA · Jan 18, 2023 · f0df8af · f0df8af
1 parent 42133ef
commit f0df8af
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 14 deletions.
diff --git a/docs/appendix/serialization.ipynb b/docs/appendix/serialization.ipynb
@@ -38,7 +38,7 @@
     "import numpy as np\n",
     "import pandas as pd\n",
     "from IPython.display import Markdown\n",
-    "from scipy.interpolate import griddata, interp2d\n",
+    "from scipy.interpolate import RegularGridInterpolator, griddata\n",
     "from tqdm.auto import tqdm\n",
     "\n",
     "from polarimetry import formulate_polarimetry\n",
@@ -257,11 +257,15 @@
    "source": [
     "## Import and interpolate\n",
     "\n",
-    "::::{margin}\n",
-    ":::{warning}\n",
-    "{obj}`~numpy.nan` values have to be replaced with `0.0` using {func}`numpy.nan_to_num`.\n",
-    ":::\n",
-    "::::"
+    "The arrays in the {ref}`exported JSON files <exported-distributions>` can be used to create a {class}`~scipy.interpolate.RegularGridInterpolator` for the intensity and for each components of $\\vec\\alpha$.\n",
+    "\n",
+    ":::{margin}\n",
+    "\n",
+    "{func}`.import_polarimetry_field` returns JAX arrays, which are read-only, {class}`~scipy.interpolate.RegularGridInterpolator` requires modifiable arrays, so we convert them to NumPy.\n",
+    "\n",
+    "Also note that the `values` array needs to be **transposed**!\n",
+    "\n",
+    ":::"
    ]
   },
   {
@@ -279,10 +283,24 @@
     "    field_definition[\"alpha_y\"],\n",
     "    field_definition[\"alpha_z\"],\n",
     ")\n",
-    "interpolated_funcs = tuple(\n",
-    "    interp2d(imported_sigma1, imported_sigma2, np.nan_to_num(z), kind=\"linear\")\n",
+    "interpolated_funcs = [\n",
+    "    RegularGridInterpolator(\n",
+    "        points=(\n",
+    "            np.array(imported_sigma1),\n",
+    "            np.array(imported_sigma2),\n",
+    "        ),\n",
+    "        values=np.array(z).transpose(),\n",
+    "        method=\"linear\",\n",
+    "    )\n",
     "    for z in imported_arrays\n",
-    ")"
+    "]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This is a function that can compute an interpolated value of each of these observables for a random point on the Dalitz plane."
    ]
   },
   {
@@ -291,14 +309,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "interpolated_funcs[1](0.8, 3.6)"
+    "interpolated_funcs[1]([0.8, 3.6])"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Use {obj}`numpy.vectorize` to compute the interpolated values over a random phase space sample:"
+    "As opposed to SciPy's deprecated {obj}`~scipy.interpolate.interp2d`, {class}`~scipy.interpolate.RegularGridInterpolator` is already in vectorized form, so there is no need to {obj}`~numpy.vectorize` it."
    ]
   },
   {
@@ -314,9 +332,7 @@
     "mini_sample = transformer(mini_sample)\n",
     "x = mini_sample[\"sigma1\"]\n",
     "y = mini_sample[\"sigma2\"]\n",
-    "z_interpolated = [\n",
-    "    np.vectorize(func)(x, y) for func in tqdm(interpolated_funcs, disable=NO_TQDM)\n",
-    "]\n",
+    "z_interpolated = [func((x, y)) for func in tqdm(interpolated_funcs, disable=NO_TQDM)]\n",
     "z_interpolated[0]"
    ]
   },

diff --git a/setup.cfg b/setup.cfg
@@ -34,6 +34,7 @@ doc =
     kaleido  # plotly static figure export
     matplotlib-inline !=0.1.5  # problem with plt.ioff()
     myst-nb >=0.14  # nb_render_markdown_format for Markdown tables
+    myst-parser >0.16.0 # field lists
     pandas
     plotly
     psutil