diff --git a/pyaml/diagnostics/chromaticity_monitor.py b/pyaml/diagnostics/chromaticity_monitor.py index 4af00902..ab60b6a6 100644 --- a/pyaml/diagnostics/chromaticity_monitor.py +++ b/pyaml/diagnostics/chromaticity_monitor.py @@ -147,7 +147,16 @@ def measure( callback: callable = None, ): """ - Main function for chromaticity measurment + Main function for chromaticity measurment. + :py:attr:`~pyaml.tuning_tools.measurement_tool.MeasurementTool.latest_measurement` contains: + + .. code-block:: python + + chromaticity:np.array # First order chromaticity, Array of [q'x,q'y] + dispersion:np.array # First order dispersion, [[dx'0,..,dx'n],[dy'0,..,dy'n]] + chromaticity_fit:np.array # Array of [[qx,qy],[q'x,q'y],[q''x,q''y],...] + dispersionx_fit:np.array # Array of [[dx0],..,[dxn],[dx'0],..,[dx'n],...] + dispersiony_fit:np.array # Array of [[dy0,..,dyn],[dy'0,..,dy'n],...] Parameters ---------- @@ -176,10 +185,22 @@ def measure( fit_dispersion : bool, optionnal Fit dispersion, [default: from config] do_plot : bool - Do you want to plot the fittinf results ? + Do you want to plot the fitting results ? callback: Callable, optional Callback is executed after each measurement or setting. If the callback return false, then the process is aborted. + callback_data dict contains: + + .. code-block:: python + + source:MeasurementTool # Tool that triggered the callback + step:int # The current step + avg_step:int # The current averaging step + rf:float # RF frequency used for the current step + tune:np.array # The measured tune (on Action.MEASURE) + orbit:np_array # The measured orbit, if fit_dispersion is True, (on Action.MEASURE) + dtune:np.array # The tune variation (on Action.RESTORE) + """ n_step = n_step if n_step is not None else self._cfg.n_step alphac = alphac if alphac is not None else self._alphac @@ -198,7 +219,8 @@ def measure( if alphac is None: raise PyAMLException("Moment compaction factor is not defined") - self.register_callback(callback) + self._init_measure("") + self._register_callback(callback) # Get devices self.check_peer() @@ -276,7 +298,8 @@ def measure( def fit(self, deltas, Q, order, orbit=None, fit_disp_order=None, do_plot=False): """ - Compute chromaticity (and dispersion) from measurement data. + Compute chromaticity (and dispersion) from input data and update + :py:attr:`~pyaml.tuning_tools.measurement_tool.MeasurementTool.latest_measurement`. Parameters ---------- @@ -292,14 +315,8 @@ def fit(self, deltas, Q, order, orbit=None, fit_disp_order=None, do_plot=False): plot : bool, optional If True, plot the fit. - Returns - ------- - dict - Dict with horizontal and veritical chromaticity and dispersion - """ chroma = np.polynomial.polynomial.polyfit(deltas, Q, order).T - self.latest_measurement = dict() self.latest_measurement["chromaticity_fit"] = chroma self.latest_measurement["chromaticity"] = chroma[:, 1] # First order chroma if orbit is not None: diff --git a/pyaml/tuning_tools/chromaticity_response_matrix.py b/pyaml/tuning_tools/chromaticity_response_matrix.py index 5c1a9db2..1f00f611 100644 --- a/pyaml/tuning_tools/chromaticity_response_matrix.py +++ b/pyaml/tuning_tools/chromaticity_response_matrix.py @@ -52,7 +52,14 @@ def measure( callback: Optional[Callable] = None, ): """ - Measure tune response matrix + Measure chromaticity response matrix. + :py:attr:`~pyaml.tuning_tools.measurement_tool.MeasurementTool.latest_measurement` contains: + + .. code-block:: python + + matrix:list[list[float] # The response matrix + variable_names:list[str] # Variable names + observable_names:list[str] # Observables names **Example** @@ -96,15 +103,16 @@ def callback(action: Action, data:dict): If the callback return false, then the scan is aborted and strength restored. callback_data dict contains: - .. code-block:: + .. code-block:: python - source:str Object that triggered the source - step:int The current step - avg_step:int The current avg step - magnet:Magnet The magnet being excited - strength:Magnet strength - chroma:np.array The measured chroma (on Action.MEASURE) - dchroma:np.array The chroma variation (on Action.RESTORE) + source:MeasurementTool # Tool that triggered the callback + idx:int # The index in the element array being processed + step:int # The current step + avg_step:int # The current averaging step + magnet:str # The magnet being excited + strength:float # Magnet strength + chroma:np.array # The measured chroma (on Action.MEASURE) + dchroma:np.array # The chroma variation (on Action.RESTORE) """ # Get devices @@ -112,7 +120,8 @@ def callback(action: Action, data:dict): sextus = self._peer.get_magnets(self._cfg.sextu_array_name) cm = self._peer.get_chromaticity_monitor(self._cfg.chromaticity_name) - self.register_callback(callback) + self._register_callback(callback) + self._init_measure("pyaml.tuning_tools.response_matrix_data") chromamat = np.zeros((len(sextus), 2)) @@ -140,7 +149,9 @@ def callback(action: Action, data:dict): # apply strength m.strength.set(str + d) - self.send_callback(Action.APPLY, {"step": qidx, "magnet": m.get_name(), "strength": float(str + d)}) + self.send_callback( + Action.APPLY, {"idx": qidx, "step": step, "magnet": m.get_name(), "strength": float(str + d)} + ) time.sleep(sleep_step) @@ -153,7 +164,8 @@ def callback(action: Action, data:dict): Qp[step] += chroma self.send_callback( - Action.MEASURE, {"step": qidx, "avg_step": avg, "magnet": m.get_name(), "chroma": chroma} + Action.MEASURE, + {"idx": qidx, "step": step, "avg_step": avg, "magnet": m.get_name(), "chroma": chroma}, ) if avg < nb_meas - 1: @@ -171,7 +183,7 @@ def callback(action: Action, data:dict): m.strength.set(str) self.send_callback( Action.RESTORE, - {"step": qidx, "magnet": m.get_name(), "strength": float(str), "dchroma": chromamat[qidx]}, + {"idx": qidx, "magnet": m.get_name(), "strength": float(str), "dchroma": chromamat[qidx]}, ) except Exception as ex: @@ -179,8 +191,8 @@ def callback(action: Action, data:dict): except KeyboardInterrupt as ex: aborted = True finally: - # Restore - m.strength.set(str) # restore strength + # Restore strength + m.strength.set(str) self.send_callback( Action.RESTORE, {"step": qidx, "magnet": m.get_name(), "strength": float(str), "dchroma": chromamat[qidx]}, @@ -194,11 +206,11 @@ def callback(action: Action, data:dict): logger.warning(f"{self.get_name()} : measurement aborted") return False - self.latest_measurement = ResponseMatrixDataConfigModel( + mat = ResponseMatrixDataConfigModel( matrix=chromamat.T.tolist(), variable_names=sextus.names(), observable_names=[cm.get_name() + ".x", cm.get_name() + ".y"], - ).model_dump() - self.latest_measurement["type"] = "pyaml.tuning_tools.response_matrix_data" + ) + self.latest_measurement.update(mat.model_dump()) return True diff --git a/pyaml/tuning_tools/dispersion.py b/pyaml/tuning_tools/dispersion.py index ef165b87..3f094eec 100644 --- a/pyaml/tuning_tools/dispersion.py +++ b/pyaml/tuning_tools/dispersion.py @@ -65,7 +65,8 @@ def measure( skip_save=True, ) - self.register_callback(callback) + self._register_callback(callback) + self._init_measure() aborted = False for code, measurement in generator: @@ -93,7 +94,7 @@ def measure( # dispersion_data.output_names = self.element_holder.get_bpms( # self.bpm_array_name # ).names() - self.latest_measurement = dispersion_data.model_dump() + self.latest_measurement.update(dispersion_data.model_dump()) def get(self): return self.latest_measurement diff --git a/pyaml/tuning_tools/measurement_tool.py b/pyaml/tuning_tools/measurement_tool.py index 0b0b6151..f75c0d6c 100644 --- a/pyaml/tuning_tools/measurement_tool.py +++ b/pyaml/tuning_tools/measurement_tool.py @@ -50,10 +50,18 @@ class MeasurementTool(Element, metaclass=ABCMeta): def __init__(self, name): super().__init__(name) - self.latest_measurement: dict = None + self._latest_measurement: dict = None + """ + """ self._peer: "ElementHolder" = None # Peer: ControlSystem or Simulator self._callback: Callable = None + def _init_measure(self, measurement_type: str | None = None): + # Initialize measurement data + self._latest_measurement = {} + if measurement_type is not None: + self._latest_measurement["type"] = measurement_type + @abstractmethod def measure(self) -> bool: """ @@ -65,28 +73,30 @@ def measure(self) -> bool: """ raise NotImplementedError() - def get(self) -> dict: + @property + def latest_measurement(self) -> dict: """ - Return last measurement data + Return last measurement data, a dictionary containing last measurement data. + See sub class of MeasurementTool to get description. Returns ------- - ResponseMatrixData + dict Return latest measurement or None """ - return self.latest_measurement + return self._latest_measurement - # TODO: Abstract this method - def load(self, load_path: Path): + def get(self) -> dict: """ - Load measurement data + Return last measurement data, a dictionary containing last measurement data. + See sub class of MeasurementTool to get description. - Parameters - ---------- - load_path: Path - Matrix filename + Returns + ------- + dict + Return latest measurement or None """ - raise NotImplementedError() + return self._latest_measurement def save(self, save_path: Path, with_type: str = "json"): """ @@ -142,16 +152,16 @@ def callback(action: Action, data: dict): """ ok = True if self._callback is not None: - # Add source + # Add source and peer cb_data["mode"] = f"{self.get_peer()}" - cb_data["source_name"] = f"{self.get_name()}" + cb_data["source"] = self ok = self._callback(action, cb_data) if not ok and raiseException: # Abort, same as ctrl+C raise KeyboardInterrupt return ok - def register_callback(self, callback: Callable): + def _register_callback(self, callback: Callable): self._callback = callback def attach(self, peer: "ElementHolder") -> Self: diff --git a/pyaml/tuning_tools/orbit_response_matrix.py b/pyaml/tuning_tools/orbit_response_matrix.py index 9f2bf17d..18511f34 100644 --- a/pyaml/tuning_tools/orbit_response_matrix.py +++ b/pyaml/tuning_tools/orbit_response_matrix.py @@ -123,9 +123,10 @@ def measure( pySC.disable_pySC_rich() aborted = False err = None - step = 0 + idx = 0 try: - self.register_callback(callback) + self._register_callback(callback) + self._init_measure() for code, measurement in generator: callback_data = measurement.response_data # to be defined better if code is ResponseCode.AFTER_SET: @@ -135,7 +136,7 @@ def measure( elif code is ResponseCode.AFTER_RESTORE: logger.info(f"Measured response of {measurement.last_input}.") self.send_callback(Action.RESTORE, callback_data) - step += 1 + idx += 1 except Exception as ex: err = ex except KeyboardInterrupt as ex: @@ -145,7 +146,7 @@ def measure( # TODO self.send_callback( Action.RESTORE, - {"step": step}, + {"idx": idx}, raiseException=False, ) @@ -157,7 +158,7 @@ def measure( return False orm_data = self._pySC_response_data_to_ORMData(measurement.response_data.model_dump()) - self.latest_measurement = orm_data.model_dump() + self.latest_measurement.update(orm_data.model_dump()) def _pySC_response_data_to_ORMData(self, data: dict) -> OrbitResponseMatrixDataConfigModel: # all metadata is discarded here. Should we keep something? diff --git a/pyaml/tuning_tools/response_matrix_data.py b/pyaml/tuning_tools/response_matrix_data.py index 1c5c2ef6..44abd184 100644 --- a/pyaml/tuning_tools/response_matrix_data.py +++ b/pyaml/tuning_tools/response_matrix_data.py @@ -26,8 +26,6 @@ class ConfigModel(BaseModel): Observable names, basically the measurements """ - model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid") - matrix: list[list[float]] variable_names: Optional[list[str]] observable_names: list[str] diff --git a/pyaml/tuning_tools/tune_response_matrix.py b/pyaml/tuning_tools/tune_response_matrix.py index 1c8491c6..c2c74681 100644 --- a/pyaml/tuning_tools/tune_response_matrix.py +++ b/pyaml/tuning_tools/tune_response_matrix.py @@ -50,7 +50,14 @@ def measure( callback: Optional[Callable] = None, ): """ - Measure tune response matrix + Measure tune response matrix. + :py:attr:`~pyaml.tuning_tools.measurement_tool.MeasurementTool.latest_measurement` contains: + + .. code-block:: python + + matrix:list[list[float] # The response matrix + variable_names:list[str] # Variable names + observable_names:list[str] # Observables names **Example** @@ -93,15 +100,16 @@ def callback(action: Action, data:dict): If the callback return false, then the scan is aborted and strength restored. callback_data dict contains: - .. code-block:: + .. code-block:: python - source:str Object that triggered the source - step:int The current step - avg_step:int The current avg step - magnet:Magnet The magnet being excited - strength:float Magnet strength - tune:np.array The measured tune (on Action.MEASURE) - dtune:np.array The tune variation (on Action.RESTORE) + source:MeasurementTool # Tool that triggered the callback + idx:int # The index in the element array being processed + step:int # The current step + avg_step:int # The current avg step + magnet:str # The magnet being excited + strength:float # Magnet strength + tune:np.array # The measured tune (on Action.MEASURE) + dtune:np.array # The tune variation (on Action.RESTORE) """ # Get devices @@ -117,7 +125,9 @@ def callback(action: Action, data:dict): sleep_step = sleep_between_step if sleep_between_step is not None else self._cfg.sleep_between_step sleep_meas = sleep_between_meas if sleep_between_meas is not None else self._cfg.sleep_between_meas - self.register_callback(callback) + self._register_callback(callback) + self._init_measure("pyaml.tuning_tools.response_matrix_data") + aborted = False err = None try: @@ -130,7 +140,9 @@ def callback(action: Action, data:dict): # apply strength m.strength.set(str + d) - self.send_callback(Action.APPLY, {"step": qidx, "magnet": m.get_name(), "strength": float(str + d)}) + self.send_callback( + Action.APPLY, {"idx": qidx, "step": step, "magnet": m.get_name(), "strength": float(str + d)} + ) sleep(sleep_step) @@ -140,7 +152,8 @@ def callback(action: Action, data:dict): tune = tm.tune.get() Q[step] += tune self.send_callback( - Action.MEASURE, {"step": qidx, "avg_step": avg, "magnet": m.get_name(), "tune": tune} + Action.MEASURE, + {"idx": qidx, "step": step, "avg_step": avg, "magnet": m.get_name(), "tune": tune}, ) if avg < nb_meas - 1: sleep(sleep_meas) @@ -157,7 +170,7 @@ def callback(action: Action, data:dict): m.strength.set(str) self.send_callback( Action.RESTORE, - {"step": qidx, "magnet": m.get_name(), "strength": float(str), "dtune": tunemat[qidx]}, + {"idx": qidx, "magnet": m.get_name(), "strength": float(str), "dtune": tunemat[qidx]}, ) except Exception as ex: @@ -180,11 +193,11 @@ def callback(action: Action, data:dict): logger.warning(f"{self.get_name()} : measurement aborted") return False - self.latest_measurement = ResponseMatrixDataConfigModel( + mat = ResponseMatrixDataConfigModel( matrix=tunemat.T.tolist(), variable_names=quads.names(), observable_names=[tm.get_name() + ".x", tm.get_name() + ".y"], - ).model_dump() - self.latest_measurement["type"] = "pyaml.tuning_tools.response_matrix_data" + ) + self.latest_measurement.update(mat.model_dump()) return True