app.py

import matplotlib.pyplot as plt
import numpy as np
import streamlit as st
import tensorflow as tf
from sklearn import datasets
from sklearn.model_selection import train_test_split
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Sequential


# Function to create datasets
def create_datasets(n_samples=1500):
    noisy_circles = datasets.make_circles(
        n_samples=n_samples, factor=0.5, noise=0.05, random_state=170
    )
    noisy_moons = datasets.make_moons(n_samples=n_samples, noise=0.05, random_state=170)
    blobs = datasets.make_blobs(n_samples=n_samples, random_state=170)
    rng = np.random.RandomState(170)
    no_structure = rng.rand(n_samples, 2), np.zeros(n_samples)
    X, y = datasets.make_blobs(n_samples=n_samples, random_state=170)
    transformation = [[0.6, -0.6], [-0.4, 0.8]]
    X_aniso = np.dot(X, transformation)
    aniso = (X_aniso, y)
    return [noisy_circles, noisy_moons, blobs, no_structure, aniso]


# Function to create a simple one-layer neural network
def create_model(input_shape, layers):
    model = Sequential()
    # Add layers based on the selection
    if layers == 'one-layer neural network':
        model.add(Dense(1, input_shape=input_shape, activation="sigmoid"))
    elif layers == 'three-layer neural network':
        model.add(Dense(128, input_shape=input_shape, activation="relu"))
        model.add(Dense(128, activation="relu"))
        model.add(Dense(128, activation="relu"))
        model.add(Dense(1, activation="sigmoid"))
    elif layers == 'five-layer neural network':
        model.add(Dense(128, input_shape=input_shape, activation="relu"))
        model.add(Dense(128, activation="relu"))
        model.add(Dense(128, activation="relu"))
        model.add(Dense(128, activation="relu"))
        model.add(Dense(128, activation="relu"))
        model.add(Dense(1, activation="sigmoid"))
    return model


def soft_quantized_influence_measure_sigmoid(
    y_true: tf.Tensor, y_pred: tf.Tensor, threshold: float = 0.5
) -> tf.Tensor:
    # Ensure y_true and y_pred are of type float32
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)

    # Calculate the global mean and standard deviation of the true values
    y_global_mean = tf.reduce_mean(y_true)
    y_std = tf.math.reduce_std(y_true)

    # Compute the absolute error between predictions and true values
    error = y_true - y_pred
    sigmoid_error = tf.sigmoid(error)

    # Determine if each error is small (within the threshold) or not
    is_small_error = sigmoid_error <= threshold

    # Count the number of small and large errors, converting boolean masks to float for sum operation
    n1 = tf.reduce_sum(tf.cast(is_small_error, tf.float32))  # Count of small errors
    n2 = tf.reduce_sum(tf.cast(~is_small_error, tf.float32))  # Count of large errors

    # Calculate weighted error losses based on the counts of small and large errors
    true_error_loss = tf.square(error) * tf.square(n1)
    false_error_loss = tf.square(error) * tf.square(n2)

    # Ensure the losses are of the same type as the condition in tf.where
    true_error_loss = tf.cast(true_error_loss, tf.float32)
    false_error_loss = tf.cast(false_error_loss, tf.float32)

    # Apply weights to errors based on whether they are small or large
    final_loss = tf.where(is_small_error, true_error_loss, false_error_loss)

    # Normalize the final loss - ensure division is properly handled with float types
    num_elements = tf.cast(tf.size(y_true), tf.float32)  # Cast number of elements to float
    final_loss = tf.reduce_mean(final_loss) / tf.square(y_std) ** 2 / num_elements

    return final_loss


def soft_loss_wrapper(threshold: float) -> callable:
    """
    Creates and returns a loss function based on a given threshold.

    The returned loss function applies a sigmoid transformation to measure influence in a soft, quantized manner.

    Parameters
    ----------
    threshold : float
        A threshold value used within the returned loss function for calculating the influence.

    Returns
    -------
    callable
        A loss function that can be used in optimization problems where a soft measurement of influence is desired.
    
    Example
    -------
    >>> loss_function = soft_loss_wrapper(0.5)
    >>> # Now `loss_function` can be used as a parameter in optimization routines.
    """
    # Assuming that `soft_quantized_influence_measure_sigmoid` is defined elsewhere,
    # and this line links the given threshold with that function.
    return soft_quantized_influence_measure_sigmoid


