initial test

Dockerfile

@@ -17,4 +17,4 @@ EXPOSE 8501
 
 HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
 
-ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]
+ENTRYPOINT ["streamlit", "run", "src/app.py", "--server.port=8501", "--server.address=0.0.0.0"]

requirements.txt

@@ -1,3 +1,6 @@
 altair
 pandas
-streamlit
+ssm-simulators
+streamlit>1.30.0
+matplotlib
+seaborn

src/app.py

@@ -0,0 +1,256 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+import streamlit as st
+
+# from hssm import simulate_data
+from ssms.config import model_config
+from ssms.basic_simulators.simulator import simulator
+import pandas as pd
+import utils
+
+
+# Function to create input select widgets
+def create_param_selectors(model_name: str, model_num: int = 1):
+    d_config = model_config[model_name]
+    params = d_config["params"]
+    param_bounds_low = d_config["param_bounds"][0]
+    param_bounds_high = d_config["param_bounds"][1]
+    param_defaults = d_config["default_params"]
+
+    d_param_slider = {}
+    for i, (name, low, high, default) in enumerate(
+        zip(
+            params,
+            param_bounds_low,
+            param_bounds_high,
+            param_defaults,
+        )
+    ):
+        d_param_slider[i] = st.slider(
+            label=name,
+            min_value=float(low),
+            max_value=float(high),
+            value=float(default),
+            key=f"param{i}"
+            f"_{model_name}"
+            f"_{model_num}"
+            f'_{st.session_state["slider_version"]}',
+        )
+    return d_param_slider
+
+
+def add_model():
+    pass
+
+
+def reset_sliders():
+    st.session_state["slider_version"] += 1
+
+
+# Initialize a slider version attribute state. Is used for resetting values
+if "slider_version" not in st.session_state:
+    st.session_state["slider_version"] = 1
+
+
+def draw_model_configurator(model_num=1):
+    # Create widgets for the sidebar
+    # st.markdown("<h2 style='text-align: center; color: black;'>Model Configurator</h1>",
+    #             unsafe_allow_html=True)
+    # Dropdown selection of model name
+    model_select = st.selectbox(
+        "Model " + str(model_num), l_model_names, key="modelname" + str(model_num)
+    )
+    # Sliders for param values
+    d_slider = create_param_selectors(model_select, model_num=model_num)
+    # Number of data points to simulate
+    nsamples = st.number_input("NSamples", value=5000, key="size" + str(model_num))
+    # Number of trajectories to show
+    ntrajectories = st.number_input(
+        "NTrajectories", value=0, key="ntraj" + str(model_num)
+    )
+    return model_select, d_slider, nsamples, ntrajectories
+
+
+st.set_page_config(layout="wide")
+
+# Get list of model names
+l_model_names = list(model_config.keys())
+
+with st.sidebar:
+    col1, col2 = st.columns(2)
+    with col1:
+        model_select_1, d_slider_1, nsamples_1, ntrajectories_1 = (
+            draw_model_configurator(model_num=1)
+        )
+    with col2:
+        model_select_2, d_slider_2, nsamples_2, ntrajectories_2 = (
+            draw_model_configurator(model_num=2)
+        )
+
+    # Button to reset sliders to default values
+    randomseed = st.number_input("RandomSeed", value=0, key="seed")
+    st.button(
+        "Reset",
+        help="Reset parameters to defaults",
+        key="reset",
+        on_click=reset_sliders,
+    )
+
+# st.title("SSM Model Plots", )
+st.markdown(
+    "<h1 style='text-align: center; color: black;'>SSM Model Plots</h1>",
+    unsafe_allow_html=True,
+)
+
+# Display components for main panel
+fig1, ax1 = plt.subplots()
+if model_config[model_select_1]["nchoices"] == 2 and not ("race" in model_select_1):
+    ax1 = utils.utils.plot_func_model(
+        model_name=model_select_1,
+        theta=[list(d_slider_1.values())],
+        axis=ax1,
+        value_range=[-0.1, 5],
+        n_samples=nsamples_1,
+        ylim=5,
+        data_color="blue",
+        add_trajectories=True,
+        n_trajectories=ntrajectories_1,
+        linewidth_model=1,
+        linewidth_histogram=1,
+        random_state=randomseed,
+    )
+else:
+    ax1 = utils.utils.plot_func_model_n(
+        model_name=model_select_1,
+        theta=[list(d_slider_1.values())],
+        axis=ax1,
+        value_range=[-0.1, 5],
+        n_samples=nsamples_1,
+        data_color="blue",
+        add_trajectories=True,
+        n_trajectories=ntrajectories_1,
+        linewidth_model=1,
+        linewidth_histogram=1,
+        random_state=randomseed,
+    )
+ax1.set_title("Model 1")
+ax1.set_xlabel("RT in seconds")
+
+fig2, ax2 = plt.subplots()
+if model_config[model_select_2]["nchoices"] == 2 and not ("race" in model_select_2):
+    ax2 = utils.utils.plot_func_model(
+        model_name=model_select_2,
+        theta=[list(d_slider_2.values())],
+        axis=ax2,
+        value_range=[-0.1, 5],
+        n_samples=nsamples_2,
+        ylim=5,
+        data_color="red",
+        add_trajectories=True,
+        n_trajectories=ntrajectories_2,
+        linewidth_model=1,
+        linewidth_histogram=1,
+        random_state=randomseed,
+    )
+else:
+    ax2 = utils.utils.plot_func_model_n(
+        model_name=model_select_2,
+        theta=[list(d_slider_2.values())],
+        axis=ax2,
+        value_range=[-0.1, 5],
+        n_samples=nsamples_2,
+        data_color="red",
+        add_trajectories=True,
+        n_trajectories=ntrajectories_2,
+        linewidth_model=1,
+        linewidth_histogram=1,
+        random_state=randomseed,
+    )
+
+ax2.set_title("Model 2")
+ax2.set_xlabel("RT in seconds")
+
+# Place figure in placeholder
+col1, col2 = st.columns(2)
+with col1:
+    figure_placeholder_1 = st.empty()  # Placeholder for figure render
+    figure_placeholder_1.pyplot(fig1)
+with col2:
+    figure_placeholder_2 = st.empty()  # Placeholder for figure render
+    figure_placeholder_2.pyplot(fig2)
+
+# Simulate two datasets:
+sim_output_1 = simulator(
+    model=model_select_1,
+    theta=[list(d_slider_1.values())],
+    n_samples=nsamples_1,
+    random_state=randomseed,
+)
+sim_output_2 = simulator(
+    model=model_select_2,
+    theta=[list(d_slider_2.values())],
+    n_samples=nsamples_2,
+    random_state=randomseed,
+)
+
+# Make metadata dataframe
+metadata = pd.DataFrame(
+    {
+        "Model 1": [
+            sim_output_1["metadata"]["model"],
+            sim_output_1["choice_p"][0, 0],
+            sim_output_1["rts"].mean(),
+            sim_output_1["metadata"]["s"],
+        ],
+        "Model 2": [
+            sim_output_2["metadata"]["model"],
+            sim_output_2["choice_p"][0, 0],
+            sim_output_2["rts"].mean(),
+            sim_output_2["metadata"]["s"],
+        ],
+    },
+    index=["Model", "Choice Probability", "Mean RT", "Noise SD"],
+)
+
+col3, col4 = st.columns(2)
+with col3:
+    if (
+        len(sim_output_1["metadata"]["possible_choices"])
+        == 2 | len(sim_output_2["metadata"]["possible_choices"])
+        == 2
+    ):
+        figure_placeholder_3 = st.empty()
+
+        # Plot the simulated data
+        fig3, ax3 = plt.subplots()
+        ax3.hist(
+            sim_output_1["rts"][np.abs(sim_output_1["rts"]) != 999]
+            * sim_output_1["choices"][np.abs(sim_output_1["rts"] != 999)],
+            histtype="step",
+            bins=50,
+            density=True,
+            color="blue",
+            fill=None,
+            label="Model 1",
+        )
+        ax3.hist(
+            sim_output_2["rts"][np.abs(sim_output_2["rts"]) != 999]
+            * sim_output_2["choices"][np.abs(sim_output_2["rts"] != 999)],
+            histtype="step",
+            bins=50,
+            density=True,
+            color="red",
+            fill=None,
+            label="Model 2",
+        )
+        ax3.legend()
+        ax3.set_xlabel("RT")
+        ax3.set_xlim(-5, 5)
+        figure_placeholder_3.pyplot(fig3)
+    else:
+        # TODO: Implement better comparison plot
+        # for models with more than 2 choice options
+        pass
+with col4:
+    st.dataframe(metadata)

src/utils/__init__.py

@@ -0,0 +1,3 @@
+from . import utils
+
+__all__ = ["utils"]

src/utils/old_plots.py

