jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def plot_entropy_vs_position(timed_pred_folder, models, pdb_code):
    all_results = []
    all_models = []
    for model in models:
        model_pred_path = timed_pred_folder / (model + ".csv")
        model_map_path = timed_pred_folder / (model + ".txt")
        pdb_to_entropy = extract_prediction_entropy_to_dict(model_pred_path, model_map_path, rotamer_mode=True if "rot" in model else False)
        pdb_code_and_chain = 0
        # Go through keys and identify the correct chain:
        for pc in pdb_to_entropy.keys():
            if pdb_code in pc:
                pdb_code_and_chain = pc
                # End loop once found:
                break
        assert pdb_code_and_chain != 0, f"PDB Code {pdb_code} not found in {list(pdb_to_entropy.keys())}"
        curr_entropy = pdb_to_entropy[pdb_code_and_chain]
        all_results.append(curr_entropy)
        all_models.append(model)

    all_results = pd.DataFrame(np.array(all_results).T, columns=all_models)
    # TODO Fix:
    fig, axs = plt.subplots(1, figsize=(10, 4.8))
    g = sns.tsplot(data=all_results, ax=axs)
    g.set(ylim=(0, 5))

    axs.yaxis.get_major_formatter().set_scientific(False)
    axs.xaxis.get_major_formatter().set_scientific(False)
    plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0)
    fig.tight_layout()
    plt.show()
    return
//  [python]
def plot_entropy_vs_position(timed_pred_folder, models, pdb_code):
    all_results = []
    all_models = []
    for model in models:
        model_pred_path = timed_pred_folder / (model + ".csv")
        model_map_path = timed_pred_folder / (model + ".txt")
        pdb_to_entropy = extract_prediction_entropy_to_dict(model_pred_path, model_map_path, rotamer_mode=True if "rot" in model else False)
        pdb_code_and_chain = 0
        # Go through keys and identify the correct chain:
        for pc in pdb_to_entropy.keys():
            if pdb_code in pc:
                pdb_code_and_chain = pc
                # End loop once found:
                break
        assert pdb_code_and_chain != 0, f"PDB Code {pdb_code} not found in {list(pdb_to_entropy.keys())}"
        curr_entropy = pdb_to_entropy[pdb_code_and_chain]
        all_results.append(curr_entropy)
        all_models.append(model)

    all_results = pd.DataFrame(np.array(all_results).T, columns=all_models)
    # TODO Fix:
    fig, axs = plt.subplots(1, figsize=(10, 4.8))
    g = sns.lineplot(data=all_results, ax=axs)
    g.set(ylim=(0, 5))

    axs.yaxis.get_major_formatter().set_scientific(False)
    axs.xaxis.get_major_formatter().set_scientific(False)
    plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0)
    fig.tight_layout()
    plt.show()
    return
//  [python]
def plot_effective_samples_per_rollout(self, tree=False):
        plt.figure()
        df = self.df.loc[self.df['Sampling Strategy'] == tree]
        nr_samples = len(df["nr_rollouts"].unique())
        sns.tsplot(data=df,
                     time="index",
                     value="effective_samples",
                     unit="nr_rollouts",
                     ci="sd",
                     palette=sns.color_palette("tab10", n_colors=nr_samples))
//  [python]
def plot_effective_samples_per_rollout(self, tree=False):
        plt.figure()
        df = self.df.loc[self.df['Sampling Strategy'] == tree]
        nr_samples = len(df["nr_rollouts"].unique())
        sns.lineplot(data=df,
                     x="index",
                     y="effective_samples",
                     hue="nr_rollouts",
                     ci="sd",
                     palette=sns.color_palette("tab10", n_colors=nr_samples))
//  [python]
def plot_cost(self, aggregator):
        fig, ax = plt.subplots()
        sns.lineplot(data=self.df,
                     x="index",
                     y="stage_cost",
                     hue=aggregator,
                     legend=True,
                     ci="sd")
//  [python]
def plot_average_cost_samples_comparison(self):
        different_samples = len(self.df["nr_rollouts"].unique())
        ax = sns.tsplot(data=self.df,
                          x="index",
                          y="stage_cost",
                          unit="nr_rollouts",
                          legend=True,
                          ci="sd",
                          err_style="band",
                          palette=sns.color_palette(
                              "tab10", n_colors=different_samples))
        df = self.df.loc[self.df['index'] < 300]
        #ax = sns.barplot(x="nr_rollouts", y="stage_cost", data=df, capsize=.2)
        #ax.set_title("Whole Body Differential Base", fontsize=20)
        #ax.set_xticklabels([5, 10, 50, 100, 200, 500, 1000, 2000], size = 30)
        #ax.set
        ax.set_yticklabels(ax.get_yticks(), size=30)
        ax.set_xlabel("samples", fontsize=40)
        #ax.set_yscale("log")
        ax.set_ylabel("average stage cost", fontsize=40)
//  [python]
def plot_average_cost_samples_comparison(self):
        different_samples = len(self.df["nr_rollouts"].unique())
        ax = sns.lineplot(data=self.df,
                          x="index",
                          y="stage_cost",
                          hue="nr_rollouts",
                          legend=True,
                          ci="sd",
                          palette=sns.color_palette(
                              "tab10", n_colors=different_samples))
        df = self.df.loc[self.df['index'] < 300]
        #ax = sns.barplot(x="nr_rollouts", y="stage_cost", data=df, capsize=.2)
        #ax.set_title("Whole Body Differential Base", fontsize=20)
        #ax.set_xticklabels([5, 10, 50, 100, 200, 500, 1000, 2000], size = 30)
        #ax.set
        ax.set_yticklabels(ax.get_yticks(), size=30)
        ax.set_xlabel("samples", fontsize=40)
        #ax.set_yscale("log")
        ax.set_ylabel("average stage cost", fontsize=40)