# Main Streamlit app
def main():
    # App mode and title
    st.set_page_config(layout="wide")
    st.title("Classification Data and Neural Network Training Visualization")

    # Sidebar configuration
    st.sidebar.header('Configuration')
    n_samples = st.sidebar.slider('Number of Samples', min_value=500, max_value=3000, value=1500, step=100)
    threshold = st.sidebar.slider('Threshold', min_value=0.0, max_value=1.0, value=0.1, step=0.01)

    # Sidebar: model selection
    model_option = st.sidebar.selectbox(
        'Choose the neural network model:',
        ('one-layer neural network', 'three-layer neural network', 'five-layer neural network')
    )
    epochs = st.sidebar.slider('Number of Epochs', min_value=10, max_value=100, value=20, step=5)

    # Generate datasets
    dataset_list = create_datasets(n_samples=n_samples)

    # Containers for the first row (scatter plots of datasets)
    st.header("Datasets")
    scatter_cols = st.columns(5)

    st.header("Training and Validation Loss with MSE")
    mse_loss_cols = st.columns(5)

    st.header("Training and Validation Loss with Custom Loss")
    custom_loss_cols = st.columns(5)

    dataset_names = [
        "Noisy Circles",
        "Noisy Moons",
        "Blobs",
        "No Structure",
        "Anisotropic",
    ]

    # Loop over each dataset
    for col_scatter, col_mse_loss, col_custom_loss, dataset, name in zip(
        scatter_cols, mse_loss_cols, custom_loss_cols, dataset_list, dataset_names
    ):
        X, y = dataset

        # Scatter plot of datasets
        with col_scatter:
            fig, ax = plt.subplots()
            if y is not None:
                ax.scatter(X[:, 0], X[:, 1], c=y, cmap="viridis", edgecolor="k", s=20)
            else:
                ax.scatter(X[:, 0], X[:, 1], cmap="viridis", edgecolor="k", s=20)
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_title(name)
            st.pyplot(fig)

        if y is not None:
            # Split the dataset
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=0.2, random_state=42
            )

            # Train with MSE loss
            model_mse = create_model((2,), model_option)
            model_mse.compile(
                optimizer="sgd", loss="binary_crossentropy", metrics=["accuracy"]
            )
            history_mse = model_mse.fit(
                X_train, y_train, validation_split=0.2, epochs=epochs, verbose=0
            )

            # Evaluate and plot MSE loss
            with col_mse_loss:
                test_loss_mse, test_acc_mse = model_mse.evaluate(
                    X_test, y_test, verbose=0
                )
                fig, ax = plt.subplots()
                ax.plot(history_mse.history["loss"], label="Training Loss")
                ax.plot(history_mse.history["val_loss"], label="Validation Loss")
                ax.set_title(f"MSE\nTest Loss: {test_loss_mse:.4f}")
                ax.legend()
                st.pyplot(fig)

            # Train with custom loss
            model_custom = create_model((2,), model_option)
            model_custom.compile(
                optimizer="sgd",
                loss=soft_loss_wrapper(threshold),
                metrics=["accuracy"],
            )
            history_custom = model_custom.fit(
                X_train, y_train, validation_split=0.2, epochs=epochs, verbose=0
            )

            # Evaluate and plot custom loss
            with col_custom_loss:
                test_loss_custom, test_acc_custom = model_custom.evaluate(
                    X_test, y_test, verbose=0
                )
                fig, ax = plt.subplots()
                ax.plot(history_custom.history["loss"], label="Training Loss")
                ax.plot(history_custom.history["val_loss"], label="Validation Loss")
                ax.set_title(f"Custom Loss\nTest Loss: {test_loss_custom:.4f}")
                ax.legend()
                st.pyplot(fig)


if __name__ == "__main__":
    main()