Merge pull request #71 from scipp/Qxy

Add support for computing I(Q_x, Q_y)

docs/examples/zoom.ipynb

@@ -107,6 +107,12 @@
     ")\n",
     "\n",
     "params[QBins] = sc.geomspace(dim='Q', start=0.004, stop=0.8, num=141, unit='1/angstrom')\n",
+    "qxy = {\n",
+    "    'Qx': sc.linspace(dim='Qx', start=-0.5, stop=0.5, num=101, unit='1/angstrom'),\n",
+    "    'Qy': sc.linspace(dim='Qy', start=-0.8, stop=0.8, num=101, unit='1/angstrom'),\n",
+    "}\n",
+    "# Uncomment to compute I(Qx, Qy) instead of I(Q)\n",
+    "# params[QxyBins] = qxy\n",
     "params[NonBackgroundWavelengthRange] = sc.array(\n",
     "    dims=['wavelength'], values=[0.7, 17.1], unit='angstrom'\n",
     ")\n",
@@ -211,7 +217,7 @@
    "outputs": [],
    "source": [
     "da = iofq.compute()\n",
-    "da.plot(norm='log', scale={'Q': 'log'})"
+    "da.plot(norm='log', scale={'Q': 'log'}, aspect='equal')"
    ]
   },
   {
@@ -264,7 +270,7 @@
     "\n",
     "iofqs = {str(key): results[key] for key in iofqs}\n",
     "iofqs = {key: val if val.bins is None else val.hist() for key, val in iofqs.items()}\n",
-    "display(sc.plot(iofqs, norm='log', scale={'Q': 'log'}))"
+    "display(sc.plot(iofqs, norm='log', scale={'Q': 'log'}, aspect='equal'))"
    ]
   }
  ],

src/esssans/beam_center_finder.py

@@ -12,9 +12,9 @@
 from .conversions import (
     ElasticCoordTransformGraph,
     calibrate_positions,
+    compute_Q,
     detector_to_wavelength,
     mask_wavelength,
-    to_Q,
 )
 from .i_of_q import merge_spectra
 from .logging import get_logger
