1. SNV Preprocessing

Import SNVScaler.

SNV standardizes each spectrum to have zero mean and unit variance. This removes baseline and scaling effects while preserving spectral shape. After SNV, all spectra have the same L2 norm:

\lVert x_{\mathrm{snv}} \rVert_2 = \sqrt{L - 1}

For example, for 256-dimensional spectra,

\lVert x_{\mathrm{snv}} \rVert_2 = \sqrt{255} \approx 15.97

Hence, SNV maps spectra onto a constant-radius hypersphere.

from chemomae.preprocessing import SNVScaler

# X_*: reflectance data (np.ndarray)
# Expected shape: (N, 256)
preprocessed = []
for X in [X_train, X_val, X_test]:
    sc = SNVScaler()
    X_snv = sc.transform(X)
    preprocessed.append(X_snv)

X_train_snv, X_val_snv, X_test_snv = preprocessed

2. Dataset and DataLoader Preparation

Convert NumPy arrays into PyTorch tensors and build DataLoaders.

from chemomae.utils import set_global_seed
import torch
from torch.utils.data import DataLoader, TensorDataset

set_global_seed(42)

train_ds = TensorDataset(torch.as_tensor(X_train_snv, dtype=torch.float32))
val_ds   = TensorDataset(torch.as_tensor(X_val_snv,   dtype=torch.float32))
test_ds  = TensorDataset(torch.as_tensor(X_test_snv,  dtype=torch.float32))

train_loader = DataLoader(train_ds, batch_size=1024, shuffle=True,  drop_last=False)
val_loader   = DataLoader(val_ds,   batch_size=1024, shuffle=False, drop_last=False)
test_loader  = DataLoader(test_ds,  batch_size=1024, shuffle=False, drop_last=False)

3. Model, Optimizer, and Scheduler Setup

Define ChemoMAE and a standard optimization pipeline.

from chemomae.models import ChemoMAE
from chemomae.training import build_optimizer, build_scheduler

model = ChemoMAE(
    seq_len=256,
    d_model=256,
    nhead=4,
    num_layers=4,
    dim_feedforward=1024,
    dropout=0.1,
    latent_dim=16,
    latent_normalize=True,
    decoder_num_layers=2,
    n_patches=32,
    n_mask=16,
)

opt = build_optimizer(
    model,
    lr=1.0e-3,
    weight_decay=0.05,
    betas=(0.9, 0.95),
)

sched = build_scheduler(
    opt,
    steps_per_epoch=max(1, len(train_loader)),
    epochs=500,
    warmup_epochs=10,
    min_lr_scale=0.1,
)

4. Optional Spectral Augmentation

Define a hypersphere-aware augmenter for SNV-normalized spectra.

from chemomae.training import SpectraAugmenter, SpectraAugmenterConfig

aug_cfg = SpectraAugmenterConfig(
    shift_prob=0.5,
    shift_delta_range=(-2.0, 2.0),
    noise_prob=0.5,
    noise_angle_deg_range=(0.5, 3.0),
    shuffle_order_per_batch=True,
    recenter_after_each_op=True,
    renorm_to_input_norm=True,
)

augmenter = SpectraAugmenter(aug_cfg)

The model input is augmented, but the reconstruction target remains the original spectrum.

This provides weak denoising-style regularization while preserving the SNV-compatible geometry of the input spectra.

5. Training Setup (Trainer + Config)

Trainer orchestrates the fixed-budget self-supervised training loop with:

• AMP (Automatic Mixed Precision)
• EMA (Exponential Moving Average of model weights)
• optional SpectraAugmenter
• selectable masked-only or full-spectrum reconstruction loss
• gradient clipping
• checkpointing / resume
• JSON logging
• final weights export

ChemoMAE does not use validation-loss-based early stopping or best-checkpoint selection. Training is controlled by a predefined epoch / step budget, and the final model is selected by an explicit rule such as EMA-last weights.

