diff --git a/CHANGELOG.md b/CHANGELOG.md index d8de761..2839f31 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,15 @@ # pylhc Changelog +## Version 0.9.0 + +Added: + +- `bpm_resync`: script to synchronize BPM data from SuperKEKB rings. + + ## Version 0.8.3 -Explicitely require the `pytz` package as a dependency since it is no longer a hard dependency of `pandas 3.x`, which we relied on.g +Explicitely require the `pytz` package as a dependency since it is no longer a hard dependency of `pandas 3.x`, which we relied on. ## Version 0.8.2 diff --git a/README.md b/README.md index 43711e6..750c5c9 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,7 @@ To use those, you should make sure to install the relevant extra dependencies wi - `KickGroup Information` - Get information about KickGroups. ([**kickgroups.py**](pylhc/kickgroups.py)) - `BSRT Logger` and `BSRT Analysis` - Saves data coming straight from LHC BSRT FESA class and allows subsequent analysis. ([**bsrt_logger.py**](pylhc/bsrt_logger.py) & [**bsrt_analysis.py**](pylhc/bsrt_analysis.py) ) - `BPM Calibration Factors` - Compute the BPM calibration factors using ballistic optics. Two methods are available: using the beta function and using the dispersion. ([**bpm_calibration.py**](pylhc/bpm_calibration.py)) +- `BPM Resynchronization` - Compute the turn offset of SuperKEKB BPMs and resyncs them to a new turn by turn file. ([**bpm_resync.py**](pylhc/bpm_resync.py)) ## License diff --git a/doc/entrypoints/bpm_resync.rst b/doc/entrypoints/bpm_resync.rst new file mode 100644 index 0000000..eb703a5 --- /dev/null +++ b/doc/entrypoints/bpm_resync.rst @@ -0,0 +1,4 @@ +BPM Resynchronization +********************** + +.. automodule:: pylhc.bpm_resync diff --git a/doc/modules/constants.rst b/doc/modules/constants.rst index 21feaa6..20a95c4 100644 --- a/doc/modules/constants.rst +++ b/doc/modules/constants.rst @@ -11,3 +11,5 @@ Constants Definitions .. automodule:: pylhc.constants.machine_settings_info .. automodule:: pylhc.constants.calibration + +.. automodule:: pylhc.constants.bpm_resync diff --git a/pylhc/__init__.py b/pylhc/__init__.py index 7e469e2..80b2159 100644 --- a/pylhc/__init__.py +++ b/pylhc/__init__.py @@ -11,7 +11,7 @@ __title__ = "pylhc" __description__ = "An accelerator physics script collection for the OMC team at CERN." __url__ = "https://github.com/pylhc/pylhc" -__version__ = "0.8.3" +__version__ = "0.9.0" __author__ = "pylhc" __author_email__ = "pylhc@github.com" __license__ = "MIT" diff --git a/pylhc/bpm_resync.py b/pylhc/bpm_resync.py new file mode 100644 index 0000000..78a086f --- /dev/null +++ b/pylhc/bpm_resync.py @@ -0,0 +1,217 @@ +""" +BPM Synchronization +------------------- + +This script resyncs the BPMs from the `LER` and `HER` rings of `SuperKEKB`. +Those BPMs are often not aligned time-wise and the values can be off by a few turns. +The resynchronization is done by looking up the phase advance of each BPM to retrieve the turn +offset. +This requires a frequency and an optics analysis of the unsynchronized turn by turn data. + +The script takes as input the original turn by turn data, the optics directory containing the +results of the optics analysis, as well as the output filename where the turn by turn data will be +written, in ASCII SDDS format. + + +Arguments: + +*--Required--* + +- **input** *(Path,str,TbtData)*: + + Input turn by turn data to be resynchronized. + Can take the form of a `Path` to a file or directly a `TbtData` object. + + flags: **['--input']** + +- **optics_dir** *(Path,str)*: + + Optics analysis path of the unsynchronized data, must contain the `total_phase_{x,y}.tfs` files. + + flags: **['--optics_dir']** + +- **output_file** *(Path,str)*: + + Output file path where to write the synchronized turn by turn data. + The directory will be created if necessary. + + flags: **['--output_file']** + +- **ring** *(str)*: + + Ring name, either `LER` or `HER`. + + flags: **['--ring']** + + choices: ``('LER', 'HER')`` + +*--Optional--* + +- **tbt_datatype** *(str)*: + Datatype of the TurnByTurn data provided as input. + + flags: **['--tbt_datatype']** + + choices: list of datatypes supported by `turn_by_turn`, in `turn_by_turn.io.TBT_MODULES` + + default: ``lhc`` + +- ** overwrite** *(bool)*: + Whether to overwrite the output file if it already exists. + + flags: **['--overwrite']** + + default: ``False`` +""" + +from copy import deepcopy +from pathlib import Path + +import numpy as np +import pandas as pd +import tfs +import turn_by_turn as tbt +from generic_parser import EntryPointParameters, entrypoint +from generic_parser.dict_parser import ArgumentError +from omc3.optics_measurements.constants import DELTA, NAME, PHASE, TUNE +from omc3.utils import logging_tools +from omc3.utils.iotools import PathOrStr + +from pylhc.constants.bpm_resync import DEFAULT_DATATYPE, PHASE_FILE, RINGS + +LOGGER = logging_tools.get_logger(__name__) + + +def _get_params() -> dict: + """ + Parse Commandline Arguments and return them as options. + + Returns: + dict + """ + + return EntryPointParameters( + input={ + "required": True, + "help": "Input turn by turn data to be resynchronized. Can take the form of a `Path` to" + "a file or directly a `TbtData` object.", + }, + optics_dir={ + "type": PathOrStr, + "required": True, + "help": "Optics path, must contain the `total_phase_{x,y}.tfs` files.", + }, + output_file={ + "type": PathOrStr, + "required": True, + "help": "Output file path where to write the turn by turn data. The directory will be" + "created if necessary.", + }, + ring={ + "type": str, + "required": True, + "choices": RINGS, + "help": (f"Ring name, from {RINGS}"), + }, + tbt_datatype={ + "type": str, + "required": False, + "choices": list(tbt.io.TBT_MODULES.keys()), + "default": DEFAULT_DATATYPE, + "help": "Datatype of the TurnByTurn data", + }, + overwrite={ + "type": bool, + "required": False, + "default": False, + "help": "Whether to overwrite the output file if it already exists.", + }, + ) + + +def sync_tbt(original_tbt: tbt.TbtData, optics_dir: Path, ring: str) -> tbt.TbtData: + """Resynchronize the BPMS in the the turn by turn data based on the phase advance. + Args: + original_tbt (tbt.TbtData): Original turn by turn data to be synchronized. + optics_dir (Path): Original optics directory containing the phase advance files. + ring (str): Ring name, either `LER` or `HER`. + Returns: + tbt.TbtData: Resynchronized turn by turn data. + """ + # Copy the original turn by turn data + synced_tbt = deepcopy(original_tbt) + + # HER and LER are in opposite direction for the phase + ring_dir = 1 if ring == "HER" else -1 + + # Some BPMs can exist in a plane but not the other, we need to check both planes to be sure + already_processed = set() + for plane in ("x", "y"): + phase_df = tfs.read(optics_dir / PHASE_FILE.format(plane=plane)) + qx = phase_df.headers[f"{TUNE}1"] + qy = phase_df.headers[f"{TUNE}2"] + bpms: pd.Series = phase_df[NAME] # using omc3 constants + dphase: pd.Series = phase_df[f"{DELTA}{PHASE}{plane.upper()}"] + + tune = (1 - qx) if plane == "x" else (1 - qy) + + # The phase advance divided by the tune will tell us how off the BPM is + ntune: pd.Series = dphase / tune + abs_ntune: pd.Series = ntune.abs() + + # If the ratio ntune is close to 1, that's one turn, otherwise, it's likely -2 turns + mag: np.ndarray = np.select([abs_ntune >= 0.8, abs_ntune >= 0.1], [1, -2], default=0) + # The final number of turns also depends on the sign of the phase + final_correction: np.ndarray = (mag * np.sign(ntune) * ring_dir).astype(int) + + # Iterate through all the BPMs and check their phase advance + for idx, bpm in enumerate(bpms): + # Check if we've seen that BPM before in the other plane + if bpm in already_processed: + continue + already_processed.add(bpm) + + if (bpm_correction := final_correction[idx]) != 0: + LOGGER.info( + f" {bpm:15s} -> turn correction of {bpm_correction} (ntune={ntune[idx]:.2f})" + ) + + # Shift the data + for plane in ("X", "Y"): + matrix = synced_tbt.matrices[0][plane] + orig_row = original_tbt.matrices[0][plane].loc[bpm] + matrix.loc[bpm] = orig_row.shift(bpm_correction, fill_value=0) + + return synced_tbt + + +@entrypoint(_get_params(), strict=True) +def main(opt): + # Open the TbT file if needed + if isinstance(opt.input, (Path, str)): + original_tbt = tbt.read(opt.input, datatype=opt.tbt_datatype) + elif isinstance(opt.input, tbt.TbtData): + original_tbt = opt.input + else: + raise ArgumentError("input must be either a Path, str or a TbtData object") + + opt.optics_dir = Path(opt.optics_dir) + opt.output_file = Path(opt.output_file) + + # Synchronise TbT + LOGGER.info(f"Resynchronizing {opt.optics_dir.name}...") + synced_tbt = sync_tbt(original_tbt, opt.optics_dir, opt.ring) # type: ignore + + # Save the resynced turn by turn data + if (opt.output_file).exists() and not opt.overwrite: + LOGGER.warning(f"File {opt.output_file} already exists, aborting.") + raise FileExistsError(f"File {opt.output_file} already exists, aborting.") + if (opt.output_file).exists() and opt.overwrite: + LOGGER.warning(f"Overwriting file {opt.output_file}.") + + opt.output_file.parent.mkdir(exist_ok=True) + tbt.write(opt.output_file, synced_tbt) + + +if __name__ == "__main__": + main() diff --git a/pylhc/constants/bpm_resync.py b/pylhc/constants/bpm_resync.py new file mode 100644 index 0000000..79a98f0 --- /dev/null +++ b/pylhc/constants/bpm_resync.py @@ -0,0 +1,16 @@ +""" +Constants: BPM Resynchronization +-------------------------------- + +Specific constants related to the BPM resynchronization script in ``PyLHC``. +""" + +from typing import Final, Literal + +# Available rings +RINGS: Final[set[Literal['LER', 'HER']]] = {'LER', 'HER'} + +# Phase file containing the phase advance of the BPMs +PHASE_FILE: Final[str] = "total_phase_{plane}.tfs" + +DEFAULT_DATATYPE: Final[Literal['lhc']] = "lhc" diff --git a/tests/inputs/bpm_resync/synced.sdds b/tests/inputs/bpm_resync/synced.sdds new file mode 100644 index 0000000..cf2b19c Binary files /dev/null and b/tests/inputs/bpm_resync/synced.sdds differ diff --git a/tests/inputs/bpm_resync/total_phase_x.tfs b/tests/inputs/bpm_resync/total_phase_x.tfs new file mode 100644 index 0000000..4a057c5 --- /dev/null +++ b/tests/inputs/bpm_resync/total_phase_x.tfs @@ -0,0 +1,53 @@ +@ Measure_optics:version %s "0.27.0" +@ Date %s "11. February 2026, 18:56:48" +@ Compensation %s "none" +@ Q1 %le 0.545561396557 +@ Q2 %le 0.590880577381 +* S MUXMDL NAME S2 NAME2 PHASEX ERRPHASEX PHASEXMDL DELTAPHASEX ERRDELTAPHASEX +$ %le %le %s %le %s %le %le %le %le %le + 2.2500008 0.2321053 "MQC2LE" 2.2500008 "MQC2LE" 0 0.000733985330348 0 0 0.000733985330348 + 27.2720861 0.4859659 "MQLC7LE" 2.2500008 "MQC2LE" 0.256293937291 0.00291312042024 0.2538606 0.002433337291 0.00291312042024 + 77.1787307 1.3487054 "MQLB4LE" 2.2500008 "MQC2LE" 0.106984972657 0.00338876555354 0.1166001 -0.009615127343 0.00338876555354 + 119.0986538 2.3883387 "MQLA5LE" 2.2500008 "MQC2LE" 0.160726344998 0.00220870794024 0.1562334 0.004492944998 0.00220870794024 + 164.5639584 3.0658979 "MQD3E1" 2.2500008 "MQC2LE" 0.866894782895 0.00346375144436 0.8337926 0.0331021828951 0.00346375144436 + 230.8020191 4.13945 "MQEAE4" 2.2500008 "MQC2LE" 0.900857436455 0.0015240457914 0.9073447 -0.0064872635448 0.0015240457914 + 270.0136808 4.7448669 "MQD3E4" 2.2500008 "MQC2LE" 0.50139549147 0.00308079948379 0.5127616 -0.01136610853 0.00308079948379 + 306.4724249 5.2875858 "MQEAE6" 2.2500008 "MQC2LE" 0.049485696792 0.00164029807706 0.0554805 -0.005994803208 0.00164029807706 + 354.4092937 5.9879131 "MQTATNE1" 2.2500008 "MQC2LE" 0.194458960987 0.00158648098661 0.7558078 0.438651160987 0.00158648098661 + 397.7476034 6.4602091 "MQTATNE2" 2.2500008 "MQC2LE" 0.675587885323 0.00148601847999 0.2281038 0.447484085323 0.00148601847999 + 472.7108711 7.6286716 "MQEAE8" 2.2500008 "MQC2LE" 0.391819883042 0.0014916135612 0.3965663 -0.004746416958 0.0014916135612 + 548.3812769 8.7768074 "MQEAE10" 2.2500008 "MQC2LE" 0.526235491844 0.00170499807942 0.5447021 -0.018466608156 0.00170499807942 + 597.0252838 9.456808 "MQEAE11" 2.2500008 "MQC2LE" 0.22046885356 0.00175282051509 0.2247027 -0.00423384644 0.00175282051509 + 696.7698054 10.9793008 "MQR2NE1" 2.2500008 "MQC2LE" 0.728825016834 0.00296818349792 0.7471955 -0.018370483166 0.00296818349792 + 881.197528 13.2607736 "MQEAE13" 2.2500008 "MQC2LE" 0.0337180564701 0.0015699715775 0.0286683 0.0050497564701 0.0015699715775 + 947.4355886 14.3343257 "MQD3E12" 2.2500008 "MQC2LE" 0.136506648725 0.00386880500401 0.1022204 0.034286248725 0.00386880500401 + 983.8943327 14.8770446 "MQEAE16" 2.2500008 "MQC2LE" 0.639640121621 0.00148885652833 0.6449393 -0.005299178379 0.00148885652833 + 1023.1059944 15.4824615 "MQD3E14" 2.2500008 "MQC2LE" 0.266329667043 0.00347174272354 0.2503562 0.015973467043 0.00347174272354 + 1059.5647385 16.0251804 "MQEAE18" 2.2500008 "MQC2LE" 0.789406306452 0.00138825429604 0.7930751 -0.003668793548 0.00138825429604 + 1107.5016074 16.7255077 "MQTANFE1" 2.2500008 "MQC2LE" 0.492659512911 0.00151747177794 0.4934024 -0.000742887089001 0.00151747177794 + 1189.3444407 17.8235473 "MQD3E16" 2.2500008 "MQC2LE" 0.614782246359 0.00371486183198 0.591442 0.023340246359 0.00371486183198 + 1301.4735906 19.5144019 "MQEAE22" 2.2500008 "MQC2LE" 0.270785621186 0.00152642618843 0.2822966 -0.011510978814 0.00152642618843 + 1350.1175975 20.1944026 "MQEAE23" 2.2500008 "MQC2LE" 0.976272243168 0.00191120446419 0.9622973 0.0139749431679 0.00191120446419 + 1404.3012987 21.2105542 "MQI6E" 2.2500008 "MQC2LE" 0.977700472626 0.00225740145544 0.9784489 -0.0007484273738 0.00225740145544 + 1418.7440064 21.2913712 "MQI5E" 2.2500008 "MQC2LE" 0.064197099179 0.00144704517394 0.0592659 0.004931199179 0.00144704517394 + 1424.1361014 21.3069978 "MQI4E" 2.2500008 "MQC2LE" 0.091309691209 0.0012886002698 0.0748925 0.016417191209 0.0012886002698 + 1494.6206604 22.6992479 "MQX2RE" 2.2500008 "MQC2LE" 0.465169163917 0.00121668867702 0.4671426 -0.001973436083 0.00121668867702 + 1564.3492168 23.8089926 "MQM2E" 2.2500008 "MQC2LE" 0.591453258231 0.00325473069156 0.5768873 0.014565958231 0.00325473069156 + 1608.8759522 24.5204202 "MQM7E" 2.2500008 "MQC2LE" 0.297975857062 0.00269383162364 0.2883149 0.009660957062 0.00269383162364 + 1710.4692283 26.1148873 "MQEAE27" 2.2500008 "MQC2LE" 0.882293284639 0.0015095987124 0.882782 -0.000488715360799 0.0015095987124 + 1746.9279724 26.6576062 "MQD3E23" 2.2500008 "MQC2LE" 0.469819035015 0.0033193234245 0.4255009 0.044318135015 0.0033193234245 + 1786.1396341 27.2630231 "MQEAE29" 2.2500008 "MQC2LE" 0.0286904119644 0.00152753560153 0.0309178 -0.0022273880356 0.00152753560153 + 1813.166033 27.7311583 "MQEAE30" 2.2500008 "MQC2LE" 0.500736844416 0.00157831659097 0.499053 0.001683844416 0.00157831659097 + 1861.1278611 28.4218222 "MQTAFOE1" 2.2500008 "MQC2LE" 0.636146490598 0.00177489315603 0.1897169 0.446429590598 0.00177489315603 + 1905.3942301 28.8967378 "MQTAFOE2" 2.2500008 "MQC2LE" 0.122229765942 0.00163590880957 0.6646325 0.457597265942 0.00163590880957 + 1980.3824571 30.0555369 "MQEAE32" 2.2500008 "MQC2LE" 0.826542423411 0.00152909928067 0.8234316 0.003110823411 0.00152909928067 + 2465.4769137 37.0395836 "MQEAE39" 2.2500008 "MQC2LE" 0.266300040741 0.00158127047893 0.8074783 0.458821740741 0.00158127047893 + 2501.9356578 37.5823024 "MQD3E33" 2.2500008 "MQC2LE" 0.776516948261 0.00325074344111 0.3501971 0.426319848261 0.00325074344111 + 2541.1473195 38.1877194 "MQEAE41" 2.2500008 "MQC2LE" 0.408528089977 0.00140886797915 0.9556141 0.452913989977 0.00140886797915 + 2616.1355465 39.3465184 "MQTAOTE1" 2.2500008 "MQC2LE" 0.575607153823 0.0015129180727 0.1144131 0.461194053823 0.0015129180727 + 2660.4019155 39.821434 "MQTAOTE2" 2.2500008 "MQC2LE" 0.0311686850482 0.00159251723576 0.5893287 0.441839985048 0.00159251723576 + 2735.3901425 40.9802331 "MQEAE44" 2.2500008 "MQC2LE" 0.220665075012 0.00142938002831 0.7481278 0.472537275012 0.00142938002831 + 2784.0341494 41.6602337 "MQEAE45" 2.2500008 "MQC2LE" 0.876616688885 0.00146588124857 0.4281284 0.448488288885 0.00146588124857 + 2850.27221 42.7337858 "MQD3E40" 2.2500008 "MQC2LE" 0.956004212917 0.00467405934326 0.5016805 0.454323712917 0.00467405934326 + 2960.9918734 44.543746 "MQLB1RE" 2.2500008 "MQC2LE" 0.217017468206 0.00201667986156 0.3116407 -0.094623231794 0.00201667986156 + 3013.8146988 45.3142773 "MQC2RE" 2.2500008 "MQC2LE" 0.976447362224 0.000781106382587 0.082172 -0.105724637776 0.000781106382587 diff --git a/tests/inputs/bpm_resync/total_phase_y.tfs b/tests/inputs/bpm_resync/total_phase_y.tfs new file mode 100644 index 0000000..54a6ef8 --- /dev/null +++ b/tests/inputs/bpm_resync/total_phase_y.tfs @@ -0,0 +1,59 @@ +@ Measure_optics:version %s "0.27.0" +@ Date %s "11. February 2026, 18:56:48" +@ Compensation %s "none" +@ Q1 %le 0.545561396557 +@ Q2 %le 0.590880577381 +* S MUYMDL NAME S2 NAME2 PHASEY ERRPHASEY PHASEYMDL DELTAPHASEY ERRDELTAPHASEY +$ %le %le %s %le %s %le %le %le %le %le + 2.2500008 0.2499407 "MQC2LE" 2.2500008 "MQC2LE" 0 0.000140459862357 0 0 0.000140459862357 + 27.2720861 0.7481189 "MQLC7LE" 2.2500008 "MQC2LE" 0.500950070873 0.000368701426063 0.4981782 0.002771870873 0.000368701426063 + 77.1787307 1.9812277 "MQLB4LE" 2.2500008 "MQC2LE" 0.718143261092 0.000486508299902 0.731287 -0.013143738908 0.000486508299902 + 119.0986538 2.9272725 "MQLA5LE" 2.2500008 "MQC2LE" 0.65825064453 0.00018669802142 0.6773318 -0.01908115547 0.00018669802142 + 164.5639584 3.4454173 "MQD3E1" 2.2500008 "MQC2LE" 0.179200966897 0.000657046353065 0.1954766 -0.016275633103 0.000657046353065 + 230.8020191 4.4204939 "MQEAE4" 2.2500008 "MQC2LE" 0.148586726077 0.000840811079963 0.1705532 -0.021966473923 0.000840811079963 + 270.0136808 5.0801102 "MQD3E4" 2.2500008 "MQC2LE" 0.82935726516 0.00055077223981 0.8301695 -0.00081223484 0.00055077223981 + 306.4724249 5.5437581 "MQEAE6" 2.2500008 "MQC2LE" 0.294624715161 0.00100492235483 0.2938174 0.000807315161 0.00100492235483 + 354.4092937 6.2657723 "MQTATNE1" 2.2500008 "MQC2LE" 0.430703957261 0.000660062289871 0.0158316 0.414872357261 0.000660062289871 + 397.7476034 6.7075129 "MQTATNE2" 2.2500008 "MQC2LE" 0.873691920878 0.000644037899246 0.4575722 0.416119720878 0.000644037899246 + 472.7108711 7.7449873 "MQEAE8" 2.2500008 "MQC2LE" 0.497354893485 0.00080249255807 0.4950466 0.0023082934853 0.00080249255807 + 548.3812769 8.8682516 "MQEAE10" 2.2500008 "MQC2LE" 0.616539395922 0.00100145861799 0.6183109 -0.001771504078 0.00100145861799 + 597.0252838 9.6760556 "MQEAE11" 2.2500008 "MQC2LE" 0.42716362641 0.00100986513193 0.4261149 0.00104872641011 0.00100986513193 + 696.7698054 10.897872 "MQR2NE1" 2.2500008 "MQC2LE" 0.638669913862 0.00182193578716 0.6479313 -0.009261386138 0.00182193578716 + 753.3246054 11.4727461 "MQFRNE3" 2.2500008 "MQC2LE" 0.217118240217 0.000911081971189 0.2228054 -0.005687159783 0.000911081971189 + 781.6020054 11.7512072 "MQDRNE5" 2.2500008 "MQC2LE" 0.487792922362 0.000957409291143 0.5012665 -0.013473577638 0.000957409291143 + 881.197528 12.9288998 "MQEAE13" 2.2500008 "MQC2LE" 0.670658135871 0.000940655129629 0.6789591 -0.008300964129 0.000940655129629 + 947.4355886 13.9039764 "MQD3E12" 2.2500008 "MQC2LE" 0.645412761035 0.000578789626078 0.6540357 -0.008622938965 0.000578789626078 + 983.8943327 14.3676243 "MQEAE16" 2.2500008 "MQC2LE" 0.107475010637 0.000820337516602 0.1176836 -0.010208589363 0.000820337516602 + 1023.1059944 15.0272406 "MQD3E14" 2.2500008 "MQC2LE" 0.779595724333 0.000564394995327 0.7772999 0.002295824333 0.000564394995327 + 1059.5647385 15.4908885 "MQEAE18" 2.2500008 "MQC2LE" 0.239113136326 0.000869874440924 0.2409478 -0.001834663674 0.000869874440924 + 1107.5016074 16.2129028 "MQTANFE1" 2.2500008 "MQC2LE" 0.968390394617 0.000582160192025 0.9629621 0.005428294617 0.000582160192025 + 1189.3444407 17.2284698 "MQD3E16" 2.2500008 "MQC2LE" 0.983151389465 0.000516904412099 0.9785291 0.004622289465 0.000516904412099 + 1301.4735906 18.815382 "MQEAE22" 2.2500008 "MQC2LE" 0.557846204515 0.000901360845498 0.5654413 -0.007595095485 0.000901360845498 + 1350.1175975 19.623186 "MQEAE23" 2.2500008 "MQC2LE" 0.393081753435 0.0010875544354 0.3732453 0.0198364534347 0.0010875544354 + 1404.3012987 20.3222869 "MQI6E" 2.2500008 "MQC2LE" 0.073688766787 0.00055896134311 0.0723462 0.001342566787 0.00055896134311 + 1418.7440064 20.3689101 "MQI5E" 2.2500008 "MQC2LE" 0.117513532365 0.000248521541781 0.1189694 -0.001455867635 0.000248521541781 + 1424.1361014 20.393492 "MQI4E" 2.2500008 "MQC2LE" 0.148709843894 0.000603364040657 0.1435513 0.005158543894 0.000603364040657 + 1494.6206604 21.7858689 "MQX2RE" 2.2500008 "MQC2LE" 0.536624917016 0.000427902157117 0.5359282 0.000696717015999 0.000427902157117 + 1564.3492168 22.4371781 "MQM2E" 2.2500008 "MQC2LE" 0.188608567769 0.000328723480924 0.1872374 0.001371167769 0.000328723480924 + 1608.8759522 22.7994754 "MQM7E" 2.2500008 "MQC2LE" 0.541942996651 0.000290994442989 0.5495347 -0.007591703349 0.000290994442989 + 1671.2575666 23.5081062 "MQD3E21" 2.2500008 "MQC2LE" 0.265060630665 0.000652031907436 0.2581655 0.006895130665 0.000652031907436 + 1710.4692283 24.1677225 "MQEAE27" 2.2500008 "MQC2LE" 0.930646540816 0.000953318310252 0.9177818 0.012864740816 0.000953318310252 + 1746.9279724 24.6313705 "MQD3E23" 2.2500008 "MQC2LE" 0.391129646078 0.000589823634343 0.3814298 0.0096998460778 0.000589823634343 + 1786.1396341 25.2909868 "MQEAE29" 2.2500008 "MQC2LE" 0.044920177187 0.000852764028394 0.0410461 0.003874077187 0.000852764028394 + 1813.166033 25.606447 "MQEAE30" 2.2500008 "MQC2LE" 0.370577760576 0.00119017485206 0.3565063 0.0140714605763 0.00119017485206 + 1861.1278611 26.3397749 "MQTAFOE1" 2.2500008 "MQC2LE" 0.502285343743 0.000842486985159 0.0898342 0.412451143743 0.000842486985159 + 1905.3942301 26.7715844 "MQTAFOE2" 2.2500008 "MQC2LE" 0.941498333633 0.000700189298474 0.5216437 0.419854633633 0.000700189298474 + 1980.3824571 27.8203725 "MQEAE32" 2.2500008 "MQC2LE" 0.570273057228 0.000875129902667 0.5704318 -0.000158742772002 0.000875129902667 + 2288.3421939 32.153925 "MQDROE5" 2.2500008 "MQC2LE" 0.318044544133 0.000991444608826 0.9039843 0.414060244133 0.000991444608826 + 2465.4769137 34.5321608 "MQEAE39" 2.2500008 "MQC2LE" 0.693865422397 0.00094482224479 0.2822201 0.411645322397 0.00094482224479 + 2501.9356578 34.9958087 "MQD3E33" 2.2500008 "MQC2LE" 0.147265369336 0.000574922387259 0.745868 0.401397369336 0.000574922387259 + 2541.1473195 35.655425 "MQEAE41" 2.2500008 "MQC2LE" 0.822349129674 0.000816832629853 0.4054843 0.416864829674 0.000816832629853 + 2616.1355465 36.7042131 "MQTAOTE1" 2.2500008 "MQC2LE" 0.874633567806 0.000706905571609 0.4542724 0.420361167806 0.000706905571609 + 2660.4019155 37.1360226 "MQTAOTE2" 2.2500008 "MQC2LE" 0.293780660695 0.000674641551279 0.8860819 0.407698760695 0.000674641551279 + 2735.3901425 38.1848107 "MQEAE44" 2.2500008 "MQC2LE" 0.350065894552 0.000882059552326 0.93487 0.415195894552 0.000882059552326 + 2784.0341494 38.9926147 "MQEAE45" 2.2500008 "MQC2LE" 0.151183769152 0.000846248962338 0.742674 0.408509769152 0.000846248962338 + 2850.27221 39.9676913 "MQD3E40" 2.2500008 "MQC2LE" 0.127040548569 0.000760020797301 0.7177506 0.409289948569 0.000760020797301 + 2960.9918734 42.343982 "MQLB1RE" 2.2500008 "MQC2LE" 0.909536785808 0.000344447017352 0.0940413 -0.184504514192 0.000344447017352 + 2989.1433603 42.8457946 "MQLC7RE" 2.2500008 "MQC2LE" 0.413955881137 0.000334805097703 0.5958539 -0.181898018863 0.000334805097703 + 3013.8146988 43.3437608 "MQC2RE" 2.2500008 "MQC2LE" 0.909315281695 0.000175129121093 0.0938201 -0.184504818305 0.000175129121093 + 3015.7846995 43.3464019 "MQC1RE" 2.2500008 "MQC2LE" 0.504719934566 7.80678610913e-05 0.0964612 0.408258734567 7.80678610913e-05 diff --git a/tests/inputs/bpm_resync/unsynced.sdds b/tests/inputs/bpm_resync/unsynced.sdds new file mode 100644 index 0000000..86ab65e Binary files /dev/null and b/tests/inputs/bpm_resync/unsynced.sdds differ diff --git a/tests/unit/test_bpm_resync.py b/tests/unit/test_bpm_resync.py new file mode 100644 index 0000000..cb8db22 --- /dev/null +++ b/tests/unit/test_bpm_resync.py @@ -0,0 +1,101 @@ +import pathlib +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import tfs +import turn_by_turn as tbt +from generic_parser.dict_parser import ArgumentError + +from pylhc import bpm_resync as resync + +INPUTS_DIR = Path(__file__).parent.parent / "inputs" / "bpm_resync" +OPTICS_DIR = Path(__file__).parent.parent / "inputs" / "bpm_resync" + + +def test_bad_arg_optics_type(): + with pytest.raises(ArgumentError) as e: + resync.main(input=Path("yeah ok"), + optics_dir=42, + output_file="yay", + ring="HER") # type: ignore + assert "optics_dir' is not of type Path" in str(e.value) + + +def test_bad_arg_output_file_type(): + with pytest.raises(ArgumentError) as e: + resync.main(input=Path("yeah ok"), + optics_dir=Path("yay"), + output_file=42, + ring="HER") # type: ignore + assert "output_file' is not of type Path" in str(e.value) + + +def test_bad_arg_ring(): + with pytest.raises(ArgumentError) as e: + resync.main(input="yeah ok", + optics_dir=Path("yay"), + output_file=Path("wat"), + ring="MOON_COLLIDER") # type: ignore + assert "ring' needs to be one of" in str(e.value) + +def test_bad_arg_tbt_datatype(): + with pytest.raises(ArgumentError) as e: + resync.main(input="yeah ok", + optics_dir=Path("yay"), + output_file=Path("wat"), + ring="HER", + tbt_datatype="quantum_sdds") # type: ignore + assert "tbt_datatype' needs to be one of" in str(e.value) + + +def test_resync(tmp_path): + # Synchronize the BPMs and check against the control + resync.main(input=INPUTS_DIR / "unsynced.sdds", + optics_dir=OPTICS_DIR, + output_file=tmp_path / "output.sdds", + ring="HER") # type: ignore + + assert Path(tmp_path / "output.sdds").exists() + + synced_tbt = tbt.read(INPUTS_DIR / "synced.sdds") + output_tbt = tbt.read(tmp_path / "output.sdds") + + assert np.all(synced_tbt.matrices[0].X == output_tbt.matrices[0].X) + assert np.all(synced_tbt.matrices[0].Y == output_tbt.matrices[0].Y) + + +def test_overwrite_ok(tmp_path): + # Write an output file to create a conflict + (tmp_path / "output.sdds").write_text("This file already exists.") + + # Synchronize the BPMs and check against the control + resync.main(input=INPUTS_DIR / "unsynced.sdds", + optics_dir=OPTICS_DIR, + output_file=tmp_path / "output.sdds", + ring="HER", + overwrite=True) # type: ignore + + assert Path(tmp_path / "output.sdds").exists() + + synced_tbt = tbt.read(INPUTS_DIR / "synced.sdds") + output_tbt = tbt.read(tmp_path / "output.sdds") + + assert np.all(synced_tbt.matrices[0].X == output_tbt.matrices[0].X) + assert np.all(synced_tbt.matrices[0].Y == output_tbt.matrices[0].Y) + + +def test_overwrite_raise(tmp_path): + # Write an output file to create a conflict + (tmp_path / "output.sdds").write_text("This file already exists.") + + # Synchronize the BPMs and check against the control + with pytest.raises(FileExistsError) as e: + resync.main(input=INPUTS_DIR / "unsynced.sdds", + optics_dir=OPTICS_DIR, + output_file=tmp_path / "output.sdds", + ring="HER", + overwrite=False) # type: ignore + + assert "output.sdds already exists, aborting." in str(e.value) \ No newline at end of file