chesskit 0.1.0

Information-geometric frame selection toolkit for atomistic datasets

CHESSKIT

CHESSKIT (Covering Histogram-based atomic Environments through Sampling on Statistical manifolds toolKIT) provides a modern, modular Python implementation for information-geometric sampling of atomistic data based on the compositional structure of local atomic environments in molecular dynamics trajectories.

Features

• Modular architecture: Easily extend with new featurizers, metrics, and algorithms
• Pluggable featurizers: SevenNet and wACSF with extensible registry pattern
• Frame selection: k-center (farthest-first greedy algorithm)

Installation

From PyPI

pip install chesskit

# Optional feature sets
pip install "chesskit[torch]"
pip install "chesskit[torch,sevennet]"

Install name and import name are intentionally different:

• install from PyPI with chesskit
• import the Python API from chess
• run the CLI as chesskit or python -m chesskit

From source (development)

git clone https://github.com/ecshin94/CHESS-toolkit.git chesskit
cd chesskit

# Install in editable mode with optional dependencies
pip install -e ".[torch,sevennet,dev]"

Dependencies

Core (required):

• numpy, scikit-learn, pyyaml, matplotlib, ase

Optional:

• torch: For GPU acceleration (CUDA/MPS support) and Fisher-Rao / isotropy post-processing
• sevennet: For SevenNet molecular featurizer
• dev: For development (pytest, black, isort, flake8, mypy)

Quick Start

1. Create configuration

# Build template (wACSF)
python -m chesskit --preset wacsf-build > config_build_wacsf.yaml

# Sample template (CHESS/Fisher-Rao)
python -m chesskit --preset chess-sample > config_sample.yaml

# Edit config_sample.yaml so paths.cache matches the cache written in build
# Example: cache/wacsf_embeddings.h5 or cache/all_embeddings.h5

2. Run the pipelines

# Step 1: Build embeddings cache from extxyz (one-time)
python -m chesskit --config config_build_wacsf.yaml --mode build

# Step 2: Sample frames from cache (fast)
python -m chesskit --config config_sample.yaml --mode sample

2.5 Quick sanity check (pytest-based)

# Fast end-to-end smoke check on tiny synthetic data
python -m chesskit --quickcheck cpu

# Run GPU smoke check (skips automatically if CUDA is unavailable)
python -m chesskit --quickcheck gpu

# Run both CPU + GPU checks
python -m chesskit --quickcheck both

Run on Google Colab with Your Own Dataset

Large trajectory files (for example, .extxyz) are intentionally not included in this repository. Use your own file in Colab, and set the filename in the provided templates under examples/.

1. Prepare files in Google Drive

In Colab, mount Drive and move to your working folder. Use the notebook template:

• examples/colab_chess_workflow.ipynb

Put these files in the same folder (or let the notebook copy them automatically):

• your dataset (for example, my_train.extxyz)
• examples/config_build_sevennet.yaml
• examples/config_sample_chess.yaml

YAML roles:

• examples/config_*.yaml: shared templates for local CLI and Colab workflows.
• examples/colab_chess_workflow.ipynb: Colab notebook that can generate runtime build.yaml and sample.yaml from the templates above.

2. Update dataset-related paths in YAML

Edit both YAML files so the dataset/cache names match your file.

In build.yaml (generated from config_build_sevennet.yaml):

• paths.input: ./my_train.extxyz
• paths.cache: cache/my_train.h5

In sample.yaml (generated from config_sample_chess.yaml):

• paths.cache: cache/my_train.h5
• paths.input is optional (set it only when you also want selected/remaining .extxyz export)

3. Run build and sample

python -m chesskit --config build.yaml --mode build
python -m chesskit --config sample.yaml --mode sample

4. Outputs

• Build cache: cache/my_train.h5
• Selection report: report/result.json (includes selected frame indices)
• Optional structure export (if paths.input is set): selected/remaining .extxyz

Note:

• .extxyz files are excluded from version control by default (.gitignore), so users can keep large private datasets locally or in Drive.

Use in Python

The published distribution is chesskit, but the import namespace remains chess for API compatibility.

from chess.pipelines import build, sample
from chess.config import load_config

# Load configs
build_cfg = load_config("config_build_wacsf.yaml")
sample_cfg = load_config("config_sample.yaml")

# Build embeddings (one-time, caches to HDF5)
build_result = build(build_cfg)
print(f"Cache: {build_result['cache_path']}")

# Select frames (fast, reuses cache)
sample_result = sample(sample_cfg)
selected_frames = sample_result["selected_frames"]
print(f"Selected {len(selected_frames)} frames")
print(f"Plots saved to: {sample_result['report_dir']}")

Mode and Featurizer Rules

Stage	Required keys	Featurizer used?	Output
mode: build	paths.input, featurizer.name	Yes (sevennet or wacsf)	writes HDF5 cache (paths.cache)
mode: sample	paths.cache, chess	No (cache-only)	selected frame indices + plots

Notes:

• Featurizer choice happens only in build stage.
• Sample stage does not run SevenNet/wACSF again; it reads cache.
• paths.cache in sample config must match the cache path produced in build.
• --mode overrides mode inside YAML when both are provided.

For the full list of configuration keys actually used by the current code, see CONFIGURATION_MANUAL.md.

Config Summary

