bunch of visual improvements and first round of exposing styling options to users

src/app.py

@@ -9,9 +9,9 @@ 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]
@@ -35,67 +35,379 @@ def create_param_selectors(model_name: str, model_num: int = 1):
             key=f"param{i}"
             f"_{model_name}"
             f"_{model_num}"
-            f'_{st.session_state["slider_version"]}',
+            f'_{st.session_state["param_version"]}',
         )
     return d_param_slider
 
 
+def create_styling_selectors(model_num: int = 1):
+    """
+    Create styling configuration widgets for plot customization.
+    
+    This function creates Streamlit widgets that allow users to customize
+    various visual aspects of the plots including colors, line widths,
+    alpha, and which model components to display.
+    
+    Note: This version is designed to work in the sidebar without using st.columns()
+    
+    Args:
+        model_num: Integer identifier for the model (1 or 2)
+        
+    Returns:
+        dict: Dictionary containing all styling parameters with their user-selected values
+    """
+    
+    # Color options for different plot elements
+    color_options = ["blue", "red", "green", "orange", "purple", "black", "gray", "brown"]
+    
+    # Legend location options (matplotlib standard locations)
+    legend_locations = ["upper right", "upper left", "lower left", "lower right", 
+                        "center", "upper center", "lower center", "center left", "center right"]
+    
+    # Marker type options for trajectories
+    marker_options = { "Diamond": "D", "Square": "s", "Line": 0, "Circle": "o", "Star": "*",  "Triangle": "^",
+                     "Plus": "+", "X": "x"}
+    
+    styling_config = {}
+    
+    # Create an expander for styling options to keep the interface clean
+    with st.expander(f"🎨 Styling", expanded=False):
+        
+        # Color Settings Section
+        st.markdown("**Colors**")
+        styling_config["data_color"] = st.selectbox(
+            "Data Color",
+            color_options,
+            index=color_options.index("blue" if model_num == 1 else "red"),
+            key=f"data_color_{model_num}_{st.session_state['styling_version']}"
+        )
+        
+        styling_config["posterior_uncertainty_color"] = st.selectbox(
+            "Model Color", 
+            color_options,
+            index=color_options.index("black"),
+            key=f"model_color_{model_num}_{st.session_state['styling_version']}"
+        )
+        
+        # Line Width Settings Section
+        st.markdown("**Lines**")
+        styling_config["linewidth_histogram"] = st.slider(
+            "Histogram Line Width",
+            min_value=0.1,
+            max_value=3.0,
+            value=1.0,
+            step=0.1,
+            key=f"hist_lw_{model_num}_{st.session_state['styling_version']}"
+        )
+        
+        styling_config["linewidth_model"] = st.slider(
+            "Model Line Width",
+            min_value=0.1,
+            max_value=3.0,
+            value=1.0,
+            step=0.1,
+            key=f"model_lw_{model_num}_{st.session_state['styling_version']}"
+        )
+        
+        # Histogram Settings Section
+        st.markdown("**Histograms**")
+        styling_config["bin_size"] = st.slider(
+            "Bin Size",
+            min_value=0.01,
+            max_value=0.2,
+            value=0.05,
+            step=0.01,
+            key=f"bin_size_{model_num}_{st.session_state['styling_version']}"
+        )
+        
+        styling_config["alpha"] = st.slider(
+            "alpha",
+            min_value=0.0,
+            max_value=1.0,
+            value=1.0,
+            step=0.05,
+            key=f"alpha_{model_num}_{st.session_state['styling_version']}"
+        )
+        
+        # Model Components Section - Toggle which parts of the model to show
+        st.markdown("**Model Components**")
+        styling_config["add_data_model_keep_boundary"] = st.checkbox(
+            "Show Boundaries", 
+            value=True,
+            key=f"show_boundary_{model_num}_{st.session_state['styling_version']}"
+        )
+        styling_config["add_data_model_keep_slope"] = st.checkbox(
+            "Show Slope/Trajectory", 
+            value=True,
+            key=f"show_slope_{model_num}_{st.session_state['styling_version']}"
+        )
+        
+        styling_config["add_data_model_keep_ndt"] = st.checkbox(
+            "Show Non-Decision Time", 
+            value=True,
+            key=f"show_ndt_{model_num}_{st.session_state['styling_version']}"
+        )
+        styling_config["add_data_model_keep_starting_point"] = st.checkbox(
+            "Show Starting Point", 
+            value=True,
+            key=f"show_start_{model_num}_{st.session_state['styling_version']}"
+        )
+
+        # Axis Limits Section
+        st.markdown("**Axis Limits**")
+        styling_config["xlim_min"] = st.number_input(
+            "x-axis min",
+            value=-0.1,
+            step=0.1,
+            key=f"xlim_min_{model_num}_{st.session_state['styling_version']}"
+        )
+        styling_config["xlim_max"] = st.number_input(
+            "x-axis max", 
+            value=5.0,
+            step=0.1,
+            key=f"xlim_max_{model_num}_{st.session_state['styling_version']}"
+        )
+        styling_config["ylim_max"] = st.number_input(
+            "y-axis max",
+            value=3.75,
+            step=0.25,
+            key=f"ylim_max_{model_num}_{st.session_state['styling_version']}"
+        )
+        
+        # Starting Point Marker Settings (only show if starting point is enabled)
+        if styling_config["add_data_model_keep_starting_point"]:
+            st.markdown("**Starting Point**")
+            styling_config["add_data_model_markersize_starting_point"] = st.slider(
+                "Marker Size",
+                min_value=10,
+                max_value=100,
+                value=35,
+                key=f"marker_size_{model_num}_{st.session_state['styling_version']}"
+            )
+            
+            styling_config["add_data_model_markertype_starting_point"] = st.selectbox(
+                "Marker Type",
+                list(marker_options.keys()),
+                index=0, # defaulting to first entry in marker options dictionary
+                key=f"marker_type_{model_num}_{st.session_state['styling_version']}"
+            )
+        else:
+            # Set defaults when starting point is not shown
+            styling_config["add_data_model_markersize_starting_point"] = 35
+            styling_config["add_data_model_markertype_starting_point"] = "Diamond"
+        
+        # AF-TODO: Legend settings weren't working yet
+
+        # # Legend Settings Section
+        # st.markdown("**Legend Settings**")
+        # styling_config["add_legend"] = st.checkbox(
+        #     "Show Legend", 
+        #     value=True,
+        #     key=f"show_legend_{model_num}_{st.session_state['styling_version']}"
+        # )
+        
+        # if styling_config["add_legend"]:
+        #     styling_config["legend_fontsize"] = st.slider(
+        #         "Legend Font Size",
+        #         min_value=6,
+        #         max_value=20,
+        #         value=12,
+        #         key=f"legend_font_{model_num}_{st.session_state['styling_version']}"
+        #     )
+            
+        #     styling_config["legend_location"] = st.selectbox(
+        #         "Legend Location",
+        #         legend_locations,
+        #         index=0,  # "upper right"
+        #         key=f"legend_loc_{model_num}_{st.session_state['styling_version']}"
+        #     )
+            
+        #     styling_config["legend_shadow"] = st.checkbox(
+        #         "Legend Shadow", 
+        #         value=True,
+        #         key=f"legend_shadow_{model_num}_{st.session_state['styling_version']}"
+        #     )
+        # else:
+        #     # Set defaults when legend is not shown
+        #     styling_config["legend_fontsize"] = 12
+        #     styling_config["legend_location"] = "upper right"
+        #     styling_config["legend_shadow"] = True
+    
+    # Convert marker type from display name to matplotlib code
+    marker_type_map = {k: v for k, v in marker_options.items()}
+    styling_config["add_data_model_markertype_starting_point"] = marker_type_map.get(
+        styling_config["add_data_model_markertype_starting_point"], "D"
+    )
+
+    return styling_config
+
+
+def get_filtered_styling_config(styling_config, plot_type="plot_func_model"):
+    """
+    Filter styling configuration based on plot type compatibility.
+    
+    Different plotting functions accept different parameters, so this function
+    filters the styling configuration to only include parameters that are
+    relevant for the specific plot type.
+    
+    Args:
+        styling_config: Dictionary of styling parameters
+        plot_type: String indicating which plot function will be used
+                  ("plot_func_model" or "plot_func_model_n")
+    
+    Returns:
+        dict: Filtered styling configuration appropriate for the plot type
+    """
+
+    if plot_type == "plot_func_model":
+        # plot_func_model accepts all styling parameters
+        return styling_config
+
+    elif plot_type == "plot_func_model_n":
+        # plot_func_model_n only accepts a subset of parameters
+        allowed_params = {
+            'linewidth_histogram', 'linewidth_model', 'bin_size', 
+            'alpha', 'legend_fontsize', 'legend_location', 'legend_shadow', 
+            'add_legend', 'add_data_model_markersize_starting_point',
+            'add_data_model_markertype_starting_point',
+            'add_data_model_keep_starting_point',
+            'add_data_model_keep_boundary',
+            'add_data_model_keep_slope', 
+            'add_data_model_keep_ndt'
+        }
+        return {k: v for k, v in styling_config.items() if k in allowed_params}
+    
+    else:
+        # Default: return all parameters
+        return styling_config
+
 def add_model():
     pass
 
