def twiss_response_creation_test():
    with logging_tools.DebugMode():
        create_response(create_twiss_resp_args + variable_cat_args +
                        optics_param_arg)
    LOG.info("\n\n\n")
    LOG.info("Creating Twiss response.")
    with logging_tools.DebugMode():
        create_response(create_twiss_resp_args + variable_cat_args)
Пример #2
0
def create_response(opt, other_opt):
    """ Entry point for creating pandas-based response matrices.

    The response matrices can be either created by response_madx or TwissResponse.

    Keyword Args:
        Required
        model_dir (str): Path to the model directory.
                         **Flags**: ['-m', '--model_dir']
        outfile_path (str): Name of fullresponse file.
                            **Flags**: ['-o', '--outfile']
        Optional
        creator (str): Create either with madx or analytically from twiss file.
                       **Flags**: --creator
                       **Choices**: ('madx', 'twiss')
                       **Default**: ``madx``
        debug: Print debug information.
               **Flags**: --debug
               **Action**: ``store_true``
        delta_k (float): Delta K1L to be applied to quads for sensitivity matrix (madx-only).
                         **Flags**: ['-k', '--deltak']
                         **Default**: ``2e-05``
        optics_params (str): List of parameters to correct upon (e.g. BBX BBY; twiss-only).
                             **Flags**: --optics_params
        variable_categories: List of the variables classes to use.
                             **Flags**: --variables
                             **Default**: ``['MQM', 'MQT', 'MQTL', 'MQY']``
    """
    with logging_tools.DebugMode(active=opt.debug,
                                 log_file=os.path.join(
                                     opt.model_dir,
                                     "generate_fullresponse.log")):
        LOG.info("Creating response.")
        accel_cls = manager.get_accel_class(other_opt)
        accel_inst = accel_cls(model_dir=opt.model_dir)
        if opt.optics_file is not None:
            accel_inst.optics_file = opt.optics_file

        if opt.creator == "madx":
            fullresponse = response_madx.generate_fullresponse(
                accel_inst, opt.variable_categories, delta_k=opt.delta_k)

        elif opt.creator == "twiss":
            fullresponse = response_twiss.create_response(
                accel_inst, opt.variable_categories, opt.optics_params)

        LOG.debug("Saving Response into file '{:s}'".format(opt.outfile_path))
        with open(opt.outfile_path, 'wb') as dump_file:
            pickle.Pickler(dump_file, -1).dump(fullresponse)
