Remove BoltzmannCombination references.

mtenn/config.py

@@ -128,14 +128,10 @@ class CombinationConfig(StringEnum):
     * mean: :py:class:`MeanCombination <mtenn.combination.MeanCombination>`
 
     * max: :py:class:`MaxCombination <mtenn.combination.MaxCombination>`
-
-    * boltzmann:
-      :py:class:`BoltzmannCombination <mtenn.combination.BoltzmannCombination>`
     """
 
     mean = "mean"
     max = "max"
-    boltzmann = "boltzmann"
 
 
 class ModelConfigBase(BaseModel):
@@ -273,8 +269,6 @@ def build(self) -> mtenn.model.Model:
                 mtenn_combination = mtenn.combination.MaxCombination(
                     negate_preds=self.max_comb_neg, pred_scale=self.max_comb_scale
                 )
-            case CombinationConfig.boltzmann:
-                mtenn_combination = mtenn.combination.BoltzmannCombination()
             case None:
                 mtenn_combination = None
 

mtenn/tests/test_combination.py

@@ -3,7 +3,7 @@
 import pytest
 import torch
 
-from mtenn.combination import MeanCombination, MaxCombination, BoltzmannCombination
+from mtenn.combination import MeanCombination, MaxCombination
 from mtenn.conversion_utils.schnet import SchNet
 
 
@@ -89,33 +89,3 @@ def test_max_combination(models_and_inputs):
             for n, p in model_test.named_parameters()
         ]
     )
-
-
-def test_boltzmann_combination(models_and_inputs):
-    model_test, model_ref, inp_list, target, loss_func = models_and_inputs
-
-    # Ref calc
-    pred_list = torch.stack([model_ref(X)[0] for X in inp_list])
-    w = torch.exp(-pred_list - torch.logsumexp(-pred_list, axis=0))
-    pred_ref = torch.dot(w.flatten(), pred_list.flatten())
-    loss = loss_func(pred_ref, target)
-    loss.backward()
-
-    # Finish setting up GroupedModel
-    model_test = SchNet.get_model(
-        model_test, grouped=True, strategy="complex", combination=BoltzmannCombination()
-    )
-
-    # Test GroupedModel
-    pred_test, _ = model_test(inp_list)
-    loss = loss_func(pred_test, target)
-    loss.backward()
-
-    # Compare
-    ref_param_dict = dict(model_ref.named_parameters())
-    assert all(
-        [
-            np.allclose(p.grad, ref_param_dict[n].grad, atol=5e-7)
-            for n, p in model_test.named_parameters()
-        ]
-    )