add function in tools to set xtype to tth or q

Author: yucongalicechen

src/diffpy/labpdfproc/functions.py

@@ -5,12 +5,7 @@
 import pandas as pd
 from scipy.interpolate import interp1d
 
-from diffpy.utils.scattering_objects.diffraction_objects import (
-    ANGLEQUANTITIES,
-    DQUANTITIES,
-    XQUANTITIES,
-    Diffraction_object,
-)
+from diffpy.utils.scattering_objects.diffraction_objects import XQUANTITIES, Diffraction_object
 
 RADIUS_MM = 1
 N_POINTS_ON_DIAMETER = 300
@@ -251,7 +246,7 @@ def _cve_method(method):
 
 def interpolate_to_xtype_grid(cve_do, xtype):
     f"""
-    interpolates the cve grid to the xtype user specifies, raise an error if xtype is invalid
+    interpolates the cve grid to the xtype user specified
 
     Parameters
     ----------
@@ -264,10 +259,7 @@ def interpolate_to_xtype_grid(cve_do, xtype):
     -------
     the new diffraction object with interpolated cve curves
     """
-
-    if xtype.lower() not in XQUANTITIES:
-        raise ValueError(f"Unknown xtype: {xtype}. Allowed xtypes are {*XQUANTITIES, }.")
-    if xtype.lower() in ANGLEQUANTITIES or xtype.lower() in DQUANTITIES:
+    if xtype == "tth":
         return cve_do
 
     orig_grid, orig_cve = cve_do.on_tth[0], cve_do.on_tth[1]

src/diffpy/labpdfproc/tools.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 
 from diffpy.labpdfproc.mud_calculator import compute_mud
+from diffpy.utils.scattering_objects.diffraction_objects import QQUANTITIES, XQUANTITIES
 from diffpy.utils.tools import get_package_info, get_user_info
 
 WAVELENGTHS = {"Mo": 0.71, "Ag": 0.59, "Cu": 1.54}
@@ -135,6 +136,25 @@ def set_wavelength(args):
     return args
 
 
+def set_xtype(args):
+    f"""
+    Set the xtype based on the given input arguments, raise an error if xtype is not one of {*XQUANTITIES, }
+
+    Parameters
+    ----------
+    args argparse.Namespace
+        the arguments from the parser
+
+    Returns
+    -------
+    args argparse.Namespace
+    """
+    if args.xtype.lower() not in XQUANTITIES:
+        raise ValueError(f"Unknown xtype: {args.xtype}. Allowed xtypes are {*XQUANTITIES, }.")
+    args.xtype = "q" if args.xtype.lower() in QQUANTITIES else "tth"
+    return args
+
+
 def set_mud(args):
     """
     Set the mud based on the given input arguments
@@ -257,6 +277,7 @@ def preprocessing_args(args):
     args = set_input_lists(args)
     args.output_directory = set_output_directory(args)
     args = set_wavelength(args)
+    args = set_xtype(args)
     args = set_mud(args)
     args = load_user_metadata(args)
     return args

tests/test_functions.py

@@ -10,7 +10,7 @@
     compute_cve,
     interpolate_to_xtype_grid,
 )
-from diffpy.utils.scattering_objects.diffraction_objects import XQUANTITIES, Diffraction_object
+from diffpy.utils.scattering_objects.diffraction_objects import Diffraction_object
 
 params1 = [
     ([0.5, 3, 1], {(0.0, -0.5), (0.0, 0.0), (0.5, 0.0), (-0.5, 0.0), (0.0, 0.5)}),
@@ -106,20 +106,6 @@ def test_interpolate_xtype(inputs, expected, mocker):
     assert actual_cve_do == expected_cve_do
 
 
-def test_interpolate_xtype_bad():
-    input_cve_do = _instantiate_test_do(
-        np.array([30, 60, 90]),
-        np.array([1, 2, 3]),
-        xtype="tth",
-        name="absorption correction, cve, for test",
-        scat_quantity="cve",
-    )
-    with pytest.raises(
-        ValueError, match=re.escape(f"Unknown xtype: invalid. Allowed xtypes are {*XQUANTITIES, }.")
-    ):
-        interpolate_to_xtype_grid(input_cve_do, xtype="invalid")
-
-
 def test_compute_cve(mocker):
     xarray, yarray = np.array([90, 90.1, 90.2]), np.array([2, 2, 2])
     expected_cve = np.array([0.5, 0.5, 0.5])

tests/test_tools.py

@@ -16,7 +16,9 @@
     set_mud,
     set_output_directory,
     set_wavelength,
+    set_xtype,
 )
+from diffpy.utils.scattering_objects.diffraction_objects import XQUANTITIES
 
 # Use cases can be found here: https://github.com/diffpy/diffpy.labpdfproc/issues/48
 
@@ -189,6 +191,31 @@ def test_set_wavelength_bad(inputs, msg):
         actual_args = set_wavelength(actual_args)
 
 
+params4 = [
+    ([], ["tth"]),
+    (["--xtype", "2theta"], ["tth"]),
+    (["--xtype", "d"], ["tth"]),
+    (["--xtype", "q"], ["q"]),
+]
+
+
+@pytest.mark.parametrize("inputs, expected", params4)
+def test_set_xtype(inputs, expected):
+    cli_inputs = ["2.5", "data.xy"] + inputs
+    actual_args = get_args(cli_inputs)
+    actual_args = set_xtype(actual_args)
+    assert actual_args.xtype == expected[0]
+
+
+def test_set_xtype_bad():
+    cli_inputs = ["2.5", "data.xy", "--xtype", "invalid"]
+    actual_args = get_args(cli_inputs)
+    with pytest.raises(
+        ValueError, match=re.escape(f"Unknown xtype: invalid. Allowed xtypes are {*XQUANTITIES, }.")
+    ):
+        actual_args = set_xtype(actual_args)
+
+
 def test_set_mud(user_filesystem):
     cli_inputs = ["2.5", "data.xy"]
     actual_args = get_args(cli_inputs)