smart_microfluidics.py

import streamlit as st
import hashlib
import requests
import json
import pandas as pd
import numpy as np
import pickle
import umap
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.graph_objects as go
import scipy.cluster.hierarchy as sch
from scipy.stats import gaussian_kde
from sklearn.cluster import KMeans
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from config import *


################################################
# AUXILIARY FUNCTIONS
################################################

def categorize_size(size):
    if size < 100:
        return "S"
    elif size < 200:
        return "M"
    else:
        return "L"


def categorize_pdi(pdi):
    if pdi < 0.1:
        return "HMD"
    elif 0.1 < pdi < 0.25:
        return "MD"
    else:
        return "PLD"


@st.cache_data
def fetch_csv_data(file_path):
    data = pd.read_csv(file_path, encoding="latin1").drop(columns=["ID", "FR-O", "FR-W"])
    data.BUFFER = data.BUFFER.astype(str).str.strip().replace({"PBS\xa0": "PBS"})
    data.CHIP = data.CHIP.replace({"Micromixer\xa0": "Micromixer"})
    data = data.map(lambda x: str(x) if isinstance(x, str) else x)
    numeric_data = data.select_dtypes(include=["float64", "int64"]).dropna()
    return data, numeric_data


@st.cache_data
def fetch_xlsx_data(file_path):
    data = pd.read_excel(file_path)
    return data


def load_credentials(url):
    response = requests.get(url)
    if response.status_code == 200:
        return json.loads(response.text)
    return None


def hash_data(data):
    return hashlib.sha256(data.encode()).hexdigest()


def authenticate(user_id, password, credentials):
    hashed_user_id = hash_data(user_id)
    hashed_password = hash_data(password)
    return hashed_user_id in credentials and credentials[hashed_user_id] == hashed_password


def display_output(size, pdi):
    if size > 500 or pdi > 0.5:
        st.write("The system doesn't form")
        st.write("`OUTPUT`: 0")
    else:
        st.write(f"`SIZE`: {size:.2f}")
        st.write(f"`PDI`: {pdi:.2f}")


def plot_model_performance(metrics, before, after):
    improvements = [
        ((after[i] - before[i]) / before[i]) * 100
        if metrics[i] == "R-squared"
        else ((before[i] - after[i]) / before[i]) * 100
        for i in range(len(metrics))
    ]

    fig, ax1 = plt.subplots(figsize=(8, 5))
    bar_width = 0.35
    x = np.arange(len(metrics))

    bars_before = ax1.bar(
        x - bar_width / 2,
        before,
        bar_width,
        label="Before",
        color="red",
        alpha=0.7,
    )
    bars_after = ax1.bar(
        x + bar_width / 2,
        after,
        bar_width,
        label="After",
        color="green",
        alpha=0.7,
    )

    for bar in bars_before:
        ax1.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height(),
            f"{bar.get_height():.2f}",
            ha="center",
            va="bottom",
            fontsize=10,
            color="black",
        )

    for bar in bars_after:
        ax1.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height(),
            f"{bar.get_height():.2f}",
            ha="center",
            va="bottom",
            fontsize=10,
            color="black",
        )

    ax2 = ax1.twinx()
    ax2.plot(
        x,
        improvements,
        marker="o",
        linestyle="--",
        color="blue",
        label="Improvement (%)",
        linewidth=2,
    )

    for i, val in enumerate(improvements):
        ax2.text(
            x[i],
            val,
            f"{val:.1f}%",
            ha="center",
            va="bottom",
            fontsize=12,
            color="blue",
        )

    ax1.set_xlabel("Metrics")
    ax1.set_ylabel("Values (Before & After)")
    ax2.set_ylabel("Improvement (%)")
    ax1.set_title("Model Performance Improvement")
    ax1.set_xticks(x)
    ax1.set_xticklabels(metrics)
    ax1.legend(loc="upper right", bbox_to_anchor=(1, 1), ncol=1, frameon=False)
    ax2.legend(loc="upper right", bbox_to_anchor=(1, 0.85), ncol=1, frameon=False)

    fig.tight_layout()
    st.pyplot(fig)
    plt.close(fig)


