jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def plot_layer_decompositions():
    path = 'expt/pure_time_series/wiki_univariate/08_radflow_no_network/serialization/evaluate-metrics.json'
    with open(path) as f:
        o3 = json.load(f)

    col1_dict = defaultdict(list)
    col2_dict = defaultdict(list)

    for k_parts in o3['f_parts']:
        for i, k_layer in enumerate(k_parts):
            col1_dict[i] += list(range(28))
            col2_dict[i] += k_layer
            # len(k_layer) == 28

    fig = plt.figure(figsize=(6, 3))

    layer_pred_out = 1
    order = [1, 2, 3, 4, 5, 6, 7, 8]
    for i, o in enumerate(order):
        ax = plt.subplot(2, 4, o)
        df = pd.DataFrame({'time': col1_dict[i], 'obs': col2_dict[i]})
        plt.plot(df["time"], df["obs"])
        ax.set_title(o)
        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_xlabel('')
        ax.set_ylabel('')

    fig.tight_layout()
    fig.savefig('figures/decomposition.pdf')
//  [python]
def plot_layer_decompositions():
    path = 'expt/pure_time_series/wiki_univariate/08_radflow_no_network/serialization/evaluate-metrics.json'
    with open(path) as f:
        o3 = json.load(f)

    col1_dict = defaultdict(list)
    col2_dict = defaultdict(list)

    for k_parts in o3['f_parts']:
        for i, k_layer in enumerate(k_parts):
            col1_dict[i] += list(range(28))
            col2_dict[i] += k_layer
            # len(k_layer) == 28

    fig = plt.figure(figsize=(6, 3))

    layer_pred_out = 1
    order = [1, 2, 3, 4, 5, 6, 7, 8]
    for i, o in enumerate(order):
        ax = plt.subplot(2, 4, o)
        df = pd.DataFrame({'time': col1_dict[i], 'obs': col2_dict[i]})
        sns.lineplot(x="time", y="obs", data=df, ax=ax)
        ax.set_title(o)
        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_xlabel('')
        ax.set_ylabel('')

    fig.tight_layout()
    fig.savefig('figures/decomposition.pdf')
//  [python]
def plot_topics_over_time(topic_distributions, topic_keys, times, topic_index, output_path=None):
    
    data_dicts = []
    for j, _distribution in enumerate(topic_distributions):        
        for _topic, _probability in enumerate(_distribution):
            if _topic == topic_index:
                data_dicts.append({'Probability': _probability,
                                   'Time': times[j]})
    data_df = pd.DataFrame(data_dicts)

    sns.set(style='ticks', font_scale=1.2)
    plt.figure(figsize=(7,2.5))
    plt.plot(data_df['Time'], data_df['Probability'], color='cornflowerblue')
    plt.xlabel('Time')
    plt.ylabel('Topic Probability')
    plt.title(' '.join(topic_keys[topic_index][:5]))
    plt.tight_layout()
    sns.despine()
    if output_path:
        plt.savefig(output_path)
    plt.show()
//  [python]
def plot_topics_over_time(topic_distributions, topic_keys, times, topic_index, output_path=None):
    
    data_dicts = []
    for j, _distribution in enumerate(topic_distributions):        
        for _topic, _probability in enumerate(_distribution):
            if _topic == topic_index:
                data_dicts.append({'Probability': _probability,
                                   'Time': times[j]})
    data_df = pd.DataFrame(data_dicts)

    sns.set(style='ticks', font_scale=1.2)
    plt.figure(figsize=(7,2.5))
    sns.lineplot(data=data_df, 
                 x='Time', 
                 y='Probability', 
                 color='cornflowerblue') 
    plt.xlabel('Time')
    plt.ylabel('Topic Probability')
    plt.title(' '.join(topic_keys[topic_index][:5]))
    plt.tight_layout()
    sns.despine()
    if output_path:
        plt.savefig(output_path)
    plt.show()
//  [python]
def get_plots(ax, data, experiment_name):
    plt.rcParams["hatch.linewidth"] = 1.5

    sns.tsplot(
        data=data, time='epoch', value='result', unit='model', ax=ax, color=PALETTE,
        err_style='band', alpha=0.75, linewidth=4
    )

    ax.set(xlabel=None, ylabel=None)
    ax.legend(title="Models", bbox_to_anchor=(1, 1))
    ax.set_title(experiment_name, fontweight="bold")
    ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))
//  [python]
def get_plots(ax, data, experiment_name):
    plt.rcParams["hatch.linewidth"] = 1.5

    sns.lineplot(
        data=data, x='epoch', y='result', ax=ax, hue='model', palette=PALETTE,
        errorbar='pi', err_style='band', alpha=0.75, linewidth=4
    )

    ax.set(xlabel=None, ylabel=None)
    ax.legend(title="Models", bbox_to_anchor=(1, 1))
    ax.set_title(experiment_name, fontweight="bold")
    ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))
//  [python]
def plot(
    qrels: Qrels,
    runs: Union[Run, List[Run]],
    graph: str = "iprec_at_recall",
    return_graph: bool = True,
    figsize: tuple = None,
    seaborn_kwargs: Dict = None,
):
    if figsize is None:
        figsize = (8, 5)

    if seaborn_kwargs is None:
        seaborn_kwargs = {"linewidth": 2.5}

    _qrels = qrels.to_typed_list()

    if isinstance(runs, list):
        _runs = [run.to_typed_list() for run in runs]
        names = [
            run.name if run.name is not None else f"run_{i+1}"
            for i, run in enumerate(runs)
        ]
    else:
        _runs = [runs.to_typed_list()]
        names = [runs.name if runs.name is not None else "run"]

    results = [
        interpolated_precision_at_recall(_qrels, run).mean(axis=0) for run in _runs
    ]

    recall = np.arange(0, 1.1, 0.1)

    df = pd.DataFrame(
        {
            "Model": [name for name in names for _ in range(11)],
            "recall": np.concatenate([recall] * len(_runs)),
            "precision": np.concatenate(results),
        }
    )

    if not return_graph:
        return df

    plt.figure(figsize=figsize)

    ax = sns.pointplot(data=df, x="recall", y="precision", hue="Model", **seaborn_kwargs)

    ax.set(
        title="Precision-Recall Curve",
        xlabel="Recall",
        ylabel="Precision",
        xticks=np.arange(0.0, 1.1, 0.1),
        yticks=np.arange(0.0, 1.1, 0.1),
    )

    plt.show()
