Moe vae tss

Mixture of Experts Variational Autoencoder (MoE-VAE) for TSS¶

This example demonstrates how to use the Mixture of Experts Variational Autoencoder (MoE-VAE) to predict Total Suspended Solids (TSS) using PACE data.

In [ ]:

# %pip install "hypercoast[all]"

# %pip install "hypercoast[all]"

Import libraries¶

In [ ]:

import os
import pickle
import numpy as np
import pandas as pd
import rasterio
import torch

import os import pickle import numpy as np import pandas as pd import rasterio import torch

In [ ]:

from hypercoast import download_file
from hypercoast.emit_utils.plot_and_save import plot_results_with_density
from hypercoast.moe_vae import (
    load_real_data_Robust,
    load_real_test_Robust,
    calculate_metrics,
    plot_results,
    save_results_to_excel,
    save_results_from_excel_for_test,
    preprocess_pace_data_Robust,
    infer_and_visualize_single_model_Robust,
    npy_to_tif,
    MoE_VAE,
    train,
    evaluate,
)

from hypercoast import download_file from hypercoast.emit_utils.plot_and_save import plot_results_with_density from hypercoast.moe_vae import ( load_real_data_Robust, load_real_test_Robust, calculate_metrics, plot_results, save_results_to_excel, save_results_from_excel_for_test, preprocess_pace_data_Robust, infer_and_visualize_single_model_Robust, npy_to_tif, MoE_VAE, train, evaluate, )

Download data¶

In [ ]:

url = "https://github.com/opengeos/datasets/releases/download/hypercoast/pace_moe_data.zip"

url = "https://github.com/opengeos/datasets/releases/download/hypercoast/pace_moe_data.zip"

In [ ]:

download_file(url, quiet=False)

download_file(url, quiet=False)

Set data paths¶

In [ ]:

nc_path = "./data/PACE_OCI.20240929T185124.L2.OC_AOP.V3_0.nc"
pace_rgb_path = "./data/snapshot-2024-08-10T00_00_00Z.tif"
wavelength_filepath = "./data/pace_wavelengths.csv"
# === Dataset paths ===
excel_path_train = "./data/Gloria_updated_07242025.xlsx"
test_files = [
    "./data/GOA_insitu_data_07242025updated.xlsx",
]
base_save_dir = "./test"
save_dir = os.path.join(base_save_dir, "TSS")
os.makedirs(save_dir, exist_ok=True)

nc_path = "./data/PACE_OCI.20240929T185124.L2.OC_AOP.V3_0.nc" pace_rgb_path = "./data/snapshot-2024-08-10T00_00_00Z.tif" wavelength_filepath = "./data/pace_wavelengths.csv" # === Dataset paths === excel_path_train = "./data/Gloria_updated_07242025.xlsx" test_files = [ "./data/GOA_insitu_data_07242025updated.xlsx", ] base_save_dir = "./test" save_dir = os.path.join(base_save_dir, "TSS") os.makedirs(save_dir, exist_ok=True)

Set PACE wavelengths for TSS¶

The wavelength selection for TSS is the same as aCDOM440, focusing on 150 bands from 400-720nm.

In [ ]:

wv_PACE = pd.read_csv(wavelength_filepath)["wavelength"].tolist()
selected_bands = wv_PACE

wv_PACE = pd.read_csv(wavelength_filepath)["wavelength"].tolist() selected_bands = wv_PACE

Read RGB image¶

In [ ]:

with rasterio.open(pace_rgb_path) as ds:
    R, G, B = ds.read(1), ds.read(2), ds.read(3)
    extent = [ds.bounds.left, ds.bounds.right, ds.bounds.bottom, ds.bounds.top]
    rgb_image = np.stack((R, G, B), axis=-1)

with rasterio.open(pace_rgb_path) as ds: R, G, B = ds.read(1), ds.read(2), ds.read(3) extent = [ds.bounds.left, ds.bounds.right, ds.bounds.bottom, ds.bounds.top] rgb_image = np.stack((R, G, B), axis=-1)

