feat: add gravity toggle

Author: jokasimr

src/ess/amor/__init__.py

@@ -30,6 +30,7 @@
     AngularResolution,
     ChopperFrequency,
     ChopperPhase,
+    GravityToggle,
     QThetaFigure,
     ReflectivityDiagnosticsView,
     SampleSizeResolution,
@@ -76,6 +77,7 @@ def default_parameters() -> dict:
             sc.scalar(-0.75, unit='deg'),
             sc.scalar(0.75, unit='deg'),
         ),
+        GravityToggle: True,
     }
 
 

src/ess/amor/conversions.py

@@ -11,7 +11,7 @@
     ZIndexLimits,
 )
 from .geometry import Detector
-from .types import CoordTransformationGraph
+from .types import CoordTransformationGraph, GravityToggle
 
 
 def theta(wavelength, divergence_angle, L2, sample_rotation, detector_rotation):
@@ -68,6 +68,24 @@ def theta(wavelength, divergence_angle, L2, sample_rotation, detector_rotation):
     return out
 
 
+def theta_no_gravity(wavelength, divergence_angle, sample_rotation, detector_rotation):
+    '''
+    Angle of reflection.
+
+    Computes the angle between the scattering direction of
+    the neutron and the sample surface while disregarding the
+    effect of gravity.
+    '''
+    theta = (
+        divergence_angle.to(unit='rad', copy=False)
+        + detector_rotation.to(unit='rad')
+        - sample_rotation.to(unit='rad')
+    )
+    if wavelength.bins:
+        return sc.bins_like(wavelength, theta)
+    return theta
+
+
 def angle_of_divergence(theta, sample_rotation, angle_to_center_of_beam):
     """
     Difference between the incident angle and the center of the incident beam.
@@ -113,11 +131,11 @@ def wavelength(
     return out.to(unit='angstrom', copy=False)
 
 
-def coordinate_transformation_graph() -> CoordTransformationGraph:
+def coordinate_transformation_graph(gravity: GravityToggle) -> CoordTransformationGraph:
     return {
         "divergence_angle": "pixel_divergence_angle",
         "wavelength": wavelength,
-        "theta": theta,
+        "theta": theta if gravity else theta_no_gravity,
         "angle_of_divergence": angle_of_divergence,
         "Q": reflectometry_q,
         "L1": lambda chopper_distance: sc.abs(chopper_distance),

src/ess/amor/types.py

@@ -47,3 +47,5 @@ class AngleCenterOfIncomingToHorizon(
 WavelengthZIndexFigure = NewType("WavelengthZIndexFigure", Any)
 QThetaFigure = NewType("QThetaFigure", Any)
 ReflectivityDiagnosticsView = NewType("ReflectivityDiagnosticsView", Any)
+
+GravityToggle = NewType("GravityToggle", bool)

tests/amor/pipeline_test.py

@@ -71,6 +71,18 @@ def test_has_expected_coordinates(amor_pipeline: sciline.Pipeline):
     assert "Q_resolution" in reflectivity_over_q.coords
 
 
+@pytest.mark.filterwarnings("ignore:Failed to convert .* into a transformation")
+@pytest.mark.filterwarnings("ignore:Invalid transformation, missing attribute")
+def test_pipeline_no_gravity_correction(amor_pipeline: sciline.Pipeline):
+    # The sample rotation value in the file is slightly off, so we set it manually
+    amor_pipeline[SampleRotation[SampleRun]] = sc.scalar(0.85, unit="deg")
+    amor_pipeline[Filename[SampleRun]] = amor.data.amor_sample_run(608)
+    amor_pipeline[amor.types.GravityToggle] = False
+    reflectivity_over_q = amor_pipeline.compute(ReflectivityOverQ)
+    assert "Q" in reflectivity_over_q.coords
+    assert "Q_resolution" in reflectivity_over_q.coords
+
+
 @pytest.mark.filterwarnings("ignore:Failed to convert .* into a transformation")
 @pytest.mark.filterwarnings("ignore:Invalid transformation, missing attribute")
 def test_orso_pipeline(amor_pipeline: sciline.Pipeline):