//  [python]
def plot(
    qrels: Qrels,
    runs: Union[Run, List[Run]],
    graph: str = "iprec_at_recall",
    return_graph: bool = True,
    figsize: tuple = None,
    seaborn_kwargs: Dict = None,
):
    if figsize is None:
        figsize = (8, 5)

    if seaborn_kwargs is None:
        seaborn_kwargs = {"linewidth": 2.5}

    _qrels = qrels.to_typed_list()

    if isinstance(runs, list):
        _runs = [run.to_typed_list() for run in runs]
        names = [
            run.name if run.name is not None else f"run_{i+1}"
            for i, run in enumerate(runs)
        ]
    else:
        _runs = [runs.to_typed_list()]
        names = [runs.name if runs.name is not None else "run"]

    results = [
        interpolated_precision_at_recall(_qrels, run).mean(axis=0) for run in _runs
    ]

    recall = np.arange(0, 1.1, 0.1)

    df = pd.DataFrame(
        {
            "Model": [name for name in names for _ in range(11)],
            "recall": np.concatenate([recall] * len(_runs)),
            "precision": np.concatenate(results),
        }
    )

    if not return_graph:
        return df

    plt.figure(figsize=figsize)

    ax = sns.lineplot(data=df, x="recall", y="precision", hue="Model", **seaborn_kwargs)

    ax.set(
        title="Precision-Recall Curve",
        xlabel="Recall",
        ylabel="Precision",
        xticks=np.arange(0.0, 1.1, 0.1),
        yticks=np.arange(0.0, 1.1, 0.1),
    )

    plt.show()
//  [python]
def uv_seasonalPlot(data, engine, xlabel, ylabel, afreq, aggf):
    #TODO : Original frequqncy wont work
    data = data.dropna().copy()
    data.rename(columns={'plotX1':ylabel}, inplace=True)

    if isinstance(aggf, dict):
        aggf = aggf[ylabel]
    data = data.groupby(['anfreq_label', 'anfreq_label1']).agg({ylabel: aggf}).reset_index()
    
    if engine == 'Static':
        fig , axes = plt.subplots(figsize=(9.4,5))
        _ = sns.tsplot(x = 'anfreq_label1', y = ylabel, data = data,
                         hue = 'anfreq_label', marker = 'o', ax = axes, legend='Week' not in afreq)

        axes.set_xlabel(afreq, fontsize = 15)
        axes.set_ylabel(ylabel, fontsize = 15)
        axes.grid(b=True, which='major', color='k', linewidth=0.25)
        axes.grid(b=True, which='minor', color='k', linewidth=0.125)
        plt.close()
        return pn.pane.Matplotlib(fig, tight=True)
    
    elif engine == 'Interactive':
        
        data = data.dropna()
        nearest = alt.selection(type='single', nearest=True,
                                on='mouseover', fields=['anfreq_label1'],
                                empty='none')
        _xshand = 'anfreq_label1:Q'
        _legend = alt.Legend()
        _yencoder = alt.Y('{0}:Q'.format(ylabel), scale=alt.Scale(zero=False), axis=alt.Axis(format='~s'))

        if afreq == 'AS':
            _legend = alt.Legend(type='gradient')
        elif afreq=='W':
            _xshand = 'anfreq_label1:N'
            _legend = None
        elif afreq=='A':
            _xshand = 'anfreq_label1:N'
            _legend = alt.Legend(type='gradient')

        line = alt.Chart(data)
        line = line.mark_line(interpolate='monotone')
        line = line.encode(x = alt.X(_xshand, axis=alt.Axis(labelAngle = 0.0), title='Day'),
                           y = _yencoder,
                           color=alt.Color('anfreq_label:N',
                                            legend = _legend,
                                            title = afreq))

        selectors = alt.Chart(data).mark_point()
        selectors = selectors.encode(x = alt.X(_xshand, axis=alt.Axis(labelAngle = 0.0)),
                                     opacity = alt.value(0))
        selectors = selectors.add_selection(nearest)

        points = line.mark_point()
        points = points.encode(opacity=alt.condition(nearest,alt.value(1),alt.value(0)))

        text = line.mark_text(align='left', dx=5, dy=-5)
        text = text.encode(text=alt.condition(nearest, '{0}:Q'.format(ylabel), alt.value(' ')))

        rules = alt.Chart(data).mark_rule(color='gray')
        rules = rules.encode(x = alt.X(_xshand, axis=alt.Axis(labelAngle = 0.0)))
        rules = rules.transform_filter(nearest)

        p = alt.layer(line,  selectors,  points,  rules,  text)
        p = p.configure_legend(gradientLength=280, gradientThickness=10)
        p = p.properties(width = 545, height = 350)
        
        return p
//  [python]
def uv_seasonalPlot(data, engine, xlabel, ylabel, afreq, aggf):
    #TODO : Original frequqncy wont work
    data = data.dropna().copy()
    data.rename(columns={'plotX1':ylabel}, inplace=True)

    if isinstance(aggf, dict):
        aggf = aggf[ylabel]
    data = data.groupby(['anfreq_label', 'anfreq_label1']).agg({ylabel: aggf}).reset_index()
    
    if engine == 'Static':
        fig , axes = plt.subplots(figsize=(9.4,5))
        _ = sns.lineplot(x = 'anfreq_label1', y = ylabel, data = data,
                         hue = 'anfreq_label', marker = 'o', ax = axes, legend='Week' not in afreq)

        axes.set_xlabel(afreq, fontsize = 15)
        axes.set_ylabel(ylabel, fontsize = 15)
        axes.grid(b=True, which='major', color='k', linewidth=0.25)
        axes.grid(b=True, which='minor', color='k', linewidth=0.125)
        plt.close()
        return pn.pane.Matplotlib(fig, tight=True)
    
    elif engine == 'Interactive':
        
        data = data.dropna()
        nearest = alt.selection(type='single', nearest=True,
                                on='mouseover', fields=['anfreq_label1'],
                                empty='none')
        _xshand = 'anfreq_label1:Q'
        _legend = alt.Legend()
        _yencoder = alt.Y('{0}:Q'.format(ylabel), scale=alt.Scale(zero=False), axis=alt.Axis(format='~s'))

        if afreq == 'AS':
            _legend = alt.Legend(type='gradient')
        elif afreq=='W':
            _xshand = 'anfreq_label1:N'
            _legend = None
        elif afreq=='A':
            _xshand = 'anfreq_label1:N'
            _legend = alt.Legend(type='gradient')

        line = alt.Chart(data)
        line = line.mark_line(interpolate='monotone')
        line = line.encode(x = alt.X(_xshand, axis=alt.Axis(labelAngle = 0.0), title='Day'),
                           y = _yencoder,
                           color=alt.Color('anfreq_label:N',
                                            legend = _legend,
                                            title = afreq))

        selectors = alt.Chart(data).mark_point()
        selectors = selectors.encode(x = alt.X(_xshand, axis=alt.Axis(labelAngle = 0.0)),
                                     opacity = alt.value(0))
        selectors = selectors.add_selection(nearest)

        points = line.mark_point()
        points = points.encode(opacity=alt.condition(nearest,alt.value(1),alt.value(0)))

        text = line.mark_text(align='left', dx=5, dy=-5)
        text = text.encode(text=alt.condition(nearest, '{0}:Q'.format(ylabel), alt.value(' ')))

        rules = alt.Chart(data).mark_rule(color='gray')
        rules = rules.encode(x = alt.X(_xshand, axis=alt.Axis(labelAngle = 0.0)))
        rules = rules.transform_filter(nearest)

        p = alt.layer(line,  selectors,  points,  rules,  text)
        p = p.configure_legend(gradientLength=280, gradientThickness=10)
        p = p.properties(width = 545, height = 350)
        
        return p
