jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def plot_lean_logs(full_df):
    # full_df: student logs during training/evaluation
    df = full_df.loc[full_df["training"]==False]
    fig, axs = plt.subplots(2, 2, sharex=True, sharey='row', figsize=(12, 6.7))
    axs[0, 0].set_title("Source Env. (Without Disturbance)")
    sns.pointplot(ax = axs[0, 0], x="iteration", y="success", hue="agent", data =df.loc[df["disturbance"] == False])
    axs[0, 0].set(xlabel='Num. demonstrations used for training', ylabel="Robustness (Success Rate)")

    axs[0, 1].set_title("Target Env. (With Disturbance)")
    sns.pointplot(ax = axs[0, 1], x="iteration", y="success", hue="agent", data =df.loc[df["disturbance"] == True])
    axs[0, 1].set(xlabel='Num. demonstrations used for training', ylabel='Robustness (Success Rate)')

    axs[1, 0].set_title("Source Env. (Without Disturbance)")
    sns.pointplot(ax = axs[1, 0], x="iteration", y="return", hue="agent", data =df.loc[df["disturbance"] == False])
    axs[1, 0].set(xlabel='Num. demonstrations used for training', ylabel='Performance (MPC Cost)')

    axs[1, 1].set_title("Target Env. (With Disturbance)")
    sns.pointplot(ax = axs[1, 1], x="iteration", y="return", hue="agent", data =df.loc[df["disturbance"] == True])
    axs[1, 1].set(xlabel='Num. demonstrations used for training', ylabel='Performance (MPC Cost)')
    plt.tight_layout()
//  [python]
def plot_lean_logs(full_df):
    # full_df: student logs during training/evaluation
    df = full_df.loc[full_df["training"]==False]
    fig, axs = plt.subplots(2, 2, sharex=True, sharey='row', figsize=(12, 6.7))
    axs[0, 0].set_title("Source Env. (Without Disturbance)")
    sns.lineplot(ax = axs[0, 0], x="iteration", y="success", hue="agent", estimator='mean', ci=ci, data =df.loc[df["disturbance"] == False])
    axs[0, 0].set(xlabel='Num. demonstrations used for training', ylabel="Robustness (Success Rate)")

    axs[0, 1].set_title("Target Env. (With Disturbance)")
    sns.lineplot(ax = axs[0, 1], x="iteration", y="success", hue="agent", estimator='mean', ci=ci, data =df.loc[df["disturbance"] == True])
    axs[0, 1].set(xlabel='Num. demonstrations used for training', ylabel='Robustness (Success Rate)')

    axs[1, 0].set_title("Source Env. (Without Disturbance)")
    sns.lineplot(ax = axs[1, 0], x="iteration", y="return", hue="agent", estimator='mean', ci=ci, data =df.loc[df["disturbance"] == False])
    axs[1, 0].set(xlabel='Num. demonstrations used for training', ylabel='Performance (MPC Cost)')

    axs[1, 1].set_title("Target Env. (With Disturbance)")
    sns.lineplot(ax = axs[1, 1], x="iteration", y="return", hue="agent", estimator='mean', ci=ci, data =df.loc[df["disturbance"] == True])
    axs[1, 1].set(xlabel='Num. demonstrations used for training', ylabel='Performance (MPC Cost)')
    plt.tight_layout()
//  [python]
def figure1(
        df: pd.DataFrame,
        output_path: str,
        log_y: bool = False,
        save_pdf: bool = True,
        name: str = "figure1",
        title: str = rf"ADE per visible chain length (w/ 95\% CI)",
        scale: float = 1.0,
) -> None:
    df = df.copy()
    df = df[df.n_timesteps_visible_chain > 1]

    fig, ax = plt.subplots(figsize=(6 * scale, 4 * scale))
    fig.suptitle(title)
    sns.pointplot(
        df,
        x="n_timesteps_visible_chain",
        y="ade_visible_chain",
        hue="name",
        palette=cycle(["GoldenRod", "r", "forestgreen", "yellow"]),
        linestyles="-",
        linewidth=2,
        ci=95,
        markers=True,
        dashes=False,
        err_style="band",
        # err_style="bars", err_kws={"fmt": 'o', "linewidth": 2, "capsize": 6},
        alpha=1,
        ax=ax,
    )
    ax.set_xlabel(rf"visibility chain length")
    ax.set_ylabel(rf"ADE")
    if log_y:
        plt.yscale('log')
    ax.legend_.set_title(None)
    plt.tight_layout()
    plt.savefig(os.path.join(output_path, f"{name}.png"))
    if save_pdf:
        plt.savefig(os.path.join(output_path, f"PDF__{name}.pdf"))
    plt.show()
    plt.close()
    plt.clf()
//  [python]
def figure1(
        df: pd.DataFrame,
        output_path: str,
        log_y: bool = False,
        save_pdf: bool = True,
        name: str = "figure1",
        title: str = rf"ADE per visible chain length (w/ 95\% CI)",
        scale: float = 1.0,
) -> None:
    df = df.copy()
    df = df[df.n_timesteps_visible_chain > 1]

    fig, ax = plt.subplots(figsize=(6 * scale, 4 * scale))
    fig.suptitle(title)
    sns.lineplot(
        df,
        x="n_timesteps_visible_chain",
        y="ade_visible_chain",
        hue="name",
        palette=cycle(["GoldenRod", "r", "forestgreen", "yellow"]),
        linestyle="-",
        linewidth=2,
        errorbar=("ci", 95),
        markers=True,
        dashes=False,
        err_style="band",
        # err_style="bars", err_kws={"fmt": 'o', "linewidth": 2, "capsize": 6},
        alpha=1,
        ax=ax,
    )
    ax.set_xlabel(rf"visibility chain length")
    ax.set_ylabel(rf"ADE")
    if log_y:
        plt.yscale('log')
    ax.legend_.set_title(None)
    plt.tight_layout()
    plt.savefig(os.path.join(output_path, f"{name}.png"))
    if save_pdf:
        plt.savefig(os.path.join(output_path, f"PDF__{name}.pdf"))
    plt.show()
    plt.close()
    plt.clf()
