[Feature Request]: MC Dropout uncertainty quantification for CHGNet inference

[0xSoftBoi]

Problem

CHGNet currently produces point predictions with no built-in uncertainty estimate. For applications like active learning, candidate screening, and reliability filtering, knowing how confident the model is in a prediction is as useful as the prediction itself—without it, you can't distinguish a confident low-energy prediction from a lucky one.

Proposed Solution

I've implemented an MCDropoutPredictor class that wraps an existing CHGNet model and enables Monte Carlo Dropout inference (Gal & Ghahramani, ICML 2016). The idea is straightforward: the pretrained weights already contain nn.Dropout layers in the MLP readout head (with p=0 by default). Setting p > 0 and running N stochastic forward passes gives a distribution over predictions—the mean approximates the standard deterministic output, while the std serves as an epistemic uncertainty proxy.

The backbone convolution layers remain in eval() mode throughout, so only the MLP head is stochastic. This keeps memory overhead low and avoids any change to the force/stress gradient graph.

Basic usage:

from chgnet.model import CHGNet
from chgnet.model.uncertainty import MCDropoutPredictor

model = CHGNet.load()
predictor = MCDropoutPredictor(model, dropout_p=0.1)

result = predictor.predict_structure(structure, n_passes=20)
print(result["energy_mean"])  # matches CHGNet's standard output (within ~1%)
print(result["energy_std"])   # epistemic uncertainty proxy

Both predict_structure and predict_graph are supported, with the same batching and task options ("e", "ef", "efs", "efsm") as the base model. The implementation passes ruff check and ruff format with select = ["ALL"] against CHGNet's pyproject.toml settings.

A prototype implementation (with tests) is available here if useful for reference:
https://github.com/0xSoftBoi/chgnet/tree/feat/mc-dropout-uncertainty

Question for maintainers

I noticed the CHGNet README mentions that feature development has moved to MatGL. Before investing time in a formal PR, I wanted to ask:

• Does this belong here, or would it be more appropriate to target MatGL?
• If CHGNet, is this the right scope for a PR, or would a smaller addition (e.g. a standalone utility function) be preferred?

Happy to submit a PR to either repo—just want to make sure it goes where it'll be most useful and have the best chance of being reviewed.

References

• Gal, Y. & Ghahramani, Z. Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning. ICML 2016. https://arxiv.org/abs/1506.02142

Code of Conduct

• I agree to follow this project's Code of Conduct

[bowen-bd]

Hi,

Thanks for this contribution.
Uncertainty prediction is out of scope for the support of CHGNet, as well as most MLIPs to date.
MCDropoutPredictor can be potentially useful, but I think this should be better maintained in a separate fork, or branch.
I'm happy to keep this issue open so future usage of uncertainty prediction can be referenced here.

In addition, this repo is no longer actively maintained, and the MatGL repo contains improved implementation and checkpoints of CHGNet. The syntax is also a bit different for MatGL as the code has been rewritten

Thanks!