//  [python]
def plot_lineplot(valdf, savepath, yaxis):
    font = {'size'   : 16}
    matplotlib.rc('font', **font)
    valdf = reshape_for_plotting(valdf, additional_cols = ['datasize'])
    valdf['$\it{Model}$'] = valdf['model_type']
    valdf['$\it{Data}$'] = valdf['scr']
    plt.figure(figsize=(7,4), dpi=300)
    my_pal = {"AutoKeras": "cornflowerblue", "DeepSwarm": "sandybrown", "TPOT": "grey"}
    ax = sns.pointplot(x = 'datasize', y = 'values', hue = '$\it{Model}$', style = '$\it{Data}$', data=valdf,  alpha = 1, linewidth = 2,  ci = 'sd', palette = my_pal, join=True)
    sns.scatterplot(x = 'datasize', y = 'values', hue = '$\it{Model}$', data=valdf, edgecolor='black', alpha = 1, linewidth = 0.4, s = 12, palette = my_pal)
    handles, labels = ax.get_legend_handles_labels()
    ax.legend(handles[:7], labels[:7], bbox_to_anchor=(1.01, 1), borderaxespad=0)
    plt.xlabel('Dataset Size')
    plt.xticks(rotation = 60)
    plt.ylabel(yaxis)
    ax.set_xscale('log')
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    plt.tight_layout()
    plt.savefig(savepath)
    plt.show()
//  [python]
def plot_lineplot(valdf, savepath, yaxis):
    font = {'size'   : 16}
    matplotlib.rc('font', **font)
    valdf = reshape_for_plotting(valdf, additional_cols = ['datasize'])
    valdf['$\it{Model}$'] = valdf['model_type']
    valdf['$\it{Data}$'] = valdf['scr']
    plt.figure(figsize=(7,4), dpi=300)
    my_pal = {"AutoKeras": "cornflowerblue", "DeepSwarm": "sandybrown", "TPOT": "grey"}
    ax = sns.lineplot(x = 'datasize', y = 'values', hue = '$\it{Model}$', style = '$\it{Data}$', data=valdf,  alpha = 1, linewidth = 2,  ci = 'sd', palette = my_pal)
    sns.scatterplot(x = 'datasize', y = 'values', hue = '$\it{Model}$', data=valdf, edgecolor='black', alpha = 1, linewidth = 0.4, s = 12, palette = my_pal)
    handles, labels = ax.get_legend_handles_labels()
    ax.legend(handles[:7], labels[:7], bbox_to_anchor=(1.01, 1), borderaxespad=0)
    plt.xlabel('Dataset Size')
    plt.xticks(rotation = 60)
    plt.ylabel(yaxis)
    ax.set_xscale('log')
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    plt.tight_layout()
    plt.savefig(savepath)
    plt.show()
//  [python]
def seedlim_plot2(args):
	xl = pd.ExcelFile(args[0])
	print(xl.sheet_names)
	df = pd.io.excel.ExcelFile.parse(xl, "CM")
	#df["Seed Limit"] = math.sqrt(int(df["SL"]))
	df["Seed Limit"] = df["SL"] 
	
	df.loc[df["Seed Limit"] == 100, "size"] = 1
	df.loc[df["Seed Limit"] == 500, "size"] = 2
	df.loc[df["Seed Limit"] == 1000, "size"] = 4
	df.loc[df["Seed Limit"] == 5000, "size"] = 8

	print(df)

	ax, tool_order, col_list = prepare_for_simul()
	#plt.gca().xaxis.grid(True)
	#plt.gca().yaxis.grid(True)
	
	kw = dict(edgecolor="none")
	
	#ax = sns.lineplot(x="Recall", y="Precision", hue="SL", style="Sample", markers=['*', 'D', 'o'], dashes=False, data=df)#, legend_out = True)
	#ax = sns.lineplot(x="Recall", y="Precision", hue="Sample", style="Sample", markers=['.', '.', '.'], dashes=False, data=df, alpha=0.7, lw=0.6, legend=False)
	ax = sns.pointplot(x="Recall", y="Precision", color="k", hue="Sample", data=df, alpha=1.0, lw=0.6, legend=False)
	ax = sns.scatterplot(x="Recall", y="Precision", hue="Sample", size="size", style="Sample", markers=['>', 'D', 'p'], data=df, sizes=(20,120), linewidth=0)

	#for line in range(0,df.shape[0]):
	#	#ax.text(df.Recall[line]+0.2, df.Precision[line], df.SL[line], horizontalalignment='left', size='medium', color='black', weight='semibold')
	#	yval = df.Precision[line] + 0.08 if line % 2 == 0 else df.Precision[line] - 0.1
	#	ax.text(df.Recall[line], yval, df.SL[line], horizontalalignment='left', color='black', fontsize=6)

	#plt.title("Mapping Accuracy", fontdict=font, pad=30.0)
	#ax[2].legend(bbox_to_anchor=(1.00,1.10), loc="lower right", borderaxespad=0)
	
	#plt.gca().set_xlim(right=100.0)
	#plt.gca().set_ylim(top=100.0)
	plt.gca().set_ylim(bottom=95.0)

	handles, labels = ax.get_legend_handles_labels()
	print(handles)
	del(handles[-1])
	del(labels[-1])
	labels[0] = 'Simulated Dataset'
	labels[4] = 'Seed Limit'
	labels[5] = 100
	labels[6] = 500
	labels[7] = 1000
	labels[8] = 5000
	print(labels)
	leg = plt.legend(handles=handles[:], labels=labels[:], loc='upper left', bbox_to_anchor=(0.0, 1.0))
	leg._legend_box.align = "left"
	ax.add_artist(leg)

	ax.set(xlabel='Recall (avg. over 5 replicates)', ylabel='Precision (avg. over 5 replicates)')
	
	#plt.savefig('sl_prec_recall.pdf', bbox_inches='tight')
	plt.savefig('sl_prec_recall.png', bbox_inches='tight', dpi=800)