//  [python]
def plot_cost_comparison(self):

        df_no_tree = self.df.loc[self.df['Sampling Strategy'] == "Monte Carlo"]
        sample_sizes_no_tree = df_no_tree["nr_rollouts"].unique()

        df_tree = self.df.loc[self.df['Sampling Strategy'] == "Tree"]
        sample_sizes_tree = df_tree["nr_rollouts"].unique()

        common_sample_sizes = set(sample_sizes_tree).intersection(
            set(sample_sizes_no_tree))
        time_padding = 0.1
        for samples_size in common_sample_sizes:
            fig, axs = plt.subplots(nrows=1, ncols=5, sharey=True)
            for i in range(5):
                df_ = self.df.loc[(self.df['nr_rollouts'] == samples_size)
                                  & (self.df['time'] < (i + 1) * 5) &
                                  (self.df['time'] > (i * 5 + time_padding))]
                df_['stage_cost'] = df_['stage_cost'].apply(
                    lambda x: np.exp(-x))
                sns.pointplot(data=df_,
                             x="time",
                             y="stage_cost",
                             hue="Sampling Strategy",
                             legend=True,
                             ci="sd",
                             palette=sns.color_palette("tab10", n_colors=2),
                             ax=axs[i])
                axs[i].set_title("Target {}".format(i + 1), fontsize=20)
                axs[i].set_xlabel("t [s]", fontsize=20)
                axs[i].set_ylabel("stage cost", fontsize=20)
                axs[i].set_ybound(lower=-0.01, upper=1.2)
                plt.legend(fontsize=18)
//  [python]
def plot_cost_comparison(self):

        df_no_tree = self.df.loc[self.df['Sampling Strategy'] == "Monte Carlo"]
        sample_sizes_no_tree = df_no_tree["nr_rollouts"].unique()

        df_tree = self.df.loc[self.df['Sampling Strategy'] == "Tree"]
        sample_sizes_tree = df_tree["nr_rollouts"].unique()

        common_sample_sizes = set(sample_sizes_tree).intersection(
            set(sample_sizes_no_tree))
        time_padding = 0.1
        for samples_size in common_sample_sizes:
            fig, axs = plt.subplots(nrows=1, ncols=5, sharey=True)
            for i in range(5):
                df_ = self.df.loc[(self.df['nr_rollouts'] == samples_size)
                                  & (self.df['time'] < (i + 1) * 5) &
                                  (self.df['time'] > (i * 5 + time_padding))]
                df_['stage_cost'] = df_['stage_cost'].apply(
                    lambda x: np.exp(-x))
                sns.lineplot(data=df_,
                             x="time",
                             y="stage_cost",
                             hue="Sampling Strategy",
                             legend=True,
                             ci="sd",
                             palette=sns.color_palette("tab10", n_colors=2),
                             ax=axs[i])
                axs[i].set_title("Target {}".format(i + 1), fontsize=20)
                axs[i].set_xlabel("t [s]", fontsize=20)
                axs[i].set_ylabel("stage cost", fontsize=20)
                axs[i].set_ybound(lower=-0.01, upper=1.2)
                plt.legend(fontsize=18)
//  [python]
def plot_anamoly_results(loss_df):
        """
        Anamoly Visualization 
        """
        import seaborn as sns
        import matplotlib.pyplot as plt
        plt.figure(figsize=(20,10))
        sns.set_style("darkgrid")
        ax = sns.distplot(loss_df["loss"], bins=100, label="Frequency")
        ax.set_title("Frequency Distribution | Kernel Density Estimation")
        ax.set(xlabel='Anomaly Confidence Score', ylabel='Frequency (sample)')
        plt.axvline(1.80, color="k", linestyle="--")
        plt.legend()


        plt.figure(figsize=(20,10))
        ax = sns.pointplot(x="timestamp", y="loss", data=loss_df, color='g', label="Anomaly Score")
        ax.set_title("Anomaly Confidence Score vs Timestamp")
        ax.set(ylabel="Anomaly Confidence Score", xlabel="Timestamp")
        plt.legend()
//  [python]
def plot_anamoly_results(loss_df):
        """
        Anamoly Visualization 
        """
        import seaborn as sns
        import matplotlib.pyplot as plt
        plt.figure(figsize=(20,10))
        sns.set_style("darkgrid")
        ax = sns.distplot(loss_df["loss"], bins=100, label="Frequency")
        ax.set_title("Frequency Distribution | Kernel Density Estimation")
        ax.set(xlabel='Anomaly Confidence Score', ylabel='Frequency (sample)')
        plt.axvline(1.80, color="k", linestyle="--")
        plt.legend()


        plt.figure(figsize=(20,10))
        ax = sns.lineplot(x="timestamp", y="loss", data=loss_df, color='g', label="Anomaly Score")
        ax.set_title("Anomaly Confidence Score vs Timestamp")
        ax.set(ylabel="Anomaly Confidence Score", xlabel="Timestamp")
        plt.legend()
//  [python]
def plot_momentum_comparison(self):
        plt.figure()

        df_no_tree = self.df.loc[self.df['Sampling Strategy'] == False]
        sample_sizes_no_tree = df_no_tree["nr_rollouts"].unique()

        for samples_size in sample_sizes_no_tree:
            plt.figure()
            df_ = df_no_tree.loc[df_no_tree['nr_rollouts'] == samples_size]
            ax = sns.tsplot(data=df_,
                              time="time",
                              value="effective_samples",
                              ci="sd",
                              condition="alpha",
                              unit="beta")
            ax.set_title("Effective Samples ({} samples)".format(samples_size))

        df_tree = self.df.loc[self.df['Sampling Strategy'] == True]
        sample_sizes_tree = df_tree["nr_rollouts"].unique()

        for samples_size in sample_sizes_tree:
            plt.figure()
            df_ = df_tree.loc[df_tree['nr_rollouts'] == samples_size]
            ax = sns.tsplot(data=df_,
                              time="time",
                              value="effective_samples",
                              ci="sd",
                              condition="alpha",
                              unit="beta")
            ax.set_title(
                "Effective Samples FD-MCTS ({} samples)".format(samples_size))