//  [python]
def plot_logs(input_df, show_training = False):

    # df: student logs
    if show_training:
        fig, axs = plt.subplots(3, 3, sharex=True, sharey='row', figsize=(21, 12))
    else: 
        fig, axs = plt.subplots(3, 2, sharex=True, sharey='row', figsize=(16, 12))

    eval_df = input_df.loc[input_df["training"]==False]
    training_df = input_df.loc[input_df["training"]==True]
    
    pd.set_option("display.max_rows", None)#Use this to print the whole dataframe 
    print(eval_df)
    # print(eval_df.loc)
    axs[0, 0].set_title("Iteration vs Return")
    sns.pointplot(ax = axs[0, 0], x="iteration", y="return", hue="agent", estimator='mean', ci=ci, data =eval_df)
    axs[0, 0].set(xlabel='Num. demonstrations used for training', ylabel='Return')

    axs[1, 0].set_title("Iteration vs len")
    sns.pointplot(ax = axs[1, 0], x="iteration", y="len", hue="agent", estimator='mean', ci=ci, data =eval_df)
    axs[1, 0].set(xlabel='Num. demonstrations used for training', ylabel='Episode Length')

    axs[2, 0].set_title("Iteration vs success")
    sns.pointplot(ax = axs[2, 0], x="iteration", y="success", hue="agent", estimator='mean', ci=ci, data =eval_df)
    axs[2, 0].set(xlabel='Num. demonstrations used for training', ylabel="Success Rate")


    

    if show_training: 
        axs[0, 2].set_title("Training Env.")
        sns.pointplot(ax = axs[0, 2], x="iteration", y="return", hue="agent", estimator='mean', ci=ci, data=training_df)
        axs[0, 2].set(xlabel='Num. demonstrations used for training', ylabel='Return')

        axs[1, 2].set_title("Training Env.")
        sns.pointplot(ax = axs[1, 2], x="iteration", y="len", hue="agent", estimator='mean', ci=ci, data=training_df)
        axs[1, 2].set(xlabel='Num. demonstrations used for training', ylabel='Episode Length')

        axs[2, 2].set_title("Training Env.")
        sns.pointplot(ax = axs[2, 2], x="iteration", y="success", hue="agent", estimator='mean', ci=ci, data =training_df)
        axs[2, 2].set(xlabel='Num. demonstrations used for training', ylabel="Success Rate")

    plt.show()
//  [python]
def plot_logs(input_df, show_training = False):

    # df: student logs
    if show_training:
        fig, axs = plt.subplots(3, 3, sharex=True, sharey='row', figsize=(21, 12))
    else: 
        fig, axs = plt.subplots(3, 2, sharex=True, sharey='row', figsize=(16, 12))

    eval_df = input_df.loc[input_df["training"]==False]
    training_df = input_df.loc[input_df["training"]==True]
    
    pd.set_option("display.max_rows", None)#Use this to print the whole dataframe 
    print(eval_df)
    # print(eval_df.loc)
    axs[0, 0].set_title("Iteration vs Return")
    sns.lineplot(ax = axs[0, 0], x="iteration", y="return", hue="agent", estimator='mean', ci=ci, data =eval_df)
    axs[0, 0].set(xlabel='Num. demonstrations used for training', ylabel='Return')

    axs[1, 0].set_title("Iteration vs len")
    sns.lineplot(ax = axs[1, 0], x="iteration", y="len", hue="agent", estimator='mean', ci=ci, data =eval_df)
    axs[1, 0].set(xlabel='Num. demonstrations used for training', ylabel='Episode Length')

    axs[2, 0].set_title("Iteration vs success")
    sns.lineplot(ax = axs[2, 0], x="iteration", y="success", hue="agent", estimator='mean', ci=ci, data =eval_df)
    axs[2, 0].set(xlabel='Num. demonstrations used for training', ylabel="Success Rate")


    

    if show_training: 
        axs[0, 2].set_title("Training Env.")
        sns.lineplot(ax = axs[0, 2], x="iteration", y="return", hue="agent", estimator='mean', ci=ci, data=training_df)
        axs[0, 2].set(xlabel='Num. demonstrations used for training', ylabel='Return')

        axs[1, 2].set_title("Training Env.")
        sns.lineplot(ax = axs[1, 2], x="iteration", y="len", hue="agent", estimator='mean', ci=ci, data=training_df)
        axs[1, 2].set(xlabel='Num. demonstrations used for training', ylabel='Episode Length')

        axs[2, 2].set_title("Training Env.")
        sns.lineplot(ax = axs[2, 2], x="iteration", y="success", hue="agent", estimator='mean', ci=ci, data =training_df)
        axs[2, 2].set(xlabel='Num. demonstrations used for training', ylabel="Success Rate")

    plt.show()
//  [python]
def boxplot(opts, evalConfigs, datasets, randValue, resTypeFile, resTypeName):
    
    # Set plot variables
    plt.rcParams.update({'font.size': 14})
    
    DFs = []
    for i, dataset in enumerate(datasets):
        evalConfig = evalConfigs[i]
        
        # Read output file containing AUROC values
        DF = pd.read_csv(str(evalConfig.output_settings.base_dir) + '/' \
                           + str(evalConfig.input_settings.datadir).split("inputs")[1] + '/' \
                           + str(evalConfig.output_settings.output_prefix) \
                           + '-' + resTypeFile + '.csv', header = 0, index_col = 0)
        
        DF = DF.T
        DFs.append(DF)    
    
    # Initialize plot window
    f, axes = plt.subplots(len(datasets), 1, figsize=(1.5*max([len(DF.columns) for DF in DFs]), 5*len(datasets)))
     
    for i, dataset in enumerate(datasets):
        DF = DFs[i]
        modifiedDF = pd.melt(DF, id_vars=DF.index, value_vars=DF.columns)
              
        if len(datasets) > 1:
            subax = axes[i]
        else:
            subax= axes
        
        # Indicate random predictors value
        ax = plt.plot(range(-1,len(DF.columns)+1), [randValue]*len(range(-1,len(DF.columns)+1)), ax = subax)
        ax[0].set_linestyle("--")
        ax[0].set_color("gray")
        
        # Plot the AUROC values as a box plot
        ax = sns.boxplot(y='value',x ='variable', data=modifiedDF, fliersize = 0,
                        palette=sns.color_palette("Set1"),
                        ax = subax)
        sns.swarmplot(y='value',x ='variable', data=modifiedDF,
                         alpha = 0.5, palette=['k'],
                         dodge = True,
                         ax = subax)
        
        ax.set_ylim([0.0,1])
        ax.set_ylabel(resTypeName, fontsize = 18)
        
        ax.set_title(resTypeName +' values for ' + dataset +' Network', fontsize = 18)
        plt.tight_layout()
    
    subax.set_xlabel('Algorithm', fontsize = 18)
               
    file = opts.output + '/' \
            + '-'.join([str(c.output_settings.output_prefix) for c in evalConfigs]) \
            + '-boxplot-' + resTypeFile + '.pdf'
    print("Boxplot saved to " + file)
    plt.savefig(file, dpi = 300)