+# def reset_sliders():
+#     st.session_state["slider_version"] += 1
+
+def reset_parameters():
+    """Reset only model parameters to defaults"""
+    st.session_state["param_version"] += 1
 
-def reset_sliders():
-    st.session_state["slider_version"] += 1
+def reset_styling():
+    """Reset only styling options to defaults"""
+    st.session_state["styling_version"] += 1
 
+def reset_all():
+    """Reset both parameters and styling to defaults"""
+    st.session_state["param_version"] += 1
+    st.session_state["styling_version"] += 1
+    st.session_state["slider_version"] += 1  # Keep for any remaining widgets
 
 # 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
+# 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)
+    
+    # 1. Dropdown selection of model name
+    model_select = st.selectbox("Model " + str(model_num), l_model_names, key="model_selector_" + str(model_num))
+
+    return model_select
+
+def draw_simulation_settings(model_num=1):
     # 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)
+        "NTrajectories", value=5, key="ntraj" + str(model_num)
     )
-    return model_select, d_slider, nsamples, ntrajectories
 
+    # Random seed setting
+    randomseed = st.number_input("RandomSeed", value=41 + model_num, key="seed_" + str(model_num))
+
+    return nsamples, ntrajectories, randomseed
 
 st.set_page_config(layout="wide")
 
 # Get list of model names
 l_model_names = list(model_config.keys())
 