//  [python]
def plot_momentum_comparison(self):
        plt.figure()

        df_no_tree = self.df.loc[self.df['Sampling Strategy'] == False]
        sample_sizes_no_tree = df_no_tree["nr_rollouts"].unique()

        for samples_size in sample_sizes_no_tree:
            plt.figure()
            df_ = df_no_tree.loc[df_no_tree['nr_rollouts'] == samples_size]
            ax = sns.lineplot(data=df_,
                              x="time",
                              y="effective_samples",
                              ci="sd",
                              hue="alpha",
                              size="beta")
            ax.set_title("Effective Samples ({} samples)".format(samples_size))

        df_tree = self.df.loc[self.df['Sampling Strategy'] == True]
        sample_sizes_tree = df_tree["nr_rollouts"].unique()

        for samples_size in sample_sizes_tree:
            plt.figure()
            df_ = df_tree.loc[df_tree['nr_rollouts'] == samples_size]
            ax = sns.lineplot(data=df_,
                              x="time",
                              y="effective_samples",
                              ci="sd",
                              hue="alpha",
                              size="beta")
            ax.set_title(
                "Effective Samples FD-MCTS ({} samples)".format(samples_size))
//  [python]
def plot_cost_momentum_comparison_tree(self):

        df_tree = self.df.loc[self.df['Sampling Strategy'] == "Tree"]
        sample_sizes_tree = df_tree["nr_rollouts"].unique()

        df_tree['beta'] = df_tree['beta'].apply(lambda x: "with"
                                                if x > 0 else "without")
        df_tree = df_tree.rename(columns={'beta': 'momentum'})
        time_padding = 0.15
        for samples_size in sample_sizes_tree:
            fig, axs = plt.subplots(nrows=1, ncols=5, sharey=True)
            for i in range(5):
                df_ = df_tree.loc[(self.df['nr_rollouts'] == samples_size) & (
                    df_tree['time'] < (i + 1) * 5 - time_padding) &
                                  (df_tree['time'] > (i * 5 + time_padding))]
                #df_['stage_cost'] = df_['stage_cost'].apply(lambda x : np.exp(-x))
                sns.pointplot(data=df_,
                             x="time",
                             y="stage_cost",
                             hue="momentum",
                             legend=False,
                             ci="sd",
                             palette=sns.color_palette("tab10", n_colors=2),
                             ax=axs[i], join=True)
                axs[i].set_title("Target {}".format(i + 1), fontsize=20)
                axs[i].set_xlabel("t [s]", fontsize=20)
                #axs[i].set_yscale('log')
                axs[i].set_ylabel("stage cost", fontsize=20)
//  [python]
def plot_cost_momentum_comparison_tree(self):

        df_tree = self.df.loc[self.df['Sampling Strategy'] == "Tree"]
        sample_sizes_tree = df_tree["nr_rollouts"].unique()

        df_tree['beta'] = df_tree['beta'].apply(lambda x: "with"
                                                if x > 0 else "without")
        df_tree = df_tree.rename(columns={'beta': 'momentum'})
        time_padding = 0.15
        for samples_size in sample_sizes_tree:
            fig, axs = plt.subplots(nrows=1, ncols=5, sharey=True)
            for i in range(5):
                df_ = df_tree.loc[(self.df['nr_rollouts'] == samples_size) & (
                    df_tree['time'] < (i + 1) * 5 - time_padding) &
                                  (df_tree['time'] > (i * 5 + time_padding))]
                #df_['stage_cost'] = df_['stage_cost'].apply(lambda x : np.exp(-x))
                sns.lineplot(data=df_,
                             x="time",
                             y="stage_cost",
                             hue="momentum",
                             legend=False,
                             ci="sd",
                             palette=sns.color_palette("tab10", n_colors=2),
                             ax=axs[i])
                axs[i].set_title("Target {}".format(i + 1), fontsize=20)
                axs[i].set_xlabel("t [s]", fontsize=20)
                #axs[i].set_yscale('log')
                axs[i].set_ylabel("stage cost", fontsize=20)
//  [python]
def plot_merge_process(self, plot_size=3):
        """Plot the change of accuracy during merge"""
        if len(self.step_data) == 0:
            print("No merge step to plot")
            return
        plot_data = []
        for step, (cur_y, *_, score, _) in self.step_data.items():
            cur_n_cluster = len(set(cur_y)) - 1
            plot_data.append({"Step": step, "# of clusters": cur_n_cluster, "Score": score})
        plot_data = pd.DataFrame(plot_data)

        fig, axes = plt.subplots(
            figsize=(plot_size * 2, plot_size),
            ncols=2,
            dpi=300,
            constrained_layout=True,
        )
        ax = axes[0]
        plt.plot(plot_data["Step"], plot_data["# of clusters"], color="gray")
        ax = axes[1]
        plt.plot(plot_data["Step"], plot_data["Score"], color="gray")
        ax.axhline(y=self.target_accuracy, color="red", linewidth=0.8, linestyle="--")
        ax.set(ylim=(plot_data["Score"].min() - 0.03, min(1, self.target_accuracy + 0.01)))
        return
//  [python]
def plot_merge_process(self, plot_size=3):
        """Plot the change of accuracy during merge"""
        if len(self.step_data) == 0:
            print("No merge step to plot")
            return
        plot_data = []
        for step, (cur_y, *_, score, _) in self.step_data.items():
            cur_n_cluster = len(set(cur_y)) - 1
            plot_data.append({"Step": step, "# of clusters": cur_n_cluster, "Score": score})
        plot_data = pd.DataFrame(plot_data)

        fig, axes = plt.subplots(
            figsize=(plot_size * 2, plot_size),
            ncols=2,
            dpi=300,
            constrained_layout=True,
        )
        ax = axes[0]
        sns.lineplot(ax=ax, data=plot_data, x="Step", y="# of clusters", color="gray")
        ax = axes[1]
        sns.lineplot(ax=ax, data=plot_data, x="Step", y="Score", color="gray")
        ax.axhline(y=self.target_accuracy, color="red", linewidth=0.8, linestyle="--")
        ax.set(ylim=(plot_data["Score"].min() - 0.03, min(1, self.target_accuracy + 0.01)))
        return
//  [python]
def plot_steady_state(params_dict: Dict[str, LogSvParams] = SS_PDF_PARAMS,
                      title: str = None,
                      ax: plt.Subplot = None
                      ) -> None:

    sigma = np.linspace(1e-4, 4.0, 1000)
    qs = []
    for key, params in params_dict.items():
        qs.append(pd.Series(steady_state(sigma=sigma, params=params), index=sigma, name=key))
    ss_pdf = pd.concat(qs, axis=1)

    for key in ss_pdf.columns:
        ax.plot(ss_pdf.index, ss_pdf[key], label=key)
    ax.set_xlim(left=0.0)
    ax.set_ylim(bottom=0.0)
    ax.set_title(title, color='darkblue')
    ax.set_xlabel('$\sigma$', fontsize=12)
    yvar_format = '{:.2f}'
    ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda z, _: yvar_format.format(z)))