def plot_ridgeline_matplotlib(
    df,
    x_range=(-5, 10),
    figsize=(12, 8),
    title="Ridgeline Plot",
):
    """
    JoyPy-free ridgeline plot.

    This replaces:
        joypy.joyplot(...)

    It avoids the Streamlit Cloud / Python 3.12 / JoyPy compatibility error:
        TypeError: 'generator' object is not subscriptable
    """

    clean_df = df.copy()
    clean_df = clean_df.select_dtypes(include=[np.number])

    valid_columns = []

    for col in clean_df.columns:
        values = clean_df[col].dropna().to_numpy(dtype=float)

        if len(values) <= 2:
            continue

        if not np.isfinite(values).all():
            continue

        if np.nanstd(values) == 0:
            continue

        valid_columns.append(col)

    fig, ax = plt.subplots(figsize=figsize)

    if not valid_columns:
        ax.text(
            0.5,
            0.5,
            "No valid numeric columns available for ridgeline plotting.",
            ha="center",
            va="center",
            transform=ax.transAxes,
        )
        ax.set_axis_off()
        fig.tight_layout()
        return fig

    x_grid = np.linspace(x_range[0], x_range[1], 500)
    vertical_offset = 1.0

    for i, col in enumerate(valid_columns):
        values = clean_df[col].dropna().to_numpy(dtype=float)

        try:
            kde = gaussian_kde(values)
            density = kde(x_grid)

            max_density = np.nanmax(density)
            if max_density > 0:
                density = density / max_density

            baseline = i * vertical_offset
            y_values = baseline + density

            ax.fill_between(
                x_grid,
                baseline,
                y_values,
                alpha=0.65,
            )

            ax.plot(
                x_grid,
                y_values,
                linewidth=1.0,
            )

            ax.text(
                x_range[0],
                baseline + 0.18,
                str(col),
                ha="left",
                va="center",
                fontsize=9,
            )

        except Exception:
            continue

    ax.set_title(title)
    ax.set_xlabel("Standardized value")
    ax.set_yticks([])
    ax.set_xlim(x_range)

    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_visible(False)

    fig.tight_layout()
    return fig


################################################
# COLOPHON
################################################

st.title("Smart Microfluidics: Machine Learning tools for Liposome Production Experiments")

# GITHUB_CREDENTIALS_URL = "https://raw.githubusercontent.com/leonardoLavagna/smart_microfluidics/main/_includes/credentials.json"
# credentials = load_credentials(GITHUB_CREDENTIALS_URL)
# if "authenticated" not in st.session_state:
#     st.session_state.authenticated = False
# if "section" not in st.session_state:
#     st.session_state.section = "Authenticate"
# if not st.session_state.authenticated:
#     st.session_state.section = "Authenticate"
#     user_id = st.text_input("Enter your ID:")
#     password = st.text_input("Enter the global password:", type="password")
#     if st.button("Login"):
#         if credentials and authenticate(user_id, password, credentials):
#             st.success("Authentication successful! You can proceed.")
#             st.session_state.authenticated = True
#             st.session_state.section = "Dataset"
#             st.session_state.authenticated = True
#             st.rerun()
#         else:
#             st.error("Authentication failed! Please check your credentials.")
# else:
#     st.sidebar.title("Choose an Option")
#     st.session_state.section = st.sidebar.selectbox(
#         "Data preprocessing, data modelization or data visualization:",
#         ["Dataset", "Data Modeling", "Data Exploration"],
#         index=["Dataset", "Data Modeling", "Data Exploration"].index(st.session_state.section)
#     )
# section = st.session_state.section
# st.write(f"You selected: {section}")

st.sidebar.title("Choose an Option")

st.session_state.section = st.sidebar.selectbox(
    "Data preprocessing, data modelization or data visualization:",
    ["Dataset", "Data Modeling", "Data Exploration"],
    index=["Dataset", "Data Modeling", "Data Exploration"].index("Dataset"),
)

section = st.session_state.section


################################################
# 1. DATA
################################################

data, numeric_data = fetch_csv_data("data/data.csv")
data_numeric = numeric_data.copy()
data_ = fetch_xlsx_data("_includes/data.xlsx")

