Expand docs with architecture overview

.github/workflows/deploy-doc.yml

@@ -0,0 +1,57 @@
+name: Build Documentation
+
+on:
+  push:
+    branches:
+      - main
+  # Allows you to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+# Allow one concurrent deployment
+concurrency:
+  group: "pages"
+  cancel-in-progress: true
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: true
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.11'
+    - name: Install dependencies
+      run: |
+        uv sync --all-extras --dev
+    - name: Build HTML documentation
+      run: |
+        cd docs
+        uv run sphinx-build -b html . _build/html
+
+    - name: Upload artifact
+      uses: actions/upload-pages-artifact@v3
+      with:
+        path: docs/_build/html
+
+  # Deploy job - only runs on main branch
+  deploy:
+    needs: build
+    # Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages
+    permissions:
+      contents: read
+      pages: write
+      id-token: write
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Deploy to GitHub Pages
+      id: deployment
+      uses: actions/deploy-pages@v4

.github/workflows/test-doc.yml

@@ -0,0 +1,34 @@
+name: Check Documentation Build
+
+on:
+  pull_request:
+    paths:
+      - 'docs/**'
+      - 'aion/**'
+      - '.github/workflows/docs-check.yml'
+      - 'pyproject.toml'
+
+jobs:
+  docs:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: true
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.11'
+    - name: Install dependencies
+      run: |
+        uv sync --all-extras --dev
+    - name: Build documentation
+      run: |
+        cd docs
+        uv run sphinx-build -W -b html . _build/html
+    - name: Check for broken links
+      run: |
+        cd docs
+        uv run sphinx-build -b linkcheck . _build/linkcheck || true

.github/workflows/test.yaml

@@ -4,7 +4,7 @@ on:
   pull_request:
   push:
     branches:
-      - master
+      - main
 
 jobs:
   pre-commit:

.readthedocs.yaml

@@ -0,0 +1,8 @@
+version: 2
+sphinx:
+  configuration: docs/conf.py
+  fail_on_warning: false
+python:
+  version: 3.10
+  install:
+    - requirements: docs/requirements.txt

README.md

