Graphs_integration

devtools/conda-envs/test_env.yaml

@@ -3,6 +3,7 @@ channels:
 
   - pytorch
   - conda-forge
+  - pyg
 
   - defaults
 dependencies:
@@ -24,10 +25,14 @@ dependencies:
   - matplotlib
   - scikit-learn
   - scipy
-
+  - pyg
+
     # Pip-only installs
   - pip:
     - KDEpy
     - nbmake
     - torch-scatter
+    - mdtraj
+    - matscipy
+
 

docs/api_core.rst

@@ -1,10 +1,14 @@
 Core modules
-------------
+============
 
 These are the building blocks which are used to construct the CVs.
 
-.. rubric:: NN
+NN
+--
+This module implements the architectures with learnable weights that can be used to build CV models.
 
+Descriptors-based
+^^^^^^^^^^^^^^^^
 .. currentmodule:: mlcolvar.core.nn
 
 .. autosummary::
@@ -13,7 +17,31 @@ These are the building blocks which are used to construct the CVs.
 
    FeedForward
 
-.. rubric:: Loss
+Graphs-based
+^^^^^^^^^^^^
+.. currentmodule:: mlcolvar.core.nn.graph
+
+Base class
+""""""""""
+.. autosummary::
+   :toctree: autosummary
+   :template: custom-class-template.rst
+
+   BaseGNN
+
+Architectures
+"""""""""""""
+.. autosummary::
+   :toctree: autosummary
+   :template: custom-class-template.rst
+
+   SchNetModel
+   GVPModel
+
+
+Loss
+----
+This module implements the loss fucntions that can be used to optimize CV models.
 
 .. currentmodule:: mlcolvar.core.loss
 
@@ -28,20 +56,40 @@ These are the building blocks which are used to construct the CVs.
    ReduceEigenvaluesLoss
    TDALoss
 
-.. rubric:: Stats 
 
+Stats 
+-----
+This module implements statistical methods with learnable weights that can be used in CV models.
+
+Base class
+^^^^^^^^^^
 .. currentmodule:: mlcolvar.core.stats
 
 .. autosummary::
    :toctree: autosummary
    :template: custom-class-template.rst
 
    Stats
+
+Linear methods
+^^^^^^^^^^^^^^
+.. currentmodule:: mlcolvar.core.stats
+
+.. autosummary::
+   :toctree: autosummary
+   :template: custom-class-template.rst
+
    PCA
    LDA 
    TICA
 
-.. rubric:: Transform
+
+Transform
+---------
+This module implements **non-learnable** pre/postprocessing tools 
+
+Base class
+^^^^^^^^^^
 
 .. currentmodule:: mlcolvar.core.transform
 
@@ -52,7 +100,9 @@ These are the building blocks which are used to construct the CVs.
    Transform
 
 
-.. rubric:: Transform.descriptors
+Descriptors
+^^^^^^^^^^^
+This submodule implements several descriptors that can be computed starting from atomic positions.
 
 .. currentmodule:: mlcolvar.core.transform.descriptors
 
@@ -66,7 +116,9 @@ These are the building blocks which are used to construct the CVs.
    EigsAdjMat
    MultipleDescriptors
 
-.. rubric:: Transform.tools
+Tools
+^^^^^
+This submodule implements pre/postporcessing tools.
 
 .. currentmodule:: mlcolvar.core.transform.tools
 

docs/api_data.rst

@@ -1,6 +1,9 @@
 Data
 ----
 
+General: dataset, module and loader
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
 .. currentmodule:: mlcolvar.data
 
 This module contains the classes used for handling datasets and for feeding them to the Lightning trainer.
@@ -11,4 +14,19 @@ This module contains the classes used for handling datasets and for feeding them
 
    DictDataset
    DictLoader
-   DictModule
+   DictModule
+
+Graph specific tools
+^^^^^^^^^^^^^^^^^^^^
+.. currentmodule:: mlcolvar.data.graph
+
+This module contains the classes used for handling and creating graphs.
+
+.. autosummary::
+   :toctree: autosummary
+   :template: custom-class-template.rst
+
+   AtomicNumberTable
+   Configuration
+   get_neighborhood
+   create_dataset_from_configurations

docs/api_explain.rst

@@ -1,20 +1,36 @@
 Explain
 -------
 
-.. rubric:: Sensitivity analysis
+Sensitivity analysis
+^^^^^^^^^^^^^^^^^^^^
 
 Perform sensitivity analysis to identify feature relevances
 
 .. currentmodule:: mlcolvar.explain.sensitivity
 
+Descriptors-based
+""""""""""""""""""
+
 .. autosummary::
    :toctree: autosummary
    :template: custom-class-template.rst
 
    sensitivity_analysis 
    plot_sensitivity
 
-.. rubric:: Sparse linear model
+Graph-based
+"""""""""""
+.. currentmodule:: mlcolvar.explain.graph_sensitivity
+
+.. autosummary::
+   :toctree: autosummary
+   :template: custom-class-template.rst
+
+   graph_node_sensitivity
+
+
+Sparse linear models
+^^^^^^^^^^^^^^^^^^^^
 
 Use sparse models to approximate classification/regression tasks
 

