pxii_bec/pxii_bec/macros/mx_methods.py

"""Use the methods in mx_basics to perform:
1) a go_to_peak scan, that scans a motor, finds the peak position and moves to peak
2) fits data from a bec history file
"""

from dataclasses import dataclass
import numpy as np

# from pxiii_parameters import FitDefaults, BPMScans, MirrorConfig

# from mx_basics import (
#     create_fit_parameters,
#     get_data_from_history,
#     fit,
#     plot_fitted_data_bec,
#     plot_live_data_bec,
# )


# Method functions
def calculate_step_size(start: float, stop: float, steps: int) -> float:
    """
    Provides the function to calculate the step size for dividing a specified range
    into a given number of steps.

    Args:
        start: The starting value of the range.
        stop: The stopping value of the range.
        steps: The number of steps to divide the range into. Must be at least 1.

    Raises:
        ValueError: If the steps value is less than 1.

    Returns:
        The calculated step size as a float, rounded to three decimal places.
    """
    if steps < 1:
        raise ValueError("Number of steps must be at least 1.")
    return round((stop - start) / steps, 3)


def move_to_position(motor_device, motor_name: str, position: float, data: dict):
    """
    Function to move a specified motor device to a given position.

    The function verifies if the requested position is within the scan range of the
    motor device provided. If the position is outside the range, the motor is
    moved to the center of its scan range, an error message is raised, and the
    operation is halted. If the position is valid, the motor is moved to the
    specified position.

    Parameters:
        motor_device: The motor device to be moved.
        motor_name: str
            The name of the motor as a string
        position: float
            The desired position to move the motor to. Position should be within
            the scan range of the motor determined by the provided data.
        data: dict
            A dictionary containing "x_data", which is used to determine the
            scan range of the motor.

    Raises:
        ValueError: Raised if the specified position is outside the valid scan
        range determined by "x_data" in the data dictionary. The motor will
        return to the center of its scan range in this case.
    """

    motor_min = np.min(data["x_data"])
    motor_max = np.max(data["x_data"])
    motor_centre = (motor_max + motor_min) / 2

    if not motor_min <= position <= motor_max:
        scans.umv(motor_device, motor_centre, relative=False)
        msg = (
            f"Position {position: .2f} is outside the scan range of "
            f"{motor_min: .2f} to {motor_max: .2f}. "
            f"Returning to centre of scan range {motor_centre: .3f}."
        )
        raise ValueError(msg)
    motor_position = round(position, 4)
    scans.umv(motor_device, motor_position, relative=False)
    print(f"\n Moving {motor_name} to position {motor_position: .3f}")


@dataclass(frozen=True)
class FitDefaults:
    """Default values for fitting routines"""

    # Constants for default models, baselines, and parameters
    MODEL = "Voigt"
    BASELINE = "Linear"
    SETTLE_TIME = 0.1
    RELATIVE_MODE = True


