add method to argparser and correct functions and docstrings

Author: yucongalicechen

src/diffpy/labpdfproc/functions.py

@@ -10,6 +10,7 @@
 RADIUS_MM = 1
 N_POINTS_ON_DIAMETER = 300
 TTH_GRID = np.arange(1, 180.1, 0.1)
+CVE_METHODS = ["brute_force", "polynomial_interpolation"]
 
 # pre-computed datasets for fast calculation
 MUD_LIST = [0.5, 1, 2, 3, 4, 5, 6]
@@ -172,15 +173,10 @@ def get_path_length(self, grid_point, angle):
         return total_distance, primary_distance, secondary_distance
 
 
-def _cve_brute_force(diffraction_data, mud, wavelength):
+def _cve_brute_force(mud):
     """
-    compute cve using brute-force method
-
-    it is computed as follows:
-    We first resample data and absorption correction to a more reasonable grid,
-    then calculate corresponding cve for the given mud in the resample grid
-    (since the same mu*D yields the same cve, we can assume that D/2=1, so mu=mud/2),
-    and finally interpolate cve to the original grid in diffraction_data.
+    compute cve for the given mud on a global grid using the brute-force method
+    assume mu=mud/2, given that the same mu*D yields the same cve and D/2=1
     """
 
     mu_sample_invmm = mud / 2
@@ -197,15 +193,15 @@ def _cve_brute_force(diffraction_data, mud, wavelength):
     return cve
 
 
-def _cve_interp_polynomial(diffraction_data, mud, wavelength):
+def _cve_polynomial_interpolation(mud):
     """
     compute cve using polynomial interpolation method, raise an error if mu*D is out of the range (0.5 to 6)
     """
 
     if mud > 6 or mud < 0.5:
         raise ValueError(
-            "mu*D is out of the acceptable range (0.5 to 6) for fast calculation. "
-            "Please rerun with a value within this range or use -b to enable brute-force calculation. "
+            f"mu*D is out of the acceptable range (0.5 to 6) for polynomial interpolation. "
+            f"Please rerun with a value within this range or specifying another method from {* CVE_METHODS, }."
         )
     coeff_a, coeff_b, coeff_c, coeff_d, coeff_e = [
         interpolation_function(mud) for interpolation_function in INTERPOLATION_FUNCTIONS
@@ -215,19 +211,20 @@ def _cve_interp_polynomial(diffraction_data, mud, wavelength):
     return cve
 
 
-def _cve_method(diffraction_data, mud, wavelength, brute_force=False):
+def _compute_cve(method, mud):
     """
-    selects the appropriate CVE calculation method
+    compute cve for the given mud on a global grid using the specified method
     """
-    if brute_force:
-        return _cve_brute_force(diffraction_data, mud, wavelength)
-    else:
-        return _cve_interp_polynomial(diffraction_data, mud, wavelength)
+    methods = {
+        "brute_force": _cve_brute_force,
+        "polynomial_interpolation": _cve_polynomial_interpolation,
+    }
+    return methods[method](mud)
 
 
-def compute_cve(diffraction_data, mud, wavelength, brute_force=False):
+def interpolate_cve(diffraction_data, mud, wavelength, method="polynomial_interpolation"):
     """
-    compute the cve for given diffraction data, mud, and wavelength, using the selected method
+    compute and interpolate the cve for the given diffraction data, mud, and wavelength, using the selected method
 
     Parameters
     ----------
@@ -237,14 +234,16 @@ def compute_cve(diffraction_data, mud, wavelength, brute_force=False):
       the mu*D of the diffraction object, where D is the diameter of the circle
     wavelength float
       the wavelength of the diffraction object
+    method str
+      the method used to calculate cve
 
     Returns
     -------
     the diffraction object with cve curves
 
     """
 
-    cve = _cve_method(diffraction_data, mud, wavelength, brute_force)
+    cve = _compute_cve(method, mud)
     orig_grid = diffraction_data.on_tth[0]
     newcve = np.interp(orig_grid, TTH_GRID, cve)
     abdo = Diffraction_object(wavelength=wavelength)

src/diffpy/labpdfproc/labpdfprocapp.py

@@ -1,7 +1,7 @@
 import sys
 from argparse import ArgumentParser
 
-from diffpy.labpdfproc.functions import apply_corr, compute_cve
+from diffpy.labpdfproc.functions import CVE_METHODS, apply_corr, interpolate_cve
 from diffpy.labpdfproc.tools import known_sources, load_metadata, preprocessing_args
 from diffpy.utils.parsers.loaddata import loadData
 from diffpy.utils.scattering_objects.diffraction_objects import XQUANTITIES, Diffraction_object
@@ -45,7 +45,7 @@ def get_args(override_cli_inputs=None):
         "-o",
         "--output-directory",
         help="The name of the output directory. If not specified "
-        "then corrected files will be written to the current directory."
+        "then corrected files will be written to the current directory. "
         "If the specified directory doesn't exist it will be created.",
         default=None,
     )
@@ -72,11 +72,11 @@ def get_args(override_cli_inputs=None):
         help="Outputs will not overwrite existing file unless --force is specified.",
     )
     p.add_argument(
-        "-b",
-        "--brute-force",
-        action="store_true",
-        help="The absorption correction will be computed using brute-force calculation "
-        "if this flag is set. Default is using fast calculation. ",
+        "-m",
+        "--method",
+        help=f"The method for computing absorption correction. Allowed methods: {*CVE_METHODS, }. "
+        f"Default method is polynomial interpolation if not specified. ",
+        default="polynomial_interpolation",
     )
     p.add_argument(
         "-u",
@@ -140,7 +140,7 @@ def main():
             metadata=load_metadata(args, filepath),
         )
 
-        absorption_correction = compute_cve(input_pattern, args.mud, args.wavelength, args.brute_force)
+        absorption_correction = interpolate_cve(input_pattern, args.mud, args.wavelength, args.method)
         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_functions.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 
-from diffpy.labpdfproc.functions import Gridded_circle, apply_corr, compute_cve
+from diffpy.labpdfproc.functions import Gridded_circle, apply_corr, interpolate_cve
 from diffpy.utils.scattering_objects.diffraction_objects import Diffraction_object
 
 params1 = [
@@ -69,13 +69,13 @@ def _instantiate_test_do(xarray, yarray, name="test", scat_quantity="x-ray"):
     return test_do
 
 
-def test_compute_cve(mocker):
+def test_interpolate_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])
     mocker.patch("diffpy.labpdfproc.functions.TTH_GRID", xarray)
     mocker.patch("numpy.interp", return_value=expected_cve)
     input_pattern = _instantiate_test_do(xarray, yarray)
-    actual_abdo = compute_cve(input_pattern, mud=1, wavelength=1.54)
+    actual_abdo = interpolate_cve(input_pattern, mud=1, wavelength=1.54)
     expected_abdo = _instantiate_test_do(
         xarray,
         expected_cve,