AugerNet

Machine learning for Auger-electron spectroscopy (AES) and x-ray photoelectron spectroscopy (XPS)

Includes:

1. Equivariant GNN predictions of:

• a: core-electron binding energies (CEBE)
• b: Auger-Electron spectra (AES)

2. CNN classifications of local bond environments (functional groups) from AES spectra augmented with CEBEs

The present code and data release only supports the GNN CEBE predictions of the following manuscript:

The next release will accompany a future manuscript on GNN Auger predictions and
CNN bond env classification and will include the data for these functionalities.
The present release only contains the code templates for GNN Auger prediction and
CNN bond environment classification.

AugerNet currently provides three model types:

Doc site template undergoing updates can be found at https://afouda11.github.io/AugerNet/

Requires Python >= 3.10 and conda.

# Clone the repository
git clone https://github.com/afouda11/AugerNet.git
cd AugerNet

# Create the conda environment (installs all dependencies + the package)
conda env create -f environment.yml
conda activate augernet

The environment.yml installs all required dependencies (PyTorch, PyTorch Geometric, RDKit, scikit-learn, SkipAtom, etc.) and the augernet package itself in editable mode.

Note: The provided environment.yml targets macOS (Apple Silicon / MPS). For Linux with CUDA, replace the PyTorch pip lines with the appropriate versions from the PyTorch install guide and the PyG install guide.

uv is a fast Python package manager; see their documentation here.

# Install uv - https://docs.astral.sh/uv/getting-started/installation/
curl -LsSf https://astral.sh/uv/install.sh | sh

# Install AugerNet Python package dependencies
uv sync

You can use uv to run the Python commands by prefixing them with uv run. uv run executes commands in an isolated virtual environment with all required dependencies. See the documentation.

All runs are controlled by a single YAML config file. Example configs are provided in examples/.

# Download processed graph data from https://zenodo.org/records/19688196
python scripts/prepare_data.py --from-zenodo
# CEBE GNN — train, will recreate same model in artifact and main paper result.
python -m augernet --config examples/gnn_cebe_configs/train.yml

AugerNet supports five run modes, set via mode: in the YAML config. All three model types share the same mode system.

Train one model on a single k-fold split with optional evaluation.

mode: train
model: cebe-gnn        # or auger-gnn or auger-cnn
train_fold: 3
n_folds: 5
run_evaluation: true
exp_split: 'both' # 'val' | 'eval' | 'both' (separate val and eval sets) | 'all' (combined val + eval set)

Output is written to {model_type}_{mode}_results/ (e.g. cebe_gnn_train_results/, auger_cnn_train_results/) with subdirectories models/, outputs/, and pngs/.

For run_evaluation: true 113 mols in expirmental cebe data split into:

• Validation set (val) (to assist fold and param search)
• Final evaluation set (eval)

Path to processed evaluation data set interally in AugerNet
exp_split: 'both' will have eval and val prefixes assigned to different outputs

Train one model per fold, evaluate each, and write a JSON summary.

mode: cv
model: cebe-gnn
n_folds: 5
split_method: random   # random | butina (GNN only)
run_evaluation: true
exp_split: 'val' 

Train one fold per configuration from a Cartesian-product grid.

mode: param
model: cebe-gnn
run_evaluation: true
exp_split: 'val' 
param_grid:
  feature_keys: ['035', '03', '0356']
  learning_rate: [0.0001, 0.0003, 0.001]
  hidden_channels: [48, 64]
  n_layers: [3, 4, 5]

A unique search_id is derived from the searched dimensions so that
different grid searches never overwrite each other.

Load a previously trained .pth model and evaluate it on experimental data.
Architecture fields must match the values used during training.

mode: evaluate
model: cebe-gnn
exp_split: 'both' # 'val' | 'eval' | 'both' (separate val and eval sets) | 'all' (combined val + eval set)
model_path: cebe_gnn_train_results/models/cebe_gnn_035_random_EQ3_h64_fold3.pth
feature_keys: '035'
layer_type: EQ
hidden_channels: 64
n_layers: 3

Run inference on a directory of .xyz files using a saved GNN model.
No pre-processing is needed — molecular graphs are built on the fly.

mode: predict
model: cebe-gnn
model_path: cebe_gnn_train_results/models/cebe_gnn_035_random_EQ3_h64_fold3.pth
predict_dir: my_molecules/
feature_keys: '035'
layer_type: EQ
hidden_channels: 64
n_layers: 3

Note: The GNN models are trained on carbon 1s properties. Predictions for non-carbon atoms are not meaningful and are marked with * in the output labels file.

Predicts per-atom carbon 1s core-electron binding energies using an equivariant or invariant message-passing neural network.
Input is a set of .xyz molecular geometries converted to PyG graphs with configurable node features.

model: cebe-gnn
feature_keys: '035'
layer_type: EQ
hidden_channels: 64
n_layers: 3

