WIP feat: SMACT Property Prediction Module

[ryannduma]

This issue tracks the development of a new property_prediction module within the SMACT codebase, utilising trained ROOST model checkpoints. The goal is to enable property prediction directly from chemical compositions, without requiring crystal structure information.

This work builds on efforts from the 2025 SMACTDOWN session and, once integrated, will allow users to compute properties such as band gap and bulk modulus via the SMACT interface.

The module includes a flexible scaffold designed to support additional prediction models in the future, ensuring long-term extensibility.

Module Structure

smact/property_prediction/
├── __init__.py              # Main module interface
├── base_predictor.py        # Abstract base class
├── wrapper.py           # Convenience functions (i.e. predict_bandgap(model, composition)) default Roost
├── roost/
│   ├── __init__.py         # ROOST module
│   └── predictor.py        # ROOST implementation
└── example_usage.py        # Example usage demonstration

Proposed Implementation

The proposed usage pattern would be:

from smact.property_prediction import RoostPropertyPredictor

# Check available properties
print(RoostPropertyPredictor.available_properties)  # ['band_gap', 'bulk_modulus']

# Create model with just property name and device
model = RoostPropertyPredictor(property_name="band_gap", device="cpu")

# Predict single or multiple compositions
model.predict("NaCl")                    # Single composition
model.predict(["NaCl", "TiO2", "GaN"])   # Multiple compositions

Key Features

1. Class-level properties: RoostPropertyPredictor.available_properties
2. SMACT validation: All compositions validated using smact_validity()
3. Extensible design: Easy to add new property prediction model types
4. Error handling: Clear error messages for invalid compositions/properties

Extensibility for Multiple Models

The architecture is designed for multiple models:

# Future implementations will work the same way:
from smact.property_prediction import OtherPropertyPredictor

other_model = OtherPropertyPredictor(property_name="bulk_modulus")

Abstract Base Class Architecture

In /smact/property_prediction/base_predictor.py:

from abc import ABC, abstractmethod

class BasePropertyPredictor(ABC):  # Abstract base class
    """Abstract base class for property predictors."""
    
    @property
    @abstractmethod
    def supported_properties(self) -> list[str]:  # Must be implemented
        """List of properties supported by this predictor."""
        pass
    
    @abstractmethod
    def predict(self, compositions: str | list[str]) -> np.ndarray:  # Must be implemented
        """Predict property values for given compositions."""
        pass

Concrete Implementation

The RoostPropertyPredictor inherits from and implements the abstract base class:

class RoostPropertyPredictor(BasePropertyPredictor):  # Concrete implementation
    
    @property
    def supported_properties(self) -> list[str]:  # Implements abstract method
        return self.available_properties
    
    def predict(self, compositions: str | list[str]) -> np.ndarray:  # Implements abstract method
        compositions = self._validate_compositions(compositions)
        # Implementation here...
        return np.zeros(len(compositions))  # Placeholder for now

Extensible Pattern

This allows for multiple concrete implementations:

# Future models will also inherit from BasePropertyPredictor
class OtherPropertyPredictor(BasePropertyPredictor):
    def supported_properties(self) -> list[str]:
        return ["band_gap", "formation_energy"]
    
    def predict(self, compositions: str | list[str]) -> np.ndarray:
        # Other-specific implementation
        pass

Current Status

The module is ready for integration with actual ROOST model checkpoints once training is complete.

Contributors

Hyunsoo Park
Jamie Swaine
Matthew Walker
Yuan Li
Zibo Zhou
Ry Nduma

Acknowledgments

With thanks to the broader SMACT community, and especially to participants of SMACTDOWN 2025 for their feedback and contributions.

Citations:

roost - Predicting materials properties without crystal structure: Deep representation learning from stoichiometry. [Paper] [arXiv]

@article{goodall_2020_predicting,
  title={Predicting materials properties without crystal structure: Deep representation learning from stoichiometry},
  author={Goodall, Rhys EA and Lee, Alpha A},
  journal={Nature Communications},
  volume={11},
  number={1},
  pages={1--9},
  year={2020},
  publisher={Nature Publishing Group}
}

[AntObi]

