diff --git a/docs/_templates/custom-class-template.rst b/docs/_templates/custom-class-template.rst new file mode 100644 index 0000000..2ae7050 --- /dev/null +++ b/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 %} diff --git a/docs/_templates/custom-module-template.rst b/docs/_templates/custom-module-template.rst index 7e55723..8c1f5e1 100644 --- a/docs/_templates/custom-module-template.rst +++ b/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 %} diff --git a/docs/api.md b/docs/api.md index d30b2bc..f982e82 100644 --- a/docs/api.md +++ b/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 ``` diff --git a/docs/conf.py b/docs/conf.py index c517aac..10f1e39 100644 --- a/docs/conf.py +++ b/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 diff --git a/docs/QUICK_START.md b/docs/guides/fine-tuning.md similarity index 97% rename from docs/QUICK_START.md rename to docs/guides/fine-tuning.md index ff61dbe..cb678f4 100644 --- a/docs/QUICK_START.md +++ b/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 diff --git a/docs/LIGHTNING.md b/docs/guides/lightning.md similarity index 99% rename from docs/LIGHTNING.md rename to docs/guides/lightning.md index 931a22a..4425ece 100644 --- a/docs/LIGHTNING.md +++ b/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. diff --git a/docs/index.md b/docs/index.md index 3f4bfd8..44e9629 100644 --- a/docs/index.md +++ b/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} diff --git a/docs/MOTIVATION.md b/docs/motivation.md similarity index 98% rename from docs/MOTIVATION.md rename to docs/motivation.md index 08c0b35..24bc8c4 100644 --- a/docs/MOTIVATION.md +++ b/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: diff --git a/docs/requirements.txt b/docs/requirements.txt index 4092e60..aeba7f7 100644 --- a/docs/requirements.txt +++ b/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==* diff --git a/src/mattertune/registry.py b/src/mattertune/registry.py index 6c8cd5c..8d1db06 100644 --- a/src/mattertune/registry.py +++ b/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", +]