Merge pull request #36 from yucongalicechen/wavelength

closes #25: overload wavelength

Author: sbillinge

src/diffpy/labpdfproc/labpdfprocapp.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 
 from diffpy.labpdfproc.functions import apply_corr, compute_cve
+from diffpy.labpdfproc.tools import set_wavelength
 from diffpy.utils.parsers.loaddata import loadData
 from diffpy.utils.scattering_objects.diffraction_objects import XQUANTITIES, Diffraction_object
 
@@ -62,7 +63,7 @@ def get_args():
 
 def main():
     args = get_args()
-    wavelength = WAVELENGTHS[args.anode_type]
+    args.wavelength = set_wavelength(args)
     filepath = Path(args.input_file)
     outfilestem = filepath.stem + "_corrected"
     corrfilestem = filepath.stem + "_cve"
@@ -80,7 +81,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,
@@ -91,7 +92,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

@@ -0,0 +1,34 @@
+import argparse
+
+import pytest
+
+from diffpy.labpdfproc.tools import set_wavelength
+
+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

@@ -0,0 +1,36 @@
+WAVELENGTHS = {"Mo": 0.71, "Ag": 0.59, "Cu": 1.54}
+known_sources = [key for key in WAVELENGTHS.keys()]
+
+
+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"]