ChemoMAE

ChemoMAE is a PyTorch library for learning representations of one-dimensional spectra, clustering directional features, and evaluating spatial label maps. Its preprocessing, model, training, clustering, and metric components can be used independently or combined in a pipeline.

Example application: unsupervised spatial mapping of hyperspectral images. (a) A Transformer encoder projects spectra into L2-normalized representations. (b) In the illustrated training setup, augmented visible patches are used to reconstruct the masked regions of the original spectra. (c) The trained encoder receives complete spectra; their representations are clustered, and the cluster labels are mapped back to the original pixel locations.

These guides document ChemoMAE v0.2.4. README links lead to the maintained repository documentation. See the v0.2.4 release notes for changes and compatibility notes.

Choose the components you need

SNV is optional: use it when removing each spectrum's mean and scale suits your data. Clustering accepts feature matrices without requiring a ChemoMAE model. Spatial evaluation requires a label map and validity mask; ordinary spectral learning and feature extraction require no pixel coordinates.

Quick Start

ChemoMAE requires Python 3.10 or newer. Install v0.2.4 with pip:

pip install chemomae==0.2.4

The example below runs on CPU and covers training, saving and reloading a model, feature extraction, and clustering.

ChemoMAE Example

This complete CPU example trains a small model, reloads its inference artifact, and clusters its features. SNV and clustering illustrate optional composition. The synthetic data, two epochs, and K=3 demonstrate APIs, not recommended tuning values. Outputs go to a new temporary directory whose path is printed.

import tempfile
from pathlib import Path

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

from chemomae.preprocessing import snv
from chemomae.models import ChemoMAE
from chemomae.training import Trainer, TrainerConfig, Extractor, ExtractorConfig
from chemomae.clustering import CosineKMeans

torch.manual_seed(42)
data_stream = torch.Generator().manual_seed(7)
spectra = snv(torch.randn(80, 64, generator=data_stream))
train_x, test_x = spectra[:64], spectra[64:]

def loader(x: torch.Tensor) -> DataLoader:
    return DataLoader(TensorDataset(x), batch_size=16, shuffle=False)

device = torch.device("cpu")
run_dir = Path(tempfile.mkdtemp(prefix="chemomae-"))
model = ChemoMAE(
    seq_len=64, n_patches=8, n_mask=4, d_model=16,
    nhead=4, num_layers=1, latent_dim=8,
).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
trainer = Trainer(
    model, optimizer, loader(train_x),
    cfg=TrainerConfig(
        out_dir=run_dir, device=device, resume_from=None, use_ema=False,
        progress=False, verbose=False,
    ),
)
result = trainer.fit(epochs=2)

assert result["final_artifact"] is not None
inference_model = ChemoMAE.load(result["final_artifact"], device=device)
extractor = Extractor(
    inference_model,
    ExtractorConfig(representation="normalized_latent", output_device="cpu"),
)
train_features = extractor(loader(train_x))
test_features = extractor(loader(test_x))

clusterer = CosineKMeans(n_components=3, max_iter=30, device=device, random_state=42)
clusterer.fit(train_features)
test_labels = clusterer.predict(test_features)
clusterer.save_centroids(run_dir / "clusters.pt")
assert train_features.shape == (64, 8) and test_features.shape == (16, 8)
assert test_labels.shape == (16,)
print(test_labels, clusterer.converged_, clusterer.stop_reason_)
print("Outputs:", run_dir)

Choices to keep explicit

The example selects a fresh output directory, raw final weights, and CPU. Keep these choices explicit when adapting it:

See starting or resuming training and the output-file guide for the exact defaults and raw/EMA file choices. CPU use requires explicitly passing device="cpu" to clustering; CUDA use requires an available CUDA device.

Continue with the staged workflow tutorial for reconstruction evaluation and optional augmentation, resume, clustering, spatial analysis, and reporting. The documentation index also provides direct routes to each API. For configuration comparisons and memory estimates, see planning a first experiment.

Model and workflow choices

ChemoMAE splits a spectrum into contiguous patches. Visible patch tokens and a CLS token enter a Transformer encoder; the projected CLS vector is the shared bottleneck from which a decoder reconstructs every spectral channel. The model guide describes masks, representations, and boundary behavior.

• Reconstruction objective: Trainer defaults to hidden-channel loss. For a full-spectrum autoencoder, choose n_mask=0 and loss_region="all" explicitly.
• Augmentation: SpectraAugmenter offers fractional channel shifts and tangent noise. Trainer keeps the pre-augmentation input as the target when one is supplied.
• Decoder and latent: one decoder layer is affine; the default two-layer decoder is an MLP. Latents are L2-normalized by default; set latent_normalize=False when unconstrained latents suit the task.
• Features: model.encode and Extractor expose CLS, raw projected, and normalized representations with all patches visible. Use model.encode in a custom gradient-enabled loop for end-to-end fine-tuning.
• Training control: Trainer provides reconstruction training, hooks, EMA, and completed-epoch resume. Model selection, downstream objectives, and data splits belong to the application. Model, augmentation, and loss primitives also work in ordinary PyTorch loops.

When evaluating generalization to unseen groups, fit learned components on the training groups and retain them for held-out prediction. Descriptive analysis of an existing dataset may use that dataset for fitting. Reconstruction loss, feature separation, and spatial coherence measure different properties; select the diagnostics that address your intended use.

Research background

ChemoMAE was extracted into a reusable library from spectral representation learning research. The associated research repository is WoodDegradationMap. A manuscript with an NIR-HSI case study is in preparation; its publication link will be added when available. Library defaults and examples are documented independently of that study's settings.

License

ChemoMAE is distributed under the Apache License 2.0. See NOTICE for notices.