//  [python]
def seedlim_plot2(args):
	xl = pd.ExcelFile(args[0])
	print(xl.sheet_names)
	df = pd.io.excel.ExcelFile.parse(xl, "CM")
	#df["Seed Limit"] = math.sqrt(int(df["SL"]))
	df["Seed Limit"] = df["SL"] 
	
	df.loc[df["Seed Limit"] == 100, "size"] = 1
	df.loc[df["Seed Limit"] == 500, "size"] = 2
	df.loc[df["Seed Limit"] == 1000, "size"] = 4
	df.loc[df["Seed Limit"] == 5000, "size"] = 8

	print(df)

	ax, tool_order, col_list = prepare_for_simul()
	#plt.gca().xaxis.grid(True)
	#plt.gca().yaxis.grid(True)
	
	kw = dict(edgecolor="none")
	
	#ax = sns.lineplot(x="Recall", y="Precision", hue="SL", style="Sample", markers=['*', 'D', 'o'], dashes=False, data=df)#, legend_out = True)
	#ax = sns.lineplot(x="Recall", y="Precision", hue="Sample", style="Sample", markers=['.', '.', '.'], dashes=False, data=df, alpha=0.7, lw=0.6, legend=False)
	ax = sns.lineplot(x="Recall", y="Precision", color="k", style="Sample", markers=[',', ',', ','], dashes=False, data=df, alpha=1.0, lw=0.6, legend=False)
	ax = sns.scatterplot(x="Recall", y="Precision", hue="Sample", size="size", style="Sample", markers=['>', 'D', 'p'], data=df, sizes=(20,120), linewidth=0)

	#for line in range(0,df.shape[0]):
	#	#ax.text(df.Recall[line]+0.2, df.Precision[line], df.SL[line], horizontalalignment='left', size='medium', color='black', weight='semibold')
	#	yval = df.Precision[line] + 0.08 if line % 2 == 0 else df.Precision[line] - 0.1
	#	ax.text(df.Recall[line], yval, df.SL[line], horizontalalignment='left', color='black', fontsize=6)

	#plt.title("Mapping Accuracy", fontdict=font, pad=30.0)
	#ax[2].legend(bbox_to_anchor=(1.00,1.10), loc="lower right", borderaxespad=0)
	
	#plt.gca().set_xlim(right=100.0)
	#plt.gca().set_ylim(top=100.0)
	plt.gca().set_ylim(bottom=95.0)

	handles, labels = ax.get_legend_handles_labels()
	print(handles)
	del(handles[-1])
	del(labels[-1])
	labels[0] = 'Simulated Dataset'
	labels[4] = 'Seed Limit'
	labels[5] = 100
	labels[6] = 500
	labels[7] = 1000
	labels[8] = 5000
	print(labels)
	leg = plt.legend(handles=handles[:], labels=labels[:], loc='upper left', bbox_to_anchor=(0.0, 1.0))
	leg._legend_box.align = "left"
	ax.add_artist(leg)

	ax.set(xlabel='Recall (avg. over 5 replicates)', ylabel='Precision (avg. over 5 replicates)')
	
	#plt.savefig('sl_prec_recall.pdf', bbox_inches='tight')
	plt.savefig('sl_prec_recall.png', bbox_inches='tight', dpi=800)
//  [python]
def seaborn_lineplot(df, ax, split_by, legend='brief', xlabel='time (s)', ylabel='$\Delta$F/F', n_boot=1000):  # NOQA W605
    cmap = colormap()

    df = df[['eventlocked_timestamps', 'eventlocked_traces', split_by]].dropna()
    values_in_split = np.sort(df[split_by].unique())

    sns.tsplot(
        time='eventlocked_timestamps',
        value='eventlocked_traces',
        data=df,
        condition=split_by,
        unit=None,
        order=np.sort(values_in_split),
        color=[cmap[split_by][state] for state in values_in_split],
        ax=ax,
        legend=legend,
        n_boot=n_boot,
    )
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel, rotation=0, ha='right')
//  [python]
def seaborn_lineplot(df, ax, split_by, legend='brief', xlabel='time (s)', ylabel='$\Delta$F/F', n_boot=1000):  # NOQA W605
    cmap = colormap()

    df = df[['eventlocked_timestamps', 'eventlocked_traces', split_by]].dropna()
    values_in_split = np.sort(df[split_by].unique())

    sns.lineplot(
        x='eventlocked_timestamps',
        y='eventlocked_traces',
        data=df,
        hue=split_by,
        hue_order=np.sort(values_in_split),
        palette=[cmap[split_by][state] for state in values_in_split],
        ax=ax,
        legend=legend,
        n_boot=n_boot,
    )
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel, rotation=0, ha='right')
//  [python]
def colocation(
    data: AnnData,
    ranks: List[int],
    fname: str,
    iterations: int = 3,
    plot_error: bool = True,
    copy: bool = False,
):
    """Decompose a tensor of pairwise colocalization quotients into signatures.

    Parameters
    ----------
    adata : AnnData
        Spatial formatted AnnData object.
    ranks : list
        List of ranks to decompose the tensor.
    iterations : int
        Number of iterations to run the decomposition.
    plot_error : bool
        Whether to plot the error of the decomposition.
    copy : bool
        Whether to return a copy of the AnnData object. Default False.
    Returns
    -------
    adata : AnnData
        .uns['factors']: Decomposed tensor factors.
        .uns['factors_error']: Decomposition error.
    """
    adata = data.copy() if copy else data

    print("Preparing tensor...")
    _colocation_tensor(adata, copy=copy)

    tensor = adata.uns["tensor"]

    print(emoji.emojize(":running: Decomposing tensor..."))
    factors, errors = decompose(tensor, ranks, iterations=iterations)

    if plot_error and errors.shape[0] > 1:
        kl = KneeLocator(errors["rank"], errors["rmse"], direction="decreasing", curve="convex")
        if kl.knee is None:
            warning('No knee found, please extend the ranks range.\nCurrent ranks range: [{ranks[0]},{ranks[-1]}]')
        else:
            info(f'Knee found at rank {kl.knee}')
            fig, ax = plt.subplots(figsize=(6, 4))
            sns.pointplot(data=errors, x="rank", y="rmse", ci=95, ax=ax)  # type: ignore
            plt.axvline(kl.knee, linestyle="--")
            plt.savefig(fname)
            plt.close(fig)
            info(f"Saved to {fname}")

    adata.uns["factors"] = factors
    adata.uns["factors_error"] = errors

    print(emoji.emojize(":heavy_check_mark: Done."))
    return adata if copy else None