+# Initialize separate version attributes for parameters and styling
+if "param_version" not in st.session_state:
+    st.session_state["param_version"] = 1
+
+if "styling_version" not in st.session_state:
+    st.session_state["styling_version"] = 1
+
+
 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)
+    st.empty()
+    st.markdown("**Model Selection**")
+    with st.container():
+        st.markdown('<div style="margin-top: -1rem;">', unsafe_allow_html=True)
+
+        col1, col2 = st.columns(2)
+        with col1:
+            model_select_1 = draw_model_configurator(model_num=1)
+
+            # Styling configuration
+            styling_config_1 = create_styling_selectors(model_num=1)
+
+        with col2:
+            model_select_2 = draw_model_configurator(model_num=2)
+
+            # Styling configuration
+            styling_config_2 = create_styling_selectors(model_num=2)
+
+        st.markdown('</div>', unsafe_allow_html=True)
+
+    st.markdown("---")
+    st.markdown("**Parameters**")
+    col1_2, col2_2 = st.columns(2)
+    with col1_2:
+        d_slider_1 = create_param_selectors(model_select_1, model_num=1)
+    with col2_2:
+        d_slider_2 = create_param_selectors(model_select_2, model_num=2)
+
+    st.markdown("---")
+    st.markdown("**Simulation Settings**")
+    col1_3, col2_3 = st.columns(2)
+    with col1_3:
+        nsamples_1, ntrajectories_1, randomseed_1 = draw_simulation_settings(model_num=1)
+    with col2_3:
+        nsamples_2, ntrajectories_2, randomseed_2 = draw_simulation_settings(model_num=2)
+
+    # Button to reset sliders to default values
+    st.markdown("---")
+    st.markdown("**Reset Options**")
+
+    # Create three columns for the reset buttons
+    reset_col1, reset_col2, reset_col3 = st.columns(3)
+
+    with reset_col1:
+        st.button(
+            "Reset Params",
+            help="Reset model parameters to defaults",
+            key="reset_params",
+            on_click=reset_parameters,
         )