//  [python]
def plot_steady_state(params_dict: Dict[str, LogSvParams] = SS_PDF_PARAMS,
                      title: str = None,
                      ax: plt.Subplot = None
                      ) -> None:

    sigma = np.linspace(1e-4, 4.0, 1000)
    qs = []
    for key, params in params_dict.items():
        qs.append(pd.Series(steady_state(sigma=sigma, params=params), index=sigma, name=key))
    ss_pdf = pd.concat(qs, axis=1)

    sns.lineplot(data=ss_pdf, dashes=False, ax=ax)
    ax.set_xlim(left=0.0)
    ax.set_ylim(bottom=0.0)
    ax.set_title(title, color='darkblue')
    ax.set_xlabel('$\sigma$', fontsize=12)
    yvar_format = '{:.2f}'
    ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda z, _: yvar_format.format(z)))
//  [python]
def plot_experiment(self, project_id: str, experiment_id: str, x='step', y: Union[str, List[str]] = 'loss', **plot_kwargs):
        """
        Generates a line plot and uploads the result to `<WRANGL_S3_URL>/<WRANGL_S3_BUCKET>/<project_id>/<experiment_id>/<x>_vs_<y>.pdf`.
        You must have `seaborn` installed for this to work.

        Args:
            project_id: name of project.
            experiment_id: name of experiment.
            x: x axis to plot.
            y: y axis (or a list of axes) to plot.
            plot_kwargs: keyword arguments to `seaborn.lineplot`.
        """
        import seaborn as sns
        from matplotlib import pyplot as plt
        config, df = self.get_experiment(project_id=project_id, experiment_id=experiment_id)

        if not len(df):
            return

        sns.set(font_scale=1.5)

        if isinstance(y, str):
            y = [y]

        for yi in y:
            yi = yi.strip()
            fig, ax = plt.subplots(figsize=(10, 5))
            sns.tsplot(x=x, y=yi, data=df, ax=ax, **plot_kwargs)

            fname = '{}_vs_{}.pdf'.format(x.replace('/', '--'), yi.replace('/', '--'))
            with tempfile.TemporaryDirectory() as tempdir_path:
                tempdir = pathlib.Path(tempdir_path)
                ffig = tempdir.joinpath(fname)
                fig.savefig(ffig, bbox_inches='tight')
                self.upload_file(project_id, experiment_id, fname, ffig, content_type='application/pdf')
//  [python]
def plot_experiment(self, project_id: str, experiment_id: str, x='step', y: Union[str, List[str]] = 'loss', **plot_kwargs):
        """
        Generates a line plot and uploads the result to `<WRANGL_S3_URL>/<WRANGL_S3_BUCKET>/<project_id>/<experiment_id>/<x>_vs_<y>.pdf`.
        You must have `seaborn` installed for this to work.

        Args:
            project_id: name of project.
            experiment_id: name of experiment.
            x: x axis to plot.
            y: y axis (or a list of axes) to plot.
            plot_kwargs: keyword arguments to `seaborn.lineplot`.
        """
        import seaborn as sns
        from matplotlib import pyplot as plt
        config, df = self.get_experiment(project_id=project_id, experiment_id=experiment_id)

        if not len(df):
            return

        sns.set(font_scale=1.5)

        if isinstance(y, str):
            y = [y]

        for yi in y:
            yi = yi.strip()
            fig, ax = plt.subplots(figsize=(10, 5))
            sns.lineplot(x=x, y=yi, data=df, ax=ax, **plot_kwargs)

            fname = '{}_vs_{}.pdf'.format(x.replace('/', '--'), yi.replace('/', '--'))
            with tempfile.TemporaryDirectory() as tempdir_path:
                tempdir = pathlib.Path(tempdir_path)
                ffig = tempdir.joinpath(fname)
                fig.savefig(ffig, bbox_inches='tight')
                self.upload_file(project_id, experiment_id, fname, ffig, content_type='application/pdf')
//  [python]
def plot_ts(da, key):
    """Plot a times series for a given xarray dataarray.

    Parameters
    ----------
    da : xr.DataArray
        one-dimensional data array
    key : str
        parameter name / ylabel
    """
    p = sns.tsplot(data=da.to_pandas(), linewidth=2)
    p.set_xlabel('time')
    p.set_ylabel(key)
//  [python]
def plot_ts(da, key):
    """Plot a times series for a given xarray dataarray.

    Parameters
    ----------
    da : xr.DataArray
        one-dimensional data array
    key : str
        parameter name / ylabel
    """
    p = sns.lineplot(data=da.to_pandas(), linewidth=2)
    p.set_xlabel('time')
    p.set_ylabel(key)
//  [python]
def plot_steady_state_x(params_dict: Dict[str, LogSvParams] = SS_PDF_PARAMS,
                        title: str = None,
                        ax: plt.Subplot = None
                        ) -> None:

    sigma = np.linspace(1e-4, 5.0, 1000)
    x = np.linspace(-5.0, 5.0, 200)
    qs = []
    for key, params in params_dict.items():
        g_sigma = steady_state(sigma=sigma, params=params)
        x_pdf = integral_x_over_sigma(x_grid=x, sigma=sigma, g_sigma=g_sigma)
        qs.append(pd.Series(x_pdf, index=x, name=key))

    ss_pdf = pd.concat(qs, axis=1)

    if ax is None:
        with sns.axes_style("darkgrid"):
            fig, ax = plt.subplots(1, 1, figsize=(10, 4.0), tight_layout=True)
    for col in ss_pdf.columns:
        ax.plot(ss_pdf.index, ss_pdf[col])
    #ax.set_xlim(left=0.0)
    ax.set_ylim(bottom=0.0)
    ax.set_title(title, color='darkblue')