def global_correction(opt, accel_opt):
    """ Do the global correction. Iteratively.

    Keyword Args:
        Required
        meas_dir: Path to the directory containing the measurement files.
                       **Flags**: --meas_dir
        model_dir: Path to the dir containing the model (twiss.dat or twiss_elements.dat) to use.
                             **Flags**: --model_dir
        Optional
        beta_file_name: Prefix of the beta file to use. E.g.: getkmodbeta
                        **Flags**: --beta_file_name
                        **Default**: ``getbeta``
        debug: Print debug information.
               **Flags**: --debug
               **Action**: ``store_true``
        eps (float): (Not implemented yet) Convergence criterion.
                     If :math:`<|\Delta(PARAM \cdot WEIGHT)|> < \epsilon`, stop iteration.
                     **Flags**: --eps
                     **Default**: ``None``
        errorcut (float): Reject BPMs whose error bar is higher than the corresponding input.
                          Input in order of optics_params.
                          **Flags**: --error_cut
        fullresponse_path: Path to the fullresponse binary file.
                           If not given, calculates the response analytically.
                           **Flags**: --fullresponse
        max_iter (int): Maximum number of correction re-iterations to perform.
                        A value of `0` means the correction is calculated once
                        (like in the old days).
                        **Flags**: --max_iter
                        **Default**: ``3``
        method (str): Optimization method to use. (Not implemented yet)
                      **Flags**: --method
                      **Choices**: ['pinv']
                      **Default**: ``pinv``
        modelcut (float): Reject BPMs whose deviation to the model is higher than the
                          correspoding input. Input in order of optics_params.
                          **Flags**: --model_cut
        optics_file: Path to the optics file to use, usually modifiers.madx.
                     If not present will default to model_path/modifiers.madx
                     **Flags**: --optics_file
        optics_params (str): List of parameters to correct upon (e.g. BBX BBY)
                             **Flags**: --optics_params
                             **Default**: ``['MUX', 'MUY', 'BBX', 'BBY', 'NDX', 'Q']``
        output_path: Path to the directory where to write the output files,
                     will default to the --meas input path.
                     **Flags**: --output_dir
                     **Default**: ``None``
        svd_cut (float): Cutoff for small singular values of the pseudo inverse. (Method: 'pinv')
                         Singular values smaller than
                         :math:`rcond \cdot largest_singular_value` are set to zero
                         **Flags**: --svd_cut
                         **Default**: ``0.01``
        use_errorbars: If True, it will take into account the measured errorbars in the correction.
                       **Flags**: --use_errorbars
                       **Action**: ``store_true``
        variable_categories: List of names of the variables classes to use.
                             **Flags**: --variables
                             **Default**: ``['MQM', 'MQT', 'MQTL', 'MQY']``
        virt_flag: If true, it will use virtual correctors.
                   **Flags**: --virt_flag
                   **Action**: ``store_true``
        weights (float): Weights to apply to each measured quantity. Input in order of optics_params.
                         **Flags**: --weights
    """

    LOG.info("Starting Iterative Global Correction.")
    with logging_tools.DebugMode(active=opt.debug,
                                 log_file=os.path.join(opt.model_dir, "iterative_correction.log")):
        not_implemented_params = [k for k in opt.optics_params
                                  if k not in _get_measurement_filters()]
        if any(not_implemented_params):
            raise NotImplementedError("Correct iterative is not equipped for parameters:"
                                      "'{:s}'".format(not_implemented_params))

        # ######### Preparations ######### #

        # check on opt
        opt = _check_opt(opt)
        meth_opt = _get_method_opt(opt)

        # get accelerator class
        accel_cls = manager.get_accel_class(accel_opt)
        accel_inst = accel_cls(model_dir=opt.model_dir)
        if opt.optics_file is not None:
            accel_inst.optics_file = opt.optics_file

        # convert numbers to dictionaries
        w_dict = dict(zip(opt.optics_params, opt.weights))
        mcut_dict = dict(zip(opt.optics_params, opt.modelcut))
        ecut_dict = dict(zip(opt.optics_params, opt.errorcut))

        # read data from files
        vars_list = _get_varlist(accel_cls, opt.variable_categories, opt.virt_flag)
        optics_params, meas_dict = _get_measurment_data(
            opt.optics_params,
            opt.meas_dir, opt.beta_file_name,
            w_dict,
        )
        mcut_dict = _automate_modelcut(mcut_dict, meas_dict, opt.variable_categories)

        if opt.fullresponse_path is not None:
            resp_dict = _load_fullresponse(opt.fullresponse_path, vars_list)
        else:
            resp_dict = response_twiss.create_response(
                accel_inst, opt.variable_categories, optics_params
            )

        # the model in accel_inst is modified later, so save nominal model here to variables
        nominal_model = _maybe_add_coupling_to_model(accel_inst.get_model_tfs(), optics_params)

        # apply filters to data
        meas_dict = _filter_measurement(
            optics_params, meas_dict, nominal_model,
            opt.use_errorbars, w_dict, ecut_dict, mcut_dict
        )
        meas_dict = _append_model_to_measurement(nominal_model, meas_dict, optics_params)
        resp_dict = _filter_response_index(resp_dict, meas_dict, optics_params)
        resp_matrix = _join_responses(resp_dict, optics_params, vars_list)

        # _dump(os.path.join(opt.output_path, "measurement_dict.bin"), meas_dict)
        delta = tfs.TfsDataFrame(0, index=vars_list, columns=["DELTA"])

        # ######### Iteration Phase ######### #

        for iteration in range(opt.max_iter + 1):
            LOG.info("Correction Iteration {:d} of {:d}.".format(iteration, opt.max_iter))

            # ######### Update Model and Response ######### #
            if iteration > 0:
                LOG.debug("Updating model via MADX.")
                corr_model_path = os.path.join(opt.output_path, "twiss_" + str(iteration) + ".dat")

                _create_corrected_model(corr_model_path, opt.change_params_path,
                                        accel_inst, opt.debug)

                corr_model_elements = tfs.read_tfs(corr_model_path, index="NAME")
                corr_model_elements = _maybe_add_coupling_to_model(
                    corr_model_elements, optics_params
                )
                corr_model = corr_model_elements.loc[tfs.get_bpms(corr_model_elements), :]

                meas_dict = _append_model_to_measurement(corr_model, meas_dict, optics_params)
                if opt.update_response:
                    LOG.debug("Updating response.")
                    # please look away for the next two lines.
                    accel_inst._model = corr_model
                    accel_inst._elements = corr_model_elements
                    resp_dict = response_twiss.create_response(
                        accel_inst, opt.variable_categories, optics_params
                    )
                    resp_dict = _filter_response_index(resp_dict, meas_dict, optics_params)
                    resp_matrix = _join_responses(resp_dict, optics_params, vars_list)

            # ######### Actual optimization ######### #
            delta += _calculate_delta(
                resp_matrix, meas_dict, optics_params, vars_list, opt.method, meth_opt)

            writeparams(opt.change_params_path, delta)
            writeparams(opt.change_params_correct_path, -delta)
            LOG.debug("Cumulative delta: {:.5e}".format(
                np.sum(np.abs(delta.loc[:, "DELTA"].values))))

        write_knob(opt.knob_path, delta)
    LOG.info("Finished Iterative Global Correction.")