@ryannduma @hspark1212 , might be worth raising an issue in https://github.com/CompRhys/aviary to see if there is any plan for any PyPI releases

[ryannduma]

okay will do

[AntObi]

Something that you should probably take into account is how you would handle different model versions of the same property. Easiest case to consider would be training separate models on the Matbench dataset and the latest version of the Materials Project; different versions of the same dataset; models trained on different subsets of the same data-source. Which model would be your "band gap" model? What if you wanted band gap models for different fidelities i.e. PBE, r2SCAN, HSE06, GLLB-SC, experimental, etc?

I would recommend looking at the MatGL library of a good example of maintaining different model versions, and use that as inspiration for how this would be developed: https://github.com/materialsvirtuallab/matgl/tree/main/pretrained_models

[ryannduma]

Rhys has been a great help setting up the PyPi version up and it's now live at https://pypi.org/project/aviary-models/1.2.0/

He'll push some updates later this week to tidy things up but should be quite useful in its current state right now 🌟

[ryannduma]

@AntObi and yeah no great point about model versioning!

I’m thinking we extend the constructor to include an optional model_version parameter:

# Default (latest/recommended model)
model = RoostPropertyPredictor(property_name="band_gap")

# Specific version/fidelity
model = RoostPropertyPredictor(
    property_name="band_gap", 
    model_version="band_gap_pbe"  # or "band_gap_hse06", "band_gap_experimental"
)

we’d maintain a model registry (similar to MatGL’s approach) that tracks:

• different fidelities (PBE, r2SCAN, HSE06, experimental)
• data sets versions (MP2024, Matbench, etc.)
• metrics and metadata

the wrapper functions would keep things simple for common use cases:

predict_bandgap("NaCl")  # Uses default/best model
predict_bandgap("NaCl", fidelity="HSE06")  # Specific fidelity

I think this would keep things tidy, but this is just an early scaffold we can build on moving forward what do you think @hspark1212 ?

[ryannduma]

Property Prediction Module Update

Note

🎄 Holiday Note: I know it's 19th December and everyone's winding down for the holidays, happy holidays to all! No need to feel obliged to review this over the break. I'll continue testing and potentially add more features over the next few weeks, then submit a comprehensive PR for review in January. Enjoy the festivities! 🎅

Summary

I have implemented a property prediction module for SMACT using the ROOST (Representation Learning from Stoichiometry) architecture. The module provides composition-only property prediction, starting with band gap prediction trained on 103,644 Materials Project DFT calculations.

Key Results:

• Test MAE: 0.283 eV (competitive with state-of-the-art composition-only methods)
• Test RMSE: 0.599 eV
• Training data: 103,644 compositions from Materials Project
• Model architecture: ROOST with MatScholar200 embeddings

---

Development Journey

Phase 1: Matbench Benchmark Validation

Before training on Materials Project data, I validated our ROOST implementation using the matbench_expt_gap benchmark (4,604 experimental band gaps). Results showed consistent ~0.98 eV MAE across three configurations:

Version	Configuration	Mean MAE (eV)
V1	Default hyperparameters	0.977 +/- 0.05
V2	Larger model, ensemble, higher LR	0.977 +/- 0.06
V3	Validation-based early stopping	0.978 +/- 0.05

This confirmed the training pipeline works correctly, though ROOST might not be competitive(im saying might here as I would like to find out more from @CompRhys ... for now this is just my naive hypothesis see further below) on this particular small experimental dataset (top methods achieve 0.29-0.35 eV).

The ROOST paper (Goodall & Lee, 2020) provides good reasons explaining that may explain the performance gap:

1. Fundamental polymorph limitation: As stated in the paper, composition-only methods "give up the ability to handle polymorphs" - different crystal structures of the same composition can have vastly different band gaps. Experimental datasets contain real-world measurements where structure matters significantly.

2. Dataset characteristics: The paper's ablation study explicitly notes that experimental band gap datasets have a "lack of diversity within different chemical sub-spaces," making them "perhaps a less discriminative benchmark." The matbench_expt_gap dataset (4,604 samples) might fall into this category.

3. DFT vs experimental domain gap: The paper demonstrates that even with transfer learning from DFT to experimental data, there remains an irreducible gap. Their EX dataset achieved 0.243 eV MAE (ensemble), showing composition-only methods struggle with experimental measurements regardless of training approach.