//  [python]
def plot_steady_state_x(params_dict: Dict[str, LogSvParams] = SS_PDF_PARAMS,
                        title: str = None,
                        ax: plt.Subplot = None
                        ) -> None:

    sigma = np.linspace(1e-4, 5.0, 1000)
    x = np.linspace(-5.0, 5.0, 200)
    qs = []
    for key, params in params_dict.items():
        g_sigma = steady_state(sigma=sigma, params=params)
        x_pdf = integral_x_over_sigma(x_grid=x, sigma=sigma, g_sigma=g_sigma)
        qs.append(pd.Series(x_pdf, index=x, name=key))

    ss_pdf = pd.concat(qs, axis=1)

    if ax is None:
        with sns.axes_style("darkgrid"):
            fig, ax = plt.subplots(1, 1, figsize=(10, 4.0), tight_layout=True)
    sns.lineplot(data=ss_pdf, ax=ax)
    #ax.set_xlim(left=0.0)
    ax.set_ylim(bottom=0.0)
    ax.set_title(title, color='darkblue')
//  [python]
def main():
    plt.figure()
    num_metric = len(metrics)
    i = 1

    for metric in metrics.keys():
        # clear data for each metric
        data = {"Site": [], "Epoch": [], "Metric": []}
        for site in sites_fl:
            record_path = os.path.join(client_results_root, "simulate_job", client_pre + site, "events.*")
            eventfile = glob.glob(record_path, recursive=True)
            assert len(eventfile) == 1, "No unique event file found!"
            eventfile = eventfile[0]
            print("adding", eventfile)
            add_eventdata(data, site, eventfile, tag=metrics[metric])

        ax = plt.subplot(1, int(num_metric), i)
        ax.set_title(metric)
        plt.plot( 'Epoch', 'Metric', data=data, linestyle='-', marker='o')
        if metric == "selector_model":
            plt.ylabel("Accuracy (%)")
        else:
            plt.ylabel("Dice")
        i = i + 1
    plt.subplots_adjust(hspace=0.3)
    plt.show()
//  [python]
def main():
    plt.figure()
    num_metric = len(metrics)
    i = 1

    for metric in metrics.keys():
        # clear data for each metric
        data = {"Site": [], "Epoch": [], "Metric": []}
        for site in sites_fl:
            record_path = os.path.join(client_results_root, "simulate_job", client_pre + site, "events.*")
            eventfile = glob.glob(record_path, recursive=True)
            assert len(eventfile) == 1, "No unique event file found!"
            eventfile = eventfile[0]
            print("adding", eventfile)
            add_eventdata(data, site, eventfile, tag=metrics[metric])

        ax = plt.subplot(1, int(num_metric), i)
        ax.set_title(metric)
        sns.lineplot(x="Epoch", y="Metric", hue="Site", data=data)
        if metric == "selector_model":
            plt.ylabel("Accuracy (%)")
        else:
            plt.ylabel("Dice")
        i = i + 1
    plt.subplots_adjust(hspace=0.3)
    plt.show()
//  [python]
def plot_ss_kurtosis(params_dict=SS_PARAMS,
                     title: str = f'Excess kurtosis of log-returns as function of $\kappa_{2}$',
                     ax: plt.Subplot = None
                     ) -> None:

    kappa2s = np.linspace(0.5, 10.0, 100)
    qs = []

    def kurtosys(params: LogSvParams) -> np.ndarray:
        kurt = np.zeros_like(kappa2s)
        for idx, kappa2 in enumerate(kappa2s):
            nu = 2.0 * (kappa2 * params.theta-params.kappa1) / params.vartheta2 - 1.0
            q = 2.0 * params.kappa1 * params.theta / params.vartheta2
            b = 2.0 * kappa2 / params.vartheta2
            arg = 2.0*np.sqrt(q*b)
            kurt[idx] = 3.0*sps.kv(nu+4.0, arg)*sps.kv(nu, arg)/np.square(sps.kv(nu+2, arg)) - 3.0
        return kurt

    for key, params in params_dict.items():
        qs.append(pd.Series(kurtosys(params=params), index=kappa2s, name=key))

    ss_pdf = pd.concat(qs, axis=1)

    for key in ss_pdf.columns:
        ax.plot(ss_pdf.index, ss_pdf[key], label=key)
    ax.set_xlabel(f'$\kappa_{2}$')
    if title is not None:
        ax.set_title(title, fontsize=12, color='darkblue')
//  [python]
def plot_ss_kurtosis(params_dict=SS_PARAMS,
                     title: str = f'Excess kurtosis of log-returns as function of $\kappa_{2}$',
                     ax: plt.Subplot = None
                     ) -> None:

    kappa2s = np.linspace(0.5, 10.0, 100)
    qs = []

    def kurtosys(params: LogSvParams) -> np.ndarray:
        kurt = np.zeros_like(kappa2s)
        for idx, kappa2 in enumerate(kappa2s):
            nu = 2.0 * (kappa2 * params.theta-params.kappa1) / params.vartheta2 - 1.0
            q = 2.0 * params.kappa1 * params.theta / params.vartheta2
            b = 2.0 * kappa2 / params.vartheta2
            arg = 2.0*np.sqrt(q*b)
            kurt[idx] = 3.0*sps.kv(nu+4.0, arg)*sps.kv(nu, arg)/np.square(sps.kv(nu+2, arg)) - 3.0
        return kurt

    for key, params in params_dict.items():
        qs.append(pd.Series(kurtosys(params=params), index=kappa2s, name=key))

    ss_pdf = pd.concat(qs, axis=1)

    sns.lineplot(data=ss_pdf, dashes=False, ax=ax)
    ax.set_xlabel(f'$\kappa_{2}$')
    if title is not None:
        ax.set_title(title, fontsize=12, color='darkblue')
//  [python]
def plot_vol_skew(params_dict=SS_PARAMS,
                  title: str = f'Skeweness of volatility as function of $\kappa_{2}$',
                  ax: plt.Subplot = None
                  ) -> None:

    kappa2s = np.linspace(0.5, 10.0, 100)
    qs = []

    def skewness(params: LogSvParams) -> np.ndarray:
        skew = np.zeros_like(kappa2s)
        for idx, kappa2 in enumerate(kappa2s):
            params.kappa2 = kappa2
            m3_r = vol_moment(params=params, r=3)
            m2_r = vol_moment(params=params, r=2)
            m1_r = vol_moment(params=params, r=1)
            m2 = m2_r - m1_r * m1_r
            # y = (m3_r-3*m1_r*m2_r-m1_r*m1_r*m1_r)/np.power(m2_r, 1.5)
            skew[idx] = (m3_r - 3 * m1_r * m2 - m1_r * m1_r * m1_r) / np.power(m2, 1.5)
        return skew

    for key, params in params_dict.items():
        qs.append(pd.Series(skewness(params=params), index=kappa2s, name=key))

    ss_pdf = pd.concat(qs, axis=1)

    for key in ss_pdf.columns:
        ax.plot(ss_pdf.index, ss_pdf[key], label=key)
    ax.set_xlabel(f'$\kappa_{2}$')
    if title is not None:
        ax.set_title(title, fontsize=12, color='darkblue')
