From dd3b3539ffa1453aa37698ca5d030d9b198a57b5 Mon Sep 17 00:00:00 2001
From: EnricoTrizio <enricotrizio@gmail.com>
Date: Mon, 20 May 2024 11:54:58 +0200
Subject: [PATCH 1/2] Made last layer activation explicit block

---
 mlcolvar/cvs/committor/committor.py | 26 +++++++++++---------------
 1 file changed, 11 insertions(+), 15 deletions(-)

diff --git a/mlcolvar/cvs/committor/committor.py b/mlcolvar/cvs/committor/committor.py
index 4a672f28..bba62370 100644
--- a/mlcolvar/cvs/committor/committor.py
+++ b/mlcolvar/cvs/committor/committor.py
@@ -3,6 +3,7 @@
 from mlcolvar.cvs import BaseCV
 from mlcolvar.core import FeedForward
 from mlcolvar.core.loss import CommittorLoss
+from mlcolvar.core.nn.utils import Custom_Sigmoid
 
 __all__ = ["Committor"]
 
@@ -31,7 +32,7 @@ class Committor(BaseCV, lightning.LightningModule):
         Utils to initialize the masses tensor for the training
     """
 
-    BLOCKS = ["nn"]
+    BLOCKS = ["nn", "sigmoid"]
 
     def __init__(
         self, 
@@ -79,23 +80,16 @@ def __init__(
         # ======= OPTIONS =======
         # parse and sanitize
         options = self.parse_options(options)
-        
-        # add the relevant nn options, set tanh for hidden layers and sharp sigmoid for output layer
-        activ_list = ["tanh" for i in range( len(layers) - 2 )]
-        activ_list.append("custom_sigmoid")
-        
-        # update options dict for activations if not already set
-        if not "activation" in options["nn"]:
-            options["nn"]["activation"] = activ_list
-
-        # ======= CHECKS =======
-        # should be empty in this case
-
 
         # ======= BLOCKS =======
-        # initialize NN turning on last layer activation
+        # initialize NN turning
         o = "nn"
-        self.nn = FeedForward(layers, last_layer_activation=True, **options[o])
+        self.nn = FeedForward(layers, **options[o])
+
+        # separately add sigmoid activation on last layer, this way it can be deactived
+        o = "sigmoid"
+        if (options[o] is not False) and (options[o] is not None):
+            self.sigmoid = Custom_Sigmoid(**options[o])
 
 
     def training_step(self, train_batch, batch_idx):
@@ -153,5 +147,7 @@ def test_committor():
     trainer = lightning.Trainer(max_epochs=5, logger=None, enable_checkpointing=False, limit_val_batches=0, num_sanity_val_steps=0)
     trainer.fit(model, datamodule)
 
+    model(X).sum().backward()
+
 if __name__ == "__main__":
     test_committor()
\ No newline at end of file

From 0d85fde8d0a65e9391a62d9318a6b6340b83ba89 Mon Sep 17 00:00:00 2001
From: EnricoTrizio <enricotrizio@gmail.com>
Date: Mon, 20 May 2024 12:06:28 +0200
Subject: [PATCH 2/2] Added doc entry for committor

---
 docs/api_cvs.rst                            | 12 ++++-
 docs/autosummary/mlcolvar.cvs.Committor.rst | 60 +++++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 docs/autosummary/mlcolvar.cvs.Committor.rst

diff --git a/docs/api_cvs.rst b/docs/api_cvs.rst
index 937b8f19..23bb4936 100644
--- a/docs/api_cvs.rst
+++ b/docs/api_cvs.rst
@@ -67,4 +67,14 @@ General framework which allows to optimize a single model with different loss fu
    :toctree: autosummary
    :template: custom-class-template.rst
 
-   MultiTaskCV
\ No newline at end of file
+   MultiTaskCV
+
+Framework for the numerical determination of the committor function based on its variational principle.
+
+.. currentmodule:: mlcolvar.cvs
+
+.. autosummary::
+   :toctree: autosummary
+   :template: custom-class-template.rst
+
+   Committor
\ No newline at end of file
diff --git a/docs/autosummary/mlcolvar.cvs.Committor.rst b/docs/autosummary/mlcolvar.cvs.Committor.rst
new file mode 100644
index 00000000..4bda53d1
--- /dev/null
+++ b/docs/autosummary/mlcolvar.cvs.Committor.rst
@@ -0,0 +1,60 @@
+mlcolvar.cvs.Committor
+====================
+
+.. currentmodule:: mlcolvar.cvs
+
+.. autoclass:: Committor
+   :members:                                   
+   :show-inheritance:                           
+   :inherited-members: Module,LightningModule                       
+
+   
+   .. automethod:: __init__
+
+   
+   .. rubric:: Methods
+
+   .. autosummary::
+   
+      ~Committor.__init__
+      ~Committor.training_step
+   
+   
+
+
+..
+   
+   
+   .. rubric:: Attributes
+
+   .. autosummary::
+   
+      ~Committor.BLOCKS
+      ~Committor.CHECKPOINT_HYPER_PARAMS_KEY
+      ~Committor.CHECKPOINT_HYPER_PARAMS_NAME
+      ~Committor.CHECKPOINT_HYPER_PARAMS_TYPE
+      ~Committor.T_destination
+      ~Committor.automatic_optimization
+      ~Committor.call_super_init
+      ~Committor.current_epoch
+      ~Committor.device
+      ~Committor.dtype
+      ~Committor.dump_patches
+      ~Committor.example_input_array
+      ~Committor.fabric
+      ~Committor.global_rank
+      ~Committor.global_step
+      ~Committor.hparams
+      ~Committor.hparams_initial
+      ~Committor.local_rank
+      ~Committor.logger
+      ~Committor.loggers
+      ~Committor.n_cvs
+      ~Committor.on_gpu
+      ~Committor.optimizer_name
+      ~Committor.trainer
+      ~Committor.training
+   
+   
+
+   
\ No newline at end of file