Improved test coverage

mlcolvar/core/loss/tda_loss.py

@@ -149,10 +149,8 @@ def tda_loss(
     """
     if not isinstance(target_centers, torch.Tensor):
         target_centers = torch.tensor(target_centers, dtype=H.dtype)
-        target_centers.requires_grad_(True)
     if not isinstance(target_sigmas, torch.Tensor):
         target_sigmas = torch.tensor(target_sigmas, dtype=H.dtype)
-        target_sigmas.requires_grad_(True)
 
     device = H.device
     target_centers = target_centers.to(device)
@@ -191,3 +189,18 @@ def tda_loss(
     if return_loss_terms:
         return loss, loss_centers, loss_sigmas
     return loss
+
+def test_tda_loss():
+    H = torch.randn(100)
+    H.requires_grad = True
+    labels = torch.zeros_like(H)
+    labels[-50:] = 1
+
+    Loss = TDALoss(n_states=2, target_centers=[-1, 1], target_sigmas=[0.1, 0.1])
+
+    loss = Loss(H=H, labels=labels, return_loss_terms=True)
+
+    loss[0].backward()
+
+if __name__ == '__main__':
+    test_tda_loss()

mlcolvar/cvs/committor/utils.py

@@ -186,4 +186,32 @@ def get_descriptors_and_derivatives(dataset,
                                 'labels': torch.clone(dataset['labels']), 
                                 'weights' : torch.clone(dataset['weights'])})
 
-    return smart_dataset, smart_derivatives
+    return smart_dataset, smart_derivatives
+
+def test_Kolmogorov_bias():
+    from mlcolvar.core.nn import FeedForward
+    model = FeedForward(layers=[4,2,1], activation='tanh')
+    inp = torch.randn((10, 4))
+    model_bias = KolmogorovBias(input_model=model, beta=1.0)
+    model_bias(inp)
+
+def test_compute_committor_weights():
+    # descriptors
+    # create dataset
+    samples = 50
+    X = torch.randn((3*samples, 6))
+
+    # create labels, bias and weights
+    y = torch.zeros(X.shape[0])
+    y[samples:] += 1
+    y[int(2*samples):] += 1
+    bias = torch.ones_like(X)
+    w = torch.zeros(X.shape[0])
+
+    # create and edit dataset
+    dataset = DictDataset({"data": X, "labels": y, "weights": w})
+    dataset = compute_committor_weights(dataset=dataset, bias=bias, data_groups=[0,1,2], beta=1.0)
+
+if __name__ == '__main__':
+    test_Kolmogorov_bias()
+    test_compute_committor_weights()