//  [python]
def plot_vol_skew(params_dict=SS_PARAMS,
                  title: str = f'Skeweness of volatility as function of $\kappa_{2}$',
                  ax: plt.Subplot = None
                  ) -> None:

    kappa2s = np.linspace(0.5, 10.0, 100)
    qs = []

    def skewness(params: LogSvParams) -> np.ndarray:
        skew = np.zeros_like(kappa2s)
        for idx, kappa2 in enumerate(kappa2s):
            params.kappa2 = kappa2
            m3_r = vol_moment(params=params, r=3)
            m2_r = vol_moment(params=params, r=2)
            m1_r = vol_moment(params=params, r=1)
            m2 = m2_r - m1_r * m1_r
            # y = (m3_r-3*m1_r*m2_r-m1_r*m1_r*m1_r)/np.power(m2_r, 1.5)
            skew[idx] = (m3_r - 3 * m1_r * m2 - m1_r * m1_r * m1_r) / np.power(m2, 1.5)
        return skew

    for key, params in params_dict.items():
        qs.append(pd.Series(skewness(params=params), index=kappa2s, name=key))

    ss_pdf = pd.concat(qs, axis=1)

    sns.lineplot(data=ss_pdf, dashes=False, ax=ax)
    ax.set_xlabel(f'$\kappa_{2}$')
    if title is not None:
        ax.set_title(title, fontsize=12, color='darkblue')
//  [python]
def line_plot():
    method_name = "gpt4v_visual_revision_prompting"

    ## load the test scores of GPT-4V 
    with open("../metrics/prediction_file_name_list_part1_new.json", "r") as f:
        file_name_list = json.load(f)
    with open("../metrics/res_dict_part1_new.json", "r") as f:
        res_dict = json.load(f)
    res_lst = res_dict[method_name]

    x_variables = []
    for filename in file_name_list:
        with open("../../testset_final/" + filename, "r") as f:
            html_content = f.read() 
            variable = count_total_nodes(html_content)
            x_variables.append(variable)

    y_variables = []
    for i, res in enumerate(res_lst):
        y_variables.append({"performance": res[1], "tags": x_variables[i], "metric": "block"})
        # y_variables.append({"performance": res[2], "tags": x_variables[i], "metric": "text"})
        y_variables.append({"performance": res[3], "tags": x_variables[i], "metric": "position"})
        # y_variables.append({"performance": res[4], "tags": x_variables[i], "metric": "color"})
        y_variables.append({"performance": res[5], "tags": x_variables[i], "metric": "CLIP"})

    # Convert list of dictionaries to DataFrame
    df = pd.DataFrame(y_variables)

    # Plotting
    plt.figure(figsize=(10, 6))
    for metric in df['metric'].unique():
        subset = df[df['metric'] == metric]
        plt.plot(subset['tags'], subset['performance'], label=metric)

    plt.title('Curves for Different Categories')
    plt.xlabel('Variable')
    plt.ylabel('Value')
    plt.legend(title='Category')
    plt.grid(True)
    plt.show()
//  [python]
def line_plot():
    method_name = "gpt4v_visual_revision_prompting"

    ## load the test scores of GPT-4V 
    with open("../metrics/prediction_file_name_list_part1_new.json", "r") as f:
        file_name_list = json.load(f)
    with open("../metrics/res_dict_part1_new.json", "r") as f:
        res_dict = json.load(f)
    res_lst = res_dict[method_name]

    x_variables = []
    for filename in file_name_list:
        with open("../../testset_final/" + filename, "r") as f:
            html_content = f.read() 
            variable = count_total_nodes(html_content)
            x_variables.append(variable)

    y_variables = []
    for i, res in enumerate(res_lst):
        y_variables.append({"performance": res[1], "tags": x_variables[i], "metric": "block"})
        # y_variables.append({"performance": res[2], "tags": x_variables[i], "metric": "text"})
        y_variables.append({"performance": res[3], "tags": x_variables[i], "metric": "position"})
        # y_variables.append({"performance": res[4], "tags": x_variables[i], "metric": "color"})
        y_variables.append({"performance": res[5], "tags": x_variables[i], "metric": "CLIP"})

    # Convert list of dictionaries to DataFrame
    df = pd.DataFrame(y_variables)

    # Plotting
    plt.figure(figsize=(10, 6))
    sns.lineplot(data=df, x='tags', y='performance', hue='metric', palette='tab10')

    plt.title('Curves for Different Categories')
    plt.xlabel('Variable')
    plt.ylabel('Value')
    plt.legend(title='Category')
    plt.grid(True)
    plt.show()
//  [python]
xs = []
ys = []
for fn in files:
  with open(fn) as f:
    curxs = []
    curys = []
    for line in f:
      tokens = line.split(' ')
      if len(tokens) == 4:
        x = int(tokens[0])
        y = float(tokens[2]) / float(tokens[3])
        if x > 300:
          continue
        curxs.append(x)
        curys.append(y)
    xs.append(curxs)
    ys.append(curys)
sns.set()
for i in range(0, len(xs)):
  plt.plot(xs[i], ys[i], color = colors[i])
//  [python]
xs = []
ys = []
for fn in files:
  with open(fn) as f:
    curxs = []
    curys = []
    for line in f:
      tokens = line.split(' ')
      if len(tokens) == 4:
        x = int(tokens[0])
        y = float(tokens[2]) / float(tokens[3])
        if x > 300:
          continue
        curxs.append(x)
        curys.append(y)
    xs.append(curxs)
    ys.append(curys)
sns.set()
for i in range(0, len(xs)):
  sns.lineplot(xs[i], ys[i], color = colors[i])