4. Top matbench methods use structure: The best-performing methods on matbench_expt_gap (0.29-0.35 eV MAE) utilise crystal structure information or employ more sophisticated approaches like MODNet that can capture structure-property relationships indirectly.

Conclusion: The ~0.98 eV MAE on matbench is not a failure of the implementation but rather reflects the fundamental limitation of composition-only methods for predicting structure-sensitive properties from experimental data. ROOST excels where it was designed to: large-scale DFT datasets where consistent computational methodology reduces noise.

Phase 2: Materials Project Training

Moved to training on Materials Project DFT band gaps, which provides:

• Much larger dataset (103,644 vs 4,604 samples)
• More consistent data (DFT calculations vs experimental measurements)
• Includes both metals (gap=0) and semiconductors
• Better suited to deep learning approaches

Training Configuration:

• Dataset: Materials Project (GGA/GGA+U calculations)
• Split: 80% train / 10% validation / 10% test
• Epochs: 232 (early stopping triggered)
• Learning rate: 3e-4 with cosine annealing
• Model: ROOST with matscholar200 embeddings

Final Model Performance:

• Test MAE: 0.283 eV
• Test RMSE: 0.599 eV
• Validation MAE: 0.282 eV

The parity plots(attached below) (predicted vs DFT band gap) show:

1. Good overall performance - Most predictions fall within ±0.5 eV of true values, with the error distribution roughly Gaussian and centred near zero.

2. Metals/zero-gap materials problem - The most visible issue is a cluster of high-error points at DFT band gap ~0 eV that are predicted to have band gaps of 2-6 eV. This is the polymorph limitation: the same composition can be metallic in one structure and semiconducting in another. ROOST predicts an "average" behaviour across polymorphs. e.g Fe₃O₄ could be metallic magnetite or insulating maghemite depending on structure. ROOST sees "Fe₃O₄" and predicts somewhere in between. Additionally, Materials Project contains many metastable metallic phases of compositions that are typically semiconducting in their ground state. This may result in the model "correctly" predicting a gap for a composition that DFT reports as metallic, as it has learned from the more common semiconducting variants.

3. Slight overprediction bias - Mean error of +0.031 eV indicates the model tends to overpredict band gaps on average slightly however, a bias this small is essentially negligible compared to the MAE of 0.283 eV... the model isn't systematically over/under-predicting.

4. Larger errors at extremes - Both very low (<1 eV) and very high (>7 eV) band gap materials show more scatter from the diagonal, likely due to fewer training examples and greater structure sensitivity in these regimes.

5. Mid-range scatter - Even in the 2-5 eV range where most data exists, there is noticeable scatter due to structure-dependent variations that composition-only models cannot capture.

These observations are consistent with the fundamental limitations of composition-only prediction discussed in the ROOST paper, rather than implementation issues.

---

Proposed Module Structure (WIP)

smact/property_prediction/
    __init__.py         
    base_predictor.py    # BasePropertyPredictor ABC + PredictionResult
    config.py            # Model registry and property metadata
    convenience.py       # High-level predict_band_gap() function
    io.py                # Model caching and download utilities should we have different ckpts in remote saves
    registry.py          # Property/model lookup functions
    roost/
        __init__.py
        predictor.py     # RoostPropertyPredictor implementation
        train.py         # Training utilities
    scripts/
        __init__.py
        convert_checkpoint.py   # Convert training checkpoints to SMACT format
    tests/
        __init__.py
        test_base_predictor.py  # Tests for base classes
        test_io.py              # Tests for IO functions
        test_registry.py        # Tests for registry functions

---

Usage

Simple Prediction

from smact.property_prediction import predict_band_gap

# Single composition
result = predict_band_gap("SiO2")
print(f"Band gap: {result[0]:.2f} eV")

# Multiple compositions
compositions = ["GaAs", "Si", "TiO2", "CdTe"]
predictions = predict_band_gap(compositions)
for comp, bg in zip(compositions, predictions):
    print(f"{comp}: {bg:.2f} eV")

With Uncertainty Estimates

from smact.property_prediction import predict_band_gap

