Get Peak to Peak from SDDS files

omc3/harpy/constants.py

@@ -28,3 +28,27 @@
 COL_FREQ = "FREQ"
 COL_PHASE = "PHASE"
 COL_ERR = "ERR"
+
+# Defocussing Monitors ---------------------------------------------------------
+LIST_1 = ["BPM.20L3.B", "BPM.18L3.B",
+ "BPM.21R3.B", "BPM.19R3.B",
+ "BPM.20L7.B", "BPM.18L7.B",
+ "BPM.21R7.B", "BPM.19R7.B",
+ "BPM.20R2.B", "BPM.18R2.B",
+ "BPM.21L4.B", "BPM.19L4.B",
+ "BPM.20R6.B", "BPM.18R6.B",
+ "BPM.21L8.B", "BPM.19L8.B",]
+
+LIST_2 = ["BPM.21L3.B", "BPM.19L3.B",
+ "BPM.20R3.B", "BPM.18R3.B",
+ "BPM.21L7.B", "BPM.19L7.B",
+ "BPM.20R7.B", "BPM.18R7.B",
+ "BPM.21R2.B", "BPM.19R2.B",
+ "BPM.20L4.B", "BPM.18L4.B",
+ "BPM.21R6.B", "BPM.19R6.B",
+ "BPM.20L8.B", "BPM.18L8.B",]
+
+DEFOCUSSING_MONITORS = {
+ 1: {"X": LIST_1, "Y": LIST_2},
+ 2: {"X": LIST_2, "Y": LIST_1}
+}

omc3/scripts/check_peak_to_peak.py