-    with col2:
-        model_select_2, d_slider_2, nsamples_2, ntrajectories_2 = (
-            draw_model_configurator(model_num=2)
+
+    with reset_col2:
+        st.button(
+            "Reset Styling",
+            help="Reset styling options to defaults", 
+            key="reset_styling",
+            on_click=reset_styling,
         )
 
-    # 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,
-    )
+    with reset_col3:
+        st.button(
+            "Reset Full",
+            help="Reset both parameters and styling to defaults",
+            key="reset_all", 
+            on_click=reset_all,
+        )
 
 # st.title("SSM Model Plots", )
 st.markdown(
@@ -106,70 +418,67 @@ st.markdown(
 # Display components for main panel
 fig1, ax1 = plt.subplots()
 if model_config[model_select_1]["nchoices"] == 2 and not ("race" in model_select_1):
+    # Use filtered styling parameters for plot_func_model
+
+    filtered_styling_1 = get_filtered_styling_config(styling_config_1, "plot_func_model")
     ax1 = utils.utils.plot_func_model(
         model_name=model_select_1,
         theta=[list(d_slider_1.values())],
         axis=ax1,
-        value_range=[-0.1, 5],
+        value_range=[styling_config_1["xlim_min"], styling_config_1["xlim_max"]],
         n_samples=nsamples_1,
-        ylim=5,
-        data_color="blue",
-        add_trajectories=True,
+        ylim=styling_config_1["ylim_max"],
         n_trajectories=ntrajectories_1,
-        linewidth_model=1,
-        linewidth_histogram=1,
-        random_state=randomseed,
+        random_state=randomseed_1,
+        **filtered_styling_1
     )
 else:
+    # Use filtered styling parameters for plot_func_model_n
+    filtered_styling_1 = get_filtered_styling_config(styling_config_1, "plot_func_model_n")
     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],
+        value_range=[styling_config_1["xlim_min"], styling_config_1["xlim_max"]],
         n_samples=nsamples_1,
-        data_color="blue",
-        add_trajectories=True,
         n_trajectories=ntrajectories_1,
-        linewidth_model=1,
-        linewidth_histogram=1,
-        random_state=randomseed,
+        random_state=randomseed_1,
+        **filtered_styling_1
     )
-ax1.set_title("Model 1")
-ax1.set_xlabel("RT in seconds")
+ax1.set_title(model_select_1.upper())
+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):
+    # Use filtered styling parameters for plot_func_model
+    filtered_styling_2 = get_filtered_styling_config(styling_config_2, "plot_func_model")
     ax2 = utils.utils.plot_func_model(
         model_name=model_select_2,
         theta=[list(d_slider_2.values())],
         axis=ax2,
-        value_range=[-0.1, 5],
+        value_range=[styling_config_2["xlim_min"], styling_config_2["xlim_max"]],
         n_samples=nsamples_2,
-        ylim=5,
-        data_color="red",
-        add_trajectories=True,
+        ylim=styling_config_2["ylim_max"],
         n_trajectories=ntrajectories_2,
-        linewidth_model=1,
-        linewidth_histogram=1,
-        random_state=randomseed,
+        random_state=randomseed_2,
+        **filtered_styling_2
     )
 else:
+    # Use filtered styling parameters for plot_func_model_n
+    filtered_styling_2 = get_filtered_styling_config(styling_config_2, "plot_func_model_n")
     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],
+        value_range=[styling_config_2["xlim_min"], styling_config_2["xlim_max"]],
         n_samples=nsamples_2,
-        data_color="red",
-        add_trajectories=True,
         n_trajectories=ntrajectories_2,
-        linewidth_model=1,
-        linewidth_histogram=1,
-        random_state=randomseed,
+        random_state=randomseed_2,
+        **filtered_styling_2
     )
 
-ax2.set_title("Model 2")
-ax2.set_xlabel("RT in seconds")
+ax2.set_title(model_select_2.upper())
+ax2.set_xlabel("rt in seconds")
 
 # Place figure in placeholder
 col1, col2 = st.columns(2)