result = predict_band_gap("ZnO", return_uncertainty=True)
print(f"Band gap: {result.mean[0]:.2f} +/- {result.std[0]:.2f} eV")

Using the Predictor Class

from smact.property_prediction import RoostPropertyPredictor

predictor = RoostPropertyPredictor(
    property_name="band_gap",
    device="cuda"  # Use GPU if available
)

# Get predictions
predictions = predictor.predict(["Si", "GaAs", "TiO2"])

Model Management

from smact.property_prediction import (
    get_supported_properties,
    get_property_unit,
    list_cached_models,
    clear_cache,
)

# List available properties
print(get_supported_properties())  # ['band_gap']

# Get property unit
print(get_property_unit("band_gap"))  # 'eV'

# List cached models
print(list_cached_models())

# Clear model cache
clear_cache()

---

Installation

Add to your SMACT installation:

pip install smact[property_prediction]

This installs the required dependencies:

• aviary-models>=1.2.0 (ROOST implementation)
• torch>=2.0.0

The model weights are automatically downloaded on first use to ~/.cache/smact/models/.

---

Tests

38 tests covering:

• PredictionResult dataclass
• BasePropertyPredictor ABC
• Composition validation
• IO functions (save/load checkpoint, caching)
• Registry functions
• Config consistency

All tests pass:

============================== 38 passed in 1.62s ==============================

---

Changes to SMACT

1. New module: smact/property_prediction/
2. pyproject.toml: Added property_prediction optional dependency:

   property_prediction = [
       "aviary-models>=1.2.0",
       "torch>=2.0.0",
   ]

---

Model Details

Roost-MP-2025.12.0-band_gap

Attribute	Value
Property	Band gap
Unit	eV
Dataset	Materials Project (103K)
Test MAE	0.283 eV
Test RMSE	0.599 eV
Architecture	ROOST
Model size	~2.3 million parameters
Embeddings	matscholar200
Message passing layers	3
Element heads	3
Training epochs	232

---

Comparison with MatGL Pretrained Models

Following @AntObi 's suggestion, I modelled our implementation on the MatGL pretrained_models approach.

What's Implemented (Aligned with MatGL)

Feature	MatGL	Our Implementation
File structure	model.json, model.pt, state.pt	Same
model.json format	@class, @module, @model_version, kwargs	Same
Remote caching	~/.matgl with RemoteFile class	~/.cache/smact/models/ with RemoteFile
Model naming	MEGNet-MP-2019.4.1-BandGap-mfi	Roost-MP-2025.12.0-band_gap
Version checking	Warns on mismatch	Warns on mismatch
Convenience functions	load_model(), clear_cache()	load_checkpoint(), clear_cache(), list_cached_models()

Model Versioning and Fidelity Support

This implementation includes infrastructure for handling concerns about model versions and fidelities:

Naming convention: Roost-<dataset>-<version>-<property>[-<fidelity>]

Config structure supports multiple fidelities:

# Current (single model):
DEFAULT_MODELS = {
    "band_gap": {
        "default": "Roost-MP-2025.12.0-band_gap",
    },
}

# Future expansion (multiple fidelities):
DEFAULT_MODELS = {
    "band_gap": {
        "default": "pbe",
        "pbe": "Roost-MP-2025.12.0-band_gap-pbe",
        "hse06": "Roost-MP-2025.12.0-band_gap-hse06",
        "r2scan": "Roost-MP-2025.12.0-band_gap-r2scan",
        "gllb_sc": "Roost-MP-2025.12.0-band_gap-gllb_sc",
        "experimental": "Roost-Expt-2025.12.0-band_gap",
    },
}

Supporting fidelity selection:

# Select specific fidelity
predictor = RoostPropertyPredictor(property_name="band_gap", fidelity="hse06")

# Or use a specific model version
predictor = RoostPropertyPredictor(model_name="Roost-MP-2025.12.0-band_gap-pbe")

Registry functions:

• get_default_model(property_name, fidelity) - resolve property+fidelity to model
• get_property_fidelities(property_name) - list available fidelities
• parse_model_name(model_name) - extract dataset, version, property, fidelity

Gaps to Address (Future Work)