//  [python]
def boxplot(opts, evalConfigs, datasets, randValue, resTypeFile, resTypeName):
    
    # Set plot variables
    plt.rcParams.update({'font.size': 14})
    
    DFs = []
    for i, dataset in enumerate(datasets):
        evalConfig = evalConfigs[i]
        
        # Read output file containing AUROC values
        DF = pd.read_csv(str(evalConfig.output_settings.base_dir) + '/' \
                           + str(evalConfig.input_settings.datadir).split("inputs")[1] + '/' \
                           + str(evalConfig.output_settings.output_prefix) \
                           + '-' + resTypeFile + '.csv', header = 0, index_col = 0)
        
        DF = DF.T
        DFs.append(DF)    
    
    # Initialize plot window
    f, axes = plt.subplots(len(datasets), 1, figsize=(1.5*max([len(DF.columns) for DF in DFs]), 5*len(datasets)))
     
    for i, dataset in enumerate(datasets):
        DF = DFs[i]
        modifiedDF = pd.melt(DF, id_vars=DF.index, value_vars=DF.columns)
              
        if len(datasets) > 1:
            subax = axes[i]
        else:
            subax= axes
        
        # Indicate random predictors value
        ax = sns.lineplot(y = randValue, 
                          x=range(-1,len(DF.columns)+1), ax = subax)
        ax.lines[0].set_linestyle("--")
        ax.lines[0].set_color("gray")
        
        # Plot the AUROC values as a box plot
        ax = sns.boxplot(y='value',x ='variable', data=modifiedDF, fliersize = 0,
                        palette=sns.color_palette("Set1"),
                        ax = subax)
        sns.swarmplot(y='value',x ='variable', data=modifiedDF,
                         alpha = 0.5, palette=['k'],
                         dodge = True,
                         ax = subax)
        
        ax.set_ylim([0.0,1])
        ax.set_ylabel(resTypeName, fontsize = 18)
        
        ax.set_title(resTypeName +' values for ' + dataset +' Network', fontsize = 18)
        plt.tight_layout()
    
    subax.set_xlabel('Algorithm', fontsize = 18)
               
    file = opts.output + '/' \
            + '-'.join([str(c.output_settings.output_prefix) for c in evalConfigs]) \
            + '-boxplot-' + resTypeFile + '.pdf'
    print("Boxplot saved to " + file)
    plt.savefig(file, dpi = 300)
//  [python]
def figure4(
        df: pd.DataFrame,
        output_path: str,
        log_y: bool = False,
        save_pdf: bool = True,
        name: str = "figure4",
        title: str = rf"PCK per threshold across scenes (w/ 95\% CI)",
        scale: float = 1.0,
) -> None:
    df = df.copy()

    pts_cols = [c for c in df.columns if c.startswith("pts_within_")]
    df = df[["name"] + pts_cols]
    df = df.melt("name", var_name="threshold", value_name="pts_within_threshold")
    df.threshold = df.threshold.apply(lambda x: float(x.replace("pts_within_", "")))

    fig, ax = plt.subplots(figsize=(6 * scale, 4 * scale))
    fig.suptitle(title)
    sns.pointplot(
        df,
        x="threshold",
        y="pts_within_threshold",
        hue="name",
        palette=cycle(["GoldenRod", "r", "forestgreen", "yellow"]),
        linestyles="-",
        linewidth=2,
        ci=95,
        markers=True,
        dashes=False,
        errwidth=2,
        capsize=0.2,
        join=True,
        alpha=1,
        legend=True,
        ax=ax,
    )
    ax.set_xlabel(rf"threshold in pixels")
    ax.set_ylabel(rf"PCK")
    if log_y:
        plt.yscale('log')
    plt.tight_layout()
    plt.savefig(os.path.join(output_path, f"{name}.png"))
    if save_pdf:
        plt.savefig(os.path.join(output_path, f"PDF__{name}.pdf"))
    plt.show()
    plt.close()
    plt.clf()
//  [python]
def figure4(
        df: pd.DataFrame,
        output_path: str,
        log_y: bool = False,
        save_pdf: bool = True,
        name: str = "figure4",
        title: str = rf"PCK per threshold across scenes (w/ 95\% CI)",
        scale: float = 1.0,
) -> None:
    df = df.copy()

    pts_cols = [c for c in df.columns if c.startswith("pts_within_")]
    df = df[["name"] + pts_cols]
    df = df.melt("name", var_name="threshold", value_name="pts_within_threshold")
    df.threshold = df.threshold.apply(lambda x: float(x.replace("pts_within_", "")))

    fig, ax = plt.subplots(figsize=(6 * scale, 4 * scale))
    fig.suptitle(title)
    sns.lineplot(
        df,
        x="threshold",
        y="pts_within_threshold",
        hue="name",
        palette=cycle(["GoldenRod", "r", "forestgreen", "yellow"]),
        linestyle="-",
        linewidth=2,
        errorbar=("ci", 95),
        markers=True,
        dashes=False,
        err_style="bars",
        alpha=1,
        legend=True,
        ax=ax,
    )
    ax.set_xlabel(rf"threshold in pixels")
    ax.set_ylabel(rf"PCK")
    if log_y:
        plt.yscale('log')
    plt.tight_layout()
    plt.savefig(os.path.join(output_path, f"{name}.png"))
    if save_pdf:
        plt.savefig(os.path.join(output_path, f"PDF__{name}.pdf"))
    plt.show()
    plt.close()
    plt.clf()