@@ -0,0 +1,199 @@
+"""
+Check Peak-to-Peak
+------------------
+
+Performs a quick check on the peak-to-peak of the given
+turn-by-turn files.
+"""
+import numbers
+import re
+from pathlib import Path
+from typing import Union, Sequence, Dict
+
+import numpy as np
+import pandas as pd
+import turn_by_turn as tbt
+from generic_parser import EntryPointParameters, entrypoint
+from turn_by_turn.utils import generate_average_tbtdata
+
+from omc3.definitions.constants import PLANES, UNIT_IN_METERS
+from omc3.harpy.constants import DEFOCUSSING_MONITORS
+from omc3.hole_in_one import HARPY_DEFAULTS
+from omc3.tbt_converter import _file_name_without_sdds
+from omc3.utils.logging_tools import get_logger
+
+LOG = get_logger(__name__)
+
+
+def get_params():
+ params = EntryPointParameters()
+ params.add_parameter(name="files", required=True, nargs='+',
+ help="Files for analysis")
+ params.add_parameter(name="beam", type=int,
+ help="LHC beam number.")
+ params.add_parameter(name="tbt_datatype",
+ default=HARPY_DEFAULTS["tbt_datatype"],
+ choices=list(tbt.io.DATA_READERS.keys()),
+ help="Choose the datatype from which to import.")
+ params.add_parameter(name="unit", type=str, default="mm",
+ choices=list(UNIT_IN_METERS.keys()),
+ help="Unit in which to log the peak-to-peak values.")
+ params.add_parameter(name="input_unit", type=str, default="mm",
+ choices=list(UNIT_IN_METERS.keys()),
+ help="Unit of the tbt input data.")
+ return params
+
+
+@entrypoint(get_params(), strict=True)
+def peak_to_peak(opt):
+ """Main function to log peak-to-peak values from SDDS files.
+
+ *--Required--*
+
+ - **files**:
+
+ Files for analysis
+
+
+ *--Optional--*
+
+ - **beam** *(int)*:
+
+ LHC beam number.
+
+
+ - **input_unit** *(str)*:
+
+ Unit of the tbt input data.
+
+ choices: ``['km', 'm', 'mm', 'um', 'nm', 'pm', 'fm', 'am']``
+
+ default: ``mm``
+
+
+ - **tbt_datatype**:
+
+ Choose the datatype from which to import.
+
+ choices: ``['lhc', 'iota', 'esrf', 'ptc', 'trackone']``
+
+ default: ``lhc``
+
+
+ - **unit** *(str)*:
+
+ Unit in which to log the peak-to-peak values.
+
+ choices: ``['km', 'm', 'mm', 'um', 'nm', 'pm', 'fm', 'am']``
+
+ default: ``mm``
+
+ """
+ for input_file in opt.files:
+ LOG.debug(f"Calculating pk2pk for file: {input_file}")
+ input_file = Path(input_file)
+ name = _file_name_without_sdds(input_file)
+ beam = _get_beam(opt.beam, filename=name)
+ bpms = DEFOCUSSING_MONITORS[beam]
+ tbt_data = tbt.read_tbt(input_file, datatype=opt.tbt_datatype)
+ for bunch, positions in zip(tbt_data.bunch_ids, tbt_data.matrices):
+ LOG.debug(f"Bunch: {bunch}")
+ pk2pk = get_pk2pk(positions, bpms)
+ _log_pk2pk(pk2pk, name, bunch, opt.input_unit, opt.unit)
+
+ if tbt_data.nbunches > 1:
+ tbt_data_av = generate_average_tbtdata(tbt_data)
+ positions = tbt_data_av.matrices[0]
+ pk2pk = get_pk2pk(positions, bpms)
+ _log_pk2pk(pk2pk, name, '(average)', opt.input_unit, opt.unit)
+
+
+def get_pk2pk(data: tbt.TransverseData, bpms: Dict[str, Sequence[str]]) -> Dict[str, float]:
+ """Get the filtered peak-to-peak from the current tbt-data from the given bpms."""
+ pk2pk = {p: None for p in PLANES}
+ for plane in PLANES:
+ positions: pd.DataFrame = getattr(data, plane)
+
+ # Get only the desired BPMs
+ mask = _get_index_mask(positions.index, bpms[plane])
+ if not any(mask):
+ LOG.debug(f"None of the required pk2pk BPMs are present "
+ f"in plane {plane} of the current tbt data.")
+ continue
+ positions = positions.loc[mask, :]
+
+ # Filter BPMs with exact zeros
+ exact_zeros = get_exact_zero_mask(positions)
+ if all(exact_zeros):
+ LOG.debug(f"Exact zeros found in all BPMs "
+ f"in plane {plane} of the current tbt data.")
+ continue
+ positions = positions.loc[~exact_zeros, :]
+
+ # Calculate peak-to-peak from remaining
+ pk2pk_per_bpm = positions.max(axis='columns') - positions.min(axis='columns')
+ pk2pk[plane] = pk2pk_per_bpm.mean()
+ return pk2pk
+
+
+def _log_pk2pk(pk2pk: Dict[str, float], name: str = None, other: Union[int, str] = None, input_unit: str = "m", unit: str = "mm"):
+ other_str = other if isinstance(other, str) else ""
+ if isinstance(other, numbers.Integral):
+ other_str = f", Bunch {other: d}"
+
+ unit_scale = UNIT_IN_METERS[input_unit] / UNIT_IN_METERS[unit]
+ pk2pk_str = " ".join(f" {plane}: {p2p*unit_scale:.4f} {unit}" for plane, p2p in pk2pk.items() if p2p is not None)
+ pk2pk_warn = ", ".join(plane for plane, p2p in pk2pk.items() if p2p is None)
+
+ LOG.info(f"Peak-to-Peak values for {name}{other_str}: {pk2pk_str} ")
+ if pk2pk_warn:
+ LOG.warning(f"No Peak-to-Peak values for {name} {other_str} in planes {pk2pk_warn}")
+
+
+# Utils ------------------------------------------------------------------------
+
+def _get_beam(beam: int, filename: Union[Path, str]) -> int:
+ """Get the beam from either given option or try to get it from the filename."""
+ if beam is None:
+ try:
+ beam = infer_beam_from_filename(filename)
+ except AttributeError:
+ raise NameError(f"No beam option given and could not infer beam from filename {filename}. "
+ f"Please provide a beam number. ")
+ LOG.debug(f"Assuming Beam {beam}")
+ return beam
+
+
+def _get_index_mask(index: pd.Index, bpms: Sequence[str]) -> Sequence[bool]:
+ """Get the boolean mask for the index/columns of the desired bpms."""
+ mask = np.zeros_like(index, dtype=bool)
+ not_found_bpms = []
+ for bpm in bpms:
+ new_mask = index.str.startswith(bpm)
+ if any(new_mask):
+ mask |= new_mask
+ else:
+ not_found_bpms.append(bpm)
+
+ if len(not_found_bpms):
+ LOG.debug(f"Some BPMs were not found in current tbt data: {not_found_bpms}")
+ return mask
+
+
+def get_exact_zero_mask(positions: pd.DataFrame):
+ """Finds bpms containing exact zeros in the data."""
+ exact_zeros = (positions == 0).any(axis='columns')
+ if any(exact_zeros):
+ LOG.debug(f"Exact zeros found in bpms {positions.columns[exact_zeros]}")
+ return exact_zeros
+
+
+def infer_beam_from_filename(filename: Union[str, Path]) -> int:
+ """Regex to find 'beam\\d' in filename and return the beam."""
+ return int(re.search(r"Beam(\d)", str(filename), flags=re.IGNORECASE).group(1))
+
+
+# Script Mode ------------------------------------------------------------------
+
+if __name__ == '__main__':
+ peak_to_peak()