remove duplicated functions in fast cve

Author: yucongalicechen

src/diffpy/labpdfproc/fast_cve.py

@@ -4,32 +4,26 @@
 import pandas as pd
 from scipy.interpolate import interp1d
 
-from diffpy.utils.scattering_objects.diffraction_objects import Diffraction_object
-
-TTH_GRID = np.arange(1, 180.1, 0.1)
+FAST_TTH_GRID = np.arange(1, 180.1, 0.1)
 MUD_LIST = [0.5, 1, 2, 3, 4, 5, 6]
 CWD = os.path.dirname(os.path.abspath(__file__))
 INVERSE_CVE_DATA = 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]
 
 
-def fast_compute_cve(diffraction_data, mud, wavelength):
+def fast_compute_cve(mud):
     """
-    use precomputed datasets to compute the cve for given diffraction data, mud and wavelength
+    use precomputed datasets to compute the cve for given mud
 
     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
+    the array of tth grid and the corresponding cve
     """
 
     coefficient_a, coefficient_b, coefficient_c, coefficient_d, coefficient_e = [
@@ -43,39 +37,4 @@ def fast_compute_cve(diffraction_data, mud, wavelength):
         + coefficient_e
     )
     cve = 1 / np.array(inverse_cve)
-
-    orig_grid = diffraction_data.on_tth[0]
-    newcve = np.interp(orig_grid, TTH_GRID, cve)
-    abdo = Diffraction_object(wavelength=wavelength)
-    abdo.insert_scattering_quantity(
-        orig_grid,
-        newcve,
-        "tth",
-        metadata=diffraction_data.metadata,
-        name=f"absorption correction, cve, for {diffraction_data.name}",
-        wavelength=diffraction_data.wavelength,
-        scat_quantity="cve",
-    )
-
-    return abdo
-
-
-def apply_fast_corr(diffraction_pattern, absorption_correction):
-    """
-    Apply absorption correction to the given diffraction object modo with the correction diffraction object abdo
-
-    Parameters
-    ----------
-    diffraction_pattern Diffraction_object
-      the input diffraction object to which the cve will be applied
-    absorption_correction Diffraction_object
-      the diffraction object that contains the cve to be applied
-
-    Returns
-    -------
-    a corrected diffraction object with the correction applied through multiplication
-
-    """
-
-    corrected_pattern = diffraction_pattern * absorption_correction
-    return corrected_pattern
+    return FAST_TTH_GRID, cve

src/diffpy/labpdfproc/functions.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 
+from diffpy.labpdfproc.fast_cve import fast_compute_cve
 from diffpy.utils.scattering_objects.diffraction_objects import Diffraction_object
 
 RADIUS_MM = 1
@@ -27,7 +28,7 @@ def _get_grid_points(self):
         self.grid = {(x, y) for x in xs for y in ys if x**2 + y**2 <= self.radius**2}
         self.total_points_in_grid = len(self.grid)
 
-    # def get_coordinate_index(self, coordinate):  # I think we probably dont need this function?
+    # def get_coordinate_index(self, coordinate):
     #     count = 0
     #     for i, target in enumerate(self.grid):
     #         if coordinate == target:
@@ -172,9 +173,9 @@ def get_path_length(self, grid_point, angle):
         return total_distance, primary_distance, secondary_distance
 
 
-def compute_cve(diffraction_data, mud, wavelength):
+def compute_cve(diffraction_data, mud, wavelength, brute_force=False):
     """
-    compute the cve for given diffraction data, mud and wavelength
+    compute the cve for given diffraction data, mud and wavelength, and a boolean to determine the way to compute
 
     Parameters
     ----------
@@ -185,31 +186,41 @@ def compute_cve(diffraction_data, mud, wavelength):
     wavelength float
       the wavelength of the diffraction object
 
-    Returns
-    -------
-    the diffraction object with cve curves
-
-    it is computed as follows:
+    the brute-force method 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.
+
+    Returns
+    -------
+    the diffraction object with cve curves
+
     """
 
-    mu_sample_invmm = mud / 2
-    abs_correction = Gridded_circle(mu=mu_sample_invmm)
-    distances, muls = [], []
-    for angle in TTH_GRID:
-        abs_correction.set_distances_at_angle(angle)
-        abs_correction.set_muls_at_angle(angle)
-        distances.append(sum(abs_correction.distances))
-        muls.append(sum(abs_correction.muls))
-    distances = np.array(distances) / abs_correction.total_points_in_grid
-    muls = np.array(muls) / abs_correction.total_points_in_grid
-    cve = 1 / muls
+    if brute_force:
+        tth_grid = TTH_GRID
+        mu_sample_invmm = mud / 2
+        abs_correction = Gridded_circle(mu=mu_sample_invmm)
+        distances, muls = [], []
+        for angle in TTH_GRID:
+            abs_correction.set_distances_at_angle(angle)
+            abs_correction.set_muls_at_angle(angle)
+            distances.append(sum(abs_correction.distances))
+            muls.append(sum(abs_correction.muls))
+        distances = np.array(distances) / abs_correction.total_points_in_grid
+        muls = np.array(muls) / abs_correction.total_points_in_grid
+        cve = 1 / muls
+    else:
+        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. "
+            )
+        tth_grid, cve = fast_compute_cve(mud)
 
     orig_grid = diffraction_data.on_tth[0]
-    newcve = np.interp(orig_grid, TTH_GRID, cve)
+    newcve = np.interp(orig_grid, tth_grid, cve)
     abdo = Diffraction_object(wavelength=wavelength)
     abdo.insert_scattering_quantity(
         orig_grid,

src/diffpy/labpdfproc/labpdfprocapp.py

@@ -1,7 +1,6 @@
 import sys
 from argparse import ArgumentParser
 
-from diffpy.labpdfproc.fast_cve import apply_fast_corr, fast_compute_cve
 from diffpy.labpdfproc.functions import apply_corr, compute_cve
 from diffpy.labpdfproc.tools import known_sources, load_metadata, preprocessing_args
 from diffpy.utils.parsers.loaddata import loadData
@@ -141,20 +140,11 @@ def main():
             metadata=load_metadata(args, filepath),
         )
 
-        if args.brute_force:
-            absorption_correction = compute_cve(input_pattern, args.mud, args.wavelength)
-            corrected_data = apply_corr(input_pattern, absorption_correction)
-        else:
-            if args.mud > 6 or args.mud < 0.5:
-                sys.exit(
-                    "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 enable brute-force calculation. "
-                )
-            absorption_correction = fast_compute_cve(input_pattern, args.mud, args.wavelength)
-            corrected_data = apply_fast_corr(input_pattern, absorption_correction)
-
+        absorption_correction = compute_cve(input_pattern, args.mud, args.wavelength, args.brute_force)
+        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")
+
         if args.output_correction:
             absorption_correction.dump(f"{corrfile}", xtype="tth")