ENH: compute decay rates from sub-intensities (#393)

* DOC: do not reverse resonance names
* DOC: rotate decay rate matrix visualization
* DX: test values of decay rate matrix
* ENH: compute sum of rate matrix with numpy.tril

Author: redeboer

.cspell/python-extra.txt

@@ -26,6 +26,7 @@ svgutils
 sympify
 tensorwaves
 tolist
+tril
 viewboxing
 viridis
 vmax

docs/intensity.ipynb

@@ -533,11 +533,20 @@
    },
    "outputs": [],
    "source": [
-    "def compute_decay_rates(func, integration_sample: DataSample):\n",
+    "def compute_decay_rates(rates: np.ndarray) -> np.ndarray:\n",
+    "    m, n = rates.shape\n",
+    "    assert m == n\n",
+    "    d = rates.diagonal()\n",
+    "    D = d * np.identity(n)\n",
+    "    X = d[None] + d[None].T\n",
+    "    return rates - X + 2 * D\n",
+    "\n",
+    "\n",
+    "def compute_sub_intensities(func, integration_sample: DataSample):\n",
     "    decay_rates = np.zeros(shape=(n_resonances, n_resonances))\n",
     "    combinations = list(product(enumerate(resonances), enumerate(resonances)))\n",
     "    progress_bar = tqdm(\n",
-    "        desc=\"Calculating rate matrix\",\n",
+    "        desc=\"Calculating sub-intensities\",\n",
     "        disable=NO_LOG,\n",
     "        total=(len(combinations) + n_resonances) // 2,\n",
     "    )\n",
@@ -548,10 +557,7 @@
     "        progress_bar.postfix = f\"{resonance1.name} × {resonance2.name}\"\n",
     "        res1 = resonance1.latex\n",
     "        res2 = resonance2.latex\n",
-    "        if res1 == res2:\n",
-    "            I_sub = sub_intensity(func, integration_sample, non_zero_couplings=[res1])\n",
-    "        else:\n",
-    "            I_sub = interference_intensity(func, integration_sample, [res1], [res2])\n",
+    "        I_sub = sub_intensity(func, integration_sample, non_zero_couplings=[res1, res2])\n",
     "        decay_rates[i, j] = I_sub / I_tot\n",
     "        if i != j:\n",
     "            decay_rates[j, i] = decay_rates[i, j]\n",
@@ -568,7 +574,6 @@
     "resonances = sorted(\n",
     "    (chain.resonance for chain in model.decay.chains),\n",
     "    key=sort_resonances,\n",
-    "    reverse=True,\n",
     ")\n",
     "n_resonances = len(resonances)"
    ]
@@ -597,70 +602,60 @@
     "    fig, ax = plt.subplots(figsize=(9, 9))\n",
     "    fig.patch.set_color(\"none\")\n",
     "    ax.set_title(title)\n",
-    "    ax.matshow(jnp.rot90(decay_rates).T, cmap=plt.cm.coolwarm, vmin=-vmax, vmax=+vmax)\n",
+    "    ax.matshow(decay_rates, cmap=plt.cm.coolwarm, vmin=-vmax, vmax=+vmax)\n",
     "\n",
     "    resonance_latex = [f\"${p.latex}$\" for p in resonances]\n",
     "    ax.set_xticks(range(n_resonances))\n",
-    "    ax.set_xticklabels(reversed(resonance_latex), rotation=90)\n",
+    "    ax.set_xticklabels(resonance_latex, rotation=90)\n",
+    "    ax.xaxis.tick_bottom()\n",
     "    ax.set_yticks(range(n_resonances))\n",
     "    ax.set_yticklabels(resonance_latex)\n",
     "    for i in range(n_resonances):\n",
     "        for j in range(n_resonances):\n",
-    "            if j < i:\n",
+    "            if i > j:\n",
     "                continue\n",
     "            rate = decay_rates[i, j]\n",
-    "            ax.text(\n",
-    "                n_resonances - j - 1,\n",
-    "                i,\n",
-    "                f\"{100 * rate:.2f}\",\n",
-    "                va=\"center\",\n",
-    "                ha=\"center\",\n",
-    "            )\n",
+    "            ax.text(i, j, f\"{100 * rate:.2f}\", ha=\"center\", va=\"center\")\n",
     "    fig.tight_layout()\n",
     "    return fig\n",
     "\n",
     "\n",
-    "decay_rates = compute_decay_rates(intensity_func, integration_sample)\n",
+    "sub_intensities = compute_sub_intensities(intensity_func, integration_sample)\n",
+    "decay_rates = compute_decay_rates(sub_intensities)\n",
     "fig = visualize_decay_rates(decay_rates)\n",
     "output_path = \"_images/rate-matrix.svg\"\n",
     "fig.savefig(output_path, bbox_inches=\"tight\")\n",
     "reduce_svg_size(output_path)\n",
     "plt.show(fig)"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    ":::{only} latex\n",
-    "{{ FIG_RATE_MATRIX }}\n",
-    ":::"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
     "jupyter": {
      "source_hidden": true
     },
-    "mystnb": {
-     "code_prompt_show": "Function for computing the total over all rates"
-    },
     "tags": [
-     "hide-cell"
+     "hide-input"
     ]
    },
    "outputs": [],
    "source": [
-    "def compute_sum_over_decay_rates(decay_rate_matrix) -> float:\n",
-    "    decay_rate_sum = 0.0\n",
-    "    for i in range(len(resonances)):\n",
-    "        for j in range(len(resonances)):\n",
-    "            if j < i:\n",
-    "                continue\n",
-    "            decay_rate_sum += decay_rate_matrix[i, j]\n",
-    "    return decay_rate_sum"
+    "np.testing.assert_array_almost_equal(\n",
+    "    100 * decay_rates[-1],\n",
+    "    [4.78, 0.16, -1.68, 0.04, 0.03, -7.82, 5.09, 1.96, -1.15, -1.95, 0.04, 14.7],\n",
+    "    decimal=2,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    ":::{only} latex\n",
+    "{{ FIG_RATE_MATRIX }}\n",
+    ":::"
    ]
   },
   {
@@ -669,6 +664,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "def compute_sum_over_decay_rates(rate_matrix: np.ndarray) -> float:\n",
+    "    return rate_matrix.diagonal().sum() + np.tril(rate_matrix, k=-1).sum()\n",
+    "\n",
+    "\n",
     "np.testing.assert_almost_equal(compute_sum_over_decay_rates(decay_rates), 1.0)"
    ]
   },
@@ -811,7 +810,8 @@
     "    y_range=sub_region_y_range,\n",
     ")\n",
     "sub_sample = transformer(sub_sample)\n",
-    "sub_decay_rates = compute_decay_rates(intensity_func, sub_sample)\n",
+    "sub_decay_intensities = compute_sub_intensities(intensity_func, sub_sample)\n",
+    "sub_decay_rates = compute_decay_rates(sub_decay_intensities)\n",
     "fig = visualize_decay_rates(sub_decay_rates, title=\"Rate matrix over sub-region\")\n",
     "output_path = \"_images/rate-matrix-sub-region.svg\"\n",
     "fig.savefig(output_path, bbox_inches=\"tight\")\n",