//  [python]
def plot_single_line(this_fig, spec_config, input_data, step_size=0.5, ent_max=5,
                     show_x_ticks=False, show_y_ticks=True, data_name="", model_name="", ymin=2, ymax=5):
    segs = np.arange(0, ent_max, step_size).tolist()
    # colorblind = sns.color_palette("coolwarm", 10)[::-1]
    last_inp_mean, cur_inp_mean, cur_pred_mean, next_pred_mean, atte_ent_mean = input_data

    axes = this_fig.add_subplot(spec_config)
    plt.plot(list(np.arange(0, 5, step_size)), atte_ent_mean, marker='o', linestyle='-')

    # axes = sns.boxplot(x=x, y=y, palette=colorblind, showfliers=False)

    axes.xaxis.set_major_locator(MultipleLocator(1))
    axes.xaxis.set_major_formatter(FormatStrFormatter('%d'))

    # For the minor ticks, use no labels; default NullFormatter.
    axes.xaxis.set_minor_locator(MultipleLocator(0.5))
    # axes.get_ylim()
    # axes.set_ylim(ymin, ymax)

    if show_x_ticks:
        # x_vals = [m * step_size for m in xticks]
        # axes.set_xticklabels(x_vals, rotation='vertical')center
        pass
    else:
        plt.setp(axes.get_xticklabels(), visible=False)

    # if not show_y_ticks:
    #     plt.setp(axes.get_yticklabels(), visible=False)
    if data_name != "":
        axes.set_ylabel(data_name)
    else:
        axes.set_ylabel("")
    if model_name != "":
        axes.set_title(model_name)
    return axes
//  [python]
def plot_single_line(this_fig, spec_config, input_data, step_size=0.5, ent_max=5,
                     show_x_ticks=False, show_y_ticks=True, data_name="", model_name="", ymin=2, ymax=5):
    segs = np.arange(0, ent_max, step_size).tolist()
    # colorblind = sns.color_palette("coolwarm", 10)[::-1]
    last_inp_mean, cur_inp_mean, cur_pred_mean, next_pred_mean, atte_ent_mean = input_data

    axes = this_fig.add_subplot(spec_config)
    sns.lineplot(x=list(np.arange(0, 5, step_size)), y=atte_ent_mean, markers=True, dashes=False)

    # axes = sns.boxplot(x=x, y=y, palette=colorblind, showfliers=False)

    axes.xaxis.set_major_locator(MultipleLocator(1))
    axes.xaxis.set_major_formatter(FormatStrFormatter('%d'))

    # For the minor ticks, use no labels; default NullFormatter.
    axes.xaxis.set_minor_locator(MultipleLocator(0.5))
    # axes.get_ylim()
    # axes.set_ylim(ymin, ymax)

    if show_x_ticks:
        # x_vals = [m * step_size for m in xticks]
        # axes.set_xticklabels(x_vals, rotation='vertical')center
        pass
    else:
        plt.setp(axes.get_xticklabels(), visible=False)

    # if not show_y_ticks:
    #     plt.setp(axes.get_yticklabels(), visible=False)
    if data_name != "":
        axes.set_ylabel(data_name)
    else:
        axes.set_ylabel("")
    if model_name != "":
        axes.set_title(model_name)
    return axes
//  [python]
def figure3(
        df: pd.DataFrame,
        output_path: str,
        log_y: bool = False,
        save_pdf: bool = True,
        name: str = "figure3",
        title: str = rf"ADE per number of visible points (w/ 95\% CI)",
        scale: float = 1.0,
) -> None:
    df = df.copy()
    fig, ax = plt.subplots(figsize=(6 * scale, 4 * scale))
    fig.suptitle(title)
    MetricLabelStyle = namedtuple("MetricLabelStle", ["metric", "label", "style"])
    metric_label_style_list = [MetricLabelStyle("ade", "ADE", "-"), MetricLabelStyle("ade_visible", "ADE Visible", ":")]
    for metric_label_style in metric_label_style_list:
        sns.tsplot(
            df,
            time="n_timesteps_visible",
            value=metric_label_style.metric,
            unit="name",
            palette=cycle(["GoldenRod", "r", "forestgreen", "yellow"]),
            linestyle=metric_label_style.style,
            linewidth=2,
            ci=95,
            markers=True,
            dashes=False,
            err_style="band",
            alpha=1,
            legend=metric_label_style.metric == "ade",
            ax=ax,
        )
    texts = [t.get_text() for t in ax.get_legend().get_texts()] + [mls.label for mls in metric_label_style_list]
    lines = ax.get_legend().get_lines() + [Line2D([0, 10], [0, 10], linewidth=2, color="black", linestyle=mls.style)
                                           for mls in metric_label_style_list]
    new_legend = plt.legend(lines, texts, loc="center left")
    ax.add_artist(new_legend)
    ax.set_xlabel(rf"number of visible points")
    ax.set_ylabel(rf"ADE")
    if log_y:
        plt.yscale('log')
    plt.tight_layout()
    plt.savefig(os.path.join(output_path, f"{name}.png"))
    if save_pdf:
        plt.savefig(os.path.join(output_path, f"PDF__{name}.pdf"))
    plt.show()
    plt.close()
    plt.clf()