@@ -0,0 +1,1157 @@
+def _plot_func_model(
+    bottom_node,
+    axis,
+    value_range=None,
+    samples=10,
+    bin_size=0.05,
+    add_data_rts=True,
+    add_data_model=True,
+    add_data_model_keep_slope=True,
+    add_data_model_keep_boundary=True,
+    add_data_model_keep_ndt=True,
+    add_data_model_keep_starting_point=True,
+    add_data_model_markersize_starting_point=50,
+    add_data_model_markertype_starting_point=0,
+    add_data_model_markershift_starting_point=0,
+    add_posterior_uncertainty_model=False,
+    add_posterior_uncertainty_rts=False,
+    add_posterior_mean_model=True,
+    add_posterior_mean_rts=True,
+    add_trajectories=False,
+    data_label="Data",
+    secondary_data=None,
+    secondary_data_label=None,
+    secondary_data_color="blue",
+    linewidth_histogram=0.5,
+    linewidth_model=0.5,
+    legend_fontsize=12,
+    legend_shadow=True,
+    legend_location="upper right",
+    data_color="blue",
+    posterior_mean_color="red",
+    posterior_uncertainty_color="black",
+    alpha=0.05,
+    delta_t_model=0.01,
+    add_legend=True,  # keep_frame=False,
+    **kwargs,
+):
+    """Calculate posterior predictive for a certain bottom node.
+
+    Arguments:
+        bottom_node: pymc.stochastic
+            Bottom node to compute posterior over.
+
+        axis: matplotlib.axis
+            Axis to plot into.
+
+        value_range: numpy.ndarray
+            Range over which to evaluate the likelihood.
+
+    Optional:
+        samples: int <default=10>
+            Number of posterior samples to use.
+
+        bin_size: float <default=0.05>
+            Size of bins used for histograms
+
+        alpha: float <default=0.05>
+            alpha (transparency) level for the sample-wise elements of the plot
+
+        add_data_rts: bool <default=True>
+            Add data histogram of rts ?
+
+        add_data_model: bool <default=True>
+            Add model cartoon for data
+
+        add_posterior_uncertainty_rts: bool <default=True>
+            Add sample by sample histograms?
+
+        add_posterior_mean_rts: bool <default=True>
+            Add a mean posterior?
+
+        add_model: bool <default=True>
+            Whether to add model cartoons to the plot.
+
+        linewidth_histogram: float <default=0.5>
+            linewdith of histrogram plot elements.
+
+        linewidth_model: float <default=0.5>
+            linewidth of plot elements concerning the model cartoons.
+
+        legend_location: str <default='upper right'>
+            string defining legend position. Find the rest of the options in the matplotlib documentation.
+
+        legend_shadow: bool <default=True>
+            Add shadow to legend box?
+
+        legend_fontsize: float <default=12>
+            Fontsize of legend.
+
+        data_color : str <default="blue">
+            Color for the data part of the plot.
+
+        posterior_mean_color : str <default="red">
+            Color for the posterior mean part of the plot.
+
+        posterior_uncertainty_color : str <default="black">
+            Color for the posterior uncertainty part of the plot.
+
+        delta_t_model:
+            specifies plotting intervals for model cartoon elements of the graphs.
+    """
+
+    # AF-TODO: Add a mean version of this!
+    if value_range is None:
+        # Infer from data by finding the min and max from the nodes
+        raise NotImplementedError("value_range keyword argument must be supplied.")
+
+    if len(value_range) > 2:
+        value_range = (value_range[0], value_range[-1])
+
+    # Extract some parameters from kwargs
+    bins = np.arange(value_range[0], value_range[-1], bin_size)
+
+    # If bottom_node is a DataFrame we know that we are just plotting real data
+    if type(bottom_node) == pd.DataFrame:
+        samples_tmp = [bottom_node]
+        data_tmp = None
+    else:
+        samples_tmp = _post_pred_generate(
+            bottom_node,
+            samples=samples,
+            data=None,
+            append_data=False,
+            add_model_parameters=True,
+        )
+        data_tmp = bottom_node.value.copy()
+
+    # Relevant for recovery mode
+    node_data_full = kwargs.pop("node_data", None)
+
+    tmp_model = kwargs.pop("model_", "angle")
+    if len(model_config[tmp_model]["choices"]) > 2:
+        raise ValueError("The model plot works only for 2 choice models at the moment")
+
+    # ---------------------------
+
+    ylim = kwargs.pop("ylim", 3)
+    hist_bottom = kwargs.pop("hist_bottom", 2)
+    hist_histtype = kwargs.pop("hist_histtype", "step")
+
+    if ("ylim_high" in kwargs) and ("ylim_low" in kwargs):
+        ylim_high = kwargs["ylim_high"]
+        ylim_low = kwargs["ylim_low"]
+    else:
+        ylim_high = ylim
+        ylim_low = -ylim
+
+    if ("hist_bottom_high" in kwargs) and ("hist_bottom_low" in kwargs):
+        hist_bottom_high = kwargs["hist_bottom_high"]
+        hist_bottom_low = kwargs["hist_bottom_low"]
+    else:
+        hist_bottom_high = hist_bottom
+        hist_bottom_low = hist_bottom
+
+    axis.set_xlim(value_range[0], value_range[-1])
+    axis.set_ylim(ylim_low, ylim_high)
+    axis_twin_up = axis.twinx()
+    axis_twin_down = axis.twinx()
+    axis_twin_up.set_ylim(ylim_low, ylim_high)
+    axis_twin_up.set_yticks([])
+    axis_twin_down.set_ylim(ylim_high, ylim_low)
+    axis_twin_down.set_yticks([])
+    axis_twin_down.set_axis_off()
+    axis_twin_up.set_axis_off()
+
+    # ADD HISTOGRAMS
+    # -------------------------------
+    # POSTERIOR BASED HISTOGRAM
+    if add_posterior_uncertainty_rts:  # add_uc_rts:
+        j = 0
+        for sample in samples_tmp:
+            tmp_label = None
+
+            if add_legend and j == 0:
+                tmp_label = "PostPred"
+
+            weights_up = np.tile(
+                (1 / bin_size) / sample.shape[0],
+                reps=sample.loc[sample.response == 1, :].shape[0],
+            )
+            weights_down = np.tile(
+                (1 / bin_size) / sample.shape[0],
+                reps=sample.loc[(sample.response != 1), :].shape[0],
+            )
+
+            axis_twin_up.hist(
+                np.abs(sample.rt[sample.response == 1]),
+                bins=bins,
+                weights=weights_up,
+                histtype=hist_histtype,
+                bottom=hist_bottom_high,
+                alpha=alpha,
+                color=posterior_uncertainty_color,
+                edgecolor=posterior_uncertainty_color,
+                zorder=-1,
+                label=tmp_label,
+                linewidth=linewidth_histogram,
+            )
+
+            axis_twin_down.hist(
+                np.abs(sample.loc[(sample.response != 1), :].rt),
+                bins=bins,
+                weights=weights_down,
+                histtype=hist_histtype,
+                bottom=hist_bottom_low,
+                alpha=alpha,
+                color=posterior_uncertainty_color,
+                edgecolor=posterior_uncertainty_color,
+                linewidth=linewidth_histogram,
+                zorder=-1,
+            )
+            j += 1
+
+    if add_posterior_mean_rts:  # add_mean_rts:
+        concat_data = pd.concat(samples_tmp)
+        tmp_label = None
+
+        if add_legend:
+            tmp_label = "PostPred Mean"
+
+        weights_up = np.tile(
+            (1 / bin_size) / concat_data.shape[0],
+            reps=concat_data.loc[concat_data.response == 1, :].shape[0],
+        )
+        weights_down = np.tile(
+            (1 / bin_size) / concat_data.shape[0],
+            reps=concat_data.loc[(concat_data.response != 1), :].shape[0],
+        )
+
+        axis_twin_up.hist(
+            np.abs(concat_data.rt[concat_data.response == 1]),
+            bins=bins,
+            weights=weights_up,
+            histtype=hist_histtype,
+            bottom=hist_bottom_high,
+            alpha=1.0,
+            color=posterior_mean_color,
+            edgecolor=posterior_mean_color,
+            zorder=-1,
+            label=tmp_label,
+            linewidth=linewidth_histogram,
+        )
+
+        axis_twin_down.hist(
+            np.abs(concat_data.loc[(concat_data.response != 1), :].rt),
+            bins=bins,
+            weights=weights_down,
+            histtype=hist_histtype,
+            bottom=hist_bottom_low,
+            alpha=1.0,
+            color=posterior_mean_color,
+            edgecolor=posterior_mean_color,
+            linewidth=linewidth_histogram,
+            zorder=-1,
+        )
+
+    # DATA HISTOGRAM
+    if (data_tmp is not None) and add_data_rts:
+        tmp_label = None
+        if add_legend:
+            tmp_label = data_label
+
+        weights_up = np.tile(
+            (1 / bin_size) / data_tmp.shape[0],
+            reps=data_tmp[data_tmp.response == 1].shape[0],
+        )
+        weights_down = np.tile(
+            (1 / bin_size) / data_tmp.shape[0],
+            reps=data_tmp[(data_tmp.response != 1)].shape[0],
+        )
+
+        axis_twin_up.hist(
+            np.abs(data_tmp[data_tmp.response == 1].rt),
+            bins=bins,
+            weights=weights_up,
+            histtype=hist_histtype,
+            bottom=hist_bottom_high,
+            alpha=1,
+            color=data_color,
+            edgecolor=data_color,
+            label=tmp_label,
+            zorder=-1,
+            linewidth=linewidth_histogram,
+        )
+
+        axis_twin_down.hist(
+            np.abs(data_tmp[(data_tmp.response != 1)].rt),
+            bins=bins,
+            weights=weights_down,
+            histtype=hist_histtype,
+            bottom=hist_bottom_low,
+            alpha=1,
+            color=data_color,
+            edgecolor=data_color,
+            linewidth=linewidth_histogram,
+            zorder=-1,
+        )
+
+    # SECONDARY DATA HISTOGRAM
+    if secondary_data is not None:
+        tmp_label = None
+        if add_legend:
+            if secondary_data_label is not None:
+                tmp_label = secondary_data_label
+
+        weights_up = np.tile(
+            (1 / bin_size) / secondary_data.shape[0],
+            reps=secondary_data[secondary_data.response == 1].shape[0],
+        )
+        weights_down = np.tile(
+            (1 / bin_size) / secondary_data.shape[0],
+            reps=secondary_data[(secondary_data.response != 1)].shape[0],
+        )
+
+        axis_twin_up.hist(
+            np.abs(secondary_data[secondary_data.response == 1].rt),
+            bins=bins,
+            weights=weights_up,
+            histtype=hist_histtype,
+            bottom=hist_bottom_high,
+            alpha=1,
+            color=secondary_data_color,
+            edgecolor=secondary_data_color,
+            label=tmp_label,
+            zorder=-100,
+            linewidth=linewidth_histogram,
+        )
+
+        axis_twin_down.hist(
+            np.abs(secondary_data[(secondary_data.response != 1)].rt),
+            bins=bins,
+            weights=weights_down,
+            histtype=hist_histtype,
+            bottom=hist_bottom_low,
+            alpha=1,
+            color=secondary_data_color,
+            edgecolor=secondary_data_color,
+            linewidth=linewidth_histogram,
+            zorder=-100,
+        )
+    # -------------------------------
+
+    if add_legend:
+        if data_tmp is not None:
+            axis_twin_up.legend(
+                fontsize=legend_fontsize, shadow=legend_shadow, loc=legend_location
+            )
+
+    # ADD MODEL:
+    j = 0
+    t_s = np.arange(0, value_range[-1], delta_t_model)
+
+    # MAKE BOUNDS (FROM MODEL CONFIG) !
+    if add_posterior_uncertainty_model:  # add_uc_model:
+        for sample in samples_tmp:
+            _add_model_cartoon_to_ax(
+                sample=sample,
+                axis=axis,
+                tmp_model=tmp_model,
+                keep_slope=add_data_model_keep_slope,
+                keep_boundary=add_data_model_keep_boundary,
+                keep_ndt=add_data_model_keep_ndt,
+                keep_starting_point=add_data_model_keep_starting_point,
+                markersize_starting_point=add_data_model_markersize_starting_point,
+                markertype_starting_point=add_data_model_markertype_starting_point,
+                markershift_starting_point=add_data_model_markershift_starting_point,
+                delta_t_graph=delta_t_model,
+                sample_hist_alpha=alpha,
+                lw_m=linewidth_model,
+                tmp_label=tmp_label,
+                ylim_low=ylim_low,
+                ylim_high=ylim_high,
+                t_s=t_s,
+                color=posterior_uncertainty_color,
+                zorder_cnt=j,
+            )
+
+    if (node_data_full is not None) and add_data_model:
+        _add_model_cartoon_to_ax(
+            sample=node_data_full,
+            axis=axis,
+            tmp_model=tmp_model,
+            keep_slope=add_data_model_keep_slope,
+            keep_boundary=add_data_model_keep_boundary,
+            keep_ndt=add_data_model_keep_ndt,
+            keep_starting_point=add_data_model_keep_starting_point,
+            markersize_starting_point=add_data_model_markersize_starting_point,
+            markertype_starting_point=add_data_model_markertype_starting_point,
+            markershift_starting_point=add_data_model_markershift_starting_point,
+            delta_t_graph=delta_t_model,
+            sample_hist_alpha=1.0,
+            lw_m=linewidth_model + 0.5,
+            tmp_label=None,
+            ylim_low=ylim_low,
+            ylim_high=ylim_high,
+            t_s=t_s,
+            color=data_color,
+            zorder_cnt=j + 1,
+        )
+
+    if add_posterior_mean_model:  # add_mean_model:
+        tmp_label = None
+        if add_legend:
+            tmp_label = "PostPred Mean"
+
+        _add_model_cartoon_to_ax(
+            sample=pd.DataFrame(pd.concat(samples_tmp).mean().astype(np.float32)).T,
+            axis=axis,
+            tmp_model=tmp_model,
+            keep_slope=add_data_model_keep_slope,
+            keep_boundary=add_data_model_keep_boundary,
+            keep_ndt=add_data_model_keep_ndt,
+            keep_starting_point=add_data_model_keep_starting_point,
+            markersize_starting_point=add_data_model_markersize_starting_point,
+            markertype_starting_point=add_data_model_markertype_starting_point,
+            markershift_starting_point=add_data_model_markershift_starting_point,
+            delta_t_graph=delta_t_model,
+            sample_hist_alpha=1.0,
+            lw_m=linewidth_model + 0.5,
+            tmp_label=None,
+            ylim_low=ylim_low,
+            ylim_high=ylim_high,
+            t_s=t_s,
+            color=posterior_mean_color,
+            zorder_cnt=j + 2,
+        )
+
+    if add_trajectories:
+        _add_trajectories(
+            axis=axis,
+            sample=samples_tmp[0],
+            tmp_model=tmp_model,
+            t_s=t_s,
+            delta_t_graph=delta_t_model,
+            **kwargs,
+        )
+
+
+# AF-TODO: Add documentation for this function
+def _add_trajectories(
+    axis=None,
+    sample=None,
+    t_s=None,
+    delta_t_graph=0.01,
+    tmp_model=None,
+    n_trajectories=10,
+    supplied_trajectory=None,
+    maxid_supplied_trajectory=1,  # useful for gifs
+    highlight_trajectory_rt_choice=True,
+    markersize_trajectory_rt_choice=50,
+    markertype_trajectory_rt_choice="*",
+    markercolor_trajectory_rt_choice="red",
+    linewidth_trajectories=1,
+    alpha_trajectories=0.5,
+    color_trajectories="black",
+    **kwargs,
+):
+    # Check markercolor type
+    if type(markercolor_trajectory_rt_choice) == str:
+        markercolor_trajectory_rt_choice_dict = {}
+        for value_ in model_config[tmp_model]["choices"]:
+            markercolor_trajectory_rt_choice_dict[
+                value_
+            ] = markercolor_trajectory_rt_choice
+    elif type(markercolor_trajectory_rt_choice) == list:
+        cnt = 0
+        for value_ in model_config[tmp_model]["choices"]:
+            markercolor_trajectory_rt_choice_dict[
+                value_
+            ] = markercolor_trajectory_rt_choice[cnt]
+            cnt += 1
+    elif type(markercolor_trajectory_rt_choice) == dict:
+        markercolor_trajectory_rt_choice_dict = markercolor_trajectory_rt_choice
+    else:
+        pass
+
+    # Check trajectory color type
+    if type(color_trajectories) == str:
+        color_trajectories_dict = {}
+        for value_ in model_config[tmp_model]["choices"]:
+            color_trajectories_dict[value_] = color_trajectories
+    elif type(color_trajectories) == list:
+        cnt = 0
+        for value_ in model_config[tmp_model]["choices"]:
+            color_trajectories_dict[value_] = color_trajectories[cnt]
+            cnt += 1
+    elif type(color_trajectories) == dict:
+        color_trajectories_dict = color_trajectories
+    else:
+        pass
+
+    # Make bounds
+    (b_low, b_high) = _make_bounds(
+        tmp_model=tmp_model,
+        sample=sample,
+        delta_t_graph=delta_t_graph,
+        t_s=t_s,
+        return_shifted_by_ndt=False,
+    )
+
+    # Trajectories
+    if supplied_trajectory is None:
+        for i in range(n_trajectories):
+            rand_int = np.random.choice(400000000)
+            out_traj = simulator(
+                theta=sample[model_config[tmp_model]["params"]].values[0],
+                model=tmp_model,
+                n_samples=1,
+                no_noise=False,
+                delta_t=delta_t_graph,
+                bin_dim=None,
+                random_state=rand_int,
+            )
+
+            tmp_traj = out_traj[2]["trajectory"]
+            tmp_traj_choice = float(out_traj[1].flatten())
+            maxid = np.minimum(np.argmax(np.where(tmp_traj > -999)), t_s.shape[0])
+
+            # Identify boundary value at timepoint of crossing
+            b_tmp = b_high[maxid] if tmp_traj_choice > 0 else b_low[maxid]
+
+            axis.plot(
+                t_s[:maxid] + sample.t.values[0],
+                tmp_traj[:maxid],
+                color=color_trajectories_dict[tmp_traj_choice],
+                alpha=alpha_trajectories,
+                linewidth=linewidth_trajectories,
+                zorder=2000 + i,
+            )
+
+            if highlight_trajectory_rt_choice:
+                axis.scatter(
+                    t_s[maxid] + sample.t.values[0],
+                    b_tmp,
+                    # tmp_traj[maxid],
+                    markersize_trajectory_rt_choice,
+                    color=markercolor_trajectory_rt_choice_dict[tmp_traj_choice],
+                    alpha=1,
+                    marker=markertype_trajectory_rt_choice,
+                    zorder=2000 + i,
+                )
+
+    else:
+        if len(supplied_trajectory["trajectories"].shape) == 1:
+            supplied_trajectory["trajectories"] = np.expand_dims(
+                supplied_trajectory["trajectories"], axis=0
+            )
+
+        for j in range(supplied_trajectory["trajectories"].shape[0]):
+            maxid = np.minimum(
+                np.argmax(np.where(supplied_trajectory["trajectories"][j, :] > -999)),
+                t_s.shape[0],
+            )
+            if j == (supplied_trajectory["trajectories"].shape[0] - 1):
+                maxid_traj = min(maxid, maxid_supplied_trajectory)
+            else:
+                maxid_traj = maxid
+
+            axis.plot(
+                t_s[:maxid_traj] + sample.t.values[0],
+                supplied_trajectory["trajectories"][j, :maxid_traj],
+                color=color_trajectories_dict[
+                    supplied_trajectory["trajectory_choices"][j]
+                ],  # color_trajectories,
+                alpha=alpha_trajectories,
+                linewidth=linewidth_trajectories,
+                zorder=2000 + j,
+            )
+
+            # Identify boundary value at timepoint of crossing
+            b_tmp = (
+                b_high[maxid_traj]
+                if supplied_trajectory["trajectory_choices"][j] > 0
+                else b_low[maxid_traj]
+            )
+
+            if maxid_traj == maxid:
+                if highlight_trajectory_rt_choice:
+                    axis.scatter(
+                        t_s[maxid_traj] + sample.t.values[0],
+                        b_tmp,
+                        # supplied_trajectory['trajectories'][j, maxid_traj],
+                        markersize_trajectory_rt_choice,
+                        color=markercolor_trajectory_rt_choice_dict[
+                            supplied_trajectory["trajectory_choices"][j]
+                        ],  # markercolor_trajectory_rt_choice,
+                        alpha=1,
+                        marker=markertype_trajectory_rt_choice,
+                        zorder=2000 + j,
+                    )
+
+
+# AF-TODO: Add documentation to this function
+def _add_model_cartoon_to_ax(
+    sample=None,
+    axis=None,
+    tmp_model=None,
+    keep_slope=True,
+    keep_boundary=True,
+    keep_ndt=True,
+    keep_starting_point=True,
+    markersize_starting_point=80,
+    markertype_starting_point=1,
+    markershift_starting_point=-0.05,
+    delta_t_graph=None,
+    sample_hist_alpha=None,
+    lw_m=None,
+    tmp_label=None,
+    ylim_low=None,
+    ylim_high=None,
+    t_s=None,
+    zorder_cnt=1,
+    color="black",
+):
+    # Make bounds
+    b_low, b_high = _make_bounds(
+        tmp_model=tmp_model,
+        sample=sample,
+        delta_t_graph=delta_t_graph,
+        t_s=t_s,
+        return_shifted_by_ndt=True,
+    )
+
+    # MAKE SLOPES (VIA TRAJECTORIES HERE --> RUN NOISE FREE SIMULATIONS)!
+    out = simulator(
+        theta=sample[model_config[tmp_model]["params"]].values[0],
+        model=tmp_model,
+        n_samples=1,
+        no_noise=True,
+        delta_t=delta_t_graph,
+        bin_dim=None,
+    )
+
+    tmp_traj = out[2]["trajectory"]
+    maxid = np.minimum(np.argmax(np.where(tmp_traj > -999)), t_s.shape[0])
+
+    if "hddm_base" in tmp_model:
+        a_tmp = sample.a.values[0] / 2
+        tmp_traj = tmp_traj - a_tmp
+
+    if keep_boundary:
+        # Upper bound
+        axis.plot(
+            t_s,  # + sample.t.values[0],
+            b_high,
+            color=color,
+            alpha=sample_hist_alpha,
+            zorder=1000 + zorder_cnt,
+            linewidth=lw_m,
+            label=tmp_label,
+        )
+
+        # Lower bound
+        axis.plot(
+            t_s,  # + sample.t.values[0],
+            b_low,
+            color=color,
+            alpha=sample_hist_alpha,
+            zorder=1000 + zorder_cnt,
+            linewidth=lw_m,
+        )
+
+    # Slope
+    if keep_slope:
+        axis.plot(
+            t_s[:maxid] + sample.t.values[0],
+            tmp_traj[:maxid],
+            color=color,
+            alpha=sample_hist_alpha,
+            zorder=1000 + zorder_cnt,
+            linewidth=lw_m,
+        )  # TOOK AWAY LABEL
+
+    # Non-decision time
+    if keep_ndt:
+        axis.axvline(
+            x=sample.t.values[0],
+            ymin=ylim_low,
+            ymax=ylim_high,
+            color=color,
+            linestyle="--",
+            linewidth=lw_m,
+            zorder=1000 + zorder_cnt,
+            alpha=sample_hist_alpha,
+        )
+    # Starting point
+    if keep_starting_point:
+        axis.scatter(
+            sample.t.values[0] + markershift_starting_point,
+            b_low[0] + (sample.z.values[0] * (b_high[0] - b_low[0])),
+            markersize_starting_point,
+            marker=markertype_starting_point,
+            color=color,
+            alpha=sample_hist_alpha,
+            zorder=1000 + zorder_cnt,
+        )
+
+
+def _make_bounds(
+    tmp_model=None,
+    sample=None,
+    delta_t_graph=None,
+    t_s=None,
+    return_shifted_by_ndt=True,
+):
+    # MULTIPLICATIVE BOUND
+    if tmp_model == "weibull" or tmp_model == "weibull_cdf":
+        b = np.maximum(
+            sample.a.values[0]
+            * model_config[tmp_model]["boundary"](
+                t=t_s, alpha=sample.alpha.values[0], beta=sample.beta.values[0]
+            ),
+            0,
+        )
+
+        # Move boundary forward by the non-decision time
+        b_raw_high = deepcopy(b)
+        b_raw_low = deepcopy(-b)
+        b_init_val = b[0]
+        t_shift = np.arange(0, sample.t.values[0], delta_t_graph).shape[0]
+        b = np.roll(b, t_shift)
+        b[:t_shift] = b_init_val
+
+    # ADDITIVE BOUND
+    elif tmp_model == "angle":
+        b = np.maximum(
+            sample.a.values[0]
+            + model_config[tmp_model]["boundary"](t=t_s, theta=sample.theta.values[0]),
+            0,
+        )
+
+        b_raw_high = deepcopy(b)
+        b_raw_low = deepcopy(-b)
+        # Move boundary forward by the non-decision time
+        b_init_val = b[0]
+        t_shift = np.arange(0, sample.t.values[0], delta_t_graph).shape[0]
+        b = np.roll(b, t_shift)
+        b[:t_shift] = b_init_val
+
+    # CONSTANT BOUND
+    elif (
+        tmp_model == "ddm"
+        or tmp_model == "ornstein"
+        or tmp_model == "levy"
+        or tmp_model == "full_ddm"
+        or tmp_model == "ddm_hddm_base"
+        or tmp_model == "full_ddm_hddm_base"
+    ):
+        b = sample.a.values[0] * np.ones(t_s.shape[0])
+
+        if "hddm_base" in tmp_model:
+            b = (sample.a.values[0] / 2) * np.ones(t_s.shape[0])
+
+        b_raw_high = b
+        b_raw_low = -b
+
+    # Separate out upper and lower bound:
+    b_high = b
+    b_low = -b
+
+    if return_shifted_by_ndt:
+        return (b_low, b_high)
+    else:
+        return (b_raw_low, b_raw_high)
+
+
+def _plot_func_model_n(
+    bottom_node,
+    axis,
+    value_range=None,
+    samples=10,
+    bin_size=0.05,
+    add_posterior_uncertainty_model=False,
+    add_posterior_uncertainty_rts=False,
+    add_posterior_mean_model=True,
+    add_posterior_mean_rts=True,
+    linewidth_histogram=0.5,
+    linewidth_model=0.5,
+    legend_fontsize=7,
+    legend_shadow=True,
+    legend_location="upper right",
+    delta_t_model=0.01,
+    add_legend=True,
+    alpha=0.01,
+    keep_frame=False,
+    **kwargs,
+):
+    """Calculate posterior predictive for a certain bottom node.
+
+    Arguments:
+        bottom_node: pymc.stochastic
+            Bottom node to compute posterior over.
+
+        axis: matplotlib.axis
+            Axis to plot into.
+
+        value_range: numpy.ndarray
+            Range over which to evaluate the likelihood.
+
+    Optional:
+        samples: int <default=10>
+            Number of posterior samples to use.
+
+        bin_size: float <default=0.05>
+            Size of bins used for histograms
+
+        alpha: float <default=0.05>
+            alpha (transparency) level for the sample-wise elements of the plot
+
+        add_posterior_uncertainty_rts: bool <default=True>
+            Add sample by sample histograms?
+
+        add_posterior_mean_rts: bool <default=True>
+            Add a mean posterior?
+
+        add_model: bool <default=True>
+            Whether to add model cartoons to the plot.
+
+        linewidth_histogram: float <default=0.5>
+            linewdith of histrogram plot elements.
+
+        linewidth_model: float <default=0.5>
+            linewidth of plot elements concerning the model cartoons.
+
+        legend_loc: str <default='upper right'>
+            string defining legend position. Find the rest of the options in the matplotlib documentation.
+
+        legend_shadow: bool <default=True>
+            Add shadow to legend box?
+
+        legend_fontsize: float <default=12>
+            Fontsize of legend.
+
+        data_color : str <default="blue">
+            Color for the data part of the plot.
+
+        posterior_mean_color : str <default="red">
+            Color for the posterior mean part of the plot.
+
+        posterior_uncertainty_color : str <default="black">
+            Color for the posterior uncertainty part of the plot.
+
+
+        delta_t_model:
+            specifies plotting intervals for model cartoon elements of the graphs.
+    """
+
+    color_dict = {
+        -1: "black",
+        0: "black",
+        1: "green",
+        2: "blue",
+        3: "red",
+        4: "orange",
+        5: "purple",
+        6: "brown",
+    }
+
+    # AF-TODO: Add a mean version of this !
+    if value_range is None:
+        # Infer from data by finding the min and max from the nodes
+        raise NotImplementedError("value_range keyword argument must be supplied.")
+
+    if len(value_range) > 2:
+        value_range = (value_range[0], value_range[-1])
+
+    # Extract some parameters from kwargs
+    bins = np.arange(value_range[0], value_range[-1], bin_size)
+
+    # Relevant for recovery mode
+    node_data_full = kwargs.pop("node_data", None)
+    tmp_model = kwargs.pop("model_", "angle")
+
+    bottom = 0
+    # ------------
+    ylim = kwargs.pop("ylim", 3)
+
+    choices = model_config[tmp_model]["choices"]
+
+    # If bottom_node is a DataFrame we know that we are just plotting real data
+    if type(bottom_node) == pd.DataFrame:
+        samples_tmp = [bottom_node]
+        data_tmp = None
+    else:
+        samples_tmp = _post_pred_generate(
+            bottom_node,
+            samples=samples,
+            data=None,
+            append_data=False,
+            add_model_parameters=True,
+        )
+        data_tmp = bottom_node.value.copy()
+
+    axis.set_xlim(value_range[0], value_range[-1])
+    axis.set_ylim(0, ylim)
+
+    # ADD MODEL:
+    j = 0
+    t_s = np.arange(0, value_range[-1], delta_t_model)
+
+    # # MAKE BOUNDS (FROM MODEL CONFIG) !
+    if add_posterior_uncertainty_model:  # add_uc_model:
+        for sample in samples_tmp:
+            tmp_label = None
+
+            if add_legend and (j == 0):
+                tmp_label = "PostPred"
+
+            _add_model_n_cartoon_to_ax(
+                sample=sample,
+                axis=axis,
+                tmp_model=tmp_model,
+                delta_t_graph=delta_t_model,
+                sample_hist_alpha=alpha,
+                lw_m=linewidth_model,
+                tmp_label=tmp_label,
+                linestyle="-",
+                ylim=ylim,
+                t_s=t_s,
+                color_dict=color_dict,
+                zorder_cnt=j,
+            )
+
+            j += 1
+
+    if add_posterior_mean_model:  # add_mean_model:
+        tmp_label = None
+        if add_legend:
+            tmp_label = "PostPred Mean"
+
+        bottom = _add_model_n_cartoon_to_ax(
+            sample=pd.DataFrame(pd.concat(samples_tmp).mean().astype(np.float32)).T,
+            axis=axis,
+            tmp_model=tmp_model,
+            delta_t_graph=delta_t_model,
+            sample_hist_alpha=1.0,
+            lw_m=linewidth_model + 0.5,
+            linestyle="-",
+            tmp_label=None,
+            ylim=ylim,
+            t_s=t_s,
+            color_dict=color_dict,
+            zorder_cnt=j + 2,
+        )
+
+    if node_data_full is not None:
+        _add_model_n_cartoon_to_ax(
+            sample=node_data_full,
+            axis=axis,
+            tmp_model=tmp_model,
+            delta_t_graph=delta_t_model,
+            sample_hist_alpha=1.0,
+            lw_m=linewidth_model + 0.5,
+            linestyle="dashed",
+            tmp_label=None,
+            ylim=ylim,
+            t_s=t_s,
+            color_dict=color_dict,
+            zorder_cnt=j + 1,
+        )
+
+    # ADD HISTOGRAMS
+    # -------------------------------
+
+    # POSTERIOR BASED HISTOGRAM
+    if add_posterior_uncertainty_rts:  # add_uc_rts:
+        j = 0
+        for sample in samples_tmp:
+            for choice in choices:
+                tmp_label = None
+
+                if add_legend and j == 0:
+                    tmp_label = "PostPred"
+
+                weights = np.tile(
+                    (1 / bin_size) / sample.shape[0],
+                    reps=sample.loc[sample.response == choice, :].shape[0],
+                )
+
+                axis.hist(
+                    np.abs(sample.rt[sample.response == choice]),
+                    bins=bins,
+                    bottom=bottom,
+                    weights=weights,
+                    histtype="step",
+                    alpha=alpha,
+                    color=color_dict[choice],
+                    zorder=-1,
+                    label=tmp_label,
+                    linewidth=linewidth_histogram,
+                )
+                j += 1
+
+    if add_posterior_mean_rts:
+        concat_data = pd.concat(samples_tmp)
+        for choice in choices:
+            tmp_label = None
+            if add_legend and (choice == choices[0]):
+                tmp_label = "PostPred Mean"
+
+            weights = np.tile(
+                (1 / bin_size) / concat_data.shape[0],
+                reps=concat_data.loc[concat_data.response == choice, :].shape[0],
+            )
+
+            axis.hist(
+                np.abs(concat_data.rt[concat_data.response == choice]),
+                bins=bins,
+                bottom=bottom,
+                weights=weights,
+                histtype="step",
+                alpha=1.0,
+                color=color_dict[choice],
+                zorder=-1,
+                label=tmp_label,
+                linewidth=linewidth_histogram,
+            )
+
+    # DATA HISTOGRAM
+    if data_tmp is not None:
+        for choice in choices:
+            tmp_label = None
+            if add_legend and (choice == choices[0]):
+                tmp_label = "Data"
+
+            weights = np.tile(
+                (1 / bin_size) / data_tmp.shape[0],
+                reps=data_tmp.loc[data_tmp.response == choice, :].shape[0],
+            )
+
+            axis.hist(
+                np.abs(data_tmp.rt[data_tmp.response == choice]),
+                bins=bins,
+                bottom=bottom,
+                weights=weights,
+                histtype="step",
+                linestyle="dashed",
+                alpha=1.0,
+                color=color_dict[choice],
+                edgecolor=color_dict[choice],
+                zorder=-1,
+                label=tmp_label,
+                linewidth=linewidth_histogram,
+            )
+    # -------------------------------
+
+    if add_legend:
+        if data_tmp is not None:
+            custom_elems = [
+                Line2D([0], [0], color=color_dict[choice], lw=1) for choice in choices
+            ]
+            custom_titles = ["response: " + str(choice) for choice in choices]
+
+            custom_elems.append(
+                Line2D([0], [0], color="black", lw=1.0, linestyle="dashed")
+            )
+            custom_elems.append(Line2D([0], [0], color="black", lw=1.0, linestyle="-"))
+            custom_titles.append("Data")
+            custom_titles.append("Posterior")
+
+            axis.legend(
+                custom_elems,
+                custom_titles,
+                fontsize=legend_fontsize,
+                shadow=legend_shadow,
+                loc=legend_location,
+            )
+
+    # FRAME
+    if not keep_frame:
+        axis.set_frame_on(False)
+
+
+def _add_model_n_cartoon_to_ax(
+    sample=None,
+    axis=None,
+    tmp_model=None,
+    delta_t_graph=None,
+    sample_hist_alpha=None,
+    lw_m=None,
+    linestyle="-",
+    tmp_label=None,
+    ylim=None,
+    t_s=None,
+    zorder_cnt=1,
+    color_dict=None,
+):
+    if "weibull" in tmp_model:
+        b = np.maximum(
+            sample.a.values[0]
+            * model_config[tmp_model]["boundary"](
+                t=t_s, alpha=sample.alpha.values[0], beta=sample.beta.values[0]
+            ),
+            0,
+        )
+
+    elif "angle" in tmp_model:
+        b = np.maximum(
+            sample.a.values[0]
+            + model_config[tmp_model]["boundary"](t=t_s, theta=sample.theta.values[0]),
+            0,
+        )
+
+    else:
+        b = sample.a.values[0] * np.ones(t_s.shape[0])
+
+    # Upper bound
+    axis.plot(
+        t_s + sample.t.values[0],
+        b,
+        color="black",
+        alpha=sample_hist_alpha,
+        zorder=1000 + zorder_cnt,
+        linewidth=lw_m,
+        linestyle=linestyle,
+        label=tmp_label,
+    )
+
+    # Starting point
+    axis.axvline(
+        x=sample.t.values[0],
+        ymin=-ylim,
+        ymax=ylim,
+        color="black",
+        linestyle=linestyle,
+        linewidth=lw_m,
+        alpha=sample_hist_alpha,
+    )
+
+    # # MAKE SLOPES (VIA TRAJECTORIES HERE --> RUN NOISE FREE SIMULATIONS)!
+    out = simulator(
+        theta=sample[model_config[tmp_model]["params"]].values[0],
+        model=tmp_model,
+        n_samples=1,
+        no_noise=True,
+        delta_t=delta_t_graph,
+        bin_dim=None,
+    )
+
+    # # AF-TODO: Add trajectories
+    tmp_traj = out[2]["trajectory"]
+
+    for i in range(len(model_config[tmp_model]["choices"])):
+        tmp_maxid = np.minimum(np.argmax(np.where(tmp_traj[:, i] > -999)), t_s.shape[0])
+
+        # Slope
+        axis.plot(
+            t_s[:tmp_maxid] + sample.t.values[0],
+            tmp_traj[:tmp_maxid, i],
+            color=color_dict[i],
+            linestyle=linestyle,
+            alpha=sample_hist_alpha,
+            zorder=1000 + zorder_cnt,
+            linewidth=lw_m,
+        )  # TOOK AWAY LABEL
+
+    return b[0]