//  [python]
def get_energies_from_distances(data, plot=False, start=None, end=None):
    kde = gaussian_kde(data)

    if start == None:
        start = np.min(data)
    if end == None:
        end = np.max(data)

    steps = 100
    positions = np.arange(start, end, (end - start) / int((end - start) / 0.02))

    probabilities = kde(positions)
    # probabilities = np.histogram(data,bins=int((end-start)/0.02),density=True,range=(start,end))[0]

    # print(probabilities)
    # print(reverse_boltzmann(probabilities[50]),reverse_boltzmann(probabilities[50],kcal=False))
    probabilities = probabilities / np.max(probabilities)  # normalize
    # print(reverse_boltzmann(probabilities[50]))
    # max_p = max(probabilities)
    energies = np.array([reverse_boltzmann(p) for p in probabilities])
    # print(energies[50])
    # Normalize energies
    # energies = energies - min(energies)
    # print(energies[50])
    if plot:
        plt.plot(positions, energies)
    return energies, positions
//  [python]
def get_energies_from_distances(data, plot=False, start=None, end=None):
    kde = gaussian_kde(data)

    if start == None:
        start = np.min(data)
    if end == None:
        end = np.max(data)

    steps = 100
    positions = np.arange(start, end, (end - start) / int((end - start) / 0.02))

    probabilities = kde(positions)
    # probabilities = np.histogram(data,bins=int((end-start)/0.02),density=True,range=(start,end))[0]

    # print(probabilities)
    # print(reverse_boltzmann(probabilities[50]),reverse_boltzmann(probabilities[50],kcal=False))
    probabilities = probabilities / np.max(probabilities)  # normalize
    # print(reverse_boltzmann(probabilities[50]))
    # max_p = max(probabilities)
    energies = np.array([reverse_boltzmann(p) for p in probabilities])
    # print(energies[50])
    # Normalize energies
    # energies = energies - min(energies)
    # print(energies[50])
    if plot:
        sns.lineplot(x=positions, y=energies)
    return energies, positions
//  [python]
def evidence_tuning_plots(df, x_input = "Mean Predicted Avg",
                          y_input = "Empirical Probability",
                          x_name="Mean Predicted",
                          y_name="Empirical Probability"):
    """ Plot the tuning plot at different evidence values """

    def lineplot(x, y, trials, methods, **kwargs):
        """method_lineplot.

        Args:
            y:
            methods:
            kwargs:
        """
        uniq_methods = set(methods.values)
        method_order = sorted(uniq_methods)

        method_new_names = [f"$\lambda={i:0.4f}$" for i in method_order]
        method_df = []
        for method_idx, (method, method_new_name) in enumerate(zip(method_order,
                                                                   method_new_names)):
            lines_y = y[methods == method]
            lines_x = x[methods == method]
            for index, (xx, yy,trial) in enumerate(zip(lines_x, lines_y, trials)):

                to_append = [{x_name  : x,
                              y_name: y,
                              "Method": method_new_name,
                              "Trial" : trial}
                    for i, (x,y) in enumerate(zip(xx,yy))]
                method_df.extend(to_append)
        method_df = pd.DataFrame(method_df)
        x = np.linspace(0,1,100)
        plt.plot(x, x, linestyle='--', color="black")
        sns.pointplot(x=x_name, y=y_name, hue="Method",
                     alpha=0.8,
                     hue_order=method_new_names, data=method_df,)
        # estimator=None, units = "Trial")

    df = df.copy()
    # Query methods that have evidence_new_reg_2.0
    df = df[["evidence" in i for i in
             df['method_name']]].reset_index()

    # Get the regularizer and reset coeff
    coeff = [float(i.split("evidence_new_reg_")[1]) for i in df['method_name']]
    df["method_name"] = coeff
    df["Data"] = convert_dataset_names(df["dataset"])
    df["Method"] = df["method_name"]

    g = sns.FacetGrid(df, col="Data",  height=6, sharex = False, sharey = False)
    g.map(lineplot, x_input,  y_input, "trial_number",
          methods=df["Method"]).add_legend()
//  [python]
def evidence_tuning_plots(df, x_input = "Mean Predicted Avg",
                          y_input = "Empirical Probability",
                          x_name="Mean Predicted",
                          y_name="Empirical Probability"):
    """ Plot the tuning plot at different evidence values """

    def lineplot(x, y, trials, methods, **kwargs):
        """method_lineplot.

        Args:
            y:
            methods:
            kwargs:
        """
        uniq_methods = set(methods.values)
        method_order = sorted(uniq_methods)

        method_new_names = [f"$\lambda={i:0.4f}$" for i in method_order]
        method_df = []
        for method_idx, (method, method_new_name) in enumerate(zip(method_order,
                                                                   method_new_names)):
            lines_y = y[methods == method]
            lines_x = x[methods == method]
            for index, (xx, yy,trial) in enumerate(zip(lines_x, lines_y, trials)):

                to_append = [{x_name  : x,
                              y_name: y,
                              "Method": method_new_name,
                              "Trial" : trial}
                    for i, (x,y) in enumerate(zip(xx,yy))]
                method_df.extend(to_append)
        method_df = pd.DataFrame(method_df)
        x = np.linspace(0,1,100)
        plt.plot(x, x, linestyle='--', color="black")
        sns.lineplot(x=x_name, y=y_name, hue="Method",
                     alpha=0.8,
                     hue_order=method_new_names, data=method_df,)
        # estimator=None, units = "Trial")

    df = df.copy()
    # Query methods that have evidence_new_reg_2.0
    df = df[["evidence" in i for i in
             df['method_name']]].reset_index()

    # Get the regularizer and reset coeff
    coeff = [float(i.split("evidence_new_reg_")[1]) for i in df['method_name']]
    df["method_name"] = coeff
    df["Data"] = convert_dataset_names(df["dataset"])
    df["Method"] = df["method_name"]

    g = sns.FacetGrid(df, col="Data",  height=6, sharex = False, sharey = False)
    g.map(lineplot, x_input,  y_input, "trial_number",
          methods=df["Method"]).add_legend()
