Metadata-Version: 2.4
Name: chesskit
Version: 0.1.0
Summary: Information-geometric frame selection toolkit for atomistic datasets
Author-email: Eui-Cheol Shin <shin@pdi-berlin.de>
License-Expression: MIT
Project-URL: Homepage, https://github.com/ecshin94/CHESS-toolkit
Project-URL: Repository, https://github.com/ecshin94/CHESS-toolkit
Project-URL: Issues, https://github.com/ecshin94/CHESS-toolkit/issues
Keywords: atomistic-simulation,frame-selection,materials-science,molecular-dynamics,sampling
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Topic :: Scientific/Engineering :: Chemistry
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=1.21.0
Requires-Dist: pyyaml>=5.4
Requires-Dist: scikit-learn>=1.0.0
Requires-Dist: matplotlib>=3.4.0
Requires-Dist: ase>=3.22.0
Requires-Dist: h5py>=3.3.0
Provides-Extra: torch
Requires-Dist: torch>=1.13.0; extra == "torch"
Requires-Dist: torch_geometric>=2.2.0; extra == "torch"
Provides-Extra: sevennet
Requires-Dist: sevenn>=0.12.0; extra == "sevennet"
Provides-Extra: wacsf
Requires-Dist: torch>=1.13.0; extra == "wacsf"
Requires-Dist: ase>=3.22.0; extra == "wacsf"
Requires-Dist: h5py>=3.3.0; extra == "wacsf"
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: pytest-cov>=3.0.0; extra == "dev"
Requires-Dist: black>=22.0.0; extra == "dev"
Requires-Dist: isort>=5.10.0; extra == "dev"
Requires-Dist: flake8>=4.0.0; extra == "dev"
Requires-Dist: mypy>=0.950; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx>=4.5.0; extra == "docs"
Requires-Dist: sphinx-rtd-theme>=1.0.0; extra == "docs"
Dynamic: license-file

<p align="center">
  <img src="assets/logo.png" width="425" alt="CHESSKIT logo"/>
</p>

<h1 align="left">CHESSKIT</h1>

 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

```bash
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)

```bash
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

```bash
# 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

```bash
# 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)

```bash
# 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

```bash
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.

```python
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](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`:

```yaml
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

```python
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:

```python
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:

```yaml
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](WORKFLOW_EXAMPLE.md) for complete templates.

**Build (SevenNet example)**:
```yaml
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)**:
```yaml
mode: sample
paths:
  cache: embeddings.h5
  report_dir: ./report
chess:
  metric: fisherrao
  k_select: 128
  k_boe: 64
```

## Testing

```bash
# Run all tests
pytest

# With coverage
pytest --cov=src/chess

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

## Documentation

- **[PIPELINE_GUIDE.md](PIPELINE_GUIDE.md)** - Complete API reference
- **[WORKFLOW_EXAMPLE.md](WORKFLOW_EXAMPLE.md)** - End-to-end usage examples  
- **[CONFIGURATION_MANUAL.md](CONFIGURATION_MANUAL.md)** - Full config key reference

## Citation

If you use CHESSKIT in your research, please cite:

```bibtex
@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](https://github.com/ecshin94/CHESS-toolkit/issues)
- Discussions: [GitHub Discussions](https://github.com/ecshin94/CHESS-toolkit/discussions)