src/utils/utils.py

@@ -0,0 +1,820 @@
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from ssms.config import model_config
+from ssms.basic_simulators.simulator import simulator
+from matplotlib.lines import Line2D
+
+
+def plot_func_model(
+    model_name,
+    theta,
+    axis,
+    value_range=None,
+    n_samples=10,
+    bin_size=0.05,
+    add_data_rts=True,
+    add_data_model_keep_slope=True,
+    add_data_model_keep_boundary=True,
+    add_data_model_keep_ndt=True,
+    add_data_model_keep_starting_point=True,
+    add_data_model_markersize_starting_point=50,
+    add_data_model_markertype_starting_point=0,
+    add_data_model_markershift_starting_point=0,
+    n_trajectories = 0,
+    linewidth_histogram=0.5,
+    linewidth_model=0.5,
+    legend_fontsize=12,
+    legend_shadow=True,
+    legend_location="upper right",
+    data_color="blue",
+    posterior_uncertainty_color="black",
+    alpha=0.05,
+    delta_t_model=0.001,
+    random_state=None,
+    add_legend=True,  # keep_frame=False,
+    **kwargs,
+):
+    """Calculate posterior predictive for a certain bottom node.
+
+    Arguments:
+        bottom_node: pymc.stochastic
+            Bottom node to compute posterior over.
+
+        axis: matplotlib.axis
+            Axis to plot into.
+
+        value_range: numpy.ndarray
+            Range over which to evaluate the likelihood.
+
+    Optional:
+        samples: int <default=10>
+            Number of posterior samples to use.
+
+        bin_size: float <default=0.05>
+            Size of bins used for histograms
+
+        alpha: float <default=0.05>
+            alpha (transparency) level for the sample-wise elements of the plot
+
+        add_data_rts: bool <default=True>
+            Add data histogram of rts ?
+
+        add_data_model: bool <default=True>
+            Add model cartoon for data
+
+        add_posterior_uncertainty_rts: bool <default=True>
+            Add sample by sample histograms?
+
+        add_posterior_mean_rts: bool <default=True>
+            Add a mean posterior?
+
+        add_model: bool <default=True>
+            Whether to add model cartoons to the plot.
+
+        linewidth_histogram: float <default=0.5>
+            linewdith of histrogram plot elements.
+
+        linewidth_model: float <default=0.5>
+            linewidth of plot elements concerning the model cartoons.
+
+        legend_location: str <default='upper right'>
+            string defining legend position. Find the rest of the options in the matplotlib documentation.
+
+        legend_shadow: bool <default=True>
+            Add shadow to legend box?
+
+        legend_fontsize: float <default=12>
+            Fontsize of legend.
+
+        data_color : str <default="blue">
+            Color for the data part of the plot.
+
+        posterior_mean_color : str <default="red">
+            Color for the posterior mean part of the plot.
+
+        posterior_uncertainty_color : str <default="black">
+            Color for the posterior uncertainty part of the plot.
+
+        delta_t_model:
+            specifies plotting intervals for model cartoon elements of the graphs.
+    """
+
+    if value_range is None:
+        # Infer from data by finding the min and max from the nodes
+        raise NotImplementedError("value_range keyword argument must be supplied.")
+
+    if len(value_range) > 2:
+        value_range = (value_range[0], value_range[-1])
+
+    # Extract some parameters from kwargs
+    bins = np.arange(value_range[0], value_range[-1], bin_size)
+
+    if model_config[model_name]["nchoices"] > 2:
+        raise ValueError("The model plot works only for 2 choice models at the moment")
+
+    # RUN SIMULATIONS
+    # -------------------------------
+
+    # Simulator Data
+    if random_state is not None:
+        np.random.seed(random_state)
+    
+    rand_int = np.random.choice(400000000)
+    sim_out = simulator(model = model_name, theta = theta, n_samples = n_samples, 
+              no_noise = False, delta_t = 0.001, 
+              bin_dim = None, random_state = rand_int)
+    
+    sim_out_traj = {}
+    for i in range(n_trajectories):
+        rand_int = np.random.choice(400000000)
+        sim_out_traj[i] = simulator(model = model_name, theta = theta, n_samples = 1, 
+                                    no_noise = False, delta_t = 0.001, 
+                                    bin_dim = None, random_state = rand_int, smooth_unif = False)
+
+    sim_out_no_noise = simulator(model = model_name, theta = theta, n_samples = 1, 
+                                 no_noise = True, delta_t = 0.001, 
+                                 bin_dim = None, smooth_unif = False)
+
+    # ADD DATA HISTOGRAMS
+    weights_up = np.tile(
+        (1 / bin_size) / sim_out['rts'][(sim_out['rts'] != -999)].shape[0],
+        reps=sim_out['rts'][(sim_out['rts'] != -999) & (sim_out['choices'] == 1)].shape[0],
+    )
+    weights_down = np.tile(
+        (1 / bin_size) / sim_out['rts'][(sim_out['rts'] != -999)].shape[0],
+        reps=sim_out['rts'][(sim_out['rts'] != -999) & (sim_out['choices'] != 1)].shape[0],
+    )
+
+    (b_high, b_low) = (np.maximum(sim_out['metadata']['boundary'], 0), 
+                       np.minimum((-1) * sim_out['metadata']['boundary'], 0))
+
+    # ADD HISTOGRAMS
+    # -------------------------------
+
+    ylim = kwargs.pop("ylim", 3)
+    #hist_bottom = kwargs.pop("hist_bottom", 2)
+    hist_histtype = kwargs.pop("hist_histtype", "step")
+
+    if ("ylim_high" in kwargs) and ("ylim_low" in kwargs):
+        ylim_high = kwargs["ylim_high"]
+        ylim_low = kwargs["ylim_low"]
+    else:
+        ylim_high = ylim
+        ylim_low = -ylim
+
+    if ("hist_bottom_high" in kwargs) and ("hist_bottom_low" in kwargs):
+        hist_bottom_high = kwargs["hist_bottom_high"]
+        hist_bottom_low = kwargs["hist_bottom_low"]
+    else:
+        hist_bottom_high = b_high[0] #hist_bottom
+        hist_bottom_low = -b_low[0] #hist_bottom
+
+    axis.set_xlim(value_range[0], value_range[-1])
+    axis.set_ylim(ylim_low, ylim_high)
+    axis_twin_up = axis.twinx()
+    axis_twin_down = axis.twinx()
+    axis_twin_up.set_ylim(ylim_low, ylim_high)
+    axis_twin_up.set_yticks([])
+    axis_twin_down.set_ylim(ylim_high, ylim_low)
+    axis_twin_down.set_yticks([])
+    axis_twin_down.set_axis_off()
+    axis_twin_up.set_axis_off()
+
+    axis_twin_up.hist(
+        np.abs(sim_out['rts'][(sim_out['rts'] != -999) & (sim_out['choices'] == 1)]),
+        bins=bins,
+        weights=weights_up,
+        histtype=hist_histtype,
+        bottom=hist_bottom_high,
+        alpha=1,
+        color=data_color,
+        edgecolor=data_color,
+        linewidth=linewidth_histogram,
+        zorder=-1,
+    )
+
+    axis_twin_down.hist(
+        np.abs(sim_out['rts'][(sim_out['rts'] != -999) & (sim_out['choices'] != 1)]),
+        bins=bins,
+        weights=weights_down,
+        histtype=hist_histtype,
+        bottom=hist_bottom_low,
+        alpha=1,
+        color=data_color,
+        edgecolor=data_color,
+        linewidth=linewidth_histogram,
+        zorder=-1,
+    )
+
+    # ADD MODEL:
+    j = 0
+    t_s = np.arange(0, sim_out['metadata']['max_t'], delta_t_model) #value_range[-1], delta_t_model)
+
+    _add_model_cartoon_to_ax(
+        sample=sim_out_no_noise,
+        axis=axis,
+        keep_slope=add_data_model_keep_slope,
+        keep_boundary=add_data_model_keep_boundary,
+        keep_ndt=add_data_model_keep_ndt,
+        keep_starting_point=add_data_model_keep_starting_point,
+        markersize_starting_point=add_data_model_markersize_starting_point,
+        markertype_starting_point=add_data_model_markertype_starting_point,
+        markershift_starting_point=add_data_model_markershift_starting_point,
+        delta_t_graph=delta_t_model,
+        sample_hist_alpha=alpha,
+        lw_m=linewidth_model,
+        ylim_low=ylim_low,
+        ylim_high=ylim_high,
+        t_s=t_s,
+        color=posterior_uncertainty_color,
+        zorder_cnt=j,
+    )
+
+    if n_trajectories > 0:
+        _add_trajectories(
+            axis=axis,
+            sample=sim_out_traj,
+            t_s=t_s,
+            delta_t_graph=delta_t_model,
+            n_trajectories=n_trajectories,
+            **kwargs,
+        )
+    
+    return axis
+
+# AF-TODO: Add documentation for this function
+def _add_trajectories(
+    axis=None,
+    sample=None,
+    t_s=None,
+    delta_t_graph=0.01,
+    n_trajectories=10,
+    supplied_trajectory=None,
+    maxid_supplied_trajectory=1,  # useful for gifs
+    highlight_trajectory_rt_choice=True,
+    markersize_trajectory_rt_choice=50,
+    markertype_trajectory_rt_choice="*",
+    markercolor_trajectory_rt_choice="red",
+    linewidth_trajectories=1,
+    alpha_trajectories=0.5,
+    color_trajectories="black",
+    **kwargs,
+    ):
+    """Add trajectories to a given axis."""
+    # Check markercolor type
+    if isinstance(markercolor_trajectory_rt_choice, str):
+        markercolor_trajectory_rt_choice_dict = {}
+        for value_ in sample[0]['metadata']['possible_choices']:
+            markercolor_trajectory_rt_choice_dict[
+                value_
+            ] = markercolor_trajectory_rt_choice
+    elif isinstance(markercolor_trajectory_rt_choice, list):
+        cnt = 0
+        for value_ in sample[0]['metadata']['possible_choices']:
+            markercolor_trajectory_rt_choice_dict[
+                value_
+            ] = markercolor_trajectory_rt_choice[cnt]
+            cnt += 1
+    elif isinstance(markercolor_trajectory_rt_choice, dict):
+        markercolor_trajectory_rt_choice_dict = markercolor_trajectory_rt_choice
+    else:
+        pass
+
+    # Check trajectory color type
+    if isinstance(color_trajectories, str):
+        color_trajectories_dict = {}
+        for value_ in sample[0]['metadata']['possible_choices']:
+            color_trajectories_dict[value_] = color_trajectories
+    elif isinstance(color_trajectories, list):
+        cnt = 0
+        for value_ in sample[0]['metadata']['possible_choices']:
+            color_trajectories_dict[value_] = color_trajectories[cnt]
+            cnt += 1
+    elif isinstance(color_trajectories, dict):
+        color_trajectories_dict = color_trajectories
+    else:
+        pass
+
+    # Make bounds
+    (b_high, b_low) = (np.maximum(sample[0]['metadata']['boundary'], 0), 
+                       np.minimum((-1) * sample[0]['metadata']['boundary'], 0))
+    
+    b_h_init = b_high[0]
+    b_l_init = b_low[0]
+    n_roll = int((sample[0]['metadata']['t'][0] / delta_t_graph) + 1)
+    b_high = np.roll(b_high, n_roll)
+    b_high[:n_roll] = b_h_init
+    b_low = np.roll(b_low, n_roll)
+    b_low[:n_roll] = b_l_init
+
+    print(n_trajectories)
+    # Trajectories
+    for i in range(n_trajectories):
+        tmp_traj = sample[i]['metadata']['trajectory']
+        tmp_traj_choice = float(sample[i]['choices'].flatten())
+        maxid = np.minimum(np.argmax(np.where(tmp_traj > -999)), t_s.shape[0])
+
+        # Identify boundary value at timepoint of crossing
+        b_tmp = b_high[maxid + n_roll] if tmp_traj_choice > 0 else b_low[maxid + n_roll]
+
+        axis.plot(
+            t_s[:maxid] + sample[i]['metadata']['t'][0], #sample.t.values[0],
+            tmp_traj[:maxid],
+            color=color_trajectories_dict[tmp_traj_choice],
+            alpha=alpha_trajectories,
+            linewidth=linewidth_trajectories,
+            zorder=2000 + i,
+        )
+
+        if highlight_trajectory_rt_choice:
+            axis.scatter(
+                t_s[maxid] + sample[i]['metadata']['t'][0], #sample.t.values[0],
+                b_tmp,
+                # tmp_traj[maxid],
+                markersize_trajectory_rt_choice,
+                color=markercolor_trajectory_rt_choice_dict[tmp_traj_choice],
+                alpha=1,
+                marker=markertype_trajectory_rt_choice,
+                zorder=2000 + i,
+            )
+
+# AF-TODO: Add documentation to this function
+def _add_model_cartoon_to_ax(
+    sample=None,
+    axis=None,
+    keep_slope=True,
+    keep_boundary=True,
+    keep_ndt=True,
+    keep_starting_point=True,
+    markersize_starting_point=80,
+    markertype_starting_point=1,
+    markershift_starting_point=-0.05,
+    delta_t_graph=None,
+    sample_hist_alpha=None,
+    lw_m=None,
+    tmp_label=None,
+    ylim_low=None,
+    ylim_high=None,
+    t_s=None,
+    zorder_cnt=1,
+    color="black",
+):
+    # Make bounds
+    (b_high, b_low) = (np.maximum(sample['metadata']['boundary'], 0), 
+                       np.minimum((-1) * sample['metadata']['boundary'], 0))
+
+    b_h_init = b_high[0]
+    b_l_init = b_low[0]
+    n_roll = int((sample['metadata']['t'][0] / delta_t_graph) + 1)
+    b_high = np.roll(b_high, n_roll)
+    b_high[:n_roll] = b_h_init
+    b_low = np.roll(b_low, n_roll)
+    b_low[:n_roll] = b_l_init
+
+    tmp_traj = sample["metadata"]["trajectory"]
+    maxid = np.minimum(np.argmax(np.where(tmp_traj > -999)),
+                       t_s.shape[0])
+
+    if keep_boundary:
+        # Upper bound
+        axis.plot(
+            t_s,  # + sample.t.values[0],
+            b_high[:t_s.shape[0]],
+            color=color,
+            alpha=1,
+            zorder=1000 + zorder_cnt,
+            linewidth=lw_m,
+            label=tmp_label,
+        )
+
+        # Lower bound
+        axis.plot(
+            t_s,  # + sample.t.values[0],
+            b_low[:t_s.shape[0]],
+            color=color,
+            alpha=1,
+            zorder=1000 + zorder_cnt,
+            linewidth=lw_m,
+        )
+
+    # Slope
+    if keep_slope:
+        axis.plot(
+            t_s[:maxid] + sample['metadata']['t'][0],
+            tmp_traj[:maxid],
+            color=color,
+            alpha=1,
+            zorder=1000 + zorder_cnt,
+            linewidth=lw_m,
+        )  # TOOK AWAY LABEL
+
+    # Non-decision time
+    if keep_ndt:
+        axis.axvline(
+            x=sample['metadata']['t'][0],
+            ymin=ylim_low,
+            ymax=ylim_high,
+            color=color,
+            linestyle="--",
+            linewidth=lw_m,
+            zorder=1000 + zorder_cnt,
+            alpha=1,
+        )
+    # Starting point
+    if keep_starting_point:
+        axis.scatter(
+            sample['metadata']['t'][0] + markershift_starting_point,
+            b_low[0] + (sample['metadata']['z'][0] * (b_high[0] - b_low[0])),
+            markersize_starting_point,
+            marker=markertype_starting_point,
+            color=color,
+            alpha=1,
+            zorder=1000 + zorder_cnt,
+        )
+
+def plot_func_model_n(
+    model_name,
+    theta,
+    axis,
+    n_trajectories=10,
+    value_range=None,
+    bin_size=0.05,
+    n_samples=10,
+    linewidth_histogram=0.5,
+    linewidth_model=0.5,
+    legend_fontsize=7,
+    legend_shadow=True,
+    legend_location="upper right",
+    delta_t_model=0.001,
+    add_legend=True,
+    alpha=1.0,
+    keep_frame=False,
+    random_state=None,
+    **kwargs,
+):
+    """Calculate posterior predictive for a certain bottom node.
+
+    Arguments:
+        bottom_node: pymc.stochastic
+            Bottom node to compute posterior over.
+
+        axis: matplotlib.axis
+            Axis to plot into.
+
+        value_range: numpy.ndarray
+            Range over which to evaluate the likelihood.
+
+    Optional:
+        samples: int <default=10>
+            Number of posterior samples to use.
+
+        bin_size: float <default=0.05>
+            Size of bins used for histograms
+
+        alpha: float <default=0.05>
+            alpha (transparency) level for the sample-wise elements of the plot
+
+        add_posterior_uncertainty_rts: bool <default=True>
+            Add sample by sample histograms?
+
+        add_posterior_mean_rts: bool <default=True>
+            Add a mean posterior?
+
+        add_model: bool <default=True>
+            Whether to add model cartoons to the plot.
+
+        linewidth_histogram: float <default=0.5>
+            linewdith of histrogram plot elements.
+
+        linewidth_model: float <default=0.5>
+            linewidth of plot elements concerning the model cartoons.
+
+        legend_loc: str <default='upper right'>
+            string defining legend position. Find the rest of the options in the matplotlib documentation.
+
+        legend_shadow: bool <default=True>
+            Add shadow to legend box?
+
+        legend_fontsize: float <default=12>
+            Fontsize of legend.
+
+        data_color : str <default="blue">
+            Color for the data part of the plot.
+
+        posterior_mean_color : str <default="red">
+            Color for the posterior mean part of the plot.
+
+        posterior_uncertainty_color : str <default="black">
+            Color for the posterior uncertainty part of the plot.
+
+
+        delta_t_model:
+            specifies plotting intervals for model cartoon elements of the graphs.
+    """
+
+    color_dict = {
+        -1: "black",
+        0: "black",
+        1: "green",
+        2: "blue",
+        3: "red",
+        4: "orange",
+        5: "purple",
+        6: "brown",
+    }
+
+    # AF-TODO: Add a mean version of this !
+    if value_range is None:
+        # Infer from data by finding the min and max from the nodes
+        raise NotImplementedError("value_range keyword argument must be supplied.")
+
+    if len(value_range) > 2:
+        value_range = (value_range[0], value_range[-1])
+
+    # Extract some parameters from kwargs
+    bins = np.arange(value_range[0], value_range[-1], bin_size)
+    # ------------
+    ylim = kwargs.pop("ylim", 4)
+
+    axis.set_xlim(value_range[0], value_range[-1])
+    axis.set_ylim(0, ylim)
+
+    # ADD MODEL:
+
+    # RUN SIMULATIONS
+    # -------------------------------
+
+    # Simulator Data
+    if random_state is not None:
+        np.random.seed(random_state)
+    
+    rand_int = np.random.choice(400000000)
+    sim_out = simulator(model = model_name, theta = theta, n_samples = n_samples, 
+              no_noise = False, delta_t = 0.001, 
+              bin_dim = None, random_state = rand_int)
+    
+    choices = sim_out['metadata']['possible_choices']
+    
+    sim_out_traj = {}
+    for i in range(n_trajectories):
+        rand_int = np.random.choice(400000000)
+        sim_out_traj[i] = simulator(model = model_name, theta = theta, n_samples = 1, 
+                                    no_noise = False, delta_t = 0.001, 
+                                    bin_dim = None, random_state = rand_int, smooth_unif = False)
+
+    sim_out_no_noise = simulator(model = model_name, theta = theta, n_samples = 1, 
+                                 no_noise = True, delta_t = 0.001, 
+                                 bin_dim = None, smooth_unif = False)
+
+    # ADD HISTOGRAMS
+    # -------------------------------
+
+    # POSTERIOR BASED HISTOGRAM
+    j = 0
+    b = np.maximum(sim_out['metadata']['boundary'], 0)
+    bottom = b[0]
+    for choice in choices:
+        tmp_label = None
+
+        if add_legend and j == 0:
+            tmp_label = "PostPred"
+
+        weights = np.tile(
+            (1 / bin_size) / sim_out['rts'].shape[0],
+            reps=sim_out['rts'][(sim_out['choices'] == choice) & (sim_out['rts'] != -999)].shape[0],
+        )
+
+        axis.hist(
+            np.abs(sim_out['rts'][(sim_out['choices'] == choice) & (sim_out['rts'] != -999)]),
+            bins=bins,
+            bottom=bottom,
+            weights=weights,
+            histtype="step",
+            alpha=alpha,
+            color=color_dict[choice],
+            zorder=-1,
+            label=tmp_label,
+            linewidth=linewidth_histogram,
+        )
+        j += 1
+
+    # ADD MODEL:
+    tmp_label = None
+    j = 0
+    t_s = np.arange(0, sim_out['metadata']['max_t'], delta_t_model)
+
+    if add_legend and (j == 0):
+        tmp_label = "PostPred"
+
+    _add_model_n_cartoon_to_ax(
+        sample=sim_out_no_noise,
+        axis=axis,
+        delta_t_graph=delta_t_model,
+        sample_hist_alpha=alpha,
+        lw_m=linewidth_model,
+        tmp_label=tmp_label,
+        linestyle="-",
+        ylim=ylim,
+        t_s=t_s,
+        color_dict=color_dict,
+        zorder_cnt=j,
+    )
+
+    if n_trajectories > 0:
+        _add_trajectories_n(
+            axis=axis,
+            sample=sim_out_traj,
+            t_s=t_s,
+            delta_t_graph=delta_t_model,
+            n_trajectories=n_trajectories,
+            **kwargs,
+        )
+
+    if add_legend:
+        custom_elems = [
+            Line2D([0], [0], color=color_dict[choice], lw=1) for choice in choices
+        ]
+        custom_titles = ["response: " + str(choice) for choice in choices]
+
+        custom_elems.append(
+            Line2D([0], [0], color="black", lw=1.0, linestyle="dashed")
+        )
+        # custom_elems.append(Line2D([0], [0], color="black", lw=1.0, linestyle="-"))
+        # custom_titles.append("Data")
+        # custom_titles.append("Posterior")
+
+        axis.legend(
+            custom_elems,
+            custom_titles,
+            fontsize=legend_fontsize,
+            shadow=legend_shadow,
+            loc=legend_location,
+        )
+
+    # FRAME
+    if not keep_frame:
+        axis.set_frame_on(False)
+
+    return axis
+
+def _add_trajectories_n(axis=None,
+                        sample=None,
+                        t_s=None,
+                        delta_t_graph=0.01,
+                        n_trajectories=10,
+                        highlight_trajectory_rt_choice=True,
+                        markersize_trajectory_rt_choice=50,
+                        markertype_trajectory_rt_choice="*",
+                        markercolor_trajectory_rt_choice="black",
+                        linewidth_trajectories=1,
+                        alpha_trajectories=0.5,
+                        color_trajectories="black",
+                        **kwargs,
+                        ):
+    
+    """Add trajectories to a given axis."""
+    color_dict = {
+        -1: "black",
+        0: "black",
+        1: "green",
+        2: "blue",
+        3: "red",
+        4: "orange",
+        5: "purple",
+        6: "brown",
+    }
+
+    # Check trajectory color type
+    if isinstance(color_trajectories, str):
+        color_trajectories_dict = {}
+        for value_ in sample[0]['metadata']['possible_choices']:
+            color_trajectories_dict[value_] = color_trajectories
+    elif isinstance(color_trajectories, list):
+        cnt = 0
+        for value_ in sample[0]['metadata']['possible_choices']:
+            color_trajectories_dict[value_] = color_trajectories[cnt]
+            cnt += 1
+    elif isinstance(color_trajectories, dict):
+        color_trajectories_dict = color_trajectories
+    else:
+        pass
+
+    # Make bounds
+    b = np.maximum(sample[0]['metadata']['boundary'], 0)
+    b_init = b[0]
+    n_roll = int((sample[0]['metadata']['t'][0] / delta_t_graph) + 1)
+    b = np.roll(b, n_roll)
+    b[:n_roll] = b_init
+
+    # Trajectories
+    for i in range(n_trajectories):
+        tmp_traj = sample[i]['metadata']['trajectory']
+        tmp_traj_choice = float(sample[i]['choices'].flatten())
+
+        for j in range(len(sample[i]['metadata']['possible_choices'])):
+            tmp_maxid = np.minimum(np.argmax(np.where(tmp_traj[:, j] > -999)), t_s.shape[0])
+
+            # Identify boundary value at timepoint of crossing
+            b_tmp = b[tmp_maxid + n_roll]
+
+            axis.plot(
+                t_s[:tmp_maxid] + sample[i]['metadata']['t'][0], #sample.t.values[0],
+                tmp_traj[:tmp_maxid, j],
+                color=color_dict[j],
+                alpha=alpha_trajectories,
+                linewidth=linewidth_trajectories,
+                zorder=2000 + i,
+            )
+
+            if highlight_trajectory_rt_choice and tmp_traj_choice == j:
+                axis.scatter(
+                    t_s[tmp_maxid] + sample[i]['metadata']['t'][0], #sample.t.values[0],
+                    b_tmp,
+                    # tmp_traj[maxid],
+                    markersize_trajectory_rt_choice,
+                    color=color_dict[tmp_traj_choice],
+                    alpha=1,
+                    marker=markertype_trajectory_rt_choice,
+                    zorder=2000 + i,
+                )
+            elif highlight_trajectory_rt_choice and tmp_traj_choice != j:
+                axis.scatter(
+                    t_s[tmp_maxid] + sample[i]['metadata']['t'][0] + 0.05, #sample.t.values[0],
+                    tmp_traj[tmp_maxid, j],
+                    # tmp_traj[maxid],
+                    markersize_trajectory_rt_choice,
+                    color=color_dict[j],
+                    alpha=1,
+                    marker=5,
+                    zorder=2000 + i,
+                )
+
+def _add_model_n_cartoon_to_ax(
+    sample=None,
+    axis=None,
+    delta_t_graph=None,
+    sample_hist_alpha=None,
+    keep_boundary=True,
+    keep_ndt=True,
+    keep_slope=True,
+    keep_starting_point=True,
+    lw_m=None,
+    linestyle="-",
+    tmp_label=None,
+    ylim=None,
+    t_s=None,
+    zorder_cnt=1,
+    color_dict=None,
+):
+    
+    b = np.maximum(sample['metadata']['boundary'], 0)
+    b_init = b[0]
+    n_roll = int((sample['metadata']['t'][0] / delta_t_graph) + 1)
+    b = np.roll(b, n_roll)
+    b[:n_roll] = b_init
+
+    # Upper bound
+    if keep_boundary:
+        axis.plot(
+            t_s,
+            b[:t_s.shape[0]],
+            color="black",
+            alpha=sample_hist_alpha,
+            zorder=1000 + zorder_cnt,
+            linewidth=lw_m,
+            linestyle=linestyle,
+            label=tmp_label,
+        )
+
+    # Starting point
+    if keep_starting_point:
+        axis.axvline(
+            x=sample['metadata']['t'][0],
+            ymin=-ylim,
+            ymax=ylim,
+            color="black",
+            linestyle=linestyle,
+            linewidth=lw_m,
+            alpha=sample_hist_alpha,
+        )
+
+    # # MAKE SLOPES (VIA TRAJECTORIES HERE --> RUN NOISE FREE SIMULATIONS)!
+    if keep_slope:
+        tmp_traj = sample["metadata"]["trajectory"]
+
+        for i in range(len(sample["metadata"]["possible_choices"])):
+            tmp_maxid = np.minimum(np.argmax(np.where(tmp_traj[:, i] > -999)), t_s.shape[0])
+
+            # Slope
+            axis.plot(
+                t_s[:tmp_maxid] + sample['metadata']['t'][0],
+                tmp_traj[:tmp_maxid, i],
+                color=color_dict[i],
+                linestyle=linestyle,
+                alpha=sample_hist_alpha,
+                zorder=1000 + zorder_cnt,
+                linewidth=lw_m,
+            )  # TOOK AWAY LABEL
+
+    return b[0]