Section	Key	Used in	Meaning
paths	input	build	extxyz or ASE-readable input structure file
paths	cache	build, sample	HDF5 cache output/input
paths	report_dir	sample	report and plot directory
misc	result_json	build, sample	run-summary JSON filename under paths.report_dir (null disables)
featurizer	name	build	sevennet or wacsf
featurizer.config	model, modality, enable_cueq, enable_flash, enable_oeq, cutoff, multi_gpu, gpu_ids	SevenNet build	SevenNet settings
featurizer.config	rcut, g2_params, g4_params	wACSF build	wACSF settings
chess	metric	sample	currently fisherrao only
chess	k_select	sample	selected frame count or all
chess	k_boe	sample	BoE codebook size
chess	radius_stop	sample	finite radius threshold (when k_select: all)

Note: for wACSF build, G2/G4 math runs on GPU but ASE neighbor-list generation is CPU-bound. Build throughput scales with CPU cores; use CHESS_WACSF_NL_WORKERS to set worker count. Note: wACSF build path is under active development; behavior and performance may change.

Multi-GPU Acceleration

Enable Multi-GPU in Configuration

Add multi_gpu and gpu_ids under featurizer.config:

featurizer:
  name: sevennet
  config:
    model: "7net-0"
    device: "cuda"
    enable_cueq: true       # Optional: cuEquivariance backend
    enable_flash: false     # Optional: flashTP backend
    enable_oeq: false       # Optional: OpenEquivariance backend (version dependent)
    multi_gpu: true          # Enable multi-GPU mode
    gpu_ids: [0, 1, 2, 3]    # Use specific GPUs (null = all available)

Python API

from chess.featurizers import FeaturizerRegistry

# Create multi-GPU featurizer
featurizer = FeaturizerRegistry.create(
    "sevennet",
    model="7net-0",
    device="cuda",
    multi_gpu=True,
    gpu_ids=[0, 1, 2, 3]  # or None for all GPUs
)

# Use normally - parallelization is automatic
E, frame_counts, F, metadata = featurizer.featurize(frames)

Requirements

• PyTorch with CUDA support
• 2+ CUDA-capable GPUs

Troubleshooting

If you encounter GPU memory issues:

• Reduce batch size in configuration
• Use fewer GPUs with gpu_ids
• Monitor GPU memory with nvidia-smi

Project Structure

chesskit/
├── src/chess/                    # Main package
│   ├── featurizers/              # Pluggable feature extractors
│   ├── transforms/               # Data transformation (PCA, pooling)
│   ├── metrics/                  # Distance metrics and isotropy analysis
│   ├── selection/                # Frame selection algorithms
│   ├── plotting/                 # Visualization
│   └── pipelines/                # High-level workflows
├── tests/                        # Test suite
├── examples/                     # Usage examples and config templates
├── pyproject.toml               # Modern packaging
└── README.md

Extending with New Featurizers

Adding a new featurizer is simple:

from chess.featurizers import Featurizer, FeaturizerRegistry

@FeaturizerRegistry.register("myfeaturizer")
class MyFeaturizer(Featurizer):
    @property
    def name(self) -> str:
        return "myfeaturizer"
    
    @property
    def embedding_dim(self) -> int:
        return 256
    
    def featurize(self, frames):
        # Your implementation
        E = ...  # (N, 256) embeddings
        frame_counts = ...  # atom counts per frame
        F = ...  # (N, 3) forces
        metadata = {...}
        return E, frame_counts, F, metadata
    
    def supports_batch(self) -> bool:
        return True
    
    def supports_forces(self) -> bool:
        return True

Then use in config:

featurizer:
  name: myfeaturizer
  config:
    param1: value1

Configuration

See examples/config_build_sevennet.yaml, examples/config_build_wacsf.yaml, and examples/config_sample_chess.yaml for complete configuration examples.

Full Configuration Example

See WORKFLOW_EXAMPLE.md for complete templates.

Build (SevenNet example):

mode: build
paths:
  input: trajectory.extxyz
  cache: embeddings.h5
  report_dir: ./report
featurizer:
  name: sevennet
  config:
    model: "7net-0"
    device: "cuda"
    enable_cueq: true
    enable_flash: false
    enable_oeq: false
    multi_gpu: true         # Optional: enable multi-GPU
    gpu_ids: [0, 1]         # Optional: specific GPUs
chess:
  metric: fisherrao

Sample (CHESS/Fisher-Rao):

mode: sample
paths:
  cache: embeddings.h5
  report_dir: ./report
chess:
  metric: fisherrao
  k_select: 128
  k_boe: 64

Testing

# Run all tests
pytest

# With coverage
pytest --cov=src/chess

# Specific test file
pytest tests/unit/test_core.py -v

Documentation

• PIPELINE_GUIDE.md - Complete API reference
• WORKFLOW_EXAMPLE.md - End-to-end usage examples
• CONFIGURATION_MANUAL.md - Full config key reference

Citation

If you use CHESSKIT in your research, please cite:

@article{shin2026chess,
  title={Information-Geometric Conditioning of Atomistic Datasets for Efficient Learning},
  author={Shin, Eui-Cheol and Jeon, Sang Ho and Nayir, Nadire},
  year={2026},
  doi={10.26434/chemrxiv.15003351/v2},
  url={https://doi.org/10.26434/chemrxiv.15003351/v2},
  publisher={ChemRxiv}
}

License

MIT License - see LICENSE file

Contributing

Contributions welcome! Please:

1. Fork the repository
2. Create a feature branch
3. Write tests for new features
4. Submit a pull request

Contact

• Issues: GitHub Issues
• Discussions: GitHub Discussions