Feature	MatGL	Our Status	Priority
Per-model README	Each model has README.md with aim, training data, metrics	Missing	High
Archive format	Downloads individual files	Downloads .tar.gz archives	Low (works)
Class deserialization	IOMixIn dynamically loads from @module/@class	Hardcoded to Roost	Medium
hubconf.py	PyTorch Hub entrypoints	Not implemented	Low
List all versions	Query GitHub API	Query manifest.json	Works
Multiple fidelity models	Trained and available	Config ready, models not yet trained	High

Draft Per-Model README Format (MatGL Style)

When hosting models, I think each should include a README.md:

# Aim
Band gap prediction from composition using ROOST architecture.

# Training dataset
MP-2025.12: Materials Project band gaps (PBE functional)
- Training set size: 82,915
- Validation set size: 10,364
- Test set size: 10,365

# Performance metrics
- Test MAE: 0.283 eV
- Test RMSE: 0.599 eV

---

Development Notes

During implementation, I encountered a subtle bug worth documenting:

Loss function shape mismatch: The robust_l1_loss function expects targets with shape [batch], but our training loop was passing [batch, 1]. PyTorch's broadcasting silently computed (pred - target) as a [batch, batch] tensor, resulting in incorrect gradients and a model that wouldn't learn. Symptoms: training MAE remained flat (~1.28 eV) regardless of epochs.

Fix: Remove .unsqueeze(1) when passing targets to the loss function.

This may help others debug similar issues.

---

Acknowledgements

Special thanks to:

• @hspark1212 (Hyunsoo Park) for supervising this work and providing guidance throughout its ongoing development
• @AntObi (Anthony Onwuli) for valuable feedback on architecture decisions, particularly the model versioning approach inspired by MatGL
• @CompRhys (Rhys Goodall) for making aviary-models available on PyPI, which made this integration possible

And to all the contributors from SMACTDOWN 2025 who laid the groundwork for this module.

---

References

• ROOST : Predicting materials properties without crystal structure: Deep representation learning from stoichiometry. [Paper] [arXiv]

• Materials Project : A. Jain*, S.P. Ong*, G. Hautier, W. Chen, W.D. Richards, S. Dacek, S. Cholia, D. Gunter, D. Skinner, G. Ceder, K.A. Persson (*=equal contributions) The Materials Project: A materials genome approach to accelerating materials innovation APL Materials, 2013, 1(1), 011002. doi:10.1063/1.4812323

• Matbench : Dunn, A., Wang, Q., Ganose, A., Dopp, D., Jain, A. Benchmarking Materials Property Prediction Methods: The Matbench Test Set and Automatminer Reference Algorithm. npj Computational Materials 6, 138 (2020). https://doi.org/10.1038/s41524-020-00406-3

• aviary-models : https://pypi.org/project/aviary-models/

---

Have a Merry Christmas !!

[CompRhys]

Roost Bandgap result on matbench_bandgap_expt feels high, the matbench_bandgap_expt dataset I believe is the same dataset from https://pubs.acs.org/doi/10.1021/acs.jpclett.8b00124 that I trained on in the version of the paper that got published. The "cognate TL" example I showed ended up with 0.243 eV MAE on a single split of that dataset so does feel like there's a bug somewhere, the pretraining on MP resulted in 0.210 eV.

I can maybe take a look to see if I broke something subtle at any point since the original work.

[hspark1212]

Great work on implementing the ROOST model! I have a few questions regarding the validation and baseline comparisons:

1. The matbench_expt_gap MAE (~0.98 eV) seems quite high. Could you share the training code in the PR so we can investigate potential issues in the training pipeline?

2. Could you visualize the band gap distributions for both datasets?

• matbench_expt_gap vs MP band gap histograms
• 0 eV vs non-zero band gap in the MP dataset
  I’m concerned that if the MP dataset has many samples near 0 eV, the model might achieve low loss by simply predicting values close to 0, without learning meaningful composition–property relationships.

3. MP model validation
   For the MP model result (0.283 eV MAE), we could validate that this is meaningful by comparing against:

• Simple baselines: constant prediction (mean/median of the training data) or random sampling predictions from the training distribution
• Linear models: linear regression or a small dense network using a compositional featurization/embedding such as Magpie