diff --git a/docs/examples/zoom.ipynb b/docs/examples/zoom.ipynb index e77946ee..ed363794 100644 --- a/docs/examples/zoom.ipynb +++ b/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'))" ] } ], diff --git a/src/esssans/beam_center_finder.py b/src/esssans/beam_center_finder.py index 0fba6f6a..6e08da88 100644 --- a/src/esssans/beam_center_finder.py +++ b/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, diff --git a/src/esssans/conversions.py b/src/esssans/conversions.py index b30f7648..12a147f8 100644 --- a/src/esssans/conversions.py +++ b/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, ) diff --git a/src/esssans/i_of_q.py b/src/esssans/i_of_q.py index 79938f4d..6db4b9a3 100644 --- a/src/esssans/i_of_q.py +++ b/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()) diff --git a/src/esssans/normalization.py b/src/esssans/normalization.py index 95226e22..46c1078f 100644 --- a/src/esssans/normalization.py +++ b/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 = ( diff --git a/src/esssans/types.py b/src/esssans/types.py index 65e8aed3..4ba01d91 100644 --- a/src/esssans/types.py +++ b/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""" diff --git a/src/esssans/uncertainty.py b/src/esssans/uncertainty.py index f708fe86..7e3346fe 100644 --- a/src/esssans/uncertainty.py +++ b/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) diff --git a/tests/loki/common.py b/tests/loki/common.py index 5de3fe5d..5bf2a10a 100644 --- a/tests/loki/common.py +++ b/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,7 +42,7 @@ 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') @@ -47,8 +50,14 @@ def make_params( 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 diff --git a/tests/loki/iofq_test.py b/tests/loki/iofq_test.py index 16b57474..503fd236 100644 --- a/tests/loki/iofq_test.py +++ b/tests/loki/iofq_test.py @@ -9,11 +9,18 @@ import scipp as sc import esssans as sans +from esssans.conversions import ElasticCoordTransformGraph from esssans.types import ( BackgroundSubtractedIofQ, BeamCenter, + CleanWavelengthMasked, + CorrectForGravity, DimsToKeep, + Numerator, + QBins, + QxyBins, ReturnEvents, + SampleRun, UncertaintyBroadcastMode, WavelengthBands, WavelengthBins, @@ -27,8 +34,9 @@ def loki_providers() -> List[Callable]: return list(sans.providers + sans.loki.providers) -def test_can_create_pipeline(): - pipeline = sciline.Pipeline(loki_providers(), params=make_params()) +@pytest.mark.parametrize('qxy', [False, True]) +def test_can_create_pipeline(qxy: bool): + pipeline = sciline.Pipeline(loki_providers(), params=make_params(qxy=qxy)) pipeline.get(BackgroundSubtractedIofQ) @@ -36,12 +44,21 @@ def test_can_create_pipeline(): 'uncertainties', [UncertaintyBroadcastMode.drop, UncertaintyBroadcastMode.upper_bound], ) -def test_pipeline_can_compute_IofQ(uncertainties): - params = make_params() +@pytest.mark.parametrize('qxy', [False, True]) +def test_pipeline_can_compute_IofQ(uncertainties, qxy: bool): + params = make_params(qxy=qxy) params[UncertaintyBroadcastMode] = uncertainties pipeline = sciline.Pipeline(loki_providers(), params=params) result = pipeline.compute(BackgroundSubtractedIofQ) - assert result.dims == ('Q',) + assert result.dims == ('Qy', 'Qx') if qxy else ('Q',) + if qxy: + assert sc.identical(result.coords['Qx'], params[QxyBins]['Qx']) + assert sc.identical(result.coords['Qy'], params[QxyBins]['Qy']) + assert result.sizes['Qx'] == 90 + assert result.sizes['Qy'] == 77 + else: + assert sc.identical(result.coords['Q'], params[QBins]) + assert result.sizes['Q'] == 100 @pytest.mark.parametrize( @@ -51,14 +68,16 @@ def test_pipeline_can_compute_IofQ(uncertainties): @pytest.mark.parametrize( 'target', [sans.IofQ[sans.SampleRun], sans.BackgroundSubtractedIofQ] ) -def test_pipeline_can_compute_IofQ_in_event_mode(uncertainties, target): - params = make_params() +@pytest.mark.parametrize('qxy', [False, True]) +def test_pipeline_can_compute_IofQ_in_event_mode(uncertainties, target, qxy: bool): + params = make_params(qxy=qxy) params[UncertaintyBroadcastMode] = uncertainties pipeline = sciline.Pipeline(loki_providers(), params=params) reference = pipeline.compute(target) pipeline[ReturnEvents] = True result = pipeline.compute(target) assert result.bins is not None + assert result.dims == ('Qy', 'Qx') if qxy else ('Q',) assert sc.allclose( sc.values(reference.data), sc.values(result.hist().data), @@ -68,7 +87,7 @@ def test_pipeline_can_compute_IofQ_in_event_mode(uncertainties, target): if uncertainties == UncertaintyBroadcastMode.drop: tol = sc.scalar(1e-14) else: - tol = sc.scalar(1e-9) + tol = sc.scalar(1e-8 if qxy else 1e-9) assert sc.allclose( sc.variances(reference).data, sc.variances(result.hist()).data, @@ -77,8 +96,9 @@ def test_pipeline_can_compute_IofQ_in_event_mode(uncertainties, target): ) -def test_pipeline_can_compute_IofQ_in_wavelength_bands(): - params = make_params() +@pytest.mark.parametrize('qxy', [False, True]) +def test_pipeline_can_compute_IofQ_in_wavelength_bands(qxy: bool): + params = make_params(qxy=qxy) params[WavelengthBands] = sc.linspace( 'wavelength', params[WavelengthBins].min(), @@ -87,12 +107,13 @@ def test_pipeline_can_compute_IofQ_in_wavelength_bands(): ) pipeline = sciline.Pipeline(loki_providers(), params=params) result = pipeline.compute(BackgroundSubtractedIofQ) - assert result.dims == ('band', 'Q') + assert result.dims == ('band', 'Qy', 'Qx') if qxy else ('band', 'Q') assert result.sizes['band'] == 10 -def test_pipeline_can_compute_IofQ_in_overlapping_wavelength_bands(): - params = make_params() +@pytest.mark.parametrize('qxy', [False, True]) +def test_pipeline_can_compute_IofQ_in_overlapping_wavelength_bands(qxy: bool): + params = make_params(qxy=qxy) # Bands have double the width edges = sc.linspace( 'band', params[WavelengthBins].min(), params[WavelengthBins].max(), 12 @@ -102,16 +123,17 @@ def test_pipeline_can_compute_IofQ_in_overlapping_wavelength_bands(): ).transpose() pipeline = sciline.Pipeline(loki_providers(), params=params) result = pipeline.compute(BackgroundSubtractedIofQ) - assert result.dims == ('band', 'Q') + assert result.dims == ('band', 'Qy', 'Qx') if qxy else ('band', 'Q') assert result.sizes['band'] == 10 -def test_pipeline_can_compute_IofQ_in_layers(): - params = make_params() +@pytest.mark.parametrize('qxy', [False, True]) +def test_pipeline_can_compute_IofQ_in_layers(qxy: bool): + params = make_params(qxy=qxy) params[DimsToKeep] = ['layer'] pipeline = sciline.Pipeline(loki_providers(), params=params) result = pipeline.compute(BackgroundSubtractedIofQ) - assert result.dims == ('layer', 'Q') + assert result.dims == ('layer', 'Qy', 'Qx') if qxy else ('layer', 'Q') assert result.sizes['layer'] == 4 @@ -131,3 +153,47 @@ def test_beam_center_from_center_of_mass_is_close_to_verified_result(): center = pipeline.compute(BeamCenter) reference = sc.vector([-0.0309889, -0.0168854, 0], unit='m') assert sc.allclose(center, reference) + + +def test_phi_with_gravity(): + params = make_params() + pipeline = sciline.Pipeline(loki_providers(), params=params) + pipeline[CorrectForGravity] = False + data_no_grav = pipeline.compute( + CleanWavelengthMasked[SampleRun, Numerator] + ).flatten(to='pixel') + graph_no_grav = pipeline.compute(ElasticCoordTransformGraph) + pipeline[CorrectForGravity] = True + data_with_grav = ( + pipeline.compute(CleanWavelengthMasked[SampleRun, Numerator]) + .flatten(to='pixel') + .hist(wavelength=sc.linspace('wavelength', 1.0, 12.0, 101, unit='angstrom')) + ) + graph_with_grav = pipeline.compute(ElasticCoordTransformGraph) + + no_grav = data_no_grav.transform_coords(('two_theta', 'phi'), graph_no_grav) + two_theta_no_grav = no_grav.coords['two_theta'] + phi_no_grav = no_grav.coords['phi'] + with_grav = data_with_grav.transform_coords(('two_theta', 'phi'), graph_with_grav) + phi_with_grav = with_grav.coords['phi'].mean('wavelength') + + assert not sc.identical(phi_no_grav, phi_with_grav) + + # Exclude pixels near the origin, since phi will vary a lot there. + not_near_origin = two_theta_no_grav > sc.scalar(0.1, unit='deg').to(unit='rad') + assert sc.all( + sc.isclose( + phi_no_grav[not_near_origin], + phi_with_grav[not_near_origin], + atol=sc.scalar(3.0, unit='deg').to(unit='rad'), + ) + ) + + # Phi is in [-pi, pi], measured from the X axis. + pos_x = sc.abs(phi_no_grav) < sc.scalar(90.0, unit='deg').to(unit='rad') + # Phi is larger with gravity, since it gives the position where it would have + # been detected without gravity. That is, with gravity all points are pulled + # "up" in the XY plane, so the angle is larger for positive X and smaller for + # negative X. + assert sc.all(phi_no_grav[pos_x] < phi_with_grav[pos_x]) + assert sc.all(phi_no_grav[~pos_x] > phi_with_grav[~pos_x])