def go_to_peak(
    motor_device,
    signal_device,
    start: float,
    stop: float,
    steps: int,
    relative: bool = FitDefaults.RELATIVE_MODE,
    plot: bool = True,
    settle: float = FitDefaults.SETTLE_TIME,
    confirm: bool = True,
    gomax: bool = False,
    negative: bool = False
):
    """
    Go to the peak of a signal by scanning a motor within a specified range and
    identifying the optimal position based on signal peak data.

    Parameters:
        motor_device: The motor device to be scanned.
        signal_device: The signal device to monitor during the scan.
        start (float): The starting position of the scan. Ignored if `relative` is True.
        stop (float): The ending position of the scan. Ignored if `relative` is True.
        steps (int): The number of steps to divide the scan range into.
        relative (bool, optional): If True, interpret `start` and `stop` as relative to
            the current motor position. Defaults to RELATIVE_MODE constant.
        plot (bool, optional): If True, plot the scan data and the fitted results.
            Defaults to True.
        settle (float, optional): The time in seconds to wait after each step for the
            signal to stabilize. Defaults to DEFAULT_SETTLE_TIME constant.
        confirm (bool, optional): If True, ask for user confirmation before starting
            the scan. Defaults to True.

    Raises:
        Exception: Raises exceptions potentially raised by dependent functions or
            operations such as plotting, fitting, or motor movement.

    Returns:
        None
    """
    motor_name = motor_device.name
    signal_name = signal_device.name
    # wf.plot(x_name=motor_name, y_name=signal_name)
    if plot:
        plot_live_data_bec(motor_name, signal_name)

    # Validate and calculate step size
    step_size = calculate_step_size(start, stop, steps)

    # Confirm the scan range
    # current_motor_position = motor_device.user_readback.get()
    current_motor_position = motor_device.read()[motor_name]["value"]
    if confirm:
        if relative:
            scan_start = current_motor_position + start
            scan_end = current_motor_position + stop
            print(
                f"\nScanning from {scan_start: .6g} to {scan_end: .6g} in "
                f"{steps} steps of size {step_size}"
            )
            print(f"Relative mode = {relative}")
        else:
            print(
                f"\nScanning from {start: .5g} to {stop: .5g} in {steps} steps of size {step_size}"
            )
            print(f"Relative mode = {relative}")
        input("Press Enter to continue...")
    # Perform the scan
    scan_result = scans.line_scan(
        motor_device, start, stop, steps=steps, relative=relative, settling_time=settle
    )
    motor_data = scan_result.scan.live_data[motor_name][motor_name].val
    signal_data = scan_result.scan.live_data[signal_name][signal_name].val
    scan_number = "Current"

    data = {
        "x_data": np.array(motor_data),
        "y_data": np.array(signal_data),
        "motor_name": motor_name,
        "signal_name": signal_name,
        "motor_device": motor_device,
        "scan_number": scan_number,
    }

    # Define and fit model to scan data
    fit_params = create_fit_parameters(deriv = False,
                                       negative = False, 
                                       model = FitDefaults.MODEL, 
                                       baseline = FitDefaults.BASELINE)
    fit_result = fit(data, fit_params)

    # Plot the fitted data if plot = True
    if plot:
        plot_fitted_data_bec(data, fit_result)

    # If gomax is set then move to the maximum value, rather than the fit centre
    if gomax:
        value = fit_result["x_max"]
        print(f"Max position is at {value}")
        move_to_position(data["motor_device"], data["motor_name"], fit_result["x_max"], data)
    else:
        # Safely move the motor to the peak position
        move_to_position(data["motor_device"], data["motor_name"], fit_result["centre"], data)


def fit_history(
    history_index: int,
    signal_name: str,
    deriv: bool = False,
    negative: bool = False,
    smoothing: bool = False,
    model: str = FitDefaults.MODEL,
    move_to_peak: bool = False,
):
    """
    Retrieve and analyze historical data by fitting a model, optionally moving to
    a peak position.

    Parameters:
        history_index (int): Index of the historical data set to retrieve.
        signal_name (str): Name of the signal to fit.
        deriv (bool, optional): Whether to include the derivative in the fitting
            procedure. Defaults to False.
        model (str, optional): Name of the model to use for fitting. Defaults to
            DEFAULT_MODEL.
        move_to_peak (bool, optional): Whether to move the motor to the peak position
            after fitting. Defaults to False.

    Raises:
        KeyError: If required keys are not found in the retrieved data dictionary.
        ValueError: If the fitting process fails or produces invalid results.

    Returns:
        None
    """
    # Retrieve historical data
    data = get_data_from_history(history_index, signal_name)

    # Define fitting parameters
    fit_params = create_fit_parameters(deriv, negative, model, FitDefaults.BASELINE)

    # Perform fit and plot the data
    fit_result = fit(data, fit_params)
    plot_fitted_data_bec(data, fit_result)
    

    # Optionally move the motor to the peak position
    if move_to_peak:
        move_to_position(data["motor_device"], data["motor_name"], fit_result["centre"], data)
    return fit_result


