move changes from previous PR

Author: yucongalicechen

src/diffpy/labpdfproc/fast_cve.py

@@ -1,12 +1,23 @@
 import math
+from pathlib import Path
 
 import numpy as np
+import pandas as pd
+from scipy.interpolate import interp1d
 
 from diffpy.utils.scattering_objects.diffraction_objects import Diffraction_object
 
 RADIUS_MM = 1
 N_POINTS_ON_DIAMETER = 300
-TTH_GRID = np.arange(1, 141, 1)
+TTH_GRID = np.arange(1, 180.1, 0.1)
+CVE_METHODS = ["brute_force", "polynomial_interpolation"]
+
+# pre-computed datasets for polynomial interpolation (fast calculation)
+MUD_LIST = [0.5, 1, 2, 3, 4, 5, 6]
+CWD = Path(__file__).parent.resolve()
+MULS = np.loadtxt(CWD / "data" / "inverse_cve.xy")
+COEFFICIENT_LIST = np.array(pd.read_csv(CWD / "data" / "coefficient_list.csv", header=None))
+INTERPOLATION_FUNCTIONS = [interp1d(MUD_LIST, coefficients, kind="quadratic") for coefficients in COEFFICIENT_LIST]
 
 
 class Gridded_circle:
@@ -172,28 +183,10 @@ def get_path_length(self, grid_point, angle):
         return total_distance, primary_distance, secondary_distance
 
 
-def compute_cve(diffraction_data, mud, wavelength):
+def _cve_brute_force(diffraction_data, mud):
     """
-    compute the cve for given diffraction data, mud and wavelength
-
-    Parameters
-    ----------
-    diffraction_data Diffraction_object
-      the diffraction pattern
-    mud float
-      the mu*D of the diffraction object, where D is the diameter of the circle
-    wavelength float
-      the wavelength of the diffraction object
-
-    Returns
-    -------
-    the diffraction object with cve curves
-
-    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
@@ -208,9 +201,85 @@ def compute_cve(diffraction_data, mud, wavelength):
     muls = np.array(muls) / abs_correction.total_points_in_grid
     cve = 1 / muls
 
+    abdo = Diffraction_object(wavelength=diffraction_data.wavelength)
+    abdo.insert_scattering_quantity(
+        TTH_GRID,
+        cve,
+        "tth",
+        metadata=diffraction_data.metadata,
+        name=f"absorption correction, cve, for {diffraction_data.name}",
+        wavelength=diffraction_data.wavelength,
+        scat_quantity="cve",
+    )
+    return abdo
+
+
+def _cve_polynomial_interpolation(diffraction_data, 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(
+            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
+    ]
+    muls = np.array(coeff_a * MULS**4 + coeff_b * MULS**3 + coeff_c * MULS**2 + coeff_d * MULS + coeff_e)
+    cve = 1 / muls
+
+    abdo = Diffraction_object(wavelength=diffraction_data.wavelength)
+    abdo.insert_scattering_quantity(
+        TTH_GRID,
+        cve,
+        "tth",
+        metadata=diffraction_data.metadata,
+        name=f"absorption correction, cve, for {diffraction_data.name}",
+        wavelength=diffraction_data.wavelength,
+        scat_quantity="cve",
+    )
+    return abdo
+
+
+def _cve_method(method):
+    """
+    retrieve the cve computation function for the given method
+    """
+    methods = {
+        "brute_force": _cve_brute_force,
+        "polynomial_interpolation": _cve_polynomial_interpolation,
+    }
+    if method not in CVE_METHODS:
+        raise ValueError(f"Unknown method: {method}. Allowed methods are {*CVE_METHODS, }.")
+    return methods[method]
+
+
+def compute_cve(diffraction_data, mud, method="polynomial_interpolation"):
+    """
+    compute and interpolate the cve for the given diffraction data and mud using the selected method
+    Parameters
+    ----------
+    diffraction_data Diffraction_object
+      the diffraction pattern
+    mud float
+      the mu*D of the diffraction object, where D is the diameter of the circle
+    method str
+      the method used to calculate cve
+
+    Returns
+    -------
+    the diffraction object with cve curves
+    """
+
+    cve_function = _cve_method(method)
+    abdo_on_global_tth = cve_function(diffraction_data, mud)
+    global_tth = abdo_on_global_tth.on_tth[0]
+    cve_on_global_tth = abdo_on_global_tth.on_tth[1]
     orig_grid = diffraction_data.on_tth[0]
-    newcve = np.interp(orig_grid, TTH_GRID, cve)
-    abdo = Diffraction_object(wavelength=wavelength)
+    newcve = np.interp(orig_grid, global_tth, cve_on_global_tth)
+    abdo = Diffraction_object(wavelength=diffraction_data.wavelength)
     abdo.insert_scattering_quantity(
         orig_grid,
         newcve,

src/diffpy/labpdfproc/functions.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, compute_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
@@ -72,6 +72,13 @@ def get_args(override_cli_inputs=None):
         action="store_true",
         help="Outputs will not overwrite existing file unless --force is specified.",
     )
+    p.add_argument(
+        "-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",
         "--user-metadata",
@@ -134,7 +141,7 @@ def main():
             metadata=load_metadata(args, filepath),
         )
 
-        absorption_correction = compute_cve(input_pattern, args.mud, args.wavelength)
+        absorption_correction = compute_cve(input_pattern, args.mud, 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/labpdfprocapp.py

@@ -75,7 +75,7 @@ def test_compute_cve(mocker):
     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 = compute_cve(input_pattern, mud=1)
     expected_abdo = _instantiate_test_do(
         xarray,
         expected_cve,

src/diffpy/labpdfproc/tests/test_fast_cve.py

@@ -312,6 +312,7 @@ def test_load_metadata(mocker, user_filesystem):
             "wavelength": 0.71,
             "output_directory": str(Path.cwd().resolve()),
             "xtype": "tth",
+            "method": "polynomial_interpolation",
             "key": "value",
             "username": "cli_username",
             "email": "[email protected]",