def analyse_with_bbq_corrections(opt):
    """ Create amplitude detuning analysis with BBQ correction from timber data.

    Keyword Args:
        Required
        beam (int): Which beam to use.
                    **Flags**: --beam
        kickac_path (str): Location of the kickac file
                           **Flags**: --kickac
        plane (str): Plane of the kicks. 'X' or 'Y'.
                           **Flags**: --plane
                           **Choices**: XY
        Optional
        ampdet_plot_out (str): Save the amplitude detuning plot here.
                          **Flags**: --ampdetplot
        ampdet_plot_show: Show the amplitude detuning plot.
                          **Flags**: --ampdetplotshow
                          **Action**: ``store_true``
        ampdet_plot_xmax (float): Maximum action (x-axis) in amplitude detuning plot.
                          **Flags**: --ampdetplotxmax
        ampdet_plot_xmin (float): Minimum action (x-axis) in amplitude detuning plot.
                                  **Flags**: --ampdetplotxmin
        ampdet_plot_ymax (float): Maximum tune (y-axis) in amplitude detuning plot.
                                  **Flags**: --ampdetplotymax
        ampdet_plot_ymin (float): Minimum tune (y-axis) in amplitude detuning plot.
                                  **Flags**: --ampdetplotymin
        bbq_out (str): Output location to save bbq data as tfs-file
                       **Flags**: --bbqout
        bbq_plot_full: Plot the full bqq data with interval as lines.
                       **Flags**: --bbqplotfull
                       **Action**: ``store_true``
        bbq_plot_out (str): Save the bbq plot here.
                            **Flags**: --bbqplot
        bbq_plot_show: Show the bbq plot.
                       **Flags**: --bbqplotshow
                       **Action**: ``store_true``
        bbq_plot_two: Two plots for the bbq plot.
                      **Flags**: --bbqplottwo
                      **Action**: ``store_true``
        debug: Activates Debug mode
               **Flags**: --debug
               **Action**: ``store_true``
        fine_cut (float): Cut (i.e. tolerance) of the tune for the fine cleaning.
                          **Flags**: --finecut
        fine_window (int): Length of the moving average window. (# data points)
                           **Flags**: --finewindow
        kickac_out (str): If given, writes out the modified kickac file
                          **Flags**: --kickacout
        label (str): Label to identify this run.
                     **Flags**: --label
        logfile (str): Logfile if debug mode is active.
                       **Flags**: --logfile
        timber_in: Fill number of desired data or path to presaved tfs-file
                   **Flags**: --timberin
        timber_out (str): Output location to save fill as tfs-file
                          **Flags**: --timberout
        tune_cut (float): Cuts for the tune. For BBQ cleaning.
                          **Flags**: --tunecut
        tune_x (float): Horizontal Tune. For BBQ cleaning.
                        **Flags**: --tunex
        tune_x_max (float): Horizontal Tune minimum. For BBQ cleaning.
                            **Flags**: --tunexmax
        tune_x_min (float): Horizontal Tune minimum. For BBQ cleaning.
                            **Flags**: --tunexmin
        tune_y (float): Vertical Tune. For BBQ cleaning.
                        **Flags**: --tuney
        tune_y_max (float): Vertical Tune minimum. For BBQ cleaning.
                            **Flags**: --tuneymax
        tune_y_min (float): Vertical  Tune minimum. For BBQ cleaning.
                            **Flags**: --tuneymin
        window_length (int): Length of the moving average window. (# data points)
                             **Flags**: --window
                             **Default**: ``20``
     """
    LOG.info("Starting Amplitude Detuning Analysis")
    with logging_tools.DebugMode(active=opt.debug, log_file=opt.logfile):
        opt = _check_analyse_opt(opt)
        figs = {}

        # get data
        kickac_df = tfs.read_tfs(opt.kickac_path, index=COL_TIME())
        bbq_df = _get_timber_data(opt.beam, opt.timber_in, opt.timber_out, kickac_df)
        x_interval = _get_approx_bbq_interval(bbq_df, kickac_df.index, opt.window_length)

        # add moving average to kickac
        kickac_df, bbq_df = kickac_modifiers.add_moving_average(kickac_df, bbq_df,
                                                                **opt.get_subdict([
                                                                    "window_length",
                                                                    "tune_x_min", "tune_x_max",
                                                                    "tune_y_min", "tune_y_max",
                                                                    "fine_cut", "fine_window"]
                                                                )
                                                                )

        # add corrected values to kickac
        kickac_df = kickac_modifiers.add_corrected_natural_tunes(kickac_df)
        kickac_df = kickac_modifiers.add_total_natq_std(kickac_df)

        # BBQ plots
        if opt.bbq_plot_out or opt.bbq_plot_show:
            if opt.bbq_plot_full:
                figs["bbq"] = bbq_tools.plot_bbq_data(
                    bbq_df,
                    output=opt.bbq_plot_out,
                    show=opt.bbq_plot_show,
                    two_plots=opt.bbq_plot_two,
                    interval=[str(datetime.datetime.fromtimestamp(xint, tz=TIMEZONE))
                              for xint in x_interval],
                )
            else:
                figs["bbq"] = bbq_tools.plot_bbq_data(
                    bbq_df.loc[x_interval[0]:x_interval[1]],
                    output=opt.bbq_plot_out,
                    show=opt.bbq_plot_show,
                    two_plots=opt.bbq_plot_two,
                )

        # amplitude detuning odr and plotting
        for tune_plane in PLANES:
            for corr in [False, True]:
                corr_label = "_corrected" if corr else ""

                # get the proper data
                data = kickac_modifiers.get_ampdet_data(kickac_df, opt.plane, tune_plane,
                                                        corrected=corr)

                # make the odr
                odr_fit = detuning_tools.do_linear_odr(**data)
                kickac_df = kickac_modifiers.add_odr(kickac_df, odr_fit, opt.plane, tune_plane,
                                                     corrected=corr)

                # plotting
                labels = ta_const.get_paired_lables(opt.plane, tune_plane)
                id_str = "J{:s}_Q{:s}{:s}".format(opt.plane.upper(), tune_plane.upper(), corr_label)

                try:
                    output = os.path.splitext(opt.ampdet_plot_out)
                except AttributeError:
                    output = None
                else:
                    output = "{:s}_{:s}{:s}".format(output[0], id_str, output[1])

                figs[id_str] = detuning_tools.plot_detuning(
                    odr_fit=odr_fit,
                    odr_plot=detuning_tools.plot_linear_odr,
                    labels={"x": labels[0], "y": labels[1], "line": opt.label},
                    output=output,
                    show=opt.ampdet_plot_show,
                    xmin=opt.ampdet_plot_xmin,
                    xmax=opt.ampdet_plot_xmax,
                    ymin=opt.ampdet_plot_ymin,
                    ymax=opt.ampdet_plot_ymax,
                    **data
                )

    # show plots if needed
    if opt.bbq_plot_show or opt.ampdet_plot_show:
        plt.show()

    # output kickac and bbq data
    if opt.kickac_out:
        tfs.write_tfs(opt.kickac_out, kickac_df, save_index=COL_TIME())

    if opt.bbq_out:
        tfs.write_tfs(opt.bbq_out, bbq_df.loc[x_interval[0]:x_interval[1]],
                      save_index=COL_TIME())

    return figs
