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 load_user_metadata
+from diffpy.labpdfproc.tools import load_user_metadata, set_wavelength
 from diffpy.utils.parsers.loaddata import loadData
 from diffpy.utils.scattering_objects.diffraction_objects import XQUANTITIES, Diffraction_object
 
@@ -73,7 +73,8 @@ def get_args():
 def main():
     args = get_args()
     args = load_user_metadata(args)
-    wavelength = WAVELENGTHS[args.anode_type]
+    args.wavelength = set_wavelength(args)
+
     filepath = Path(args.input_file)
     outfilestem = filepath.stem + "_corrected"
     corrfilestem = filepath.stem + "_cve"
@@ -91,7 +92,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,
@@ -102,7 +103,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

@@ -1,8 +1,9 @@
+import argparse
 from argparse import ArgumentParser
 
 import pytest
 
-from diffpy.labpdfproc.tools import load_user_metadata
+from diffpy.labpdfproc.tools import load_user_metadata, set_wavelength
 
 params5 = [
     ([[]], []),
@@ -25,3 +26,33 @@ def test_load_user_metadata(inputs, expected):
     for expected_pair in expected:
         setattr(expected_args, expected_pair[0], expected_pair[1])
     assert actual_args == expected_args
+
+
+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,3 +1,7 @@
+WAVELENGTHS = {"Mo": 0.71, "Ag": 0.59, "Cu": 1.54}
+known_sources = [key for key in WAVELENGTHS.keys()]
+
+
 def load_key_value_pair(s):
     items = s.split("=")
     key = items[0].strip()
@@ -29,3 +33,37 @@ def load_user_metadata(args):
 # question 2: more than one =?
 # if i have key is hello, value is hello=world, then i can have hello=hello=world to process okay
 # if i have = in key then it's still processing as value => would this be an issue?
+
+
+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"]