merged updates from wavelength

Author: yucongalicechen

src/diffpy/labpdfproc/labpdfprocapp.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 
 from diffpy.labpdfproc.functions import apply_corr, compute_cve
-from diffpy.labpdfproc.tools import set_output_directory
+from diffpy.labpdfproc.tools import set_output_directory, set_wavelength
 from diffpy.utils.parsers.loaddata import loadData
 from diffpy.utils.scattering_objects.diffraction_objects import XQUANTITIES, Diffraction_object
 
@@ -64,9 +64,10 @@ def get_args():
 
 def main():
     args = get_args()
-    wavelength = WAVELENGTHS[args.anode_type]
-    filepath = Path(args.input_file)
     args.output_directory = set_output_directory(args)
+    args.wavelength = set_wavelength(args)
+
+    filepath = Path(args.input_file)
     outfilestem = filepath.stem + "_corrected"
     corrfilestem = filepath.stem + "_cve"
     outfile = args.output_directory / (outfilestem + ".chi")
@@ -83,7 +84,7 @@ def main():
             f"exists. Please rerun specifying -f if you want to overwrite it"
         )
 
-    input_pattern = Diffraction_object(wavelength=wavelength)
+    input_pattern = Diffraction_object(wavelength=args.wavelength)
     xarray, yarray = loadData(args.input_file, unpack=True)
     input_pattern.insert_scattering_quantity(
         xarray,
@@ -94,7 +95,7 @@ def main():
         metadata={"muD": args.mud, "anode_type": args.anode_type},
     )
 
-    absorption_correction = compute_cve(input_pattern, args.mud, wavelength)
+    absorption_correction = compute_cve(input_pattern, args.mud, args.wavelength)
     corrected_data = apply_corr(input_pattern, absorption_correction)
     corrected_data.name = f"Absorption corrected input_data: {input_pattern.name}"
     corrected_data.dump(f"{outfile}", xtype="tth")

src/diffpy/labpdfproc/tests/test_tools.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from diffpy.labpdfproc.tools import set_output_directory
+from diffpy.labpdfproc.tools import set_output_directory, set_wavelength
 
 params1 = [
     ([None], [Path.cwd().resolve()]),
@@ -35,3 +35,33 @@ def test_set_output_directory_bad():
         actual_args.output_directory = set_output_directory(actual_args)
         assert Path(actual_args.output_directory).exists()
         assert not Path(actual_args.output_directory).is_dir()
+
+
+params2 = [
+    ([None, None], [0.71]),
+    ([None, "Ag"], [0.59]),
+    ([0.25, "Ag"], [0.25]),
+    ([0.25, None], [0.25]),
+]
+
+
+@pytest.mark.parametrize("inputs, expected", params2)
+def test_set_wavelength(inputs, expected):
+    expected_wavelength = expected[0]
+    actual_args = argparse.Namespace(wavelength=inputs[0], anode_type=inputs[1])
+    actual_wavelength = set_wavelength(actual_args)
+    assert actual_wavelength == expected_wavelength
+
+
+params3 = [
+    ([None, "invalid"]),
+    ([0, None]),
+    ([-1, "Mo"]),
+]
+
+
+@pytest.mark.parametrize("inputs", params3)
+def test_set_wavelength_bad(inputs):
+    with pytest.raises(ValueError):
+        actual_args = argparse.Namespace(wavelength=inputs[0], anode_type=inputs[1])
+        actual_args.wavelength = set_wavelength(actual_args)

src/diffpy/labpdfproc/tools.py

@@ -1,5 +1,8 @@
 from pathlib import Path
 
+WAVELENGTHS = {"Mo": 0.71, "Ag": 0.59, "Cu": 1.54}
+known_sources = [key for key in WAVELENGTHS.keys()]
+
 
 def set_output_directory(args):
     """
@@ -23,3 +26,37 @@ def set_output_directory(args):
     output_dir = Path(args.output_directory).resolve() if args.output_directory else Path.cwd().resolve()
     output_dir.mkdir(parents=True, exist_ok=True)
     return output_dir
+
+
+def set_wavelength(args):
+    """
+    Set the wavelength based on the given input arguments
+
+    Parameters
+    ----------
+    args argparse.Namespace
+        the arguments from the parser
+
+    Returns
+    -------
+        float: the wavelength value
+
+    we raise an ValueError if the input wavelength is non-positive
+    or if the input anode_type is not one of the known sources
+
+    """
+    if args.wavelength is not None and args.wavelength <= 0:
+        raise ValueError("Please rerun the program specifying a positive float number.")
+    if not args.wavelength and args.anode_type and args.anode_type not in WAVELENGTHS:
+        raise ValueError(
+            f"Invalid anode type {args.anode_type}. "
+            f"Please rerun the program to either specify a wavelength as a positive float number "
+            f"or specify anode_type as one of {known_sources}."
+        )
+
+    if args.wavelength:
+        return args.wavelength
+    elif args.anode_type:
+        return WAVELENGTHS[args.anode_type]
+    else:
+        return WAVELENGTHS["Mo"]