Add a "reference" section to the documentation

.github/workflows/build.yml

@@ -8,6 +8,8 @@ on:
     branches:
       - master
   pull_request:
+    paths-ignore:
+      - 'docs/**'
 
 permissions:
   contents: read

.github/workflows/docs.yml

@@ -9,11 +9,21 @@ jobs:
     name: Deploy docs
     runs-on: ubuntu-latest
     steps:
-      - name: Checkout main
+      - name: Checkout
         uses: actions/checkout@v4
+      - name: Set up Python 3.12
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.12
+      - name: Install pdm
+        run: pip install pdm
+      - name: Install dependencies
+        run: pdm install
 
       - name: Deploy docs
-        uses: mhausenblas/[email protected]
-        # Or use mhausenblas/mkdocs-deploy-gh-pages@nomaterial to build without the mkdocs-material theme
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: mkdocs gh-deploy --force
+        # note: Checking if we really need the one below:
+        # uses: mhausenblas/[email protected]
+        # # Or use mhausenblas/mkdocs-deploy-gh-pages@nomaterial to build without the mkdocs-material theme
+        # env:
+        #   GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

docs/generate_reference_docs.py

@@ -10,119 +10,75 @@
 
 from project.utils.env_vars import REPO_ROOTDIR
 
-nav = mkdocs_gen_files.nav.Nav()
-
-
-package = "project"
-module = "project/main.py"
-submodules = ["project.datamodules", "project.utils", "project.networks", "project.algorithms"]
+module = "project"
+modules = [
+    "project/main.py",
+    "project/experiment.py",
+]
+submodules = [
+    "project.algorithms",
+    "project.configs",
+    "project.datamodules",
+    "project.networks",
+    "project.utils",
+]
 
 
 def _get_import_path(module_path: Path) -> str:
     """Returns the path to use to import a given (internal) module."""
     return ".".join(module_path.relative_to(REPO_ROOTDIR).with_suffix("").parts)
 
 
-def add_doc_for_module(module_path: Path) -> None:
-    a = "reference" / (module_path.relative_to(REPO_ROOTDIR).with_suffix(".md"))
-    module_import_path = _get_import_path(module_path)
+def main():
+    nav = mkdocs_gen_files.nav.Nav()
 
-    with mkdocs_gen_files.open(a, "w") as f:
-        print(
-            textwrap.dedent(f"""\
-            ::: {module_import_path}
-            """),
-            file=f,
-        )
-    docs_dir = REPO_ROOTDIR / "docs"
-    module_path_relative_to_docs_dir = module_path.relative_to(docs_dir, walk_up=True)
-    mkdocs_gen_files.set_edit_path(a, module_path_relative_to_docs_dir)
+    add_doc_for_module(REPO_ROOTDIR / "project", nav)
 
+    # with mkdocs_gen_files.open("reference/SUMMARY.md", "w") as nav_file:
+    #     # assert False, "\n".join(nav.build_literate_nav())
+    #     nav_file.writelines(nav.build_literate_nav())
 
-def get_modules(package: Path) -> list[Path]:
-    return [
-        p
-        for p in package.glob("*.py")
-        if not p.name.endswith("_test.py") and not p.name == "__init__.py"
-    ]
 
+def add_doc_for_module(module_path: Path, nav: mkdocs_gen_files.nav.Nav) -> None:
+    """TODO."""
+
+    assert module_path.is_dir() and (module_path / "__init__.py").exists(), module_path
 