//  [python]
def conf_percentile_lineplot(df, error_col_name="rmse", error_title="RMSE",
                                   xlabel="Confidence Percentile",
                                   y_points = 1000,
                                   truncate_tail = False):
    """conf_percentile_lineplot.

    Args:
        df: Consolidated df of results
        error_col_name: Name of column in df that containse rror
        error_title: Renamed error in the output plot
        xlabel: xlabel
        y_points: Number of data points to plot
        truncate_tail: If true, don't plot the tail of the plot
    """

    methods = df['method_name']
    y = df[error_col_name]

    uniq_methods = set(methods.values)
    method_order = [j for j in METHOD_ORDER if j in uniq_methods]
    method_df = []

    # Truncate the last point if truncate_til is true
    slice_obj = slice(None) if not truncate_tail else slice(-1)

    lines = [line if len(line) < y_points
             else  np.array(line)[np.linspace(0, len(line) - 1, y_points).astype(int)]
             for line in y]

    method_df = [ {"Ratio": x, error_title : y, "Method" : method}
                 for line, method in zip(lines, methods)
                 for x, y in zip(np.linspace(0,1, len(line))[slice_obj] ,
                                 line[slice_obj])]
    method_df = pd.DataFrame(method_df)
    sns.pointplot(x="Ratio", y=error_title, hue="Method",
                 palette = METHOD_COLORS,
                 hue_order=method_order, data=method_df)
//  [python]
def conf_percentile_lineplot(df, error_col_name="rmse", error_title="RMSE",
                                   xlabel="Confidence Percentile",
                                   y_points = 1000,
                                   truncate_tail = False):
    """conf_percentile_lineplot.

    Args:
        df: Consolidated df of results
        error_col_name: Name of column in df that containse rror
        error_title: Renamed error in the output plot
        xlabel: xlabel
        y_points: Number of data points to plot
        truncate_tail: If true, don't plot the tail of the plot
    """

    methods = df['method_name']
    y = df[error_col_name]

    uniq_methods = set(methods.values)
    method_order = [j for j in METHOD_ORDER if j in uniq_methods]
    method_df = []

    # Truncate the last point if truncate_til is true
    slice_obj = slice(None) if not truncate_tail else slice(-1)

    lines = [line if len(line) < y_points
             else  np.array(line)[np.linspace(0, len(line) - 1, y_points).astype(int)]
             for line in y]

    method_df = [ {"Ratio": x, error_title : y, "Method" : method}
                 for line, method in zip(lines, methods)
                 for x, y in zip(np.linspace(0,1, len(line))[slice_obj] ,
                                 line[slice_obj])]
    method_df = pd.DataFrame(method_df)
    sns.lineplot(x="Ratio", y=error_title, hue="Method",
                 palette = METHOD_COLORS,
                 hue_order=method_order, data=method_df)
//  [python]
def plot_convergence(saving_folder: str | os.PathLike) -> None:
    """Plot the loss values sampled by the various methods as a function of the batch number.

    Args:
        saving_folder: the folder where the calibration results were saved

    """
    calibration_results_file = Path(saving_folder) / "calibration_results.csv"
    data_frame = pd.read_csv(calibration_results_file)

    losses_cummin = data_frame.groupby("batch_num_samp").min()["losses_samp"].cummin()

    plt.figure()
    g = sns.tsplot(
        data=data_frame,
        time="batch_num_samp",
        value="losses_samp",
        unit="method_samp",
        color="method_samp",
        palette="tab10",
    )

    g = sns.tsplot(
        time=np.arange(max(data_frame["batch_num_samp"]) + 1),
        data=losses_cummin,
        color="black",
        ls="--",
        marker="o",
        label="min loss",
    )

    ids = data_frame["method_samp"].unique()
    sampler_names = _get_samplers_names(saving_folder, ids)

    handles, labels = g.get_legend_handles_labels()
    labels = [*sampler_names, labels[-1]]

    plt.legend(handles, labels, loc="upper right")
//  [python]
def plot_convergence(saving_folder: str | os.PathLike) -> None:
    """Plot the loss values sampled by the various methods as a function of the batch number.

    Args:
        saving_folder: the folder where the calibration results were saved

    """
    calibration_results_file = Path(saving_folder) / "calibration_results.csv"
    data_frame = pd.read_csv(calibration_results_file)

    losses_cummin = data_frame.groupby("batch_num_samp").min()["losses_samp"].cummin()

    plt.figure()
    g = sns.lineplot(
        data=data_frame,
        x="batch_num_samp",
        y="losses_samp",
        hue="method_samp",
        palette="tab10",
    )

    g = sns.lineplot(
        x=np.arange(max(data_frame["batch_num_samp"]) + 1),
        y=losses_cummin,
        color="black",
        ls="--",
        marker="o",
        label="min loss",
    )

    ids = data_frame["method_samp"].unique()
    sampler_names = _get_samplers_names(saving_folder, ids)

    handles, labels = g.get_legend_handles_labels()
    labels = [*sampler_names, labels[-1]]

    plt.legend(handles, labels, loc="upper right")
//  [python]
def plot_loss_over_min_loss_vs_epoch_by_run_id(
        runs_histories_df: pd.DataFrame,
        plot_dir: str):
    plt.close()
    sns.tsplot(
        data=runs_histories_df,
        time='_step',
        value='loss_over_min_loss',
        unit='run_id',
        legend=False,
        linewidth=1.5,
    )

    plt.yscale('log')
    plt.ylabel('Loss / Min(Loss)')
    plt.xlabel('Epoch')

    plt.savefig(os.path.join(plot_dir,
                             f'loss_over_min_loss_vs_epoch_by_run_id.png'),
                bbox_inches='tight',
                dpi=300)
    # plt.show()
    plt.close()
//  [python]
def plot_loss_over_min_loss_vs_epoch_by_run_id(
        runs_histories_df: pd.DataFrame,
        plot_dir: str):
    plt.close()
    sns.lineplot(
        data=runs_histories_df,
        x='_step',
        y='loss_over_min_loss',
        hue='run_id',
        legend=False,
        linewidth=1.5,
    )

    plt.yscale('log')
    plt.ylabel('Loss / Min(Loss)')
    plt.xlabel('Epoch')

    plt.savefig(os.path.join(plot_dir,
                             f'loss_over_min_loss_vs_epoch_by_run_id.png'),
                bbox_inches='tight',
                dpi=300)
    # plt.show()
    plt.close()