fix K-mod import filtering (#471)

Author: JoschD

CHANGELOG.md

@@ -1,5 +1,10 @@
 # OMC3 Changelog
 
+#### 2024-11-14 - v0.20.1 - _jdilly_
+
+- Fixed:
+  - `kmod_import` fixed issue with too harsh filtering of BPM data.
+
 #### 2024-11-14 - v0.20.0 - _jdilly_, _awegsche_
 
 - Added:

omc3/__init__.py

@@ -11,7 +11,7 @@
 __title__ = "omc3"
 __description__ = "An accelerator physics tools package for the OMC team at CERN."
 __url__ = "https://github.com/pylhc/omc3"
-__version__ = "0.20.0"
+__version__ = "0.20.1"
 __author__ = "pylhc"
 __author_email__ = "[email protected]"
 __license__ = "MIT"

omc3/scripts/kmod_import.py

@@ -362,17 +362,25 @@ def _read_kmod_results(paths: Sequence[Path | str], beam: int
     # sort into bpm and betastar --
     bpm_results = _filter_bpm_results(all_dfs, beam=beam)
     betastar_results = _filter_betastar_results(all_dfs)
-
+    LOG.debug(
+        f"Found {len(bpm_results)} BPM results and {len(betastar_results)} betastar results for beam {beam}."
+    )
     return bpm_results, betastar_results
 
 
+def _is_bpm_df(df: tfs.TfsDataFrame, beam: int) -> bool:
+    """ Check if the given df is a BPM results file for the given beam. """
+    # bpm files must have a beta column
+    if f"{BETA}X" not in df.columns:
+        return False
+    
+    # They should also have at least one element (e.g. the BPM) matching the beam
+    elements = df.index if NAME not in df.columns else df[NAME]
+    return elements.str.match(fr".*\.B{beam}$", flags=re.IGNORECASE).any()
+
+
 def _filter_bpm_results(dfs: Sequence[tfs.TfsDataFrame], beam: int) -> list[tfs.TfsDataFrame]:
-    bpm_dfs = [df for df in dfs if f"{BETA}X" in df.columns]
-    bpm_dfs = [df for df in bpm_dfs if (
-        NAME not in df.columns or all(df[NAME].str.match(fr"(.*\.B{beam}$|IP\d$)", flags=re.IGNORECASE))
-    )
-    ]
-    return bpm_dfs
+    return [df for df in dfs if _is_bpm_df(df, beam)]
 
 
 def _filter_betastar_results(dfs: Sequence[tfs.TfsDataFrame]) -> list[tfs.TfsDataFrame]:
@@ -386,8 +394,9 @@ def _write_output(dfs: dict[str, tfs.TfsDataFrame], output_dir: Path):
             df = dfs.get(f"{id_}{plane}")
             if df is None:
                 continue
-
-            tfs.write(output_dir / f"{filename}{plane.lower()}{EXT}", df, save_index=NAME)
+            outfile = output_dir / f"{filename}{plane.lower()}{EXT}"
+            LOG.info(f"Writing output file {outfile}")
+            tfs.write(outfile, df, save_index=NAME)
 
 
 # Script Mode ------------------------------------------------------------------