model_id format: cebe_gnn_{feature_keys}_{split_method}{n_folds}_{layer}{n_layers}_h{hidden}
Example: cebe_gnn_035_random5_EQ3_h64

Predicts Auger-electron spectra from molecular graphs.
The details for running this model will be released in a future release.

Classifies carbon environments from broadened Auger spectra using a 1D CNN.
The details for running this model will be released in a future release.

See docs/configuration.md for the full reference, or see the summary tables below.

Details for this model will be released in a future release.

Details for this model will be released in a future release.

Each model type writes to its own results directory:

Each contains outputs/ and pngs/ subdirectories. Train, cv, and param modes also create a models/ subdirectory.

Details for this model will be released in a future release.

Processed and raw data files are stored at https://zenodo.org/records/19688196\

To download pre-processed data to data/processed/ run:

python scripts/prepare_data.py --from-zenodo

To dowloand both the pre-processed data to data/processed/ and
raw data to data/raw/ and unpack run:

python scripts/prepare_data.py --from-zenodo --with-raw

To dowloand both the raw data to data/raw/, unpack then process graphs locally run:

python scripts/prepare_data.py --with-raw

To just process graphs locally from pre-downloaded raw run:

python scripts/prepare_data.py

This repository contains a compressed dir of the skipatom: https://github.com/lantunes/skipatom\ files required to use the skipatom-200 and skipatom-30 vectors as atom type rep node features in: data/raw/skipatom.tar.gz
prepare_data.py will unpack this if it has not already been unpacked.

Tests are split into two tiers using pytest markers:

Currently only test_cebe_gnn_config.py is ran in the CI workflow to reduce run-time

# Essential tests only 
uv run pytest tests/ -m essential -v --tb=short

# Full suite (all tests)
uv run pytest tests/ -v --tb=short

# Single file
uv run pytest tests/test_cebe_gnn_model.py -v

Graph and model tests use a real molecule (dsgdb9nsd_133427) from tests/test_mol/ rather than synthetic data. Model symmetry tests verify that CEBE predictions are invariant to rotation and translation and equivariant to atom reordering properties required by the physics of the problem.

The artifact showcases the main result of the release (in plots) and includes the\ config file and model weights that produced it. scripts/export_best_model.py is used to copy the selected model to the artifact.
For a cv or param run identify the best fold and copy its weights, plots, and\ config into the tracked artifacts/ directory for release.
Here the train run for the main rersult in the CEBE GNN manuscript is used for the artifact.

uv run python scripts/export_best_model.py
uv run python scripts/export_best_model.py --results-dir auger_gnn_param_results
uv run python scripts/export_best_model.py --overwrite

artifacts/data_manifest.yml records the Zenodo DOI and SHA-256 checksums for all data files. To verify integrity:

shasum -a 256 data/processed/*.pt data/raw/*.tar.gz

AugerNet/
  src/
    augernet/                           # Python package (src layout)
      __init__.py                       # Sets data paths
      __main__.py                       # CLI entry point
      config.py                         # YAML to AugerNetConfig dataclass
      train_driver.py                   # Mode dispatch, CV, param search
      backend_gnn.py                    # Unified GNN backend (cebe-gnn + auger-gnn)
      backend_cnn.py                    # CNN backend (auger-cnn)
      feature_assembly.py               # Runtime feature selection and scaling
      gnn_train_utils.py                # MPNN model, train loop, unit tests
      cnn_train_utils.py                # AugerCNN1D model, CNN trainer
      carbon_dataframe.py               # CarbonDataset for CNN spectra
      carbon_environment.py             # Carbon environment patterns and labels
      class_merging.py                  # Carbon-class merging schemes
      spec_utils.py                     # Spectrum broadening and processing
      build_molecular_graphs.py         # XYZ to PyG graphs
      eneg_diff.py                      # Electronegativity scoring
      evaluation_scripts/
        evaluate_cebe_model.py          # CEBE evaluation plots and metrics
  scripts/
    prepare_data.py                     # Regenerate processed datasets from raw
    export_best_model.py                # Export best CV fold to artifacts/
  examples/
    gnn_cebe_configs/                   # Example YAML configs for CEBE GNN
  tests/
    conftest.py                         # Shared fixtures and markers
    test_mol/                           # Real molecule XYZ data for tests
    test_cebe_gnn_config.py             # Config defaults, resolution, YAML loading
    test_cebe_gnn_features.py           # Feature tags, parsing, scaling, assembly
    test_cebe_gnn_graph.py              # XYZ parsing, graph building, node/edge features
    test_cebe_gnn_model.py              # MPNN construction, forward pass, symmetry tests
  data/
    raw/                                # Raw XYZ + CEBE + atomic BE + skipatom files 
    processed/                          # Pre-built PyG datasets + CNN pickles + norm stat files
  artifacts/                            # Release artifacts (tracked in git)
    data_manifest.yml
    config/
    model_weights/
    plots/
  environment.yml
  pyproject.toml
  README.md