@@ -167,11 +167,15 @@ def _iofq_in_quadrants(
     phi = data.transform_coords(
         'phi', graph=graph, keep_intermediate=False, keep_inputs=False
     ).coords['phi']
+    if phi.bins is not None or 'wavelength' in phi.dims:
+        # If gravity-correction is enabled, phi depends on wavelength (and event).
+        # We cannot handle this below, so we approximate phi by the mean value.
+        phi = phi.mean('wavelength')
     phi_bins = sc.linspace('phi', -pi, pi, 5, unit='rad')
     quadrants = ['south-west', 'south-east', 'north-east', 'north-west']
 
     providers = [
-        to_Q,
+        compute_Q,
         merge_spectra,
         normalize,
         iofq_denominator,

src/esssans/conversions.py

@@ -19,6 +19,7 @@
     MaskedData,
     MonitorType,
     Numerator,
+    QxyBins,
     RawMonitor,
     RunType,
     WavelengthMask,
@@ -73,7 +74,7 @@ def two_theta(
     scattered_beam: sc.Variable,
     wavelength: sc.Variable,
     gravity: sc.Variable,
-) -> sc.Variable:
+) -> dict[str, sc.Variable]:
     """
     Compute the scattering angle from the incident and scattered beam vectors, taking
     into account the effects of gravity.
@@ -84,7 +85,6 @@ def two_theta(
     L2 = sc.norm(scattered_beam)
 
     x_term = cylindrical_x(cyl_x_unit_vector(gravity, incident_beam), scattered_beam)
-    x_term *= x_term
 
     y_term = sc.to_unit(wavelength, elem_unit(L2), copy=True)
     y_term *= y_term
@@ -98,6 +98,9 @@ def two_theta(
     else:
         y_term = drop * y_term
     y_term += cylindrical_y(cyl_y_unit_vector(gravity), scattered_beam)
+    phi = sc.atan2(y=y_term, x=x_term)
+
+    x_term *= x_term
     y_term *= y_term
 
     if set(x_term.dims).issubset(set(y_term.dims)):
@@ -107,17 +110,31 @@ def two_theta(
     out = sc.sqrt(y_term, out=y_term)
     out /= L2
     out = sc.asin(out, out=out)
-    return out
+    return {'two_theta': out, 'phi': phi}
 
 
-def phi(cylindrical_x: sc.Variable, cylindrical_y: sc.Variable) -> sc.Variable:
+def phi_no_gravity(
+    cylindrical_x: sc.Variable, cylindrical_y: sc.Variable
+) -> sc.Variable:
     """
-    Compute the cylindrial phi angle around the incident beam. Note that it is assumed
+    Compute the cylindrical phi angle around the incident beam. Note that it is assumed
     here that the incident beam is perpendicular to the gravity vector.
     """
     return sc.atan2(y=cylindrical_y, x=cylindrical_x)
 
 
+def Qxy(Q: sc.Variable, phi: sc.Variable) -> dict[str, sc.Variable]:
+    """
+    Compute the Qx and Qy components of the scattering vector from the scattering angle,
+    wavelength, and phi angle.
+    """
+    Qx = sc.cos(phi)
+    Qx *= Q
+    Qy = sc.sin(phi)
+    Qy *= Q
+    return {'Qx': Qx, 'Qy': Qy}
+
+
 ElasticCoordTransformGraph = NewType('ElasticCoordTransformGraph', dict)
 MonitorCoordTransformGraph = NewType('MonitorCoordTransformGraph', dict)
 
@@ -157,12 +174,15 @@ def sans_elastic(gravity: Optional[CorrectForGravity]) -> ElasticCoordTransformG
     """  # noqa: E501
     graph = {**beamline.beamline(scatter=True), **tof.elastic_Q('tof')}
     if gravity:
-        graph['two_theta'] = two_theta
+        del graph['two_theta']
+        graph[('two_theta', 'phi')] = two_theta
+    else:
+        graph['phi'] = phi_no_gravity
     graph['cyl_x_unit_vector'] = cyl_x_unit_vector
     graph['cyl_y_unit_vector'] = cyl_y_unit_vector
     graph['cylindrical_x'] = cylindrical_x
     graph['cylindrical_y'] = cylindrical_y
-    graph['phi'] = phi
+    graph[('Qx', 'Qy')] = Qxy
     return ElasticCoordTransformGraph(graph)
 
 
@@ -222,15 +242,22 @@ def mask_wavelength(
     return CleanWavelengthMasked[RunType, IofQPart](da)
 
 
-def to_Q(
-    data: CleanWavelengthMasked[RunType, IofQPart], graph: ElasticCoordTransformGraph
+def compute_Q(
+    data: CleanWavelengthMasked[RunType, IofQPart],
+    graph: ElasticCoordTransformGraph,
+    compute_Qxy: Optional[QxyBins],
 ) -> CleanQ[RunType, IofQPart]:
     """
     Convert a data array from wavelength to Q.
     """
-    # Keep naming of wavelength dim, subsequent steps use a (Q, wavelength) binning.
+    # Keep naming of wavelength dim, subsequent steps use a (Q[xy], wavelength) binning.
     return CleanQ[RunType, IofQPart](
-        data.transform_coords('Q', graph=graph, rename_dims=False)
+        data.transform_coords(
+            ('Qx', 'Qy') if compute_Qxy else 'Q',
+            graph=graph,
+            keep_intermediate=False,
+            rename_dims=False,
+        )
     )
 
 
@@ -241,5 +268,5 @@ def to_Q(
     monitor_to_wavelength,
     detector_to_wavelength,
     mask_wavelength,
-    to_Q,
+    compute_Q,
 )

src/esssans/i_of_q.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 # Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
 
+import uuid
 from typing import Optional
 
 import scipp as sc
@@ -22,6 +23,7 @@
     MonitorType,
     NonBackgroundWavelengthRange,
     QBins,
+    QxyBins,
     ReturnEvents,
     RunType,
     SampleRun,
@@ -130,7 +132,10 @@ def resample_direct_beam(
 
 
 def merge_spectra(
-    data: CleanQ[RunType, IofQPart], q_bins: QBins, dims_to_keep: Optional[DimsToKeep]
+    data: CleanQ[RunType, IofQPart],
+    q_bins: Optional[QBins],
+    qxy_bins: Optional[QxyBins],
+    dims_to_keep: Optional[DimsToKeep],
 ) -> CleanSummedQ[RunType, IofQPart]:
     """
     Merges all spectra:
@@ -157,7 +162,11 @@ def merge_spectra(
     if dims_to_keep is not None:
         dims_to_reduce -= set(dims_to_keep)
 
-    edges = {'Q': q_bins} if isinstance(q_bins, int) else {q_bins.dim: q_bins}
+    if qxy_bins:
+        # We make Qx the inner dim, such that plots naturally show Qx on the x-axis.
+        edges = {'Qy': qxy_bins['Qy'], 'Qx': qxy_bins['Qx']}
+    else:
+        edges = {'Q': q_bins}
 
     if data.bins is not None:
         q_all_pixels = data.bins.concat(dims_to_reduce)
@@ -168,7 +177,10 @@ def merge_spectra(
         # flattening more expensive.
         stripped = data.copy(deep=False)
         for name, coord in data.coords.items():
-            if name not in {'Q', 'wavelength'} and set(coord.dims) & dims_to_reduce:
+            if (
+                name not in {'Q', 'Qx', 'Qy', 'wavelength'}
+                and set(coord.dims) & dims_to_reduce
+            ):
                 del stripped.coords[name]
         to_flatten = [dim for dim in data.dims if dim in dims_to_reduce]
 
@@ -178,10 +190,26 @@ def merge_spectra(
             data_dims.remove(dim)
             data_dims.append(dim)
         stripped = stripped.transpose(data_dims)
-        # Flatten to Q such that `hist` below will turn this into the new Q dim.
-        flat = stripped.flatten(dims=to_flatten, to='Q')
+        # Flatten to helper dim such that `hist` will turn this into the new Q dim(s).
+        # For sc.hist this has to be named 'Q'.
+        helper_dim = 'Q'
+        flat = stripped.flatten(dims=to_flatten, to=helper_dim)
 
-        out = flat.hist(**edges)
+        if len(edges) == 1:
+            out = flat.hist(**edges)
+        else:
+            # sc.hist (or the underlying sc.bin) cannot deal with extra data dims,
+            # work around by flattening and regrouping.
+            for dim in flat.dims:
+                if dim == helper_dim:
+                    continue
+                if dim not in flat.coords:
+                    flat.coords[dim] = sc.arange(dim, flat.sizes[dim])
+            out = (
+                flat.flatten(to=str(uuid.uuid4()))
+                .group(*[flat.coords[dim] for dim in flat.dims if dim != helper_dim])
+                .hist(**edges)
+            )
     return CleanSummedQ[RunType, IofQPart](out.squeeze())
 
 

src/esssans/normalization.py

@@ -406,11 +406,7 @@ def _reduce(da: sc.DataArray) -> sc.DataArray:
                 )
     elif numerator.bins is not None:
         numerator = numerator.hist()
-    if numerator.bins is not None and denominator.bins is None:
-        da = numerator.bins / sc.lookup(func=denominator, dim='Q')
-    else:
-        da = numerator / denominator
-    return IofQ[RunType](da)
+    return IofQ[RunType](numerator / denominator)
 
 
 providers = (

src/esssans/types.py

@@ -114,6 +114,9 @@ class TransmissionRun(sciline.Scope[RunType, int], int):
 QBins = NewType('QBins', sc.Variable)
 """Q binning"""
 
+QxyBins = NewType('QxyBins', dict[str, sc.Variable])
+"""Binning for 'Qx' and 'Qy'. If set this overrides QBins."""
+
 NonBackgroundWavelengthRange = NewType('NonBackgroundWavelengthRange', sc.Variable)
 """Range of wavelengths that are not considered background in the monitor"""
 

src/esssans/uncertainty.py

@@ -68,14 +68,25 @@ def broadcast_to_events_with_upper_bound_variances(
     if da.variances is None:
         return da
     constituents = events.bins.constituents
-    Q = constituents['data'].coords['Q']
-    constituents['data'] = sc.DataArray(
-        sc.ones(sizes=Q.sizes, dtype='float32'), coords={'Q': Q}
-    )
+
+    if 'Q' in constituents['data'].coords:
+        Q = constituents['data'].coords['Q']
+        constituents['data'] = sc.DataArray(
+            sc.ones(sizes=Q.sizes, dtype='float32'), coords={'Q': Q}
+        )
+        edges = {'Q': da.coords['Q']}
+    else:
+        Qx = constituents['data'].coords['Qx']
+        Qy = constituents['data'].coords['Qy']
+        constituents['data'] = sc.DataArray(
+            sc.ones(sizes=Qx.sizes, dtype='float32'),
+            coords={'Qx': Qx, 'Qy': Qy},
+        )
+        edges = {'Qy': da.coords['Qy'], 'Qx': da.coords['Qx']}
     # Combine all bins of the events that correspond to the same bin in the input data
     to_concat = set(events.dims) - set(da.dims)
     binned = sc.DataArray(sc.bins(**constituents).bins.concat(to_concat))
-    counts = binned.hist(Q=da.coords['Q'])
+    counts = binned.hist(**edges)
     da = da.copy()
     da.variances *= counts.values
     da.data = sc.bins_like(events, da.data)

tests/loki/common.py

@@ -13,6 +13,7 @@
     EmptyBeamRun,
     FileList,
     QBins,
+    QxyBins,
     ReturnEvents,
     SampleRun,
     TransmissionRun,
@@ -22,7 +23,9 @@
 
 
 def make_params(
-    sample_runs: Optional[List[str]] = None, background_runs: Optional[List[str]] = None
+    sample_runs: Optional[List[str]] = None,
+    background_runs: Optional[List[str]] = None,
+    qxy: bool = False,
 ) -> dict:
     params = default_parameters.copy()
 
@@ -39,16 +42,22 @@ def make_params(
     params[FileList[EmptyBeamRun]] = ['60392-2022-02-28_2215.nxs']
 
     params[WavelengthBins] = sc.linspace(
-        'wavelength', start=1.0, stop=13.0, num=201, unit='angstrom'
+        'wavelength', start=1.0, stop=13.0, num=51, unit='angstrom'
     )
     params[BeamStopPosition] = sc.vector([-0.026, -0.022, 0.0], unit='m')
     params[BeamStopRadius] = sc.scalar(0.042, unit='m')
     params[CorrectForGravity] = True
     params[UncertaintyBroadcastMode] = UncertaintyBroadcastMode.upper_bound
     params[ReturnEvents] = False
 
-    params[QBins] = sc.linspace(
-        dim='Q', start=0.01, stop=0.3, num=101, unit='1/angstrom'
-    )
+    if qxy:
+        params[QxyBins] = {
+            'Qx': sc.linspace('Qx', -0.3, 0.3, 91, unit='1/angstrom'),
+            'Qy': sc.linspace('Qy', -0.2, 0.3, 78, unit='1/angstrom'),
+        }
+    else:
+        params[QBins] = sc.linspace(
+            dim='Q', start=0.01, stop=0.3, num=101, unit='1/angstrom'
+        )
 
     return params