@@ -185,29 +494,32 @@ sim_output_1 = simulator(
     model=model_select_1,
     theta=[list(d_slider_1.values())],
     n_samples=nsamples_1,
-    random_state=randomseed,
+    random_state=randomseed_1,
 )
 sim_output_2 = simulator(
     model=model_select_2,
     theta=[list(d_slider_2.values())],
     n_samples=nsamples_2,
-    random_state=randomseed,
+    random_state=randomseed_2,
 )
 
 # Make metadata dataframe
+# AF-TODO: Should be transposed and then resolve consequences
+# because right now the dataframe has mixed data types in each columns
+# which leads streamlit to complain about arrow incompatibility
 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"],
+            str(sim_output_1["metadata"]["model"]),
+            float(sim_output_1["choice_p"][0, 0]),
+            float(sim_output_1["rts"].mean()),
+            float(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"],
+            str(sim_output_2["metadata"]["model"]),
+            float(sim_output_2["choice_p"][0, 0]),
+            float(sim_output_2["rts"].mean()),
+            float(sim_output_2["metadata"]["s"]),
         ],
     },
     index=["Model", "Choice Probability", "Mean RT", "Noise SD"],
@@ -230,9 +542,9 @@ with col3:
             histtype="step",
             bins=50,
             density=True,
-            color="blue",
+            color=styling_config_1["data_color"],  # Use user-selected color
             fill=None,
-            label="Model 1",
+            label=model_select_1.upper(),
         )
         ax3.hist(
             sim_output_2["rts"][np.abs(sim_output_2["rts"]) != 999]
@@ -240,12 +552,12 @@ with col3:
             histtype="step",
             bins=50,
             density=True,
-            color="red",
+            color=styling_config_2["data_color"],  # Use user-selected color
             fill=None,
-            label="Model 2",
+            label=model_select_2.upper(),
         )
         ax3.legend()
-        ax3.set_xlabel("RT")
+        ax3.set_xlabel("rt")
         ax3.set_xlim(-5, 5)
         figure_placeholder_3.pyplot(fig3)
     else:
@@ -253,4 +565,4 @@ with col3:
         # for models with more than 2 choice options
         pass
 with col4:
-    st.dataframe(metadata)
+    st.dataframe(metadata)

src/streamlit_app.py

@@ -1,40 +0,0 @@
-import altair as alt
-import numpy as np
-import pandas as pd
-import streamlit as st
-
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))

src/utils/old_plots.py

@@ -1,1157 +0,0 @@
-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

@@ -187,7 +187,7 @@ def plot_func_model(
         weights=weights_up,
         histtype=hist_histtype,
         bottom=hist_bottom_high,
-        alpha=1,
+        alpha=alpha,
         color=data_color,
         edgecolor=data_color,
         linewidth=linewidth_histogram,
@@ -200,7 +200,7 @@ def plot_func_model(
         weights=weights_down,
         histtype=hist_histtype,
         bottom=hist_bottom_low,
-        alpha=1,
+        alpha=alpha,
         color=data_color,
         edgecolor=data_color,
         linewidth=linewidth_histogram,
@@ -308,7 +308,6 @@ def _add_trajectories(
     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']
@@ -426,7 +425,7 @@ def _add_model_cartoon_to_ax(
         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,
+            s=markersize_starting_point,
             marker=markertype_starting_point,
             color=color,
             alpha=1,
@@ -448,7 +447,7 @@ def plot_func_model_n(
     legend_location="upper right",
     delta_t_model=0.001,
     add_legend=True,
-    alpha=1.0,
+    alpha=1,
     keep_frame=False,
     random_state=None,
     **kwargs,
@@ -472,7 +471,7 @@ def plot_func_model_n(
         bin_size: float <default=0.05>
             Size of bins used for histograms
 
-        alpha: float <default=0.05>
+        alpha: float <default=1.0>
             alpha (transparency) level for the sample-wise elements of the plot
 
         add_posterior_uncertainty_rts: bool <default=True>
@@ -560,7 +559,7 @@ def plot_func_model_n(
     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, 
+                                    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,