ls_server.py

import h5py
from bokeh.plotting import curdoc
from bokeh.models import ColumnDataSource
from bokeh.layouts import column, row
import json
import sys
import argparse
import os
from skimage import measure
from bokeh.plotting import output_file, save
import numpy as np
from visualizer.ls_decisionboundary import LSBoundaryVisualizer


def extract_boundary_lines(xx, yy, zz):
    contours = measure.find_contours(zz, level=0.5)  # Assuming boundary at 0.5 probability
    xs, ys = [], []
    for contour in contours:
        xs.append(xx[0, 0] + contour[:, 1] * (xx[0, -1] - xx[0, 0]) / zz.shape[1])
        ys.append(yy[0, 0] + contour[:, 0] * (yy[-1, 0] - yy[0, 0]) / zz.shape[0])
    return xs, ys

parser = argparse.ArgumentParser(description="Launch the Bokeh server with an HDF5 file, this plot is to display changes in model behavior over training step.")
parser.add_argument("--file", type=str, required=True, help="Path to the HDF5 file")
parser.add_argument("--output", type=str, required=False, help="If specified filename, while running on python not bokeh serve, the html will be saved in ./output")
args = parser.parse_args()


# Load the HDF5 file
h5_file = args.file

# Check if the file has an .h5 extension
if not h5_file.lower().endswith(".h5"):
    print(f"Error: The input file '{h5_file}' is not an HDF5 (.h5) file.")
    sys.exit(1)

# Check if the file exists
if not os.path.isfile(h5_file):
    print(f"Error: The file '{h5_file}' does not exist.")
    sys.exit(1)

if args.output is not None:
    os.makedirs('./output', exist_ok=True)
    output_file(filename=f"./output/{args.output}.html", title="Static HTML file", mode="inline")

with h5py.File(h5_file, "r") as f:
    read = f["config"]["config_data"][()]
    config_json = read.decode("utf-8")
    config = json.loads(config_json)
    dataset = config.get("dataset")
    max_epoch = config.get("max_epochs")

    X_coord = np.array(f["coord/X_train"])
    y_train = np.array(f["coord/y_train"])

    sentivities = [f[f"scores/epoch_{epoch}"]["sensitivities"][()] for epoch in range(max_epoch)]
    all_epoch_noises = [f[f"scores/epoch_{epoch}"]["noise"][()] for epoch in range(max_epoch)]

    #all_induced_noises = [f[f"scores/epoch_{epoch}"]["all_noise"][()] for epoch in range(max_epoch)]

    # Extract decision boundary data
    xx = [f[f"scores/epoch_{epoch}"]["decision_boundary"]["xx"][:] for epoch in range(max_epoch)]
    yy = [f[f"scores/epoch_{epoch}"]["decision_boundary"]["yy"][:] for epoch in range(max_epoch)]
    Z = [f[f"scores/epoch_{epoch}"]["decision_boundary"]["Z"][:] for epoch in range(max_epoch)]

colors = ["white", "white"]
marker = ["circle", "star"]

xs = []
ys = []
for epoch in range(max_epoch):
    xx_step = xx[epoch]
    yy_step = yy[epoch]
    zz_step = Z[epoch]
    
    # Extract boundary for each step
    boundary_x, boundary_y = extract_boundary_lines(xx_step, yy_step, zz_step)
    xs.append(boundary_x)
    ys.append(boundary_y)

# Compute global min and max noise across all epochs
min_noise, max_noise = np.min(all_epoch_noises), np.max(all_epoch_noises)

# Avoid division by zero
if max_noise - min_noise == 0:
    normalized_noises = np.ones_like(all_epoch_noises)  # Default to uniform size
else:
    normalized_noises = (all_epoch_noises - min_noise) / (max_noise - min_noise)

# Scale to desired range (e.g., 5 to 50)
min_size, max_size = 5, 50
scaled_sizes = min_size + normalized_noises * (max_size - min_size)

# Convert to a list of lists (each epoch's sizes as a list)
scaled_sizes_list = scaled_sizes.tolist()

# Scale alpha per epoch
scaled_alphas_list = []
for epoch_noises in all_epoch_noises:
    min_noise, max_noise = np.min(epoch_noises), np.max(epoch_noises)

    if max_noise - min_noise == 0:
        normalized_noises = np.ones_like(epoch_noises)
    else:
        normalized_noises = (epoch_noises - min_noise) / (max_noise - min_noise)

    num_levels = 4
    alpha_levels = [0.05, 0.4, 0.7, 1.0]
    quantiles = np.linspace(0, 1, num_levels + 1)
    alpha_assignments = np.zeros_like(normalized_noises)

    for i in range(num_levels):
        lower_bound = quantiles[i]
        upper_bound = quantiles[i + 1]
        mask = (normalized_noises >= lower_bound) & (normalized_noises < upper_bound)
        alpha_assignments[mask] = alpha_levels[i]

    scaled_alphas_list.append(alpha_assignments.tolist())


shared_resource = ColumnDataSource(data={
    "epoch": list(range(max_epoch)),
    "xs": xs,
    "ys": ys,
    "size": scaled_sizes_list,
    "alpha": scaled_alphas_list
})

shared_source = ColumnDataSource(data={
    "x": X_coord[:, 0],
    "y": X_coord[:, 1],
    "class": y_train,
    "color": [colors[cls] for cls in y_train],
    "marker": [marker[cls] for cls in y_train],
    "size": scaled_sizes_list[0],
    "alpha": scaled_alphas_list[0]
})

boundary = LSBoundaryVisualizer(shared_source, shared_resource, max_epoch, colors)

boundary_layout = column(boundary.get_layout(), sizing_mode="scale_both")

layout = row(boundary_layout)

curdoc().add_root(layout)

if args.output is not None:
    layout.sizing_mode = "scale_both"
    save(layout)