@@ -86,6 +86,8 @@ AION-1’s tokenizers cover **39 distinct data types**, grouped by survey and da
 Start with our interactive tutorial:
 - **[Open in Google Colab](https://colab.research.google.com/github/PolymathicAI/AION/blob/main/notebooks/Tutorial.ipynb)** - Learn AION basics interactively, no local setup required!
 
+For detailed guides, see the [online documentation](https://polymathic-ai.github.io/AION/).
+
 ## 📦 Advanced Installation
 
 AION offers flexible installation options to suit your environment and requirements.

aion/codecs/config.py

@@ -1,54 +1,54 @@
 # Configuration for codecs
 
+from aion.codecs.catalog import CatalogCodec
+from aion.codecs.image import ImageCodec
+from aion.codecs.scalar import (
+    GridScalarCodec,
+    LogScalarCodec,
+    MultiScalarCodec,
+    ScalarCodec,
+)
+from aion.codecs.scalar_field import ScalarFieldCodec
+from aion.codecs.spectrum import SpectrumCodec
 from aion.modalities import (
+    HSCAG,
+    HSCAI,
+    HSCAR,
+    HSCAY,
+    HSCAZ,
+    Dec,
+    GaiaFluxBp,
+    GaiaFluxG,
+    GaiaFluxRp,
+    GaiaParallax,
+    GaiaXpBp,
+    GaiaXpRp,
+    HSCMagG,
+    HSCMagI,
+    HSCMagR,
+    HSCMagY,
+    HSCMagZ,
+    HSCShape11,
+    HSCShape12,
+    HSCShape22,
     Image,
-    Spectrum,
     LegacySurveyCatalog,
-    LegacySurveySegmentationMap,
+    LegacySurveyEBV,
     LegacySurveyFluxG,
-    LegacySurveyFluxR,
     LegacySurveyFluxI,
-    LegacySurveyFluxZ,
+    LegacySurveyFluxR,
     LegacySurveyFluxW1,
     LegacySurveyFluxW2,
     LegacySurveyFluxW3,
     LegacySurveyFluxW4,
-    LegacySurveyShapeR,
+    LegacySurveyFluxZ,
+    LegacySurveySegmentationMap,
     LegacySurveyShapeE1,
     LegacySurveyShapeE2,
-    LegacySurveyEBV,
-    Z,
-    HSCAG,
-    HSCAR,
-    HSCAI,
-    HSCAZ,
-    HSCAY,
-    HSCMagG,
-    HSCMagR,
-    HSCMagI,
-    HSCMagZ,
-    HSCMagY,
-    HSCShape11,
-    HSCShape22,
-    HSCShape12,
-    GaiaFluxG,
-    GaiaFluxBp,
-    GaiaFluxRp,
-    GaiaParallax,
+    LegacySurveyShapeR,
     Ra,
-    Dec,
-    GaiaXpBp,
-    GaiaXpRp,
-)
-from aion.codecs.image import ImageCodec
-from aion.codecs.spectrum import SpectrumCodec
-from aion.codecs.catalog import CatalogCodec
-from aion.codecs.scalar_field import ScalarFieldCodec
-from aion.codecs.scalar import (
-    ScalarCodec,
-    LogScalarCodec,
-    MultiScalarCodec,
-    GridScalarCodec,
+    Spectrum,
+    Z,
 )
 
 CODEC_CONFIG = {

aion/codecs/manager.py

@@ -3,12 +3,21 @@
 Handles dynamic loading and management of codecs for different modalities.
 """
 
-from typing import Dict, Union, Optional, Type
+from typing import Dict, Optional, Type, Union
+
 import torch
 
-from aion.modalities import Modality
 from aion.codecs.base import Codec
 from aion.codecs.config import CODEC_CONFIG
+from aion.modalities import Modality
+
+
+class ModalityTypeError(TypeError):
+    """Error raised when a modality type is not supported."""
+
+
+class TokenKeyError(ValueError):
+    """Error raised when a token key is not found in the tokens dictionary."""
 
 
 class CodecManager:
@@ -53,7 +62,7 @@ def _load_codec(self, modality_type: Type[Modality]) -> Codec:
             ):
                 config = CODEC_CONFIG[modality_type.__base__]
             else:
-                raise ValueError(
+                raise ModalityTypeError(
                     f"No codec configuration found for modality type: {modality_type.__name__}"
                 )
         else:
@@ -76,6 +85,7 @@ def _load_codec(self, modality_type: Type[Modality]) -> Codec:
 
         return codec
 
+    @torch.no_grad()
     def encode(self, *modalities: Modality) -> Dict[str, torch.Tensor]:
         """Encode multiple modalities.
 
@@ -98,14 +108,15 @@ def encode(self, *modalities: Modality) -> Dict[str, torch.Tensor]:
             if hasattr(modality, "token_key"):
                 token_key = modality.token_key
             else:
-                raise ValueError(
+                raise ModalityTypeError(
                     f"Modality {type(modality).__name__} does not have a token_key attribute"
                 )
 
             tokens[token_key] = tokenized
 
         return tokens
 
+    @torch.no_grad()
     def decode(
         self,
         tokens: Dict[str, torch.Tensor],
@@ -124,14 +135,14 @@ def decode(
             Decoded modality instance
         """
         if not hasattr(modality_type, "token_key"):
-            raise ValueError(
-                f"Modality type {modality_type.__name__} does not have a token_key attribute"
+            raise ModalityTypeError(
+                f"Modality type {modality_type} does not have a token_key attribute"
             )
 
         token_key = modality_type.token_key
         if token_key not in tokens:
-            raise ValueError(
-                f"Token key '{token_key}' for modality {modality_type.__name__} not found in tokens dictionary"
+            raise TokenKeyError(
+                f"Token key '{token_key}' for modality {modality_type} not found in tokens dictionary"
             )
 
         # Get the appropriate codec
@@ -140,7 +151,7 @@ def decode(
         # Decode using the codec with any provided metadata
         decoded_modality = codec.decode(tokens[token_key], **metadata)
 
-        # Casting the decoded modality to be the specific modality type requested
+        # Cast decoded modality to the correct type
         decoded_modality = modality_type(**decoded_modality.model_dump())
 
         return decoded_modality

docs/Makefile

@@ -0,0 +1,7 @@
+SPHINXBUILD := sphinx-build
+SOURCEDIR := .
+BUILDDIR := _build
+
+.PHONY: html
+html:
+	$(SPHINXBUILD) -M html $(SOURCEDIR) $(BUILDDIR)