-def get_subpackages(package: Path) -> list[Path]:
-    return [
+    children = list(
         p
-        for p in package.iterdir()
-        if p.is_dir() and not p.name.startswith("__") and (p / "__init__.py").exists()
-    ]
-
-
-project_nav = mkdocs_gen_files.nav.Nav()
-with mkdocs_gen_files.open("reference/project/main.md", "w") as f:
-    print(
-        textwrap.dedent("""\
-        ::: project.main
-        """),
-        file=f,
+        for p in module_path.glob("*.py")
+        if not p.name.startswith("__") and not p.name.endswith("_test.py")
     )
-nav["project", "main"] = "project/main.md"
-mkdocs_gen_files.set_edit_path("reference/project/main.md", "../project/main.py")
-
-with mkdocs_gen_files.open("reference/project/experiment.md", "w") as f:
-    print(
-        textwrap.dedent("""\
-        ::: project.experiment
-        """),
-        file=f,
-    )
-nav["project", "experiment"] = "reference/project/experiment.md"
-mkdocs_gen_files.set_edit_path("reference/project/experiment.md", "../project/experiment.py")
-
-project_utils_nav = mkdocs_gen_files.nav.Nav()
-with mkdocs_gen_files.open("reference/project/utils/types.md", "w") as f:
-    print(
-        textwrap.dedent("""\
-        ::: project.utils.types
-            options:
-                show_source: true
-        """),
-        file=f,
+    for child_module_path in children:
+        child_module_import_path = _get_import_path(child_module_path)
+        doc_file = child_module_path.relative_to(REPO_ROOTDIR).with_suffix(".md")
+        write_doc_file = f"reference/{doc_file}"
+
+        nav[tuple(child_module_import_path.split("."))] = f"{doc_file}"
+
+        with mkdocs_gen_files.open(write_doc_file, "w") as f:
+            print(
+                textwrap.dedent(f"""\
+                ::: {child_module_import_path}
+                """),
+                file=f,
+            )
+        docs_dir = REPO_ROOTDIR / "docs"
+        module_path_relative_to_docs_dir = child_module_path.relative_to(docs_dir, walk_up=True)
+        mkdocs_gen_files.set_edit_path(write_doc_file, str(module_path_relative_to_docs_dir))
+
+    submodules = list(
+        p
+        for p in module_path.iterdir()
+        if p.is_dir()
+        and (p / "__init__.py").exists()
+        and not p.name.endswith("_test")
+        and not p.name.startswith((".", "__"))
     )
-nav["project", "utils", "types"] = "reference/project/utils/types.md"
-mkdocs_gen_files.set_edit_path("reference/project/utils/types.md", "../project/utils/types.py")
-
-
-with mkdocs_gen_files.open("reference.md", "w") as nav_file:
-    # assert False, "\n".join(nav.build_literate_nav())
-    nav_file.writelines(nav.build_literate_nav())
-
-# with mkdocs_gen_files.open("reference/SUMMARY.md", "w") as project_nav_file:
-#     project_nav_file.writelines(project_nav.build_literate_nav())
-
-
-# project_root = REPO_ROOTDIR / "project"
-# for python_module_path in sorted(
-#     f
-#     # for f in project_root.glob("*.py")
-#     for f in [(project_root / "project")]
-#     if not f.name.endswith("_test.py") and not f.name == "__init__.py"
-# ):
-#     doc_path = python_module_path.relative_to(REPO_ROOTDIR).with_suffix(".md")
-
-#     full_doc_path = Path("reference") / doc_path
-
-#     nav[full_doc_path.with_suffix("").parts] = str(full_doc_path)
-
-#     with mkdocs_gen_files.open(full_doc_path, "w") as f:
-#         module_import_path = ".".join(
-#             python_module_path.relative_to(REPO_ROOTDIR).with_suffix("").parts
-#         )
-#         print(f"::: {module_import_path}", file=f)
-
-#     mkdocs_gen_files.set_edit_path(
-#         full_doc_path, python_module_path.relative_to(REPO_ROOTDIR / "docs", walk_up=True)
-#     )
+    for submodule in submodules:
+        add_doc_for_module(submodule, nav)
 
-# nav["mkdocs_autorefs", "references"] = "autorefs/references.md"
-# nav["mkdocs_autorefs", "plugin"] = "autorefs/plugin.md"
 
-# with mkdocs_gen_files.open("reference.md", "w") as nav_file:
-#     nav_file.writelines(nav.build_literate_nav())
+if __name__ in ["__main__", "<run_path>"]:
+    # Run when executed directly or by mkdocs. Seems like the __name__ is <run_path> during `mkdocs serve`
+    main()

docs/index.md

@@ -17,7 +17,7 @@ This project makes use of the following libraries:
 
 - [Hydra](https://hydra.cc/) is used to configure the project. It allows you to define configuration files and override them from the command line.
 - [PyTorch Lightning](https://lightning.ai/docs/pytorch/stable/) is used to as the training framework. It provides a high-level interface to organize ML research code.
-    - 🔥 Please note: You can also use [Jax](https://jax.readthedocs.io/en/latest/) with this repo, as is shown in the [Jax example](examples/examples.md#using-jax) 🔥
+    - 🔥 Please note: You can also use [Jax](https://jax.readthedocs.io/en/latest/) with this repo, as is shown in the [Jax example](examples/jax.md) 🔥
 - [Weights & Biases](https://wandb.ai) is used to log metrics and visualize results.
 - [pytest](https://docs.pytest.org/en/stable/) is used for testing.
 

project/utils/utils.py

@@ -35,7 +35,6 @@ def log_once(message: str, level: int) -> None:
     """Logs a message once per logger instance. The message is logged at the specified level.
 
     Args:
-        logger: The logger instance to use.
         message: The message to log.
         level: The logging level to use.
     """