Cylinder absorp corr

docs/user-guide/absorption-correction.ipynb

@@ -0,0 +1,234 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "a6de8431-db39-49a5-b5dc-1afee82bc7d2",
+   "metadata": {},
+   "source": [
+    "# Absorption correction for cylindrical samples\n",
+    "\n",
+    "The intensity per unit solid angle at a detector element located at $\\mathbf{x}$ is modeled as\n",
+    "$$\n",
+    "I(\\lambda, \\mathbf{x}) = \\int_{sample} S(\\lambda, \\theta(\\hat{\\mathbf{b}}, \\mathbf{x}-\\mathbf{x}')) T(\\lambda, \\mathbf{x}, \\mathbf{x}') \\ d\\mathbf{x}'\n",
+    "$$\n",
+    "where $S(\\lambda, \\theta)$ is the probability that a neutron with wavelength $\\lambda$ scatters with the scattering angle $\\theta$, $T(\\lambda, \\mathbf{x}, \\mathbf{x}')$ is the transmission rate for neutrons scattering at $\\mathbf{x}'$ towards $\\mathbf{x}$, and $\\hat{\\mathbf{b}}$ is the direction of the incoming beam.\n",
+    "\n",
+    "If the detector is far away from the sample relative to the size of the sample then $\\theta$ is close to constant as a function of $\\mathbf{x}'$ and the $S$ term can be moved outside of the integtral:\n",
+    "$$\n",
+    "I(\\lambda, \\mathbf{x}) \\approx S(\\lambda, \\theta(\\hat{\\mathbf{b}}, \\mathbf{x} -\\bar{\\mathbf{x}'})) \\int_{sample}  T(\\lambda, \\mathbf{x}, \\mathbf{x}') \\ d\\mathbf{x}'.\n",
+    "$$\n",
+    "To compute the scattering probabiltiy $S$ from the intensity $I$ we need to estimate the term\n",
+    "$$\n",
+    "C(\\lambda, \\mathbf{x}) = \\int_{sample}  T(\\lambda, \\mathbf{x}, \\mathbf{x}') \\ d\\mathbf{x}'.\n",
+    "$$\n",
+    "This is the \"absorption correction\" term.\n",
+    "\n",
+    "The transmission fraction is a function of path length $L$ of the neutron going through the sample\n",
+    "$$\n",
+    "T(\\lambda, \\mathbf{x}, \\mathbf{x}') = \\exp{(-\\mu(\\lambda) L(\\hat{\\mathbf{b}}, \\mathbf{x}, \\mathbf{x}'))}\n",
+    "$$\n",
+    "where $\\mu$ is material dependent.\n",
+    "\n",
+    "The path length through the sample depends on the shape of the sample, the scattering point $\\mathbf{x}'$ and the detection position $\\mathbf{x}$.\n",
+    "\n",
+    "To compute $C(\\lambda, \\mathbf{x})$ you can use\n",
+    "```python\n",
+    "scippneutron.absorption.base.compute_transmission_map(\n",
+    "    shape,\n",
+    "    material,\n",
+    "    beam_direction,\n",
+    "    wavelength,\n",
+    "    detector_position\n",
+    ")\n",
+    "```\n",
+    "where `shape` and `material` are sample properties, `shape` defines `L` and `material` defines $\\mu$, `beam_direction` is $\\mathbf{b}$, `wavelength` is $\\lambda$ and `detector_position` is $\\mathbf{x}$. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "07bd8cdb-c18a-48e4-a60a-041a0d4cd07c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import scipp as sc\n",
+    "import numpy as np\n",
+    "\n",
+    "from scippneutron.absorption import compute_transmission_map, Cylinder, Material\n",
+    "from scippneutron.atoms import ScatteringParams\n",
+    "\n",
+    "\n",
+    "sample_material = Material(scattering_params=ScatteringParams.for_isotope('V'), effective_sample_number_density=sc.scalar(1, unit='1/angstrom**3'))\n",
+    "\n",
+    "sample_shape = Cylinder(\n",
+    "    symmetry_line=sc.vector([0, 1, 0]),\n",
+    "    center_of_base=sc.vector([0, -.5, 0], unit='cm'),\n",
+    "    radius=sc.scalar(1., unit='cm'),\n",
+    "    height=sc.scalar(0.3, unit='cm')\n",
+    ")\n",
+    "\n",
+    "# Create some dummy detector positions.\n",
+    "theta = sc.linspace('theta', 0, np.pi, 60, endpoint=True, unit='rad')\n",
+    "phi = sc.linspace('phi', 0, 2 * np.pi, 120, endpoint=False, unit='rad')\n",
+    "r = sc.array(dims='r', values=[5.], unit='m')\n",
+    "dims = (*r.dims, *theta.dims, *phi.dims)\n",
+    "\n",
+    "# Detector positions in grid on a sphere around the origin\n",
+    "detector_positions = sc.vectors(\n",
+    "    dims=dims,\n",
+    "    values=sc.concat(\n",
+    "        [\n",
+    "            r * sc.sin(theta) * sc.cos(phi),\n",
+    "            r * sc.sin(theta) * sc.sin(phi),\n",
+    "            sc.broadcast(r * sc.cos(theta), sizes={**r.sizes, **theta.sizes, **phi.sizes}),\n",
+    "        ],\n",
+    "        dim='row',\n",
+    "    )\n",
+    "    .transpose([*dims, 'row'])\n",
+    "    .values,\n",
+    "    unit=r.unit,\n",
+    ")\n",
+    "\n",
+    "def transmission(quadrature_kind):\n",
+    "    'Evaluates the transmission map for different kinds of quadrature'\n",
+    "    da = compute_transmission_map(\n",
+    "        sample_shape,\n",
+    "        sample_material,\n",
+    "        beam_direction=sc.vector([0, 0, 1]),\n",
+    "        wavelength=sc.linspace('wavelength', 0.5, 2.5, 10, unit='angstrom'),\n",
+    "        detector_position=detector_positions,\n",
+    "        quadrature_kind=quadrature_kind,\n",
+    "    )\n",
+    "    da.coords['phi'] = phi\n",
+    "    da.coords['theta'] = theta\n",
+    "    return da\n",
+    "\n",
+    "\n",
+    "def show_correction_map(da):\n",
+    "    'Plotting utility'\n",
+    "    return (\n",
+    "        da['wavelength', 0]['r', 0].plot() /\n",
+    "        da['wavelength', 0]['r', 0]['theta', da.sizes['theta']//2].plot() /\n",
+    "        da['wavelength', 0]['r', 0]['theta', da.sizes['theta']//2]['phi', da.sizes['phi']//4 - da.sizes['phi']//6:da.sizes['phi']//4 + da.sizes['phi']//6].plot() /\n",
+    "        da['wavelength', 0]['r', 0]['theta', da.sizes['theta']//2]['phi', da.sizes['phi']//2 - da.sizes['phi']//6:da.sizes['phi']//2 + da.sizes['phi']//6].plot() /\n",
+    "        da['wavelength', 0]['r', 0]['phi', da.sizes['phi']//2].plot()\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6f03111b-8d1c-4c82-a9e9-b2fcde5c1003",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "transmission('cheap')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "abb0d783-9a57-4daa-b4dc-7295044beb78",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "show_correction_map(transmission('cheap'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1512c244-118e-40a0-ac88-1510504603da",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "show_correction_map(transmission('medium'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d34de43c-0341-4656-b739-656cb3aff4a1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "show_correction_map(transmission('expensive'))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "93a0af6e-dd16-4ee4-838d-df2ef014ff4a",
+   "metadata": {},
+   "source": [
+    "## Error estimate"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dbbdbf77-07b0-491e-b898-6daccc03ccd8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import plopp as pp\n",
+    "cheap = transmission('cheap')['wavelength', 0]['r', 0]\n",
+    "medium = transmission('medium')['wavelength', 0]['r', 0]\n",
+    "expensive = transmission('expensive')['wavelength', 0]['r', 0]\n",
+    "\n",
+    "t = cheap.sizes['theta'] // 2\n",
+    "pmin = cheap.sizes['phi'] // 2 - cheap.sizes['phi'] // 8\n",
+    "pmax = cheap.sizes['phi'] // 2 + cheap.sizes['phi'] // 8\n",
+    "pp.plot(\n",
+    "    {\n",
+    "        'cheap': cheap['theta', t]['phi', pmin:pmax],\n",
+    "        'medium': medium['theta', t]['phi', pmin:pmax],\n",
+    "        'expensive': expensive['theta', t]['phi', pmin:pmax],\n",
+    "    }\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "180ab825-0df5-4422-91cf-474523c63fe6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Relative difference between cheap and expensive quadrature\n",
+    "(sc.abs(cheap - expensive)/sc.abs(expensive)).mean().value"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a3fb34e-449b-4a4c-8ea6-849b0c5166cc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Relative difference between medium and expensive quadrature\n",
+    "(sc.abs(medium - expensive)/sc.abs(expensive)).mean().value"
+   ]
+  }
+ ],
+ "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",
+   "version": "3.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/user-guide/index.md

@@ -10,4 +10,5 @@ instrument-view
 recipes
 masking-tool
 algorithms-background/index
+absorption-correction
 ```

src/scippneutron/absorption/__init__.py

@@ -0,0 +1,6 @@
+from .base import compute_transmission_map
+from .cylinder import Cylinder
+from .material import Material
+
+
+__all__ = ('compute_transmission_map', 'Cylinder', 'Material')

src/scippneutron/absorption/base.py

@@ -0,0 +1,101 @@
+from functools import partial
+from typing import Any
+
+import scipp as sc
+
+from .material import Material
+from .types import SampleShape
+
+
+def compute_transmission_map(
+    sample_shape: SampleShape,
+    sample_material: Material,
+    beam_direction: sc.Variable,
+    wavelength: sc.Variable,
+    detector_position: sc.Variable,
+    quadrature_kind: Any = 'medium',
+) -> sc.DataArray:
+    points, weights = sample_shape.quadrature(quadrature_kind)
+    transmission = _integrate_transmission_fraction(
+        partial(
+            _single_scatter_distance_through_sample,
+            sample_shape,
+            points,
+            beam_direction,
+        ),
+        partial(_transmission_fraction, sample_material),
+        points,
+        weights,
+        detector_position,
+        wavelength,
+    )
+    return sc.DataArray(
+        data=transmission / sample_shape.volume,
+        coords={'detector_position': detector_position, 'wavelength': wavelength},
+    )
+
+
+def _single_scatter_distance_through_sample(
+    sample_shape, scatter_point, initial_direction, scatter_direction
+):
+    L1 = sample_shape.beam_intersection(scatter_point, -initial_direction)
+    L2 = sample_shape.beam_intersection(scatter_point, scatter_direction)
+    return L1 + L2
+
+
+def _transmission_fraction(material, distance_through_sample, wavelength):
+    return sc.exp(
+        -(material.attenuation_coefficient(wavelength) * distance_through_sample).to(
+            unit='dimensionless', copy=False
+        )
+    )
+
+
+def _integrate_transmission_fraction(
+    distance_through_sample,
+    transmission,
+    points,
+    weights,
+    detector_position,
+    wavelengths,
+):
+    # If size after broadcast is too large
+    # then don't vectorize the operation to avoid OOM
+    if points.size * detector_position.size > 20_000_000:
+        out = []
+        dim = detector_position.dims[0]
+        for i in range(detector_position.sizes[dim]):
+            out.append(  # noqa: PERF401
+                _integrate_transmission_fraction(
+                    distance_through_sample,
+                    transmission,
+                    points,
+                    weights,
+                    detector_position[dim, i],
+                    wavelengths,
+                )
+            )
+
+        return sc.concat(
+            out,
+            dim=dim,
+        )
+
+    scatter_direction = detector_position - points.to(unit=detector_position.unit)
+    scatter_direction /= sc.norm(scatter_direction)
+    Ltot = distance_through_sample(scatter_direction)
+
+    # The Ltot array is already large, to avoid OOM, don't vectorize this operation
+    return sc.concat(
+        [
+            # Instead of broadcast multiply and sum, use matvec for efficiency
+            # this becomes relevant when the number of wavelength points grows
+            sc.array(
+                dims=(tf := transmission(Ltot, w)).dims[:-1],
+                values=tf.values @ weights.values,
+                unit=tf.unit * weights.unit,
+            )
+            for w in wavelengths
+        ],
+        dim=wavelengths.dim,
+    )