Load training data¶

Using robust scaling for TSS prediction.

In [ ]:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}")

In [ ]:

(
    train_real_dl,
    test_real_dl,
    input_dim,
    output_dim,
    train_ids,
    test_ids,
    scalers_Rrs_real,
    scalers_dict,
) = load_real_data_Robust(
    excel_path=excel_path_train,
    selected_bands=selected_bands,
    target_parameter="TSS",
)

# Save scalers for later use
torch.save(scalers_dict, os.path.join(save_dir, "scaler.pt"))
with open(os.path.join(save_dir, "scalers_Rrs_real.pkl"), "wb") as f:
    pickle.dump(scalers_Rrs_real, f)

( train_real_dl, test_real_dl, input_dim, output_dim, train_ids, test_ids, scalers_Rrs_real, scalers_dict, ) = load_real_data_Robust( excel_path=excel_path_train, selected_bands=selected_bands, target_parameter="TSS", ) # Save scalers for later use torch.save(scalers_dict, os.path.join(save_dir, "scaler.pt")) with open(os.path.join(save_dir, "scalers_Rrs_real.pkl"), "wb") as f: pickle.dump(scalers_Rrs_real, f)

Load test data¶

In [ ]:

test_dls, test_ids_list, test_dates_list = [], [], []
for file in test_files:
    dl, _, _, ids, dates = load_real_test_Robust(
        excel_path=file,
        selected_bands=selected_bands,
        scaler_Rrs=scalers_Rrs_real,
        scalers_dict=scalers_dict,
        target_parameter="TSS",
    )
    test_dls.append(dl)
    test_ids_list.append(ids)
    test_dates_list.append(dates)

test_dls, test_ids_list, test_dates_list = [], [], [] for file in test_files: dl, _, _, ids, dates = load_real_test_Robust( excel_path=file, selected_bands=selected_bands, scaler_Rrs=scalers_Rrs_real, scalers_dict=scalers_dict, target_parameter="TSS", ) test_dls.append(dl) test_ids_list.append(ids) test_dates_list.append(dates)

Load PACE data¶

In [ ]:

test_loader, Rrs, mask, latitude, longitude = preprocess_pace_data_Robust(
    nc_path=nc_path,
    scaler_Rrs=scalers_Rrs_real,
    use_diff=False,
    full_band_wavelengths=np.array(selected_bands),
)

test_loader, Rrs, mask, latitude, longitude = preprocess_pace_data_Robust( nc_path=nc_path, scaler_Rrs=scalers_Rrs_real, use_diff=False, full_band_wavelengths=np.array(selected_bands), )

Initialize model¶

The model architecture for TSS uses smaller hidden dimensions compared to aCDOM440.

In [ ]:

model = MoE_VAE(
    input_dim=input_dim,
    output_dim=output_dim,
    latent_dim=16,
    encoder_hidden_dims=[64, 32],
    decoder_hidden_dims=[32, 64],
    activation="leakyrelu",
    use_norm="layer",
    use_dropout=False,
    use_softplus_output=False,
    num_experts=4,
    k=2,
    noisy_gating=True,
).to(device)

model = MoE_VAE( input_dim=input_dim, output_dim=output_dim, latent_dim=16, encoder_hidden_dims=[64, 32], decoder_hidden_dims=[32, 64], activation="leakyrelu", use_norm="layer", use_dropout=False, use_softplus_output=False, num_experts=4, k=2, noisy_gating=True, ).to(device)

Model training¶

In [ ]:

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
train_log = train(
    model,
    train_real_dl,
    device,
    epochs=200,
    optimizer=optimizer,
    save_dir=save_dir,
)
best_train_loss = train_log["best_loss"]

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) train_log = train( model, train_real_dl, device, epochs=200, optimizer=optimizer, save_dir=save_dir, ) best_train_loss = train_log["best_loss"]

Model evaluation¶

In [ ]:

