fixed vclist bug and added test

peakipy/cli/fit.py

@@ -59,6 +59,7 @@ class FitPeaksArgs:
     jack_knife_sample_errors: bool = False
     mp: bool = (True,)
     verbose: bool = (False,)
+    vclist_data: Optional[np.array] = None
 
 
 @dataclass
@@ -409,7 +410,7 @@ def rename_columns_for_compatibility(df):
 
 
 def add_vclist_to_df(fit_input: FitPeaksInput, df: pd.DataFrame):
-    vclist_data = fit_input.args.get("vclist_data")
+    vclist_data = fit_input.args.vclist_data
     df["vclist"] = df.plane.apply(lambda x: vclist_data[x])
     return df
 

test/test_cli.py

@@ -51,6 +51,16 @@ def test_fit_main_with_default(protein_L):
     peakipy.cli.main.fit(**args)
 
 
+def test_fit_main_with_vclist(protein_L):
+    args = dict(
+        peaklist_path=protein_L / Path("test.csv"),
+        data_path=protein_L / Path("test1.ft2"),
+        output_path=protein_L / Path("fits_PV.csv"),
+        vclist=protein_L / Path("vclist"),
+    )
+    peakipy.cli.main.fit(**args)
+
+
 def test_fit_main_with_gaussian(protein_L):
     args = dict(
         peaklist_path=protein_L / Path("test.csv"),

test/test_fit.py

@@ -19,6 +19,9 @@
     create_parameter_dict,
     perform_initial_lineshape_fit_on_cluster_of_peaks,
     merge_unpacked_parameters_with_metadata,
+    add_vclist_to_df,
+    FitPeaksArgs,
+    FitPeaksInput,
 )
 from peakipy.core import Lineshape, pvoigt2d
 
@@ -344,6 +347,32 @@ def test_merge_unpacked_parameters_with_metadata():
     assert expected_result.equals(actual_result)
 
 
+def test_add_vclist_to_df():
+    args = FitPeaksArgs(
+        noise=0, uc_dics={}, lineshape=Lineshape.PV, vclist_data=np.array([1, 2, 3])
+    )
+    fit_peaks_input = FitPeaksInput(
+        args=args, data=None, config=None, plane_numbers=None
+    )
+    df = pd.DataFrame(dict(plane=[0, 1, 2]))
+    expected_df = pd.DataFrame(dict(plane=[0, 1, 2], vclist=[1, 2, 3]))
+    actual_df = add_vclist_to_df(fit_peaks_input, df)
+    assert actual_df.equals(expected_df)
+
+
+def test_add_vclist_to_df_plane_order():
+    args = FitPeaksArgs(
+        noise=0, uc_dics={}, lineshape=Lineshape.PV, vclist_data=np.array([1, 2, 3])
+    )
+    fit_peaks_input = FitPeaksInput(
+        args=args, data=None, config=None, plane_numbers=None
+    )
+    df = pd.DataFrame(dict(plane=[2, 1, 0]))
+    expected_df = pd.DataFrame(dict(plane=[2, 1, 0], vclist=[3, 2, 1]))
+    actual_df = add_vclist_to_df(fit_peaks_input, df)
+    assert actual_df.equals(expected_df)
+
+
 # def test_perform_initial_lineshape_fit_on_cluster_of_peaks(pseudo_voigt_model_result):
 #     expected_result = pseudo_voigt_model_result
 #     actual_result = perform_initial_lineshape_fit_on_cluster_of_peaks()