Merge pull request #15 from arm61/amor_reflectometry

Amor reflectometry

src/ess/amor/__init__.py

@@ -0,0 +1,2 @@
+# flake8: noqa: F401
+from ess.amor.amor_data import AmorData, AmorReference, Normalisation

src/ess/reflectometry/__init__.py

@@ -0,0 +1,8 @@
+import scipp as sc
+from scipy.constants import neutron_mass, h, g
+
+h = (h * 1e20 / 1e6 *
+     (sc.units.kg * sc.units.angstrom * sc.units.angstrom / sc.units.us))
+HDM = h / (neutron_mass * sc.units.kg)
+G = -g * (sc.units.m / (sc.units.s * sc.units.s))
+G_ACC = sc.geometry.position(0. * G.unit, G, 0. * G.unit)

src/ess/reflectometry/binning.py

@@ -0,0 +1,82 @@
+# flake8: noqa: E501
+"""
+This module is focused on enabling different binning for reflectometry data.
+"""
+
+# author: Andrew R. McCluskey (arm61)
+
+import numpy as np
+import scipp as sc
+import scippneutron as scn
+
+
+def q_bin(data, bins=None, unit=None):
+    """
+    Return data that has been binned in the q-bins passed.
+
+    Args:
+        data (`ess.reflectometry.ReflData` or `ess.amor.AmorData` or `ess.amor.AmorReference`): data to be binned.
+        bins (`array_like`, optional): q-bin edges. Defaults to 200 bins between max and min q.
+        unit (`scipp._scipp.core.Unit`, optional): Unit of q bins. Defaults to 1/Å.
+
+    Returns:
+        (`scipp._scipp.core.DataArray`): Data array binned into qz with resolution.
+    """
+    if "qz" in data.event.coords and "tof" in data.event.coords:
+        if unit is None:
+            unit = data.event.coords['qz'].unit
+        erase = ['tof'] + data.data.dims
+        q_z_vector = sc.to_unit(data.event.coords["qz"], unit)
+        if bins is None:
+            bins = np.linspace(
+                sc.min(q_z_vector).values,
+                sc.max(q_z_vector).values, 200)
+        data.event.coords["qz"] = q_z_vector
+        edges = sc.array(dims=["qz"], unit=unit, values=bins)
+        binned = sc.bin(data.data, erase=erase, edges=[edges])
+        if "sigma_qz_by_qz" in data.event.coords:
+            qzr = []
+            for i in binned.data.values:
+                try:
+                    qzr.append(i.coords["sigma_qz_by_qz"].values.max())
+                except ValueError:
+                    qzr.append(0)
+            qzr = np.array(qzr)
+            binned.coords["sigma_qz_by_qz"] = sc.Variable(values=qzr,
+                                                          dims=["qz"])
+    else:
+        raise sc.NotFoundError("qz or tof coordinate cannot be found.")
+    return binned / (data.event.shape[0] * sc.units.dimensionless)
+
+
+def two_dimensional_bin(data, dims, bins=None, units=None):
+    """
+    Perform some arbitrary two-dimensional binning.
+
+    Args:
+        data (`ess.reflectometry.ReflData` or `ess.amor.AmorData` or `ess.amor.AmorReference`): data to be binned.
+        dims (`tuple` of `str`): The dimensions to be binned
+        bins (`tuple` of `array_like`, optional): Bin edges. Optional, defaults to min and max with 50 bins in each dim.
+        unit (`scipp._scipp.core.Unit`): Unit of bins. Optional, defaults to units of coord.
+
+    Returns:
+        (`scipp._scipp.core.DataArray`): Data array binned into qz with resolution.
+    """
+    for d in dims:
+        if d not in list(data.event.coords):
+            raise sc.NotFoundError(f'dim {d} not found in coords')
+    if bins is None:
+        bins = []
+        for d in dims:
+            vals = data.event.coords[d].values
+            bins.append(np.linspace(vals.min(), vals.max(), 50))
+    if units is None:
+        units = []
+        for d in dims:
+            units.append(data.event.coords[d].unit)
+    bin_edges = []
+    for i, d in enumerate(dims):
+        unit_change = sc.to_unit(data.event.coords[d], units[i])
+        data.event.coords[d] = unit_change
+        bin_edges.append(sc.array(dims=[d], unit=units[i], values=bins[i]))
+    return sc.bin(data.data.bins.concatenate('detector_id'), edges=bin_edges)

src/ess/reflectometry/corrections.py