if st.session_state.section == "Dataset":
    st.markdown(
        """This app provides machine learning and data analysis tools for laboratory operators while carrying out microfluidic liposome experiments.
               The app is in a developing phase. Current version: `v0.2`. In this version the user can:"""
    )
    st.markdown("""- Work with pretrained machine learning architectures on a default dataset;""")
    st.markdown(
        """- Simulate experiments with the tools in the Data Modeling section (drop-down menu on the left); """
    )
    st.markdown(
        """- Get in depth data analysis with the visualizations in the Data Exploration section (drop-down menu on the left). """
    )
    st.markdown(
        """Report bugs or request updates [@Leonardo Lavagna](https://leonardolavagna.github.io/) by [opening an issue in the project repository](https://github.com/leonardoLavagna/smart_microfluidics/issues).
                The training of the models was done on a publicly available dataset and in compliance with current legislation."""
    )
    st.header("Dataset")
    st.warning(
        ":warning: Current models have been trained on the default dataset which is of intermediate dimension. Higher performances require additional training data."
    )
    st.write("Get the data for subsequent processing.")
    user_choice = st.radio("Upload custom data?", ("No", "Yes"))

    if user_choice == "Yes":
        st.warning(":no_entry_sign: Customized data processing not available for current user.")
    else:
        st.write("Loading default dataset...")
        st.success("File loaded successfully!")
        st.dataframe(data_)
        st.markdown("""The default dataset contains the following features:""")
        st.markdown("""- `ID`: a unique identifier for each formulation;""")
        st.markdown("""- `ML`: the Main Lipid used;""")
        st.markdown("""- `CHIP`: the type of chip used in the microfluidic setup;""")
        st.markdown("""- `ESM`: the concentration (mg/mL) of the ESM lipid;""")
        st.markdown("""- `HSPC`: the concentration (mg/mL) of the HSPC lipid;""")
        st.markdown("""- `CHOL`: the concentration (mg/mL) of Cholesterol used;""")
        st.markdown("""- `PEG`: the concentration (mg/mL) of DSPE-PEG2000 used;""")
        st.markdown("""- `TFR`: the Total Flow Rate (mL/min) in the microfluidic considered setup;""")
        st.markdown("""- `FRR`: the Flow Rate Ratio between flux streams in the considered setup;""")
        st.markdown("""- `BUFFER`: the aqueous medium considered;""")
        st.markdown("""- `OUTPUT`: the result of the formation of the system;""")
        st.markdown("""- `SIZE`: the size (nm) of the liposomes;""")
        st.markdown("""- `PDI`: the PolyDispersity Index.""")


################################################
# 2. MODELS
################################################