//  [python]
def colocation(
    data: AnnData,
    ranks: List[int],
    fname: str,
    iterations: int = 3,
    plot_error: bool = True,
    copy: bool = False,
):
    """Decompose a tensor of pairwise colocalization quotients into signatures.

    Parameters
    ----------
    adata : AnnData
        Spatial formatted AnnData object.
    ranks : list
        List of ranks to decompose the tensor.
    iterations : int
        Number of iterations to run the decomposition.
    plot_error : bool
        Whether to plot the error of the decomposition.
    copy : bool
        Whether to return a copy of the AnnData object. Default False.
    Returns
    -------
    adata : AnnData
        .uns['factors']: Decomposed tensor factors.
        .uns['factors_error']: Decomposition error.
    """
    adata = data.copy() if copy else data

    print("Preparing tensor...")
    _colocation_tensor(adata, copy=copy)

    tensor = adata.uns["tensor"]

    print(emoji.emojize(":running: Decomposing tensor..."))
    factors, errors = decompose(tensor, ranks, iterations=iterations)

    if plot_error and errors.shape[0] > 1:
        kl = KneeLocator(errors["rank"], errors["rmse"], direction="decreasing", curve="convex")
        if kl.knee is None:
            warning('No knee found, please extend the ranks range.\nCurrent ranks range: [{ranks[0]},{ranks[-1]}]')
        else:
            info(f'Knee found at rank {kl.knee}')
            fig, ax = plt.subplots(figsize=(6, 4))
            sns.lineplot(data=errors, x="rank", y="rmse", ci=95, marker="o", ax=ax)  # type: ignore
            plt.axvline(kl.knee, linestyle="--")
            plt.savefig(fname)
            plt.close(fig)
            info(f"Saved to {fname}")

    adata.uns["factors"] = factors
    adata.uns["factors_error"] = errors

    print(emoji.emojize(":heavy_check_mark: Done."))
    return adata if copy else None
//  [python]
def single_plot(df, value_keys, config, output_dir):
    tasks = config['tasks']
    x_axis_sci_limit = config['x_axis_sci_limit']
    for task_name in tqdm(tasks):
        task_output_dir = os.path.join(output_dir, task_name)
        if not os.path.exists(task_output_dir):
            os.makedirs(task_output_dir)
        task_df = df.loc[df['task_name']==task_name]
        for value_key in value_keys:
            value_df = task_df[['timestep', 'algo_name', value_key]].dropna()
            sns.tsplot(data = value_df, time = 'timestep', value = value_key, unit = 'algo_name')
            
            plt.ticklabel_format(axis="x", style="sci", scilimits=x_axis_sci_limit)
            output_path = os.path.join(task_output_dir, value_key + ".svg")
            plt.savefig(output_path)
            plt.clf()
//  [python]
def single_plot(df, value_keys, config, output_dir):
    tasks = config['tasks']
    x_axis_sci_limit = config['x_axis_sci_limit']
    for task_name in tqdm(tasks):
        task_output_dir = os.path.join(output_dir, task_name)
        if not os.path.exists(task_output_dir):
            os.makedirs(task_output_dir)
        task_df = df.loc[df['task_name']==task_name]
        for value_key in value_keys:
            value_df = task_df[['timestep', 'algo_name', value_key]].dropna()
            sns.lineplot(data = value_df,x = 'timestep', y = value_key, hue = 'algo_name')
            
            plt.ticklabel_format(axis="x", style="sci", scilimits=x_axis_sci_limit)
            output_path = os.path.join(task_output_dir, value_key + ".svg")
            plt.savefig(output_path)
            plt.clf()
//  [python]
def plot_losses(pipeline_result, *, ax=None):
    """Plot the losses per epoch."""
    import seaborn as sns

    sns.set_style("darkgrid")

    ax = _ensure_ax(ax)
    rv = sns.pointplot(x=range(len(pipeline_result.losses)), y=pipeline_result.losses, ax=ax)

    loss_name = loss_resolver.normalize_inst(pipeline_result.model.loss)
    ax.set_ylabel(f"{loss_name} Loss")
    ax.set_xlabel("Epoch")
    ax.set_title(pipeline_result.title if pipeline_result.title is not None else "Losses Plot")
    return rv
//  [python]
def plot_losses(pipeline_result, *, ax=None):
    """Plot the losses per epoch."""
    import seaborn as sns

    sns.set_style("darkgrid")

    ax = _ensure_ax(ax)
    rv = sns.lineplot(x=range(len(pipeline_result.losses)), y=pipeline_result.losses, ax=ax)

    loss_name = loss_resolver.normalize_inst(pipeline_result.model.loss)
    ax.set_ylabel(f"{loss_name} Loss")
    ax.set_xlabel("Epoch")
    ax.set_title(pipeline_result.title if pipeline_result.title is not None else "Losses Plot")
    return rv
//  [python]
def plot_early_stopping(pipeline_result, *, ax=None, lineplot_kwargs=None):
    """Plot the evaluations during early stopping."""
    import seaborn as sns

    if not isinstance(pipeline_result.stopper, EarlyStopper) or not pipeline_result.stopper.results:
        raise ValueError

    ax = _ensure_ax(ax)

    x = [(1 + e) * pipeline_result.stopper.frequency for e in range(len(pipeline_result.stopper.results))]
    rv = sns.tsplot(x=x, y=pipeline_result.stopper.results, ax=ax, marker="o", **(lineplot_kwargs or {}))

    ax.set_ylabel(pipeline_result.stopper.metric)
    ax.set_xlabel("Epoch")
    ax.set_title(pipeline_result.title if pipeline_result.title is not None else "Early Stopper Evaluation Plot")
    return rv
//  [python]
def plot_early_stopping(pipeline_result, *, ax=None, lineplot_kwargs=None):
    """Plot the evaluations during early stopping."""
    import seaborn as sns

    if not isinstance(pipeline_result.stopper, EarlyStopper) or not pipeline_result.stopper.results:
        raise ValueError

    ax = _ensure_ax(ax)

    x = [(1 + e) * pipeline_result.stopper.frequency for e in range(len(pipeline_result.stopper.results))]
    rv = sns.lineplot(x=x, y=pipeline_result.stopper.results, ax=ax, marker="o", **(lineplot_kwargs or {}))

    ax.set_ylabel(pipeline_result.stopper.metric)
    ax.set_xlabel("Epoch")
    ax.set_title(pipeline_result.title if pipeline_result.title is not None else "Early Stopper Evaluation Plot")
    return rv