//  [python]
def figure3(
        df: pd.DataFrame,
        output_path: str,
        log_y: bool = False,
        save_pdf: bool = True,
        name: str = "figure3",
        title: str = rf"ADE per number of visible points (w/ 95\% CI)",
        scale: float = 1.0,
) -> None:
    df = df.copy()
    fig, ax = plt.subplots(figsize=(6 * scale, 4 * scale))
    fig.suptitle(title)
    MetricLabelStyle = namedtuple("MetricLabelStle", ["metric", "label", "style"])
    metric_label_style_list = [MetricLabelStyle("ade", "ADE", "-"), MetricLabelStyle("ade_visible", "ADE Visible", ":")]
    for metric_label_style in metric_label_style_list:
        sns.lineplot(
            df,
            x="n_timesteps_visible",
            y=metric_label_style.metric,
            hue="name",
            palette=cycle(["GoldenRod", "r", "forestgreen", "yellow"]),
            linestyle=metric_label_style.style,
            linewidth=2,
            errorbar=("ci", 95),
            markers=True,
            dashes=False,
            err_style="band",
            alpha=1,
            legend=metric_label_style.metric == "ade",
            ax=ax,
        )
    texts = [t.get_text() for t in ax.get_legend().get_texts()] + [mls.label for mls in metric_label_style_list]
    lines = ax.get_legend().get_lines() + [Line2D([0, 10], [0, 10], linewidth=2, color="black", linestyle=mls.style)
                                           for mls in metric_label_style_list]
    new_legend = plt.legend(lines, texts, loc="center left")
    ax.add_artist(new_legend)
    ax.set_xlabel(rf"number of visible points")
    ax.set_ylabel(rf"ADE")
    if log_y:
        plt.yscale('log')
    plt.tight_layout()
    plt.savefig(os.path.join(output_path, f"{name}.png"))
    if save_pdf:
        plt.savefig(os.path.join(output_path, f"PDF__{name}.pdf"))
    plt.show()
    plt.close()
    plt.clf()
//  [python]
def figure2(
        df: pd.DataFrame,
        output_path: str,
        ade_metric: str = "ade_visible",
        log_y: bool = False,
        save_pdf: bool = True,
        name: str = "figure2",
        title: str = rf"ADE for mostly visible (w/ 95\% CI)"
) -> None:
    df = df.copy()
    df_list = []
    for mostly_visible_threshold in range(1, int(df.n_timesteps_visible.max()) + 1):
        df_ = df.copy()
        df_["mostly_visible_threshold"] = mostly_visible_threshold
        df_["mostly_visible"] = (df_.n_timesteps_visible >= mostly_visible_threshold).apply(
            lambda x: "Mostly Visible" if x else "Mostly Occluded")
        df_list += [df_]

        df_ = df.copy()
        df_["mostly_visible_threshold"] = mostly_visible_threshold
        df_["mostly_visible"] = "All"
        df_list += [df_]

    df = pd.concat(df_list).reset_index()

    fig, ax = plt.subplots(figsize=(8, 5))
    fig.suptitle(title)
    mostly_visible_types = ["Mostly Visible", "Mostly Occluded", "All"]
    mostly_visible_linestlye = ["--", ":", "-"]
    for mostly_visible, linestyle in zip(mostly_visible_types, mostly_visible_linestlye, strict=True):
        sns.pointplot(
            data=df[df.mostly_visible == mostly_visible],
            x="mostly_visible_threshold",
            y=ade_metric,
            hue="name",
            palette=cycle(["GoldenRod", "r", "forestgreen", "yellow"]),
            linestyles=linestyle,
            linewidth=2,
            ci=95,
            markers=True,
            dashes=False,
            join=True,
            alpha=1,
            legend=mostly_visible == "All",
            ax=ax,
        )
    texts = [t.get_text() for t in ax.get_legend().get_texts()] + mostly_visible_types
    lines = ax.get_legend().get_lines() + [Line2D([0, 10], [0, 10], linewidth=2, color="black", linestyle=linestyle)
                                           for linestyle in mostly_visible_linestlye]
    new_legend = plt.legend(lines, texts, loc="center left")
    ax.add_artist(new_legend)
    ax.set_xlabel(rf"mostly visible threshold")
    ax.set_ylabel(rf"ADE")
    if log_y:
        plt.yscale('log')
    plt.tight_layout()
    plt.savefig(os.path.join(output_path, f"{name}_{ade_metric}.png"))
    if save_pdf:
        plt.savefig(os.path.join(output_path, f"PDF__{name}_{ade_metric}.pdf"))
    plt.show()
    plt.close()
    plt.clf()
//  [python]
def figure2(
        df: pd.DataFrame,
        output_path: str,
        ade_metric: str = "ade_visible",
        log_y: bool = False,
        save_pdf: bool = True,
        name: str = "figure2",
        title: str = rf"ADE for mostly visible (w/ 95\% CI)"
) -> None:
    df = df.copy()
    df_list = []
    for mostly_visible_threshold in range(1, int(df.n_timesteps_visible.max()) + 1):
        df_ = df.copy()
        df_["mostly_visible_threshold"] = mostly_visible_threshold
        df_["mostly_visible"] = (df_.n_timesteps_visible >= mostly_visible_threshold).apply(
            lambda x: "Mostly Visible" if x else "Mostly Occluded")
        df_list += [df_]

        df_ = df.copy()
        df_["mostly_visible_threshold"] = mostly_visible_threshold
        df_["mostly_visible"] = "All"
        df_list += [df_]

    df = pd.concat(df_list).reset_index()

    fig, ax = plt.subplots(figsize=(8, 5))
    fig.suptitle(title)
    mostly_visible_types = ["Mostly Visible", "Mostly Occluded", "All"]
    mostly_visible_linestlye = ["--", ":", "-"]
    for mostly_visible, linestyle in zip(mostly_visible_types, mostly_visible_linestlye, strict=True):
        sns.lineplot(
            df[df.mostly_visible == mostly_visible],
            x="mostly_visible_threshold",
            y=ade_metric,
            hue="name",
            palette=cycle(["GoldenRod", "r", "forestgreen", "yellow"]),
            linestyle=linestyle,
            linewidth=2,
            errorbar=("ci", 95),
            markers=True,
            dashes=False,
            err_style="band",
            alpha=1,
            legend=mostly_visible == "All",
            ax=ax,
        )
    texts = [t.get_text() for t in ax.get_legend().get_texts()] + mostly_visible_types
    lines = ax.get_legend().get_lines() + [Line2D([0, 10], [0, 10], linewidth=2, color="black", linestyle=linestyle)
                                           for linestyle in mostly_visible_linestlye]
    new_legend = plt.legend(lines, texts, loc="center left")
    ax.add_artist(new_legend)
    ax.set_xlabel(rf"mostly visible threshold")
    ax.set_ylabel(rf"ADE")
    if log_y:
        plt.yscale('log')
    plt.tight_layout()
    plt.savefig(os.path.join(output_path, f"{name}_{ade_metric}.png"))
    if save_pdf:
        plt.savefig(os.path.join(output_path, f"PDF__{name}_{ade_metric}.pdf"))
    plt.show()
    plt.close()
    plt.clf()