def scan_bpm(bpmname):
    """
    Runs a grid scan of a BPM in x and y,  and plots each channel
    as a heatmap.

    Parameters:
        bpmname: the name of the bpm to be scanned e.g. "fe"

    """

    # Open a dock area and set up the heatmaps
    dock_area = bec.gui.new("XBPM_Scan")
    wf5 = dock_area.new("Sum").new(bec.gui.available_widgets.Heatmap)
    wf1 = dock_area.new("Ch1", relative_to="Sum", position="bottom").new(
        bec.gui.available_widgets.Heatmap
    )
    wf3 = dock_area.new("Ch3", relative_to="Ch1", position="right").new(
        bec.gui.available_widgets.Heatmap
    )
    wf4 = dock_area.new("Ch4", relative_to="Ch3", position="bottom").new(
        bec.gui.available_widgets.Heatmap
    )
    wf2 = dock_area.new("Ch2", relative_to="Ch1", position="bottom").new(
        bec.gui.available_widgets.Heatmap
    )
    wfscan = dock_area.new("ScanControl").new(bec.gui.available_widgets.ScanControl)

    cfg = getattr(BPMScans, bpmname)

    wf1.x_label = cfg["x_name"]
    wf1.y_label = cfg["y_name"]
    wf1.plot(x_name=cfg["x_name"], y_name=cfg["y_name"], z_name=cfg["z1_name"], color_map="plasma")

    wf2.x_label = cfg["x_name"]
    wf2.y_label = cfg["y_name"]
    wf2.plot(x_name=cfg["x_name"], y_name=cfg["y_name"], z_name=cfg["z2_name"], color_map="plasma")

    wf3.x_label = cfg["x_name"]
    wf3.y_label = cfg["y_name"]
    wf3.plot(x_name=cfg["x_name"], y_name=cfg["y_name"], z_name=cfg["z3_name"], color_map="plasma")

    wf4.x_label = cfg["x_name"]
    wf4.y_label = cfg["y_name"]
    wf4.plot(x_name=cfg["x_name"], y_name=cfg["y_name"], z_name=cfg["z4_name"], color_map="plasma")

    wf5.x_label = cfg["x_name"]
    wf5.y_label = cfg["y_name"]
    wf5.plot(x_name=cfg["x_name"], y_name=cfg["y_name"], z_name=cfg["z5_name"], color_map="plasma")
    # Run the scan
    x_mot = cfg["x_device"]
    y_mot = cfg["y_device"]
    # scans.grid_scan(x_mot, -0.5, 0.5, 20, y_mot, -0.5, 0.5, 20,
    #                     exp_time=0.5, relative=False, snaked=True)


def optimise_kb(mirror):
    """
    Runs a grid scan of a the upstream and downstream benders,
    and plots a heatmap of the sample camera x or y sigma.

    Parameters:
        mirror: either "hfm" or :vfm"

    """

    # Open a dock area and set up the heatmaps
    dock_area = bec.gui.new(mirror)
    wf1 = dock_area.new("Heatmap").new(bec.gui.available_widgets.Heatmap)

    wfscan = dock_area.new("ScanControl").new(bec.gui.available_widgets.ScanControl)

    cfg = getattr(MirrorConfig, mirror)

    wf1.x_label = cfg["bu_name"]
    wf1.y_label = cfg["bd_name"]
    wf1.plot(x_name=cfg["bu_name"], y_name=cfg["bd_name"], z_name=cfg["z_name"], color_map="plasma")

    # Run the scan
    x_mot = cfg["x_device"]
    y_mot = cfg["y_device"]
    # scans.grid_scan(x_mot, -0.02, 0.02, 11, y_mot, -0.02, 0.02, 11,
    #                     exp_time=0.5, relative=True, snaked=True)