//  [python]
def vis_training_curves(history=None, output_dir: str = None) -> None:
    """
    Visualizes the training curves of the model.

    Args:
        history: The history object returned by the model.fit() method.
        output_dir (Optional[str]): The output directory to save the training curves plot.

    Returns:
        None
    """
    # Extract the accuracy and loss values for training and validation data
    acc = history.history['accuracy']
    val_acc = history.history['val_accuracy']
    loss = history.history['loss']
    val_loss = history.history['val_loss']
    epochs_range = range(len(acc))

    # Create dataframes for the training and validation data
    df_val = pd.DataFrame({'run': 'validation', 'step': epochs_range, 'epoch_accuracy': val_acc, 'epoch_loss': val_loss})
    df_train = pd.DataFrame({'run': 'train', 'step': epochs_range, 'epoch_accuracy': acc, 'epoch_loss': loss})

    # Concatenate the dataframes
    frames = [df_val, df_train]
    df = pd.concat(frames)
    df = df.reset_index()

    # Plot the training curves
    plt.figure(figsize=(16, 6))
    plt.suptitle('Training curves', fontsize=16)
    plt.subplot(1, 2, 1)
    sns.pointplot(data=df, x="step", y="epoch_accuracy", hue="run").set_title("accuracy")
    plt.grid()
    plt.subplot(1, 2, 2)
    sns.pointplot(data=df, x="step", y="epoch_loss", hue="run").set_title("loss")
    plt.grid()
    plt.savefig(os.path.join(output_dir, 'Training_curves.png'))
//  [python]
def vis_training_curves(history=None, output_dir: str = None) -> None:
    """
    Visualizes the training curves of the model.

    Args:
        history: The history object returned by the model.fit() method.
        output_dir (Optional[str]): The output directory to save the training curves plot.

    Returns:
        None
    """
    # Extract the accuracy and loss values for training and validation data
    acc = history.history['accuracy']
    val_acc = history.history['val_accuracy']
    loss = history.history['loss']
    val_loss = history.history['val_loss']
    epochs_range = range(len(acc))

    # Create dataframes for the training and validation data
    df_val = pd.DataFrame({'run': 'validation', 'step': epochs_range, 'epoch_accuracy': val_acc, 'epoch_loss': val_loss})
    df_train = pd.DataFrame({'run': 'train', 'step': epochs_range, 'epoch_accuracy': acc, 'epoch_loss': loss})

    # Concatenate the dataframes
    frames = [df_val, df_train]
    df = pd.concat(frames)
    df = df.reset_index()

    # Plot the training curves
    plt.figure(figsize=(16, 6))
    plt.suptitle('Training curves', fontsize=16)
    plt.subplot(1, 2, 1)
    sns.lineplot(data=df, x="step", y="epoch_accuracy", hue="run").set_title("accuracy")
    plt.grid()
    plt.subplot(1, 2, 2)
    sns.lineplot(data=df, x="step", y="epoch_loss", hue="run").set_title("loss")
    plt.grid()
    plt.savefig(os.path.join(output_dir, 'Training_curves.png'))
//  [python]
def save_line_plot(filename, y, maintitle, subtitle, df):
    print("Creating " + filename)
    fig, ax = plt.subplots(figsize=(8, 6))
    sns.pointplot(x="Concurrent Users", y=y, hue="Scenario Name", data=df, markers=True, linestyles=["-", "-", "-"])
    ax.yaxis.set_major_formatter(tkr.FuncFormatter(lambda y, p: "{:,}".format(y)))
    plt.suptitle(maintitle)
    plt.title(subtitle)
    plt.savefig(filename)
    plt.clf()
    plt.close(fig)
//  [python]
def save_line_plot(filename, y, maintitle, subtitle, df):
    print("Creating " + filename)
    fig, ax = plt.subplots(figsize=(8, 6))
    sns.lineplot(x="Concurrent Users", y=y, hue="Scenario Name", data=df, markers=True, style="Scenario Name", dashes=False)
    ax.yaxis.set_major_formatter(tkr.FuncFormatter(lambda y, p: "{:,}".format(y)))
    plt.suptitle(maintitle)
    plt.title(subtitle)
    plt.savefig(filename)
    plt.clf()
    plt.close(fig)
//  [python]
def plot_generalization_lines(all_res_pd, data_cate, data_cate_name):
    import seaborn as sns
    from matplotlib import pyplot as plt
    import matplotlib.pylab as pylab

    plt.clf()
    sns.set()
    fig, axes = plt.subplots(1, 3, figsize=(6, 4))
    print(all_res_pd.columns.tolist())

    plot_data = all_res_pd.loc[all_res_pd["data"].isin(data_cate)]

    plot_data = plot_data.loc[plot_data["method"] != "Global-Train"]
    plot_data = plot_data.loc[plot_data["method"] != "Isolated"]
    plot_data = plot_data.loc[plot_data["method"] != "FedOpt"]
    plot_data = plot_data.loc[plot_data["method"] != "FedOpt-FT"]
    filter_out_methods = ["Global-Train", "Isolated", "FedOpt", "FedOpt-FT"]
    for i, metric in enumerate(column_names_generalization):
        plt.clf()
        sns.set()
        fig, axes = plt.subplots(1, 1, figsize=(2, 3))
        x = "data"
        if data_cate_name == "femnist_all":
            x = "s"
        if data_cate_name == "cifar10_all":
            x = "alpha"

        ax = sns.pointplot(
            ax=axes,
            data=plot_data,
            x=x,
            y=metric,
            hue="method",
            join=True,
            markers=True,
            hue_order=[
                m for m in expected_method_names if m not in filter_out_methods
            ],
            sort=True,
        )
        ax.set(ylabel=column_names_generalization_for_plot[i])
        plt.gca().invert_xaxis()

        if data_cate_name == "cifar10_all":
            ax.set_xscale('log')

        plt.legend(bbox_to_anchor=(1, 1), loc=2, ncol=2, borderaxespad=0.)
        plt.tight_layout()
        plt.savefig(f"generalization_all_{data_cate_name}_{i}.pdf",
                    bbox_inches='tight',
                    pad_inches=0)

        plt.show()