def analyse_with_bbq_corrections(opt):
    """ Create amplitude detuning analysis with BBQ correction from timber data. """
    LOG.info("Starting Amplitude Detuning Analysis")
    with logging_tools.DebugMode(active=opt.debug, log_file=opt.logfile):
        opt = _check_analyse_opt(opt)
        figs = {}

        # get data
        bbq_df = _get_timber_data(opt.beam, opt.timber_in, opt.timber_out)
        kickac_df = tfs.read_tfs(opt.kickac_path, index=COL_TIME())
        x_interval = _get_approx_bbq_interval(bbq_df, kickac_df.index,
                                              opt.window_length)

        # add moving average to kickac
        kickac_df, bbq_df = _add_moving_average(
            kickac_df, bbq_df,
            **opt.get_subdict([
                "window_length", "tune_x_min", "tune_x_max", "tune_y_min",
                "tune_y_max", "fine_cut", "fine_window"
            ]))

        # add corrected values to kickac
        kickac_df = _add_corrected_natural_tunes(kickac_df)

        # output kickac and bbq data
        if opt.kickac_out:
            tfs.write_tfs(opt.kickac_out, kickac_df, save_index=COL_TIME())

        if opt.bbq_out:
            tfs.write_tfs(opt.bbq_out,
                          bbq_df.loc[x_interval[0]:x_interval[1]],
                          save_index=COL_TIME())

        if opt.bbq_plot_out or opt.bbq_plot_show:
            if opt.bbq_plot_full:
                figs["bbq"] = bbq_tools.plot_bbq_data(
                    bbq_df,
                    output=opt.bbq_plot_out,
                    show=opt.bbq_plot_show,
                    two_plots=opt.bbq_plot_two,
                    interval=[
                        str(datetime.datetime.fromtimestamp(xint))
                        for xint in x_interval
                    ],
                )
            else:
                figs["bbq"] = bbq_tools.plot_bbq_data(
                    bbq_df.loc[x_interval[0]:x_interval[1]],
                    output=opt.bbq_plot_out,
                    show=opt.bbq_plot_show,
                    two_plots=opt.bbq_plot_two,
                )

        # amplitude detuning analysis
        plane = ta_const.get_plane_from_orientation(opt.orientation)
        for other_plane in PLANES:
            labels = ta_const.get_paired_lables(plane, other_plane)
            id_str = "J{:s}_Q{:s}".format(plane.upper(), other_plane.upper())

            # get proper data
            columns = ta_const.get_paired_columns(plane, other_plane)
            data = {
                key: kickac_df.loc[:, columns[key]]
                for key in columns.keys()
            }

            # plotting
            try:
                output = os.path.splitext(opt.ampdet_plot_out)
                output = "{:s}_{:s}{:s}".format(output[0], id_str, output[1])
            except AttributeError:
                output = None

            figs[id_str] = detuning_tools.plot_detuning(
                odr_plot=detuning_tools.linear_odr_plot,
                labels={
                    "x": labels[0],
                    "y": labels[1],
                    "line": opt.label
                },
                output=output,
                show=opt.ampdet_plot_show,
                x_min=opt.ampdet_plot_xmin,
                x_max=opt.ampdet_plot_xmax,
                y_min=opt.ampdet_plot_ymin,
                y_max=opt.ampdet_plot_ymax,
                **data)

    if opt.bbq_plot_show or opt.ampdet_plot_show:
        plt.show()

    return figs