from chemomae.training import TrainerConfig, Trainer

trainer_cfg = TrainerConfig(
    out_dir="runs",
    device="cuda",
    amp=True,
    amp_dtype="bf16",
    enable_tf32=False,
    grad_clip=1.0,
    use_ema=True,
    ema_decay=0.999,
    loss_type="mse",
    loss_region="masked",
    reduction="mean",
    resume_from="auto",
)

trainer = Trainer(
    model,
    opt,
    train_loader,
    scheduler=sched,
    augmenter=augmenter,
    cfg=trainer_cfg,
)

result = trainer.fit(epochs=500)
print(result["final_model"])  # "ema_last_model.pt" if EMA is enabled

loss_region="masked" preserves the original masked-only objective. For an ordinary autoencoder objective, construct ChemoMAE with n_mask=0 and set loss_region="all" explicitly; this mode computes loss against the full clean target spectrum. loss_region="masked" with no masked elements raises ValueError instead of producing a silent zero loss.

During training, ChemoMAE produces the following outputs under out_dir:

runs/
├── training_history.json
│    ↳ Per-epoch records:
│       [
│         {
│           "epoch": 1,
│           "train_loss": ...,
│           "lr": ...,
│           "time_sec": ...,
│           "loss_region": "masked",
│           "n_mask": 16
│         },
│         ...
│       ]
│
├── last_model.pt
│    ↳ Final raw model weights at the end of training
│
├── ema_last_model.pt
│    ↳ Final EMA weights at the end of training
│       (saved only when EMA is enabled)
│
└── checkpoints/
     └── last.pt
          ↳ Full checkpoint for resume:
             model + optimizer + scheduler + scaler + EMA + loss region + history

6. Evaluation (Tester + Config)

The Tester evaluates masked reconstruction loss on a dataset.

It supports:

• masked-only reconstruction loss
• AMP (bf16 / fp16)
• optional fixed visible masks
• optional SpectraAugmenter
• JSON logging to a test-history file

from chemomae.training import TesterConfig, Tester

tester_cfg = TesterConfig(
    out_dir="runs",
    device="cuda",
    amp=True,
    amp_dtype="bf16",
    loss_type="mse",
    reduction="mean",
    fixed_visible=None,
    log_history=True,
    history_filename="test_history.json",
)

tester = Tester(
    model,
    tester_cfg,
    augmenter=None,
)

test_loss = tester(test_loader)
print(f"Test Loss: {test_loss:.6f}")

When augmenter is provided, the Tester applies augmentation to the model input while keeping the reconstruction target as the original spectrum:

tester = Tester(
    model,
    tester_cfg,
    augmenter=augmenter,
)

This evaluates reconstruction robustness under input perturbations.

7. Latent Extraction (Extractor + Config)

Extract latent embeddings from a trained ChemoMAE using all-visible encoding.

By default, Extractor does not use ChemoMAE masking. It directly calls the encoder with an all-visible mask, so the full spectrum is used for latent feature extraction.

from chemomae.training import ExtractorConfig, Extractor

extractor_cfg = ExtractorConfig(
    device="cuda",
    amp=True,
    amp_dtype="bf16",
    save_path=None,
    return_numpy=False,
)

extractor = Extractor(
    model,
    extractor_cfg,
    augmenter=None,
)

latent_test = extractor(test_loader)

When augmenter is provided, the Extractor applies augmentation before encoder inference:

extractor = Extractor(
    model,
    extractor_cfg,
    augmenter=augmenter,
)

latent_test_aug = extractor(test_loader)

Without an augmenter, latent extraction is deterministic with respect to ChemoMAE masking because all positions are treated as visible. With an augmenter, extracted embeddings may vary due to stochastic spectral shift/noise augmentation.

8. Clustering with Cosine K-Means

Cluster latent vectors using cosine geometry.

from chemomae.clustering import CosineKMeans, elbow_ckmeans