if section == "Data Modeling":
    st.header("Data Modeling")
    st.write(
        """Start by choosing a model in the drop-down menu, then try the selected model with your data. If you want to test multiple models (Random forest regressor, XGBoost and the Advanced models) on the same input data,
                just enter each data value once in the table below. The inverse model and the single target models work differently, so you will have to enter new inputs."""
    )

    option = st.sidebar.selectbox(
        "Select a model",
        [
            "Random forest regressor",
            "XGBoost",
            "Inverse model",
            "Single target models",
        ],
    )

    if option != "Inverse model" and option != "Single target models":
        ml = st.selectbox("ML", ["HSPC", "ESM"])
        chip = st.selectbox("CHIP", ["Micromixer", "Droplet junction"])

        esm_disabled = ml == "HSPC"
        hspc_disabled = ml == "ESM"

        esm = st.number_input(
            "ESM",
            value=0.0 if esm_disabled else 0.1,
            min_value=0.0,
            max_value=100.0,
            step=0.1,
            disabled=esm_disabled,
        )
        hspc = st.number_input(
            "HSPC",
            value=0.0 if hspc_disabled else 3.75,
            min_value=0.0,
            max_value=100.0,
            step=0.1,
            disabled=hspc_disabled,
        )
        chol = st.number_input("CHOL", value=0.0, min_value=0.0, max_value=100.0, step=0.1)
        peg = st.number_input("PEG", value=1.25, min_value=0.0, max_value=100.0, step=0.1)
        tfr = st.number_input("TFR", value=1.0, min_value=0.0, max_value=100.0, step=0.1)
        frr = st.number_input("FRR", value=3.0, min_value=0.0, max_value=100.0, step=0.1)
        buffer = st.selectbox("BUFFER", ["PBS", "MQ"])
        tlp = (hspc if hspc > 0 else esm) + chol + peg

    elif option == "Inverse model":
        size = st.number_input("SIZE", value=118.0, min_value=0.0, max_value=5000.0, step=10.0)
        pdi = st.number_input("PDI", value=0.33, min_value=0.0, max_value=1.0, step=0.1)

    elif option == "Single target models":
        esm = st.number_input("ESM", value=0.0, min_value=0.0, max_value=100.0, step=0.1, key="esm_in")
        hspc = st.number_input("HSPC", value=3.75, min_value=0.0, max_value=100.0, step=0.1, key="hspc_in")
        chol = st.number_input("CHOL", value=0.0, min_value=0.0, max_value=100.0, step=0.1, key="chol_in")
        peg = st.number_input("PEG", value=1.25, min_value=0.0, max_value=100.0, step=0.1, key="peg_in")
        tfr = st.number_input("TFR", value=1.0, min_value=0.0, max_value=100.0, step=0.1, key="tfr_in")
        frr = st.number_input("FRR", value=3.0, min_value=0.0, max_value=100.0, step=0.1, key="frr_in")
        buffer = st.selectbox("AQUEOUS", ["PBS", "MQ"], key="buf_in")
        size = st.number_input("SIZE", value=0.0, min_value=0.0, max_value=5000.0, step=10.0, key="siz_in")
        pdi = st.number_input("PDI", value=0.33, min_value=0.0, max_value=1.0, step=0.1, key="pd_in")

    # 2.1 Random forest regressor
    if option == "Random forest regressor":
        st.subheader("Random forest regressor")
        st.markdown(
            "Using multiple [decision trees](https://en.wikipedia.org/wiki/Decision_tree) in parallel and [bagging](http://en.wikipedia.org/wiki/Bootstrap_aggregating) this model based on a [random forest](https://en.wikipedia.org/wiki/Random_forest) provides joint predictions for `SIZE` and `PDI`."
        )

        with open(random_forest_model, "rb") as file:
            model = pickle.load(file)

        st.write(f"Loaded {random_forest_model} with the following performance metrics.")
        st.table(
            pd.DataFrame(
                {
                    "Metric": [
                        "R-squared",
                        "Mean Squared Error (MSE)",
                        "Mean Absolute Error (MAE)",
                    ],
                    "Value": [
                        0.36157896454981364,
                        1958.890858993266,
                        15.086741645521377,
                    ],
                }
            )
        )

        if st.button("Predict"):
            input_data = pd.DataFrame(
                {
                    "ML": [ml],
                    "CHIP": [chip],
                    "TLP": [tlp],
                    "ESM": [0.0 if esm_disabled else esm],
                    "HSPC": [0.0 if hspc_disabled else hspc],
                    "CHOL": [chol],
                    "PEG": [peg],
                    "TFR ": [tfr],
                    "FRR": [frr],
                    "BUFFER": [buffer],
                }
            )

            st.markdown("**Input data**")
            st.write(input_data)

            predictions = model.predict(input_data)
            size, pdi = predictions[0]

            st.markdown("**Model predictions**")
            display_output(size, pdi)

    # 2.2 XGBoost
    elif option == "XGBoost":
        st.subheader("XGBoost")
        st.markdown(
            "Using [eXtreme Gradient Boosting](https://en.wikipedia.org/wiki/XGBoost) the model provides joint predictions for `SIZE` and `PDI`."
        )

        with open(xgboost_model, "rb") as file:
            model = pickle.load(file)

        st.write(f"Loaded {xgboost_model} with the following performance metrics.")
        st.table(
            pd.DataFrame(
                {
                    "Metric": [
                        "R-squared",
                        "Mean Squared Error (MSE)",
                        "Mean Absolute Error (MAE)",
                    ],
                    "Value": [
                        0.32854801416397095,
                        1967.81201171875,
                        14.646432876586914,
                    ],
                }
            )
        )

        if st.button("Predict"):
            input_data = pd.DataFrame(
                {
                    "ML": [ml],
                    "CHIP": [chip],
                    "TLP": [tlp],
                    "ESM": [0.0 if esm_disabled else esm],
                    "HSPC": [0.0 if hspc_disabled else hspc],
                    "CHOL": [chol],
                    "PEG": [peg],
                    "TFR ": [tfr],
                    "FRR": [frr],
                    "BUFFER": [buffer],
                }
            )

            input_data_ = pd.DataFrame(
                {
                    "ML": [ml],
                    "CHIP": [chip],
                    "TLP": [tlp],
                    "ESM": [0.0 if esm_disabled else esm],
                    "HSPC": [0.0 if hspc_disabled else hspc],
                    "CHOL": [chol],
                    "PEG": [peg],
                    "TFR ": [tfr],
                    "FRR": [frr],
                    "BUFFER": [buffer],
                    "OUTPUT": [1],
                }
            )

            st.markdown("**Input data**")
            st.write(input_data)

            predictions = model.predict(input_data_)
            size, pdi = predictions[0]

            st.markdown("**Model predictions**")
            display_output(size, pdi)

    # 2.3 Inverse model
    elif option == "Inverse model":
        st.subheader("Inverse model")
        st.markdown(
            "This inference method pivots the pretrained XGBoost model to solve an [inverse problem](https://en.wikipedia.org/wiki/Inverse_problem): given target `SIZE` and `PDI` returns predictions for the other numerical features."
        )
        st.warning("🧑🏻‍💻 Work in progress... Only numerical predictions available.")

        with open(inverse_xgboost_model, "rb") as file:
            model = pickle.load(file)

        st.write(f"Loaded {inverse_xgboost_model} with the folloing performance metrics.")
        st.table(
            pd.DataFrame(
                {
                    "Metric": [
                        "R-squared",
                        "Mean Squared Error",
                        "Mean Absolute Error",
                    ],
                    "Value": [
                        0.11767681688070297,
                        89.26007080078125,
                        3.7459394931793213,
                    ],
                }
            )
        )

        if st.button("Predict"):
            input_data = pd.DataFrame({"SIZE": [size], "PDI": [pdi]})
            predictions = model.predict(input_data)
            predictions = np.abs(predictions)
            predictions = np.where(predictions < 0.5, 0, predictions)

            st.markdown("**Input data**")
            st.write(input_data)

            st.markdown("**Model predictions**")
            prediction_df = pd.DataFrame(
                predictions,
                columns=["ESM", "HSPC", "CHOL", "PEG", "TFR", "FRR"],
            )
            st.write(prediction_df)

    # 2.4 Advanced single target models
    elif option == "Single target models":
        st.header("Advanced single target models")
        st.markdown(
            r"""These models will predict one target variable at a time, for example `SIZE`, given all the other features. If you want to predict `SIZE` set `PDI` to zero.
                        If you want to predict `PDI` set `SIZE` to zero. These kind of advanced models can be integrated in a single mixture of experts to predict jointly the target
                        variables, as detailed in www....................."""
        )

        st.subheader("Preview of some available advanced models for predicting `SIZE` or `PDI`")

        st.write("`xgboost_one_target_pdi`")
        st.table(
            pd.DataFrame(
                {
                    "Metric": [
                        "R-squared",
                        "Mean Squared Error (MSE)",
                        "Mean Absolute Error (MAE)",
                    ],
                    "Value": [
                        0.9922573566436768,
                        0.000269298150669783355,
                        0.012760158628225327,
                    ],
                }
            )
        )

        st.write("`xgboost_one_target_size`")
        st.table(
            pd.DataFrame(
                {
                    "Metric": [
                        "R-squared",
                        "Mean Squared Error (MSE)",
                        "Mean Absolute Error (MAE)",
                    ],
                    "Value": [
                        0.9996433258056641,
                        492.2565612792969,
                        8.92746639251709,
                    ],
                }
            )
        )

        st.write("With an improvement in terms of metrics given by the following plot for `SIZE`. Similar results hold for `PDI`.")

        metrics = ["R-squared", "MSE", "MAE"]
        before = [0.3285, 1967.81, 14.65]
        after = [0.9996433258056641, 492.2565612792969, 8.92746639251709]
        plot_model_performance(metrics, before, after)

        st.write("Depending on the number of samples available for training and validation we can boost the performances even further.")
        st.subheader("Try the `best_xgboost_model_size` model")

        if st.button("Predict"):
            if pdi > 0.0 and size == 0.0:
                input_data = pd.DataFrame(
                    {
                        "ESM": [esm],
                        "HSPC": [hspc],
                        "CHOL": [chol],
                        "PEG": [peg],
                        "TFR": [tfr],
                        "FRR": [frr],
                        "AQUEOUS": [buffer],
                        "PDI": [pdi],
                    }
                )

                with open("_models/best_xgboost_model_size.pkl", "rb") as file:
                    model = pickle.load(file)

                st.markdown("**Input data**")
                st.write(input_data)

                predictions = model.predict(input_data)
                predicted_size = predictions[0]

                st.markdown("**Model predictions**")
                st.write(f"Predicted `SIZE`: {predicted_size}")

            elif size > 0.0 and pdi == 0.0:
                st.subheader("Try the `best_xgboost_model_pdi` model")

                input_data = pd.DataFrame(
                    {
                        "ESM": [esm],
                        "HSPC": [hspc],
                        "CHOL": [chol],
                        "PEG": [peg],
                        "TFR": [tfr],
                        "FRR": [frr],
                        "AQUEOUS": [buffer],
                        "SIZE": [size],
                    }
                )

                with open("_models/best_xgboost_model_pdi.pkl", "rb") as file:
                    model = pickle.load(file)

                st.markdown("**Input data**")
                st.write(input_data)

                predictions = model.predict(input_data)
                predicted_pdi = predictions[0]

                st.markdown("**Model predictions**")
                st.write(f"Predicted `PDI`: {predicted_pdi}")

            else:
                st.warning("Invalid inputs.")


