Absorption correction for cylindrical samples (#532)

* feat: absorption correction for cylindrical samples

docs/user-guide/absorption-correction.ipynb

@@ -0,0 +1,201 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "a6de8431-db39-49a5-b5dc-1afee82bc7d2",
+   "metadata": {},
+   "source": [
+    "# Absorption correction example\n",
+    "\n",
+    "Assuming single scattering, the neutron intensity per unit solid angle at the position $\\mathbf{p}$ is modeled as\n",
+    "$$\n",
+    "I(\\lambda, \\mathbf{x}) = \\int_{sample} S(\\lambda, \\widehat{\\mathbf{p}-\\mathbf{x}}) T(\\lambda, \\mathbf{p}, \\mathbf{x}) \\ d\\mathbf{x}\n",
+    "$$\n",
+    "where $S(\\lambda, \\hat{\\mathbf{s}})$ is the probability that a neutron with wavelength $\\lambda$ scatters in the direction $\\hat{\\mathbf{s}}$ and $T(\\lambda, \\mathbf{p}, \\mathbf{x})$ is the transmission rate for neutrons scattering at $\\mathbf{x}$ towards $\\mathbf{p}$. Here the $\\widehat{\\quad \\cdot\\quad}$ means that the vector is normalized to length 1.\n",
+    "\n",
+    "If the detector is far away from the sample relative to the size of the sample then $\\widehat{\\mathbf{p}-\\mathbf{x}}$ is close to constant in the integral, or if $S$ represents inhomogeneous scattering, the $S$ term does not depend on $\\mathbf{x}$ and can be moved outside of the integtral:\n",
+    "$$\n",
+    "I(\\lambda, \\mathbf{x}) \\approx S(\\lambda, \\widehat{\\mathbf{p} -\\bar{\\mathbf{x}}}) \\int_{sample}  T(\\lambda, \\mathbf{p}, \\mathbf{x}) \\ d\\mathbf{x}\n",
+    "$$\n",
+    "where $\\bar{\\mathbf{x}}$ denotes the center of the sample.\n",
+    "\n",
+    "To compute the scattering probabiltiy $S$ from the intensity $I$ we need to estimate the term\n",
+    "$$\n",
+    "C(\\lambda, \\mathbf{p}) = \\int_{sample}  T(\\lambda, \\mathbf{p}, \\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{p}, \\mathbf{x}) = \\exp{\\big(-\\mu(\\lambda) L(\\hat{\\mathbf{b}}, \\mathbf{p}, \\mathbf{x})\\big)}\n",
+    "$$\n",
+    "where $\\mu$ is material dependent and $\\hat{\\mathbf{b}}$ is the direction of the incoming neutron.\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{p}$.\n",
+    "\n",
+    "To compute $C(\\lambda, \\mathbf{p})$ 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:\n",
+    "\n",
+    "* `shape` defines `L`\n",
+    "* `material` defines $\\mu$\n",
+    "* `beam_direction` is $\\hat{\\mathbf{b}}$\n",
+    "* `wavelength` is $\\lambda$\n",
+    "* `detector_position` is $\\mathbf{p}$"
+   ]
+  },
+  {
+   "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",
+    "# Create a material with scattering parameters to represent a Vanadium sample\n",
+    "sample_material = Material(scattering_params=ScatteringParams.for_isotope('V'), effective_sample_number_density=sc.scalar(1, unit='1/angstrom**3'))\n",
+    "\n",
+    "# Create a sample shape\n",
+    "sample_shape = Cylinder(\n",
+    "    symmetry_line=sc.vector([0, 1, 0]),\n",
+    "    # center_of_base is placed slightly below the origin, so that the center of the cylinder is at the origin\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 detector positions, in this case the detector positions are placed in a sphere around the sample\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",
+    "            sc.broadcast(r * sc.cos(theta), sizes={**r.sizes, **theta.sizes, **phi.sizes}),\n",
+    "            r * sc.sin(theta) * sc.sin(phi),\n",
+    "        ],\n",
+    "        dim='row',\n",
+    "    )\n",
+    "    .transpose([*dims, 'row'])\n",
+    "    .values,\n",
+    "    unit=r.unit,\n",
+    ")\n",
+    "\n",
+    "def transmission_fraction(quadrature_kind):\n",
+    "    'Evaluate the transmission fraction using the selected quadrature kind'\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_fraction('medium')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "abb0d783-9a57-4daa-b4dc-7295044beb78",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "show_correction_map(transmission_fraction('cheap'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1512c244-118e-40a0-ac88-1510504603da",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "show_correction_map(transmission_fraction('medium'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d34de43c-0341-4656-b739-656cb3aff4a1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "show_correction_map(transmission_fraction('expensive'))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c04f6dcc-5ad9-4204-9de6-68a88650b067",
+   "metadata": {},
+   "source": [
+    "## What quadrature should I use?\n",
+    "\n",
+    "Use `medium` first, if it'd not good enough try `expensive`."
+   ]
+  }
+ ],
+ "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

@@ -11,4 +11,5 @@ recipes
 masking-tool
 wfm/index
 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,
+    )