k_list, inertias, K, idx, kappa = elbow_ckmeans(
    CosineKMeans,
    latent_test,
    device="cuda",
    k_max=50,
    chunk=5_000_000,
    random_state=42,
)

ckm = CosineKMeans(
    n_components=K,
    tol=1e-4,
    max_iter=500,
    device="cuda",
    random_state=42,
)

ckm.fit(latent_test, chunk=5_000_000)
ckm.save_centroids("runs/ckm.pt")
labels = ckm.predict(latent_test, chunk=5_000_000)

9. Clustering with vMF Mixture

Probabilistic hyperspherical clustering.

from chemomae.clustering import VMFMixture, elbow_vmf

k_list, scores, K, idx, kappa = elbow_vmf(
    VMFMixture,
    latent_test,
    device="cuda",
    k_max=50,
    chunk=5_000_000,
    random_state=42,
    criterion="bic",
)

vmf = VMFMixture(
    n_components=K,
    tol=1e-4,
    max_iter=500,
    device="cuda",
    random_state=42,
)

vmf.fit(latent_test, chunk=5_000_000)
vmf.save("runs/vmf.pt")
labels = vmf.predict(latent_test, chunk=5_000_000)

SNVScaler

• Document
• Implementation

SNVScaler performs row-wise mean subtraction and variance scaling . Each spectrum is centered and divided by its unbiased standard deviation (ddof=1). It is a stateless transformer supporting both NumPy and PyTorch , preserving the original framework, device, and dtype .

When transform_stats=True, it returns (Y, mu, sd), where sd already includes eps and can be directly used for inverse reconstruction.

After SNV, all rows have zero mean and unit variance , producing a constant L2 norm of sqrt(L - 1), thereby mapping spectra onto a constant-radius hypersphere — ideal for cosine-based clustering.

import numpy as np
from chemomae.preprocessing import SNVScaler

X = np.array([[1.0, 2.0, 3.0],
              [4.0, 5.0, 6.0]], dtype=np.float32)

scaler = SNVScaler()
Y = scaler.transform(X)

scaler = SNVScaler(transform_stats=True)
Y, mu, sd = scaler.transform(X)
X_rec = scaler.inverse_transform(Y, mu=mu, sd=sd)

Key Features

• unbiased standard deviation (ddof=1, with automatic fallback for L=1)
• numerically stable eps handling
• float64 internal computation
• Torch-compatible device and dtype preservation

When to Use

• Standard preprocessing for NIR spectra
• Before cosine-based modeling or clustering

cosine_fps_downsample

• Document
• Implementation

cosine_fps_downsample performs Farthest-Point Sampling (FPS) under hyperspherical geometry , selecting spectra that are maximally diverse in direction .

Internally, all rows are L2-normalized for selection, but the returned subset is drawn from the original-scale input X. It supports both NumPy and PyTorch inputs and automatically leverages CUDA when available.

import numpy as np
from chemomae.preprocessing import cosine_fps_downsample

X = np.random.randn(1000, 128).astype(np.float32)
X_sub = cosine_fps_downsample(X, ratio=0.1, seed=42)

Key Features

• internal normalization for cosine-based selection
• output kept in original scale
• device-aware Torch support

When to Use

• diversity-driven subsampling
• reducing redundancy in large spectral datasets

ChemoMAE

• Document
• Implementation

ChemoMAE is a Masked Autoencoder for 1D spectra.

It adopts a patch-token formulation , where contiguous spectral bands are grouped into patches and masking is performed at the patch level . The encoder processes only the visible patch tokens together with a CLS token, and the decoder reconstructs the full spectrum using a lightweight MLP decoder .

The CLS output is projected to a latent_dim vector and may be L2-normalized, yielding embeddings naturally suited to cosine similarity and hyperspherical clustering.

import torch
from chemomae.models import ChemoMAE