//  [python]
def plot_results(results, original_df, model_name):
    """Plots predictions over original data to visualize results. Saves each
    plot as a png.

    Keyword arguments:
    -- results: a dataframe with unscaled predictions
    -- original_df: the original monthly sales dataframe
    -- model_name: the name that will be used in the plot title
    """
    fig, ax = plt.subplots(figsize=(15, 5))
    plt.plot(original_df.date, original_df.sales, label='Original', color='mediumblue')
    plt.plot(results.date, results.pred_value, label='Predicted', color='red')
    ax.set(xlabel="Date",
           ylabel="Sales",
           title=f"{model_name} Sales Forecasting Prediction")
    ax.legend()
    sns.despine()

    plt.savefig(f'../model_output/{model_name}_forecast.png')
//  [python]
def plot_results(results, original_df, model_name):
    """Plots predictions over original data to visualize results. Saves each
    plot as a png.

    Keyword arguments:
    -- results: a dataframe with unscaled predictions
    -- original_df: the original monthly sales dataframe
    -- model_name: the name that will be used in the plot title
    """
    fig, ax = plt.subplots(figsize=(15, 5))
    sns.lineplot(original_df.date, original_df.sales, data=original_df, ax=ax,
                 label='Original', color='mediumblue')
    sns.lineplot(results.date, results.pred_value, data=results, ax=ax,
                 label='Predicted', color='red')
    ax.set(xlabel="Date",
           ylabel="Sales",
           title=f"{model_name} Sales Forecasting Prediction")
    ax.legend()
    sns.despine()

    plt.savefig(f'../model_output/{model_name}_forecast.png')
//  [python]
def positional_gia_plot(
    sdata,
    vars: list,
    id_var: str = "id",
    xlab: str = "Position",
    ylab: str = "Predicted Score",
    ylim: Optional[tuple] = None,
    save: Optional[PathLike] = None,
    return_axes: bool = False,
):
    """Plot a lineplot for each position of the sequence after implanting a feature.

    Assumes that the value corresponding to each seqsm_var in the sdata.uns dictionary
    has the same shape, namely (L, ) where L are the positions where a feature was implanted
    and scores were calculated using a model. Plots the scores as a line plot with a 95% CI
    corresponding to the number of sequences used to make the plot.

    Parameters
    ----------
    sdata : SeqData
        The SeqData object with sequences and scores to plot
    id_var : str
        The name of the variable in sdata.obs to use as the x-axis
    xlab : str
        The x-axis label
    ylab : str
        The y-axis label
    save : PathLike
        The path to save the figure to
    return_axes : bool
        Whether to return the matplotlib axes object
    """
    concat_df = pd.DataFrame()
    for key in vars:
        df = pd.DataFrame(index=sdata[id_var].values, data=sdata[key].values).melt(
            var_name=xlab, value_name=ylab, ignore_index=False
        )
        df["feature"] = key
        concat_df = pd.concat([concat_df, df])
    concat_df.reset_index(drop=True, inplace=True)
    g = sns.pointplot(data=concat_df, x=xlab, y=ylab, hue="feature")
    if ylim is not None:
        g.set(ylim=ylim)
    if save:
        _save_fig(save)
    if return_axes:
        return g
//  [python]
def positional_gia_plot(
    sdata,
    vars: list,
    id_var: str = "id",
    xlab: str = "Position",
    ylab: str = "Predicted Score",
    ylim: Optional[tuple] = None,
    save: Optional[PathLike] = None,
    return_axes: bool = False,
):
    """Plot a lineplot for each position of the sequence after implanting a feature.

    Assumes that the value corresponding to each seqsm_var in the sdata.uns dictionary
    has the same shape, namely (L, ) where L are the positions where a feature was implanted
    and scores were calculated using a model. Plots the scores as a line plot with a 95% CI
    corresponding to the number of sequences used to make the plot.

    Parameters
    ----------
    sdata : SeqData
        The SeqData object with sequences and scores to plot
    id_var : str
        The name of the variable in sdata.obs to use as the x-axis
    xlab : str
        The x-axis label
    ylab : str
        The y-axis label
    save : PathLike
        The path to save the figure to
    return_axes : bool
        Whether to return the matplotlib axes object
    """
    concat_df = pd.DataFrame()
    for key in vars:
        df = pd.DataFrame(index=sdata[id_var].values, data=sdata[key].values).melt(
            var_name=xlab, value_name=ylab, ignore_index=False
        )
        df["feature"] = key
        concat_df = pd.concat([concat_df, df])
    concat_df.reset_index(drop=True, inplace=True)
    g = sns.lineplot(data=concat_df, x=xlab, y=ylab, hue="feature")
    if ylim is not None:
        g.set(ylim=ylim)
    if save:
        _save_fig(save)
    if return_axes:
        return g
//  [python]
def plot(df: pd.DataFrame):
    base_values = {"batch_size": 1024, "max_iter_e_step": 1, "max_no_improvement": 5}
    for variable in base_values.keys():
        df_to_plot = df
        for other_variable in base_values.keys():
            if other_variable != variable:
                df_to_plot = df_to_plot[
                    df_to_plot[other_variable] == base_values[other_variable]
                ]
        # 2 subplots with a shared legend
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
        sns.pointplot(
            x="max_rows",
            y="score_train",
            hue="high_card_feature",
            style=variable,
            data=df_to_plot,
            ax=ax1,
            alpha=0.5,
        )
        ax1.get_legend().remove()
        sns.pointplot(
            x="max_rows",
            y="time_fit",
            hue="high_card_feature",
            style=variable,
            data=df_to_plot,
            ax=ax2,
            alpha=0.5,
        )
        # log scale
        ax1.set_yscale("log")
        ax2.set_yscale("log")
        fig.suptitle(f"Effect of {variable} on score_train and time_fit")
        # legend outside
        plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0)
        # tight layout
        plt.tight_layout()
        plt.show()