docs/api_utils.rst

@@ -1,7 +1,9 @@
 Utils
 -----
 
-.. rubric:: Input/Output
+
+Input/Output
+^^^^^^^^^^^^
 
 Helper functions for loading dataframes (incl. PLUMED files) and directly creating datasets from them.
 
@@ -14,7 +16,9 @@ Helper functions for loading dataframes (incl. PLUMED files) and directly creati
    load_dataframe
    create_dataset_from_files
 
-.. rubric:: Time-lagged datasets
+
+Time-lagged datasets
+^^^^^^^^^^^^^^^^^^^^
 
 Create a dataset of pairs of time-lagged configurations.
 
@@ -26,31 +30,35 @@ Create a dataset of pairs of time-lagged configurations.
 
    create_timelagged_dataset
 
-.. rubric:: FES
-
-.. rubric:: Trainer
 
-Functions used in conjunction with the lightning Trainer (e.g. logging, metrics...).
+FES
+^^^
+Compute (and plot) the free energy surface along the CVs.
 
-.. currentmodule:: mlcolvar.utils.trainer
+.. currentmodule:: mlcolvar.utils.fes
 
 .. autosummary::
    :toctree: autosummary
    :template: custom-class-template.rst
 
-   MetricsCallback
+   compute_fes
 
-Compute (and plot) the free energy surface along the CVs.
 
-.. currentmodule:: mlcolvar.utils.fes
+Trainer
+^^^^^^^
+Functions used in conjunction with the lightning Trainer (e.g. logging, metrics...).
+
+.. currentmodule:: mlcolvar.utils.trainer
 
 .. autosummary::
    :toctree: autosummary
    :template: custom-class-template.rst
 
-   compute_fes
+   MetricsCallback
+
 
-Plotting utils
+Plot
+^^^^
 
 .. currentmodule:: mlcolvar.utils.plot
 

docs/autosummary/mlcolvar.core.nn.FeedForward.rst

@@ -17,6 +17,7 @@
    .. autosummary::
 
       ~FeedForward.__init__
+      ~FeedForward.backward
       ~FeedForward.forward
 
 
@@ -29,27 +30,9 @@
 
    .. autosummary::
 
-      ~FeedForward.CHECKPOINT_HYPER_PARAMS_KEY
-      ~FeedForward.CHECKPOINT_HYPER_PARAMS_NAME
-      ~FeedForward.CHECKPOINT_HYPER_PARAMS_TYPE
       ~FeedForward.T_destination
-      ~FeedForward.automatic_optimization
       ~FeedForward.call_super_init
-      ~FeedForward.current_epoch
-      ~FeedForward.device
-      ~FeedForward.dtype
       ~FeedForward.dump_patches
-      ~FeedForward.example_input_array
-      ~FeedForward.fabric
-      ~FeedForward.global_rank
-      ~FeedForward.global_step
-      ~FeedForward.hparams
-      ~FeedForward.hparams_initial
-      ~FeedForward.local_rank
-      ~FeedForward.logger
-      ~FeedForward.loggers
-      ~FeedForward.on_gpu
-      ~FeedForward.trainer
       ~FeedForward.training
 
 

docs/autosummary/mlcolvar.core.nn.graph.BaseGNN.rst

@@ -0,0 +1,41 @@
+mlcolvar.core.nn.graph.BaseGNN
+==============================
+
+.. currentmodule:: mlcolvar.core.nn.graph
+
+.. autoclass:: BaseGNN
+   :members:                                   
+   :show-inheritance:                           
+   :inherited-members: Module,LightningModule                       
+
+
+   .. automethod:: __init__
+
+
+   .. rubric:: Methods
+
+   .. autosummary::
+
+      ~BaseGNN.__init__
+      ~BaseGNN.embed_edge
+
+
+
+
+..
+
+
+   .. rubric:: Attributes
+
+   .. autosummary::
+
+      ~BaseGNN.T_destination
+      ~BaseGNN.call_super_init
+      ~BaseGNN.dump_patches
+      ~BaseGNN.in_features
+      ~BaseGNN.out_features
+      ~BaseGNN.training
+
+
+
+