mae = ChemoMAE(
    seq_len=256,
    d_model=256,
    nhead=4,
    num_layers=4,
    dim_feedforward=1024,
    decoder_num_layers=2,
    latent_dim=8,
    latent_normalize=True,
    n_patches=32,
    n_mask=16,
)

x = torch.randn(8, 256)
x_rec, z, visible = mae(x)

Key Features

• patch-wise masking
• Transformer encoder over visible tokens
• lightweight MLP decoder
• optional L2-normalized latent
• cosine-friendly embeddings

When to Use

• learning geometry-aware spectral representations
• downstream clustering, visualization, or supervised fine-tuning

build_optimizer & build_scheduler

• Document
• Implementation

Utility functions for a standardized Transformer-style optimization pipeline.

• build_optimizer creates grouped AdamW
• build_scheduler creates linear warmup → cosine decay

from chemomae.models import ChemoMAE
from chemomae.training import build_optimizer, build_scheduler

model = ChemoMAE(seq_len=256)
optimizer = build_optimizer(model, lr=1.0e-3, weight_decay=0.05)
scheduler = build_scheduler(
    optimizer,
    steps_per_epoch=len(train_loader),
    epochs=100,
    warmup_epochs=5,
)

SpectraAugmenterConfig & SpectraAugmenter

• Document
• Implementation

SpectraAugmenter provides SNV-geometry-aware augmentation for SNV-normalized spectra.

After SNV preprocessing, each spectrum is approximately mean-centered and has a fixed per-spectrum norm. Instead of applying unconstrained Euclidean perturbations, SpectraAugmenter applies weak spectral perturbations and optionally projects the result back toward the SNV-compatible geometry by:

• re-centering each spectrum to mean zero
• re-normalizing each spectrum to the original per-spectrum L2 norm

The current implementation supports two augmentations:

• fractional shift A small wavelength-axis perturbation implemented by interpolation.

• tangent Gaussian noise A random local perturbation constructed by projecting Gaussian noise onto the tangent space and rotating the spectrum by a sampled angle.

Fractional shift is controlled by shift_delta_range, while tangent Gaussian noise is controlled by noise_angle_deg_range.

from chemomae.training import SpectraAugmenter, SpectraAugmenterConfig

aug_cfg = SpectraAugmenterConfig(
    shift_prob=0.5,
    shift_delta_range=(-2.0, 2.0),
    noise_prob=0.5,
    noise_angle_deg_range=(0.5, 3.0),
    shuffle_order_per_batch=True,
    recenter_after_each_op=True,
    renorm_to_input_norm=True,
    eps=1.0e-8,
)

augmenter = SpectraAugmenter(aug_cfg)

augmenter.train()
x_aug = augmenter(x)

Key Features

• SNV-compatible spectral augmentation
• fractional wavelength-axis shift
• tangent-space Gaussian perturbation
• probability-controlled operation application
• delta-controlled fractional shift
• angle-controlled tangent Gaussian noise
• optional random ordering of shift/noise operations
• optional re-centering to zero mean after each operation
• optional re-normalization to the input L2 norm after each operation
• implemented as nn.Module, so it supports .to(device)
• automatically inactive in eval() mode

Mode Behavior

SpectraAugmenter applies augmentation only in train() mode.

augmenter.train()
x_aug = augmenter(x)  # augmentation is applied

augmenter.eval()
x_same = augmenter(x)  # input is returned unchanged

This behavior is important when using the augmenter outside the Trainer. For example, Extractor and Tester temporarily set the augmenter to train() when augmentation is explicitly requested, while keeping the ChemoMAE model itself in eval() mode.

When to Use

• during ChemoMAE pretraining on SNV-normalized spectra
• when you want denoising-style pretraining, where the model receives weakly perturbed spectra but reconstructs the original spectra
• when evaluating reconstruction robustness under controlled spectral perturbations
• when extracting augmented latent representations for robustness checks or test-time augmentation style analysis
• when perturbations should remain compatible with cosine-based or hyperspherical downstream analysis