//  [python]
def plot(df: pd.DataFrame):
    base_values = {"batch_size": 1024, "max_iter_e_step": 1, "max_no_improvement": 5}
    for variable in base_values.keys():
        df_to_plot = df
        for other_variable in base_values.keys():
            if other_variable != variable:
                df_to_plot = df_to_plot[
                    df_to_plot[other_variable] == base_values[other_variable]
                ]
        # 2 subplots with a shared legend
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
        sns.lineplot(
            x="max_rows",
            y="score_train",
            hue="high_card_feature",
            style=variable,
            data=df_to_plot,
            ax=ax1,
            alpha=0.5,
        )
        ax1.get_legend().remove()
        sns.lineplot(
            x="max_rows",
            y="time_fit",
            hue="high_card_feature",
            style=variable,
            data=df_to_plot,
            ax=ax2,
            alpha=0.5,
        )
        # log scale
        ax1.set_yscale("log")
        ax2.set_yscale("log")
        fig.suptitle(f"Effect of {variable} on score_train and time_fit")
        # legend outside
        plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0)
        # tight layout
        plt.tight_layout()
        plt.show()
//  [python]
def plot(self, on_ts: bool = True) -> None:
        if on_ts:
            title = "Anomaly detection (plot on top of the time series)"
        else:
            title = "Anomaly detection (plot on top of the anomaly scores)"

        sns.set(rc={"figure.figsize": (10, 5)})
        sns.set_style("whitegrid")
        if on_ts:
            ax = sns.pointplot(x=self.ts_data_df.index, y=self.ts_data_df.values, color="midnightblue")
        else:
            ax = sns.pointplot(
                x=self.anomaly_scores_df["time"],
                y=self.anomaly_scores_df.drop(columns="time").values,
                color="midnightblue",
            )

        first = True
        if self.anomaly_interval_list is not None:
            # pyre-ignore[16]
            for ia in self.anomaly_interval_list:
                if first:
                    ax.axvspan(
                        ia.start,
                        ia.end,
                        alpha=0.98,
                        color="pink",
                        label="Anomaly",
                    )
                    first = False
                else:
                    ax.axvspan(
                        ia.start,
                        ia.end,
                        alpha=0.98,
                        color="pink",
                    )

        ax.legend(loc=0)
        ax.set_title(title)
        plt.xticks(rotation=90)
        plt.show()

        return
//  [python]
def plot(self, on_ts: bool = True) -> None:
        if on_ts:
            title = "Anomaly detection (plot on top of the time series)"
        else:
            title = "Anomaly detection (plot on top of the anomaly scores)"

        sns.set(rc={"figure.figsize": (10, 5)})
        sns.set_style("whitegrid")
        if on_ts:
            ax = sns.lineplot(self.ts_data_df, color="midnightblue", marker="o")
        else:
            ax = sns.lineplot(
                self.anomaly_scores_df.set_index("time"),
                color="midnightblue",
                marker="o",
            )

        first = True
        if self.anomaly_interval_list is not None:
            # pyre-ignore[16]
            for ia in self.anomaly_interval_list:
                if first:
                    ax.axvspan(
                        ia.start,
                        ia.end,
                        alpha=0.98,
                        color="pink",
                        label="Anomaly",
                    )
                    first = False
                else:
                    ax.axvspan(
                        ia.start,
                        ia.end,
                        alpha=0.98,
                        color="pink",
                    )

        ax.legend(loc=0)
        ax.set_title(title)
        plt.xticks(rotation=90)
        plt.show()

        return
//  [python]
def plot_frontier(frontier, all_points, writer=None, epoch=0, title="", beta=None, mixture=None, figsize=(8, 6),
                  verbose=True, clamp_qc=None):
    """
        Plot the hull of all Qc, Qr points for different (action, budget).

        If a threshold beta and corresponding mixture is provided, plot them.
    :param frontier: points of the Pareto frontier
    :param all_points: all points (Qc, Qr)
    :param SummaryWriter writer: will log the image to tensorboard if not None
    :param epoch: timestep for tensorboard log
    :param title: figure title
    :param beta: a budget threshold used at decision time
    :param mixture: the optimal mixture corresponding to this budget beta
    :param figsize: figure size, inches
    :param verbose: should the legend be displayed
    :param clamp_qc: if qc is clamped, use these values at x axis limits
    :return: the string description of the image, and its size
    """
    # Figure creation
    dfa, dfh = pd.DataFrame(all_points), pd.DataFrame(frontier)
    fig = plt.figure(figsize=figsize, tight_layout=True)
    sns.scatterplot(data=dfa, x="qc", y="qr", hue="action", legend="full")
    plt.plot(dfh["qc"], dfh["qr"], marker="x", label="hull")
    if clamp_qc:  # known limits
        plt.xlim(clamp_qc[0]-0.1, clamp_qc[1]+0.1)
    if beta is not None:
        plt.axvline(x=beta)
    if mixture:
        plt.plot([mixture.inf.qc, mixture.sup.qc], [mixture.inf.qr, mixture.sup.qr],
                     color="red", marker="o")
    plt.title(title)
    leg = plt.legend(loc='upper right')
    if not verbose:
        leg.remove()
        plt.xlabel('')
        plt.ylabel('')

    # Figure export
    fig.canvas.draw()
    data_str = fig.canvas.tostring_rgb()
    if writer:
        data = np.fromstring(data_str, dtype=np.uint8, sep='')
        data = np.rollaxis(data.reshape(fig.canvas.get_width_height()[::-1] + (3,)), 2, 0)
        writer.add_image(clean_tag(title), data, epoch)
    plt.close()
    return data_str, fig.canvas.get_width_height()