################################################
# 3. DATA EXPLORATION
################################################

elif section == "Data Exploration":
    option = st.sidebar.selectbox(
        "Choose a Visualization:",
        [
            "Ridgeline plot",
            "Correlation heatmaps",
            "Alluvial plot",
            "Feature importance",
            "PCA and clustering",
        ],
    )

    # 3.1 Ridgeline plot
    if option == "Ridgeline plot":
        st.header("Ridgeline plot")
        st.markdown(
            "Displays the distributions of individual features as "
            "[overlapping density curves](https://en.wikipedia.org/wiki/Ridgeline_plot)."
        )

        st.subheader("Numerical columns summary")
        st.write(numeric_data.describe())

        scaler = StandardScaler()
        df_standardized = pd.DataFrame(
            scaler.fit_transform(numeric_data),
            columns=numeric_data.columns,
        )

        st.subheader("Stacked densities")

        fig = plot_ridgeline_matplotlib(
            df_standardized,
            x_range=(-5, 10),
            figsize=(12, 8),
            title="Ridgeline Plot",
        )

        st.pyplot(fig)
        plt.close(fig)

    # 3.2 Correlation heatmap
    elif option == "Correlation heatmaps":
        st.header("Correlation heatmap")
        st.markdown(
            "Displays the [correlation between numerical features](https://en.wikipedia.org/wiki/Heat_map) in the dataset."
        )

        formed = st.text_input(
            "Are you interested in the dataset of formed liposomes? Answer YES (Y) or NO (N):",
            "Y",
        )
        n_ids = st.number_input(
            "Enter the number of formulations you desire to visualize:",
            min_value=9,
            max_value=15,
            value=10,
        )

        numeric_cols = data.select_dtypes(include=["number"]).columns

        if formed == "YES" or formed == "Y":
            # 3.2.1 Correlations by features
            st.subheader("Correlations by features")

            yes_data = data[data["OUTPUT"] == "YES"]
            numeric_yes_data = yes_data[numeric_cols]
            correlation_matrix = numeric_yes_data.corr()

            fig, ax = plt.subplots(figsize=(10, 8))
            sns.heatmap(correlation_matrix, annot=True, fmt=".2f", cmap="coolwarm", ax=ax)
            fig.tight_layout()
            st.pyplot(fig)
            plt.close(fig)

            # 3.2.2 Correlations by IDs
            st.subheader("Correlations by IDs")

            data_red = data[:n_ids]
            yes_data_red = data_red[data_red["OUTPUT"] == "YES"]
            numeric_yes_data_red = yes_data_red[numeric_cols]
            correlation_matrix_t = numeric_yes_data_red.T.corr()

            fig, ax = plt.subplots(figsize=(10, 8))
            sns.heatmap(correlation_matrix_t, annot=True, fmt=".2f", cmap="coolwarm", ax=ax)
            fig.tight_layout()
            st.pyplot(fig)
            plt.close(fig)

            # 3.2.3 Clustered correlations by features
            st.subheader("Clustered correlations by features")

            _ = sch.linkage(numeric_yes_data, method="ward")
            grid = sns.clustermap(
                numeric_yes_data.corr(),
                method="ward",
                cmap="coolwarm",
                annot=True,
                figsize=(10, 8),
            )
            st.pyplot(grid.fig)
            plt.close(grid.fig)

            # 3.2.4 Clustered correlations by IDs
            st.subheader("Clustered correlations by IDs")

            _ = sch.linkage(numeric_yes_data_red.T, method="ward")
            grid = sns.clustermap(
                numeric_yes_data_red.T.corr(),
                method="ward",
                cmap="coolwarm",
                annot=True,
                figsize=(10, 8),
            )
            st.pyplot(grid.fig)
            plt.close(grid.fig)

        else:
            # 3.2.5 Correlations by features
            no_data = data[data["OUTPUT"] == "NO"]
            numeric_no_data = no_data[numeric_cols]
            correlation_matrix = numeric_no_data.corr()

            fig, ax = plt.subplots(figsize=(10, 8))
            sns.heatmap(correlation_matrix, annot=True, fmt=".2f", cmap="coolwarm", ax=ax)
            fig.tight_layout()
            st.pyplot(fig)
            plt.close(fig)

            # 3.2.6 Clustered correlations by IDs
            st.subheader("Clustered correlations by IDs")

            data_red = data[:n_ids]
            no_data_red = data_red[data_red["OUTPUT"] == "NO"]
            numeric_no_data_red = no_data_red[numeric_cols]
            correlation_matrix_t = numeric_no_data_red.T.corr()

            fig, ax = plt.subplots(figsize=(10, 8))
            sns.heatmap(correlation_matrix_t, annot=True, fmt=".2f", cmap="coolwarm", ax=ax)
            fig.tight_layout()
            st.pyplot(fig)
            plt.close(fig)

    # 3.3 Alluvial plot
    elif option == "Alluvial plot":
        st.header("Alluvial plot")
        st.write(
            "Displays the [flow between categorical data](https://en.wikipedia.org/wiki/Alluvial_diagram). Here the target variables `SIZE` and `PDI` are made categorical by grupping them into the following categories:"
        )
        st.markdown("- **S** (Small) if `SIZE` is less than 100;")
        st.markdown("- **M** (Medium) if `SIZE` is between 100 and 199;")
        st.markdown("- **L** (Large) if `SIZE` is 200 or more;")
        st.markdown("- **HMD** (Highly Monodisperse) if `PDI` is less than 0.1;")
        st.markdown("- **MD** (Monodisperse) if `PDI` is between 0.1 and 0.25;")
        st.markdown("- **PLD** (Polydisperse) if `PDI` is 0.25 or more.")

        data_ = data.copy()

        for col in data_.columns:
            if data_[col].dtype == "object":
                data_[col] = data_[col].astype(str)

        data_["SIZE"] = data_["SIZE"].apply(categorize_size)
        data_["PDI"] = data_["PDI"].apply(categorize_pdi)

        # 3.3.1 Sankey diagram of two variables
        st.subheader("Sankey diagram of two variables")

        source = st.selectbox(
            "Choose the first categorical variable of interest.",
            ("ML", "CHIP", "BUFFER", "OUTPUT", "SIZE", "PDI"),
        )
        target = st.selectbox(
            "Choose the second categorical variable of interest (different from the first).",
            ("CHIP", "ML", "BUFFER", "OUTPUT", "SIZE", "PDI"),
        )

        if source == target:
            st.write(":red[INPUT ERROR. The selected variables cannot be equal.]")
        else:
            value_counts = data_.groupby([source, target]).size().reset_index(name="value")

            categories = list(set(value_counts[source]).union(set(value_counts[target])))
            category_to_index = {category: i for i, category in enumerate(categories)}

            sources = value_counts[source].map(category_to_index).tolist()
            targets = value_counts[target].map(category_to_index).tolist()
            values = value_counts["value"].tolist()

            sankey_fig = go.Figure(
                data=[
                    go.Sankey(
                        node=dict(
                            pad=15,
                            thickness=20,
                            line=dict(color="black", width=0.5),
                            label=categories,
                        ),
                        link=dict(
                            source=sources,
                            target=targets,
                            value=values,
                        ),
                    )
                ]
            )

            st.plotly_chart(sankey_fig)

        # 3.3.2 Sankey diagram of more than two flows
        st.subheader("Sankey diagram of a customizable categorical flow")

        default_selection = ["ML", "CHIP", "BUFFER", "OUTPUT"]

        selected_columns = st.multiselect(
            "Select categorical variables to visualize:",
            options=["ML", "CHIP", "BUFFER", "OUTPUT", "SIZE", "PDI"],
            default=default_selection,
        )

        if len(selected_columns) < 3:
            st.warning("Please select more then two two categorical variables.")
        else:
            value_counts = data_.groupby(selected_columns).size().reset_index(name="value")

            all_categories = []
            for col in selected_columns:
                all_categories.extend(value_counts[col].unique())

            all_categories = list(set(all_categories))
            category_to_index = {category: index for index, category in enumerate(all_categories)}

            sources = []
            targets = []
            values = []

            for _, row in value_counts.iterrows():
                for i in range(len(selected_columns) - 1):
                    sources.append(category_to_index[row[selected_columns[i]]])
                    targets.append(category_to_index[row[selected_columns[i + 1]]])
                    values.append(row["value"])

            fig = go.Figure(
                data=[
                    go.Sankey(
                        node=dict(
                            pad=15,
                            thickness=20,
                            line=dict(color="black", width=0.5),
                            label=all_categories,
                        ),
                        link=dict(
                            source=sources,
                            target=targets,
                            value=values,
                        ),
                    )
                ]
            )

            st.plotly_chart(fig, key="cust_sank")

    # 3.4 Feature importance
    elif option == "Feature importance":
        st.header("Feature importance with a Random Forest Regressor")
        st.markdown(
            "Displays the [importance of each feature](https://en.wikipedia.org/wiki/Feature_selection) for predicting a set of input targets."
        )

        # 3.4.1 Single target feature importance
        st.subheader("Single target feature importance")

        target_feature = st.selectbox("Select a target feature.", ("SIZE", "PDI"))

        numeric_data_ = numeric_data.copy()

        if target_feature == "SIZE":
            numeric_data_ = numeric_data.drop(columns=["PDI"])
        if target_feature == "PDI":
            numeric_data_ = numeric_data.drop(columns=["SIZE"])

        X = numeric_data_.drop(columns=[target_feature])

        if "SIZE" in X.columns:
            X = X.drop(columns=["SIZE"])
        if "PDI" in X.columns:
            X = X.drop(columns=["PDI"])

        y = numeric_data_[target_feature]

        rf = RandomForestRegressor(n_estimators=100, random_state=42)
        rf.fit(X, y)

        feature_importances = pd.DataFrame(
            {
                "Feature": X.columns,
                "Importance": rf.feature_importances_,
            }
        ).sort_values(by="Importance", ascending=False)

        fig, ax = plt.subplots(figsize=(10, 8))
        sns.barplot(
            data=feature_importances,
            x="Importance",
            y="Feature",
            hue="Feature",
            palette="viridis",
            dodge=False,
            ax=ax,
        )

        if ax.get_legend() is not None:
            ax.get_legend().remove()

        fig.tight_layout()
        st.pyplot(fig)
        plt.close(fig)

        # 3.4.2 Two targets feature importance
        target_feature_1 = st.selectbox("Select the first target feature.", ("SIZE", "PDI"))
        target_feature_2 = st.selectbox("Select the second target feature.", ("PDI", "SIZE"))

        if target_feature_1 == target_feature_2:
            st.error("INPUT ERROR: The selected variables cannot be the same.")
        else:
            correlation_matrix = numeric_data.corr()
            target_features = [target_feature_1, target_feature_2]

            joint_correlation = (
                correlation_matrix[target_features]
                .drop(index=target_features, errors="ignore")
                .mean(axis=1)
            )

            joint_correlation_df = joint_correlation.reset_index()
            joint_correlation_df.columns = ["Feature", "Mean Correlation"]
            joint_correlation_df = joint_correlation_df.sort_values(
                by="Mean Correlation",
                ascending=False,
            )

            fig, ax = plt.subplots(figsize=(10, 8))
            sns.barplot(
                data=joint_correlation_df,
                x="Mean Correlation",
                y="Feature",
                palette="viridis",
                dodge=False,
                ax=ax,
            )
            ax.set_xlabel("Mean Correlation with Targets")
            ax.set_ylabel("Feature")
            ax.set_title("Feature Importance for Selected Targets")

            fig.tight_layout()
            st.pyplot(fig)
            plt.close(fig)

    # 3.5 PCA and clustering
    elif option == "PCA and clustering":
        st.header("Principal component analysis and associated clusters.")
        st.markdown(
            """Displays the principal data features and their clusters using a standard [PCA](https://en.wikipedia.org/wiki/Principal_component_analysis) algorithm, a [UMAP](https://en.wikipedia.org/wiki/Nonlinear_dimensionality_reduction#Uniform_manifold_approximation_and_projection),
        and [k-means](https://en.wikipedia.org/wiki/K-means_clustering) clustering. These plots help in model design, pinpoint processing requirements and can provide guidance during training and validation."""
        )

        # 3.5.1 Standard PCA
        st.subheader("Standard PCA onto two principal components")

        scaler = StandardScaler()
        data_standardized = scaler.fit_transform(data_numeric)

        pca = PCA(n_components=2)
        pca_result = pca.fit_transform(data_standardized)

        df_pca = pd.DataFrame(pca_result, columns=["PC1", "PC2"])

        fig, ax = plt.subplots(figsize=(8, 6))
        sns.kdeplot(
            x=df_pca["PC1"],
            y=df_pca["PC2"],
            fill=True,
            cmap="Blues",
            thresh=0.05,
            ax=ax,
        )
        ax.set_xlabel("Principal Component 1")
        ax.set_ylabel("Principal Component 2")
        ax.set_title("2D Density Plot of Data Distribution (PCA Transformed)")

        fig.tight_layout()
        st.pyplot(fig)
        plt.close(fig)

        # 3.5.2 UMAP
        st.subheader("UMAP decomposition and k-means clustering")

        scaler = StandardScaler()
        data_standardized = scaler.fit_transform(data_numeric)

        n_components = st.slider(
            "Select number of UMAP components:",
            min_value=2,
            max_value=10,
            value=2,
        )

        umap_model = umap.UMAP(n_components=n_components, random_state=42)
        umap_result = umap_model.fit_transform(data_standardized)

        n_clusters = st.slider(
            "Select number of clusters for K-Means:",
            min_value=2,
            max_value=10,
            value=3,
        )

        kmeans = KMeans(n_clusters=n_clusters, random_state=42)
        clusters = kmeans.fit_predict(umap_result)

        df_umap = pd.DataFrame(
            umap_result,
            columns=[f"UMAP{i + 1}" for i in range(n_components)],
        )
        df_umap["Cluster"] = clusters

        st.write(
            f"UMAP Projection with K-Means Cluster Density "
            f"(UMAP components: {n_components}, Clusters: {n_clusters}):"
        )

        fig, ax = plt.subplots(figsize=(10, 8))
        sns.kdeplot(
            x=df_umap["UMAP1"],
            y=df_umap["UMAP2"],
            hue=df_umap["Cluster"],
            fill=True,
            palette="tab10",
            thresh=0.05,
            alpha=0.7,
            ax=ax,
        )
        ax.set_xlabel("UMAP Component 1")
        ax.set_ylabel("UMAP Component 2")
        ax.set_title(f"UMAP Projection with K-Means Cluster Density (Clusters: {n_clusters})")

        fig.tight_layout()
        st.pyplot(fig)
        plt.close(fig)