//  [python]
def plot_generalization_lines(all_res_pd, data_cate, data_cate_name):
    import seaborn as sns
    from matplotlib import pyplot as plt
    import matplotlib.pylab as pylab

    plt.clf()
    sns.set()
    fig, axes = plt.subplots(1, 3, figsize=(6, 4))
    print(all_res_pd.columns.tolist())

    plot_data = all_res_pd.loc[all_res_pd["data"].isin(data_cate)]

    plot_data = plot_data.loc[plot_data["method"] != "Global-Train"]
    plot_data = plot_data.loc[plot_data["method"] != "Isolated"]
    plot_data = plot_data.loc[plot_data["method"] != "FedOpt"]
    plot_data = plot_data.loc[plot_data["method"] != "FedOpt-FT"]
    filter_out_methods = ["Global-Train", "Isolated", "FedOpt", "FedOpt-FT"]
    for i, metric in enumerate(column_names_generalization):
        plt.clf()
        sns.set()
        fig, axes = plt.subplots(1, 1, figsize=(2, 3))
        x = "data"
        if data_cate_name == "femnist_all":
            x = "s"
        if data_cate_name == "cifar10_all":
            x = "alpha"

        ax = sns.lineplot(
            ax=axes,
            data=plot_data,
            x=x,
            y=metric,
            hue="method",
            style="method",
            markers=True,
            dashes=True,
            hue_order=[
                m for m in expected_method_names if m not in filter_out_methods
            ],
            sort=True,
        )
        ax.set(ylabel=column_names_generalization_for_plot[i])
        plt.gca().invert_xaxis()

        if data_cate_name == "cifar10_all":
            ax.set_xscale('log')

        plt.legend(bbox_to_anchor=(1, 1), loc=2, ncol=2, borderaxespad=0.)
        plt.tight_layout()
        plt.savefig(f"generalization_all_{data_cate_name}_{i}.pdf",
                    bbox_inches='tight',
                    pad_inches=0)

        plt.show()
//  [python]
def save_learning_curve_plot(
    save_img_path: str,
    training_losses: List,
    validation_losses: List,
) -> None:
    """create and save leanring curve plot

    Args:
        save_dir_path (str): Directory path to save figure
        training_losses (List): training_losses from train(...)
        validation_losses (List): validation_losses from train(...)
        validation_accuracy (Optional[List], optional): validation_accuracy from train(...). Defaults to None.

    Raises:
        ValueError: training_losses and validation_losses must be the same length
        ValueError: training_losses and validation_accuracy must be the same length

    Returns:
        str: Saved figure path.
    """
    if len(training_losses) != len(validation_losses):
        raise ValueError("train_losses and validation_losses must be the same length.")

    _, ax = plt.subplots(figsize=(8, 6))
    ax.set_xlabel("Epochs")
    x = [i for i in range(len(training_losses))]
    plt.plot(x, training_losses, label="Training Loss", color="tab:orange")
    plt.plot(x, validation_losses, label="Validation Loss", color="tab:blue")
    ax.set_ylabel("Training & Validation Loss")

    ax.legend(loc="upper center")
    plt.tight_layout()
    plt.savefig(save_img_path)
    plt.close()
//  [python]
def save_learning_curve_plot(
    save_img_path: str,
    training_losses: List,
    validation_losses: List,
) -> None:
    """create and save leanring curve plot

    Args:
        save_dir_path (str): Directory path to save figure
        training_losses (List): training_losses from train(...)
        validation_losses (List): validation_losses from train(...)
        validation_accuracy (Optional[List], optional): validation_accuracy from train(...). Defaults to None.

    Raises:
        ValueError: training_losses and validation_losses must be the same length
        ValueError: training_losses and validation_accuracy must be the same length

    Returns:
        str: Saved figure path.
    """
    if len(training_losses) != len(validation_losses):
        raise ValueError("train_losses and validation_losses must be the same length.")

    _, ax = plt.subplots(figsize=(8, 6))
    ax.set_xlabel("Epochs")
    x = [i for i in range(len(training_losses))]
    sns.lineplot(
        x=x, y=training_losses, label="Training Loss", ax=ax, color="tab:orange"
    )
    sns.lineplot(
        x=x, y=validation_losses, label="Validation Loss", ax=ax, color="tab:blue"
    )
    ax.set_ylabel("Training & Validation Loss")

    ax.legend(loc="upper center")
    plt.tight_layout()
    plt.savefig(save_img_path)
    plt.close()
//  [python]
def graph_compiled_data(compiled_data):
    compiled_data_to_graph = compiled_data[
        ["new_cases", "In hospital confirmed", "in_icu", "on_ventilator"]
    ].dropna()

    compiled_data_to_graph["date"] = compiled_data_to_graph.index

    melted_compiled = pd.melt(
        compiled_data_to_graph, id_vars=["date"], value_name="patients"
    )

    ax = seaborn.pointplot(x="date", y="patients", hue="variable", data=melted_compiled, join=True)
    handles, labels = ax.get_legend_handles_labels()
    ax.legend(handles=handles[1:], labels=labels[1:])

    plt.xticks(rotation=90)
//  [python]
def graph_compiled_data(compiled_data):
    compiled_data_to_graph = compiled_data[
        ["new_cases", "In hospital confirmed", "in_icu", "on_ventilator"]
    ].dropna()

    compiled_data_to_graph["date"] = compiled_data_to_graph.index

    melted_compiled = pd.melt(
        compiled_data_to_graph, id_vars=["date"], value_name="patients"
    )

    ax = seaborn.lineplot(x="date", y="patients", hue="variable", data=melted_compiled)
    handles, labels = ax.get_legend_handles_labels()
    ax.legend(handles=handles[1:], labels=labels[1:])

    plt.xticks(rotation=90)
//  [python]
steps = 30

percentage_cutoff = 0.1 # compute intdiv on top 10% molecules in the samples 


### Get QED distribution at different steps of CbAS

tanim_dist_all, tanim_dist_top = [],[]

for step in np.arange(1,steps+1):

    samples = pd.read_csv(f'../cbas/slurm/results/{name}/docking_results/{step}.csv')
    samples = samples.sort_values('score')
    N = int(samples.shape[0]*percentage_cutoff)
    
    smiles = samples.smile
    smiles = [s for s in smiles if Chem.MolFromSmiles(s) is not None]
    mols = [Chem.MolFromSmiles(s) for s in smiles]
    fps = [AllChem.GetMorganFingerprintAsBitVect(m , 3, nBits=2048) for m in mols]
    
    fps= np.array(fps)
    
    D= pairwise_distances(fps, metric = 'jaccard')
    D_top = D[:N,:N]
    
    tanim_dist_all.append(np.mean(D))
    tanim_dist_top.append(np.mean(D_top))
    
plt.plot(np.arange(1, step+1), tanim_dist_all, color = 'b', label = 'all samples')
//  [python]
steps = 30

percentage_cutoff = 0.1 # compute intdiv on top 10% molecules in the samples 