predictions, actuals = evaluate(
    model, test_real_dl, device, TSS_scalers_dict=scalers_dict
)
epsilon, beta, rmse, rmsle, mape, bias, mae = calculate_metrics(predictions, actuals)
test_loss = mae

save_results_to_excel(
    test_ids, actuals, predictions, os.path.join(save_dir, "test_results.xlsx")
)

predictions, actuals = evaluate( model, test_real_dl, device, TSS_scalers_dict=scalers_dict ) epsilon, beta, rmse, rmsle, mape, bias, mae = calculate_metrics(predictions, actuals) test_loss = mae save_results_to_excel( test_ids, actuals, predictions, os.path.join(save_dir, "test_results.xlsx") )

In [ ]:

plot_results_with_density(
    predictions, actuals, save_dir, mode="test_results", xlim=(-2, 4), ylim=(-2, 4)
)

plot_results_with_density( predictions, actuals, save_dir, mode="test_results", xlim=(-2, 4), ylim=(-2, 4) )

In [ ]:

plot_results(predictions, actuals, save_dir, mode="test_results")

plot_results(predictions, actuals, save_dir, mode="test_results")

In [ ]:

for dl, ids, dates, path in zip(test_dls, test_ids_list, test_dates_list, test_files):
    preds, acts = evaluate(model, dl, device, TSS_scalers_dict=scalers_dict)
    save_results_from_excel_for_test(preds, acts, ids, dates, path, save_dir)

for dl, ids, dates, path in zip(test_dls, test_ids_list, test_dates_list, test_files): preds, acts = evaluate(model, dl, device, TSS_scalers_dict=scalers_dict) save_results_from_excel_for_test(preds, acts, ids, dates, path, save_dir)

Model inference¶

In [ ]:

Output = infer_and_visualize_single_model_Robust(
    model=model,
    test_loader=test_loader,
    Rrs=Rrs,
    mask=mask,
    latitude=latitude,
    longitude=longitude,
    save_folder=save_dir,
    extent=extent,
    rgb_image=rgb_image,
    structure_name="TSS",
    TSS_scalers_dict=scalers_dict,
    vmin=0,
    vmax=50,
)

print(
    f"✅ Finished training, train loss: {best_train_loss:.4f}, test loss: {test_loss:.4f}"
)

Output = infer_and_visualize_single_model_Robust( model=model, test_loader=test_loader, Rrs=Rrs, mask=mask, latitude=latitude, longitude=longitude, save_folder=save_dir, extent=extent, rgb_image=rgb_image, structure_name="TSS", TSS_scalers_dict=scalers_dict, vmin=0, vmax=50, ) print( f"✅ Finished training, train loss: {best_train_loss:.4f}, test loss: {test_loss:.4f}" )

Save as GeoTIFF (optional)¶

In [ ]:

import matplotlib.pyplot as plt

tif_path = os.path.join(save_dir, "TSS.tif")

npy_to_tif(npy_input=Output, out_tif=tif_path, resolution_m=1000)

with rasterio.open(tif_path) as src:
    img = src.read(1)
    transform = src.transform
    bounds = src.bounds
img_masked = np.where(img < 0, np.nan, img)
extent = [bounds.left, bounds.right, bounds.bottom, bounds.top]
plt.figure(figsize=(8, 6))
im = plt.imshow(img_masked, cmap="jet", vmin=0, vmax=50, extent=extent, origin="upper")
plt.colorbar(im, label="TSS")
plt.title("TSS")
plt.show()

import matplotlib.pyplot as plt tif_path = os.path.join(save_dir, "TSS.tif") npy_to_tif(npy_input=Output, out_tif=tif_path, resolution_m=1000) with rasterio.open(tif_path) as src: img = src.read(1) transform = src.transform bounds = src.bounds img_masked = np.where(img < 0, np.nan, img) extent = [bounds.left, bounds.right, bounds.bottom, bounds.top] plt.figure(figsize=(8, 6)) im = plt.imshow(img_masked, cmap="jet", vmin=0, vmax=50, extent=extent, origin="upper") plt.colorbar(im, label="TSS") plt.title("TSS") plt.show()