Refactor documentation structure and add API reference

We have to now install some deps in the workflow, such as nshconfig, torch, ase, so that we can build the docs.

docs/_templates/custom-class-template.rst

@@ -0,0 +1,21 @@
+{{ fullname | escape | underline}}
+
+.. currentmodule:: {{ module }}
+
+.. autoclass:: {{ objname }}
+   :members:
+   :show-inheritance:
+   :inherited-members:
+   :special-members: __call__, __add__, __mul__
+
+   {% block methods %}
+   {% if methods %}
+   .. rubric:: {{ _('Methods') }}
+
+   .. autosummary::
+      :nosignatures:
+   {% for item in methods %}
+      ~{{ name }}.{{ item }}
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}

docs/_templates/custom-module-template.rst

@@ -4,16 +4,35 @@
    :members:
    :undoc-members:
    :show-inheritance:
-   :special-members: __init__
+   :variables:
+   :module-variables:
+
+   {% block attributes %}
+   {% if attributes %}
+   .. rubric:: Module Variables
+
+   .. autosummary::
+      {% for item in attributes %}
+      {{ item }}
+      {%- endfor %}
+
+   {% for item in attributes %}
+   .. autodata:: {{ item }}
+      :annotation:
+   {%- endfor %}
+   {% endif %}
+   {% endblock %}
 
    {% block modules %}
    {% if modules %}
    .. rubric:: Submodules
+
    .. autosummary::
       :toctree:
       :recursive:
-   {% for item in modules %}
+
+      {% for item in modules %}
       {{ item }}
-   {%- endfor %}
+      {%- endfor %}
    {% endif %}
    {% endblock %}

docs/api.md

@@ -1,18 +1,11 @@
-# API Documentation
+# API Reference
 
 ```{eval-rst}
+.. currentmodule:: mattertune
+
 .. autosummary::
    :toctree: _autosummary
-   :template: custom-module-template.rst
    :recursive:
 
-   mattertune.backbones
-   mattertune.callbacks
-   mattertune.configs
-   mattertune.data
-   mattertune.finetune
-   mattertune.loggers
-   mattertune.main
-   mattertune.normalization
-   mattertune.registry
+   mattertune
 ```

docs/conf.py

@@ -52,14 +52,50 @@
     "member-order": "bysource",
     "special-members": "__init__",
     "undoc-members": True,
-    "exclude-members": "__weakref__",
+    "exclude-members": "__weakref__, __pydantic_core_schema__, __pydantic_validator__, __pydantic_serializer__, \
+__pydantic_fields_set__, __pydantic_extra__, __pydantic_private__, __pydantic_post_init__, __pydantic_decorators__, \
+__pydantic_parent_namespace__, __pydantic_generic_metadata__, __pydantic_custom_init__, __pydantic_complete__, \
+__fields__, __fields_set__, model_fields, model_config, model_computed_fields, __class_vars__, __private_attributes__, \
+__signature__, __pydantic_root_model__, __slots__, __dict__, model_extra, model_fields_set, model_post_init",
+    "show-module-summary": True,
 }
 
+
+# To exclude private members (those starting with _)
+def skip_private_members(app, what, name, obj, skip, options):
+    if (
+        name.startswith("_")
+        and not name.startswith("__")
+        and not name.startswith("__pydantic")
+    ):
+        return True
+    return skip
+
+
+def setup(app):
+    app.connect("autodoc-skip-member", skip_private_members)
+
+
+autodoc_mock_imports = [
+    # "torch",
+    "pytorch_lightning",
+    "torch-geometric",
+    "torchmetrics",
+]
+
 # Type hints settings
 autodoc_typehints = "description"
 autodoc_typehints_format = "short"
 
-# Add autosummary settings
+# Enable automatic doc generation
 autosummary_generate = True
-autosummary_imported_members = True
+autodoc_member_order = "bysource"
+add_module_names = True
+
+# Custom templates
+autodoc_template_dir = ["_templates"]
 templates_path = ["_templates"]
+
+# always_use_bars_union for type hints
+always_use_bars_union = True
+always_document_param_types = True

docs/QUICK_START.md → docs/guides/fine-tuning.md

@@ -1,4 +1,4 @@
-## Quick Start: Fine-Tuning a Pre-trained Model
+# Fine-Tuning a Pre-trained Model
 
 This guide will walk you through fine-tuning a pre-trained model for predicting molecular properties. We'll use a complete example with detailed explanations.
 
@@ -89,7 +89,7 @@ predictions = property_predictor.predict([water], ["energy"])
 print(f"Predicted energy: {predictions[0]['energy']} eV")
 ```
 
-### Key Components Explained
+## Key Components Explained
 
 1. **Configuration Structure**:
    - `MatterTunerConfig`: The main configuration container
@@ -107,7 +107,7 @@ print(f"Predicted energy: {predictions[0]['energy']} eV")
    - Automatic train/validation splitting
    - Configurable batch size for memory management
 
-### Common Customizations
+## Common Customizations
 
 ```python
 # Add force prediction
@@ -143,7 +143,7 @@ trainer=mt.configs.TrainerConfig(
 )
 ```
 
-### Tips for Successful Fine-Tuning
+## Tips for Successful Fine-Tuning
 
 1. **Data Quality**:
    - Ensure your training data is clean and properly formatted

docs/LIGHTNING.md → docs/guides/lightning.md

@@ -1,4 +1,4 @@
-# Lightning Integration in MatterTune
+# Advanced: Lightning Integration
 
 MatterTune uses PyTorch Lightning as its core training framework. This document outlines how Lightning is integrated and what functionality it provides.
 

docs/index.md

@@ -7,17 +7,19 @@ MatterTune is a flexible and powerful machine learning library designed specific
 :caption: Getting Started
 
 introduction
+motivation
 installation
-quick_start
 ```
 
 ```{toctree}
 :maxdepth: 2
 :caption: User Guide
 
+guides/fine-tuning
 guides/model_usage
 guides/training_config
 guides/data_handling
+guides/lightning
 ```
 
 ```{toctree}

docs/MOTIVATION.md → docs/motivation.md

@@ -128,7 +128,7 @@ class GraphPropertyConfig:
     loss_coefficient: float = 1.0
     reduction: Literal["mean", "sum", "max"]  # How to aggregate atomic features
 
-PropertyConfig: TypeAlias = EnergyPropertyConfig | ForcesPropertyConfig | StressesPropertyConfig | GraphPropertyConfig
+PropertyConfig = TypeAliasType("PropertyConfig", EnergyPropertyConfig | ForcesPropertyConfig | StressesPropertyConfig | GraphPropertyConfig)
 ```
 
 Benefits:

docs/requirements.txt

@@ -3,3 +3,6 @@ sphinx-rtd-theme>=1.0.0
 myst-parser>=0.18.0
 sphinx-autodoc-typehints>=1.12.0
 sphinx-copybutton>=0.5.0
+nshconfig==*
+torch==*
+ase==*

src/mattertune/registry.py

@@ -5,4 +5,11 @@
 from .finetune.base import FinetuneModuleBaseConfig
 
 backbone_registry = C.Registry(FinetuneModuleBaseConfig, discriminator="name")
+"""Registry for backbone modules."""
+
 data_registry = C.Registry(C.Config, discriminator="type")
+"""Registry for data modules."""
+__all__ = [
+    "backbone_registry",
+    "data_registry",
+]