### Get QED distribution at different steps of CbAS

tanim_dist_all, tanim_dist_top = [],[]

for step in np.arange(1,steps+1):

    samples = pd.read_csv(f'../cbas/slurm/results/{name}/docking_results/{step}.csv')
    samples = samples.sort_values('score')
    N = int(samples.shape[0]*percentage_cutoff)
    
    smiles = samples.smile
    smiles = [s for s in smiles if Chem.MolFromSmiles(s) is not None]
    mols = [Chem.MolFromSmiles(s) for s in smiles]
    fps = [AllChem.GetMorganFingerprintAsBitVect(m , 3, nBits=2048) for m in mols]
    
    fps= np.array(fps)
    
    D= pairwise_distances(fps, metric = 'jaccard')
    D_top = D[:N,:N]
    
    tanim_dist_all.append(np.mean(D))
    tanim_dist_top.append(np.mean(D_top))
    
sns.lineplot(x=np.arange(1, step+1), y=tanim_dist_all, color = 'b', label = 'all samples')
//  [python]
def _plot_aggregation_evolution(df_plot, data_var, standard_name, units, mode='boxplot', ax=None, save_path=None,
                                show=False, aggregation_time='D', **kwargs):
    if ax is None:
        fig, ax = plt.subplots()
    if mode == 'boxplot':
        sns.boxplot(data=df_plot, x='aggregated_time', y=data_var, whis=2.5, ax=ax, **kwargs)
    elif mode == 'violin':
        sns.violinplot(data=df_plot, x='aggregated_time', y=data_var, ax=ax, **kwargs)
    elif mode == 'quantiles':
        if 'hue' in kwargs.keys():
            df_plot_list = df_plot.groupby(kwargs['hue'])
            kwargs.pop('hue')
        else:
            df_plot_list = [(0, df_plot)]
        for _, df_plot_i in df_plot_list:
            quantiles_plot = df_plot_i.groupby('aggregated_time').quantile([0.1, 0.5, 0.9], numeric_only=True)
            quantiles_plot = quantiles_plot.unstack()
            quantiles_plot = quantiles_plot[data_var]
            plt.plot(quantiles_plot.index, quantiles_plot[0.5], **kwargs)

            if 'color' in kwargs.keys():
                ax.fill_between(x=quantiles_plot.index,
                                y1=quantiles_plot[0.1].values,
                                y2=quantiles_plot[0.9].values,
                                alpha=0.2, color=kwargs['color'])
            else:
                ax.fill_between(x=quantiles_plot.index,
                                y1=quantiles_plot[0.1].values,
                                y2=quantiles_plot[0.9].values,
                                alpha=0.2)
    else:
        raise ValueError('mode %s is not implemented. Only boxplot, quantiles and violin' % mode)

    ax.set_xlabel('Time [%s]' % aggregation_time)
    ax.set_ylabel(r'%s [$%s$]' % (re.sub('_', ' ', standard_name).title(), units))
    ax.set_facecolor('white')
    plt.gca().xaxis.set_major_locator(mdates.AutoDateLocator())
    plt.xticks(rotation=45)
    plt.tight_layout()

    if save_path is not None:
        plt.savefig(save_path)
    if show:
        plt.show()

    return ax
//  [python]
def _plot_aggregation_evolution(df_plot, data_var, standard_name, units, mode='boxplot', ax=None, save_path=None,
                                show=False, aggregation_time='D', **kwargs):
    if ax is None:
        fig, ax = plt.subplots()
    if mode == 'boxplot':
        sns.boxplot(data=df_plot, x='aggregated_time', y=data_var, whis=2.5, ax=ax, **kwargs)
    elif mode == 'violin':
        sns.violinplot(data=df_plot, x='aggregated_time', y=data_var, ax=ax, **kwargs)
    elif mode == 'quantiles':
        if 'hue' in kwargs.keys():
            df_plot_list = df_plot.groupby(kwargs['hue'])
            kwargs.pop('hue')
        else:
            df_plot_list = [(0, df_plot)]
        for _, df_plot_i in df_plot_list:
            quantiles_plot = df_plot_i.groupby('aggregated_time').quantile([0.1, 0.5, 0.9], numeric_only=True)
            quantiles_plot = quantiles_plot.unstack()
            quantiles_plot = quantiles_plot[data_var]
            sns.lineplot(data=quantiles_plot, y=0.5, x='aggregated_time', ax=ax, **kwargs)

            if 'color' in kwargs.keys():
                ax.fill_between(x=quantiles_plot.index,
                                y1=quantiles_plot[0.1].values,
                                y2=quantiles_plot[0.9].values,
                                alpha=0.2, color=kwargs['color'])
            else:
                ax.fill_between(x=quantiles_plot.index,
                                y1=quantiles_plot[0.1].values,
                                y2=quantiles_plot[0.9].values,
                                alpha=0.2)
    else:
        raise ValueError('mode %s is not implemented. Only boxplot, quantiles and violin' % mode)

    ax.set_xlabel('Time [%s]' % aggregation_time)
    ax.set_ylabel(r'%s [$%s$]' % (re.sub('_', ' ', standard_name).title(), units))
    ax.set_facecolor('white')
    plt.gca().xaxis.set_major_locator(mdates.AutoDateLocator())
    plt.xticks(rotation=45)
    plt.tight_layout()

    if save_path is not None:
        plt.savefig(save_path)
    if show:
        plt.show()

    return ax
//  [python]
def plot_long_seq(
    data,
    ax,
    method_name,
    lineplot_kwargs=None,
    manual_average=True,
    smoothen=False,
    title="Average Success",
):
    if lineplot_kwargs is None:
        lineplot_kwargs = {}

    data = data.copy()
    if manual_average:
        cols = data.columns[data.columns.str.contains(f"test/stochastic/.*/success", regex=True)]
        data["average_success"] = data[cols].mean(axis=1)
        avg_success = data
    else:
        avg_success = data.rename(columns={"test/stochastic/average_success": "average_success"})
    if smoothen:
        avg_success = smoothen_long_seq(avg_success)

    avg_success["Average Success"] = avg_success["average_success"]
    avg_success["Method"] = avg_success["identifier"]
    avg_success = avg_success.replace({"Method": RENAMER})

    style = "half" if "half" in data.columns else None
    s.lineplot(
        data=avg_success,
        ax=ax,
        x="x",
        y="Average Success",
        hue="Method",
        style=style,
        **lineplot_kwargs,
    )

    plot_task_separators(data, steps_per_task=get_steps_per_task(data), ax=ax)
    ax.set_title(f"{title}, {method_name}")