From 648f748464d5804ad67fd0d2c8a9c8b82fb16944 Mon Sep 17 00:00:00 2001
From: Tsz Wai Ko <47970742+kenko911@users.noreply.github.com>
Date: Tue, 31 Dec 2024 11:29:57 +0800
Subject: [PATCH] Add an option for Bussi ensemble

Signed-off-by: Tsz Wai Ko <47970742+kenko911@users.noreply.github.com>
---
 src/matgl/ext/ase.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/matgl/ext/ase.py b/src/matgl/ext/ase.py
index 3aaa451a..cd16c667 100644
--- a/src/matgl/ext/ase.py
+++ b/src/matgl/ext/ase.py
@@ -379,7 +379,7 @@ def __init__(
         state_attr: torch.Tensor | None = None,
         stress_weight: float = 1 / 160.21766208,
         ensemble: Literal[
-            "nve", "nvt", "nvt_langevin", "nvt_andersen", "npt", "npt_berendsen", "npt_nose_hoover"
+            "nve", "nvt", "nvt_langevin", "nvt_andersen", "nvt_bussi", "npt", "npt_berendsen", "npt_nose_hoover"
         ] = "nvt",
         temperature: int = 300,
         timestep: float = 1.0,