TrainerConfig & Trainer

• Document
• Implementation

TrainerConfig and Trainer form the core training engine of ChemoMAE.

They provide a fixed-budget training loop with an explicit choice between masked-only and full-spectrum reconstruction loss, with support for:

• AMP (bf16 / fp16)
• optional TF32
• EMA parameter tracking
• optional SpectraAugmenter
• loss_region="masked" (default) or loss_region="all"
• gradient clipping
• checkpointing and resume
• weights-only export for final model variants
• JSON history logging

ChemoMAE does not use validation-loss-based early stopping or best-checkpoint selection. Training is controlled by a predefined epoch / step budget, and the final model is selected by an explicit rule such as EMA-last weights.

from chemomae.models import ChemoMAE
from chemomae.training import (
    Trainer,
    TrainerConfig,
    SpectraAugmenter,
    SpectraAugmenterConfig,
    build_optimizer,
    build_scheduler,
)

model = ChemoMAE(seq_len=256, latent_dim=16, n_patches=32, n_mask=24)

cfg = TrainerConfig(
    out_dir="runs",
    device="cuda",
    amp=True,
    amp_dtype="bf16",
    enable_tf32=False,
    grad_clip=1.0,
    use_ema=True,
    ema_decay=0.999,
    loss_type="mse",
    loss_region="masked",
    reduction="mean",
    resume_from="auto",
)

aug_cfg = SpectraAugmenterConfig(
    shift_prob=0.5,
    shift_delta_range=(-2.0, 2.0),
    noise_prob=0.5,
    noise_angle_deg_range=(0.5, 3.0),
    shuffle_order_per_batch=True,
    recenter_after_each_op=True,
    renorm_to_input_norm=True,
)
augmenter = SpectraAugmenter(aug_cfg)

optimizer = build_optimizer(model, lr=1.0e-3, weight_decay=0.05)

epochs = 800
scheduler = build_scheduler(
    optimizer,
    steps_per_epoch=len(train_loader),
    epochs=epochs,
    warmup_epochs=40,
)

trainer = Trainer(
    model,
    optimizer,
    train_loader,
    scheduler=scheduler,
    augmenter=augmenter,
    cfg=cfg,
)

_ = trainer.fit(epochs=epochs)

For full-spectrum autoencoder training, use n_mask=0 together with loss_region="all". The loss region is never inferred from n_mask.

Key Features

• automatic device and precision handling
• fixed epoch / fixed step SSL pretraining
• EMA tracking after each optimizer step
• EMA-consistent final export behavior:
  • final raw weights → last_model.pt
  • final EMA weights → ema_last_model.pt if EMA is enabled
• checkpoints/last.pt stores the full resumable training state
• optional train-time spectral augmentation
• batch-wise scheduler stepping
• atomic JSON history logging

Outputs

runs/
├── training_history.json
│    ↳ Per-epoch records:
│       [
│         {
│           "epoch": 1,
│           "train_loss": ...,
│           "lr": ...,
│           "time_sec": ...,
│           "loss_region": "masked",
│           "n_mask": 24
│         },
│         ...
│       ]
│
├── last_model.pt
│    ↳ Final raw model weights at the end of training
│
├── ema_last_model.pt
│    ↳ Final EMA weights at the end of training
│       (saved only when EMA is enabled)
│
└── checkpoints/
     └── last.pt
          ↳ Full checkpoint for resume:
             model + optimizer + scheduler + scaler + EMA + loss region + history

When to Use

• fixed-budget masked reconstruction pretraining for ChemoMAE
• full-spectrum autoencoder training with n_mask=0 and loss_region="all"
• validation-free SSL pretraining
• training runs where model selection is handled by an explicit final rule, such as EMA-last weights

TesterConfig & Tester

• Document
• Implementation

Tester provides a lightweight evaluation loop for trained ChemoMAE models.