//  [python]
def plot_frontier(frontier, all_points, writer=None, epoch=0, title="", beta=None, mixture=None, figsize=(8, 6),
                  verbose=True, clamp_qc=None):
    """
        Plot the hull of all Qc, Qr points for different (action, budget).

        If a threshold beta and corresponding mixture is provided, plot them.
    :param frontier: points of the Pareto frontier
    :param all_points: all points (Qc, Qr)
    :param SummaryWriter writer: will log the image to tensorboard if not None
    :param epoch: timestep for tensorboard log
    :param title: figure title
    :param beta: a budget threshold used at decision time
    :param mixture: the optimal mixture corresponding to this budget beta
    :param figsize: figure size, inches
    :param verbose: should the legend be displayed
    :param clamp_qc: if qc is clamped, use these values at x axis limits
    :return: the string description of the image, and its size
    """
    # Figure creation
    dfa, dfh = pd.DataFrame(all_points), pd.DataFrame(frontier)
    fig = plt.figure(figsize=figsize, tight_layout=True)
    sns.scatterplot(data=dfa, x="qc", y="qr", hue="action", legend="full")
    sns.lineplot(data=dfh, x="qc", y="qr", marker="x", label="hull")
    if clamp_qc:  # known limits
        plt.xlim(clamp_qc[0]-0.1, clamp_qc[1]+0.1)
    if beta is not None:
        plt.axvline(x=beta)
    if mixture:
        sns.lineplot(x=[mixture.inf.qc, mixture.sup.qc], y=[mixture.inf.qr, mixture.sup.qr],
                     color="red", marker="o")
    plt.title(title)
    leg = plt.legend(loc='upper right')
    if not verbose:
        leg.remove()
        plt.xlabel('')
        plt.ylabel('')

    # Figure export
    fig.canvas.draw()
    data_str = fig.canvas.tostring_rgb()
    if writer:
        data = np.fromstring(data_str, dtype=np.uint8, sep='')
        data = np.rollaxis(data.reshape(fig.canvas.get_width_height()[::-1] + (3,)), 2, 0)
        writer.add_image(clean_tag(title), data, epoch)
    plt.close()
    return data_str, fig.canvas.get_width_height()
//  [python]
def plot_mean_std(self, ax=None, fname=None) -> None:
        """
        Plot the means and standard deviations of each dataset.

        Args:
            ax: Axis to plot on. If none, a new figure is made.
            fname: If not none, saves the plot with this file name.
        """
        real = self.real
        fake = self.fake
        if ax is None:
            fig, ax = plt.subplots(1, 2, figsize=(10, 5))
            fig.suptitle(
                "Absolute Log Mean and STDs of numeric data\n", fontsize=16
            )

        ax[0].grid(True)
        ax[1].grid(True)
        real = real.select_dtypes(include=np.number).reset_index()
        fake = fake.select_dtypes(include=np.number).reset_index()
        real_mean = np.log(np.add(abs(real.mean()).values, 1e-5))
        fake_mean = np.log(np.add(abs(fake.mean()).values, 1e-5))
        min_mean = min(real_mean) - 1
        max_mean = max(real_mean) + 1
        line = np.arange(min_mean, max_mean)
        plt.plot(line, line, ax=ax[0])
        sns.scatterplot(x=real_mean, y=fake_mean, ax=ax[0])
        ax[0].set_title("Means of real and fake data")
        ax[0].set_xlabel("real data mean (log)")
        ax[0].set_ylabel("fake data mean (log)")

        real_std = np.log(np.add(real.std().values, 1e-5))
        fake_std = np.log(np.add(fake.std().values, 1e-5))
        min_std = min(real_std) - 1
        max_std = max(real_std) + 1
        line = np.arange(min_std, max_std)
        plt.plot(line, line, ax=ax[1])
        sns.scatterplot(x=real_std, y=fake_std, ax=ax[1])
        ax[1].set_title("Stds of real and fake data")
        ax[1].set_xlabel("real data std (log)")
        ax[1].set_ylabel("fake data std (log)")

        if fname is not None:
            plt.savefig(fname)

        if ax is None:
            plt.show()
//  [python]
def plot_mean_std(self, ax=None, fname=None) -> None:
        """
        Plot the means and standard deviations of each dataset.

        Args:
            ax: Axis to plot on. If none, a new figure is made.
            fname: If not none, saves the plot with this file name.
        """
        real = self.real
        fake = self.fake
        if ax is None:
            fig, ax = plt.subplots(1, 2, figsize=(10, 5))
            fig.suptitle(
                "Absolute Log Mean and STDs of numeric data\n", fontsize=16
            )

        ax[0].grid(True)
        ax[1].grid(True)
        real = real.select_dtypes(include=np.number).reset_index()
        fake = fake.select_dtypes(include=np.number).reset_index()
        real_mean = np.log(np.add(abs(real.mean()).values, 1e-5))
        fake_mean = np.log(np.add(abs(fake.mean()).values, 1e-5))
        min_mean = min(real_mean) - 1
        max_mean = max(real_mean) + 1
        line = np.arange(min_mean, max_mean)
        sns.lineplot(x=line, y=line, ax=ax[0])
        sns.scatterplot(x=real_mean, y=fake_mean, ax=ax[0])
        ax[0].set_title("Means of real and fake data")
        ax[0].set_xlabel("real data mean (log)")
        ax[0].set_ylabel("fake data mean (log)")

        real_std = np.log(np.add(real.std().values, 1e-5))
        fake_std = np.log(np.add(fake.std().values, 1e-5))
        min_std = min(real_std) - 1
        max_std = max(real_std) + 1
        line = np.arange(min_std, max_std)
        sns.lineplot(x=line, y=line, ax=ax[1])
        sns.scatterplot(x=real_std, y=fake_std, ax=ax[1])
        ax[1].set_title("Stds of real and fake data")
        ax[1].set_xlabel("real data std (log)")
        ax[1].set_ylabel("fake data std (log)")

        if fname is not None:
            plt.savefig(fname)

        if ax is None:
            plt.show()
//  [python]
def exp_plot():
    x = np.linspace(1, 100, num=10)
    y = np.ones(len(x))
    plt.plot(x, y)
    plt.savefig("./paper_plots/fig1.eps")
//  [python]
def exp_plot():
    x = np.linspace(1, 100, num=10)
    y = np.ones(len(x))
    sns.lineplot(x=x, y=y)
    plt.savefig("./paper_plots/fig1.eps")