@@ -0,0 +1,104 @@
+# flake8: noqa: E501
+"""
+Corrections to be used for neutron reflectometry reduction processes.
+"""
+import numpy as np
+import scipp as sc
+from scipy.special import erf
+from ess.reflectometry import HDM, G_ACC
+
+
+def angle_with_gravity(data, pixel_position, sample_position):
+    """
+    Find the angle of reflection when accounting for the presence of gravity.
+
+    Args:
+        data (`scipp._scipp.core.DataArray`): Reduction data array.
+        pixel_position (`scipp._scipp.core.VariableView`): Detector pixel positions, should be a `vector_3_float64`-type object.
+        sample_position (`scipp._scipp.core.VariableView`): Scattered neutron origin position.
+
+    Returns:
+        (`scipp._scipp.core.Variable`): Gravity corrected angle values.
+    """
+    # This is a workaround until scipp #1819 is resolved, at which time the following should be used instead
+    # velocity = sc.to_unit(HDM / wavelength, 'm/s')
+    # At which point the args can be changed to wavelength
+    # (where this is data.bins.constituents['data'].coords['wavelength'].astype(sc.dtype.float64) or similar)
+    # instead of data
+    velocity = sc.to_unit(
+        HDM / data.bins.constituents["data"].coords["wavelength"].astype(
+            sc.dtype.float64),
+        "m/s",
+    )
+    data.bins.constituents["data"].coords["velocity"] = velocity
+    velocity = data.bins.coords["velocity"]
+    velocity.unit = sc.units.m / sc.units.s
+    y_measured = sc.geometry.y(pixel_position)
+    z_measured = sc.geometry.z(pixel_position)
+    z_origin = sc.geometry.z(sample_position)
+    y_origin = sc.geometry.y(sample_position)
+    y_dash = y_dash0(velocity, z_origin, y_origin, z_measured, y_measured)
+    intercept = y_origin - y_dash * z_origin
+    y_true = z_measured * y_dash + intercept
+    angle = sc.to_unit(
+        sc.atan(y_true / z_measured).bins.constituents["data"], 'deg')
+    return angle
+
+
+def y_dash0(velocity, z_origin, y_origin, z_measured, y_measured):
+    """
+    Evaluation of the first dervative of the kinematic equations for for the trajectory of a neutron reflected from a surface.
+
+    Args:
+        velocity (`scipp._scipp.core.VariableView`): Neutron velocity.
+        z_origin (`scipp._scipp.core.Variable`): The z-origin position for the reflected neutron.
+        y_origin (`scipp._scipp.core.Variable`): The y-origin position for the reflected neutron.
+        z_origin (`scipp._scipp.core.Variable`): The z-measured position for the reflected neutron.
+        y_origin (`scipp._scipp.core.Variable`): The y-measured position for the reflected neutron.
+
+    Returns:
+        (`scipp._scipp.core.VariableView` or `array_like`): The gradient of the trajectory of the neutron at the origin position.
+    """
+    velocity2 = velocity * velocity
+    z_diff = z_measured - z_origin
+    y_diff = y_measured - y_origin
+    return -0.5 * sc.norm(G_ACC) * z_diff / velocity2 + y_diff / z_diff
+
+
+def illumination_correction(beam_size, sample_size, theta):
+    """
+    The factor by which the intensity should be multiplied to account for the
+    scattering geometry, where the beam is Gaussian in shape.
+
+    Args:
+        beam_size (:py:attr:`sc.Variable`): Width of incident beam.
+        sample_size (:py:attr:`sc.Variable`): Width of sample in the dimension of the beam.
+        theta (:py:attr:`sc.Variable`): Incident angle.
+
+    Returns:
+        (:py:attr:`array_like`): Correction factor.
+    """
+    beam_on_sample = beam_size / sc.sin(theta)
+    fwhm_to_std = 2 * np.sqrt(2 * np.log(2))
+    scale_factor = erf((sample_size / beam_on_sample * fwhm_to_std).values)
+    return sc.Variable(values=scale_factor, dims=theta.dims)
+
+
+def illumination_of_sample(beam_size, sample_size, theta):
+    """
+    Determine the illumination of the sample by the beam and therefore the size of this illuminated length.
+
+    Args:
+        beam_size (:py:attr:`sc.Variable`): Width of incident beam, in metres.
+        sample_size (:py:attr:`sc.Variable`): Width of sample in the dimension of the beam, in metres.
+        theta (:py:attr:`sc.Variable`): Incident angle.
+
+    Returns:
+        (`sc.Variable`): The size of the beam, for each theta, on the sample.
+    """
+    beam_on_sample = beam_size / sc.sin(theta)
+    if ((sc.mean(beam_on_sample)) > sample_size).value:
+        beam_on_sample = sc.broadcast(sample_size,
+                                      shape=theta.shape,
+                                      dims=theta.dims)
+    return beam_on_sample