It computes masked reconstruction loss (SSE/MSE) over a DataLoader, with AMP support, optional fixed visible masks, optional SpectraAugmenter, and JSON logging.

from chemomae.training import Tester, TesterConfig

cfg = TesterConfig(
    out_dir="runs",
    device="cuda",
    amp=True,
    amp_dtype="bf16",
    loss_type="mse",
    reduction="mean",
    fixed_visible=None,
    log_history=True,
    history_filename="test_history.json",
)

tester = Tester(
    model,
    cfg,
    augmenter=None,
)

avg_loss = tester(test_loader)
print("Test loss:", avg_loss)

When augmenter is provided, Tester evaluates reconstruction from augmented inputs while keeping the reconstruction target as the original spectrum.

tester = Tester(
    model,
    cfg,
    augmenter=augmenter,
)

When to Use

• evaluating masked reconstruction loss
• comparing reconstruction behavior under fixed visible masks
• evaluating robustness to weak spectral perturbations
• logging test losses separately from training history

ExtractorConfig & Extractor

• Document
• Implementation

Extractor provides a latent extraction pipeline from trained ChemoMAE models in all-visible mode.

It directly calls the encoder with an all-visible mask, so ChemoMAE masking is not used during extraction. Without an augmenter, this gives deterministic latent embeddings with respect to masking.

It supports AMP inference, Torch/NumPy return types, optional saving, and optional SpectraAugmenter.

from chemomae.training import Extractor, ExtractorConfig

cfg = ExtractorConfig(
    device="cuda",
    amp=True,
    amp_dtype="bf16",
    save_path=None,
    return_numpy=True,
)

extractor = Extractor(
    model,
    cfg,
    augmenter=None,
)

Z = extractor(loader)

When augmenter is provided, Extractor applies augmentation before encoder inference.

extractor = Extractor(
    model,
    cfg,
    augmenter=augmenter,
)

Z_aug = extractor(loader)

With an augmenter, extracted embeddings may vary across calls because spectral shift/noise augmentation can be stochastic.

When to Use

• extracting latents for clustering
• visualization
• downstream analysis
• robustness checks using augmented latent embeddings

CosineKMeans & elbow_ckmeans

• Document
• Implementation

CosineKMeans implements hyperspherical k-means with cosine similarity.

import torch
from chemomae.clustering import CosineKMeans, elbow_ckmeans

X = torch.randn(10_000, 64)
ckm = CosineKMeans(n_components=12, device="cuda", random_state=42)
ckm.fit(X)
labels = ckm.predict(X)

When to Use

• clustering unit-sphere embeddings
• model selection of K under cosine geometry

VMFMixture & elbow_vmf

• Document
• Implementation

VMFMixture fits a von Mises–Fisher mixture model on the unit hypersphere.

import torch
from chemomae.clustering import VMFMixture, elbow_vmf

X = torch.randn(10000, 64, device="cuda")
vmf = VMFMixture(n_components=16, device="cuda", random_state=42)
vmf.fit(X)
labels = vmf.predict(X)

When to Use

• probabilistic clustering of unit-sphere embeddings
• BIC / NLL-based model selection under cosine geometry

silhouette_samples_cosine_gpu & silhouette_score_cosine_gpu

• Document
• Implementation

Cosine-based GPU-accelerated silhouette metrics for clustering evaluation.

import numpy as np
from chemomae.clustering import silhouette_score_cosine_gpu

X = np.random.randn(100, 16).astype(np.float32)
labels = np.random.randint(0, 4, size=100)

score = silhouette_score_cosine_gpu(X, labels, device="cpu")
print(score)

set_global_seed

• Document
• Implementation

Unified seeding for Python, NumPy, and PyTorch, with optional CuDNN determinism.

from chemomae.utils import set_global_seed

set_global_seed(42)

When to Use

• at the start of any experiment
• before training, testing, clustering, or extraction