Merge pull request #33 from alan-turing-institute/add-base

Update base classes

Author: sgreenbury

docs/DEVELOPER_GUIDE.md

@@ -0,0 +1,32 @@
+# Developer Guide
+TODO
+
+## Overview
+TODO
+
+## API notes
+
+
+### Trainer and Model Integration
+Example usage with `lightning` Trainer:
+
+```python
+model = EncoderDecoder()  # Anything that inherits for L.LightningModule
+trainer = L.Trainer()
+trainer.fit(model, train_dataloader)  # train_dataloader should output a batch of data from an iterable.
+
+model = EncoderProcessorDecoder()
+trainer = L.Trainer()
+trainer.fit(model, train_dataloader)
+```
+
+### Model API
+Subclasses of `LightningModule` from `lightning` aim to have API:
+```python
+def training_step(self, batch: Batch, batch_idx: int) -> Tensor: ...
+def forward(self, x: Tensor) -> Tensor: ...
+```
+
+Direct subclasses of `nn.Module` from `torch` aim to have API:
+```python
+def forward(self, x: Tensor) -> Tensor: ...

pyproject.toml

@@ -6,8 +6,14 @@ readme = "README.md"
 authors = [
     { name = "AI for Physical Systems Team at The Alan Turing Institute", email = "ai4physics@turing.ac.uk" },
 ]
-requires-python = ">=3.10,<3.13"
-dependencies = []
+requires-python = ">=3.11,<3.13"
+dependencies = [
+    "einops>=0.8.1",
+    "h5py>=3.15.1",
+    "lightning>=2.5.6",
+    "the-well>=1.1.0",
+    "torch>=2.9.1",
+]
 
 [project.optional-dependencies]
 dev = [
@@ -69,6 +75,7 @@ ignore = [
     "ISC001",  # Conflicts with formatter
     # "D417",    # Missing trailing new line in docstring
     "D100",    # Missing docstring in public module
+    "D102",    # Missing docstring in public class, TODO: remove in future
     "D104",    # Missing docstring in public package
     "PLR0913", # too many arguments
 ]

src/auto_cast/decoders/__init__.py

@@ -0,0 +1,3 @@
+from .base import Decoder
+
+__all__ = ["Decoder"]

src/auto_cast/decoders/base.py

@@ -0,0 +1,9 @@
+from typing import Any
+
+from torch import Tensor, nn
+
+
+class Decoder(nn.Module):
+    """Base Decoder."""
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any: ...

src/auto_cast/encoders/__init__.py

@@ -0,0 +1,3 @@
+from .base import Encoder
+
+__all__ = ["Encoder"]

src/auto_cast/encoders/base.py

@@ -0,0 +1,12 @@
+from typing import Any
+
+from torch import nn
+
+
+class Encoder(nn.Module):
+    """Base encoder."""
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
+        """Forward Pass through the Encoder."""
+        msg = "To implement."
+        raise NotImplementedError(msg)

src/auto_cast/models/encoder_decoder.py

@@ -0,0 +1,59 @@
+from typing import Any
+
+import lightning as L
+import torch
+from torch import nn
+
+from auto_cast.decoders import Decoder
+from auto_cast.encoders import Encoder
+from auto_cast.processors.base import Preprocessor
+from auto_cast.types import Batch, Tensor
+
+
+class EncoderDecoder(L.LightningModule):
+    """Encoder-Decoder Model."""
+
+    encoder: Encoder
+    decoder: Decoder
+    preprocessor: Preprocessor
+    loss_func: nn.Module
+
+    def __init__(self):
+        pass
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
+        return self.decoder(self.encoder(*args, **kwargs))
+
+    def training_step(self, batch: Batch, batch_idx: int) -> Tensor:  # noqa: ARG002
+        x = self.preprocessor(batch)
+        output = self(x)
+        loss = self.loss_func(output, batch["output_fields"])
+        return loss  # noqa: RET504
+
+    def validation_step(self, batch: Batch, batch_idx: int) -> Tensor: ...
+
+    def test_step(self, batch: Batch, batch_idx: int) -> Tensor: ...
+
+    def predict_step(self, batch: Batch, batch_idx: int) -> Tensor: ...
+
+    def encode(self, x: Batch) -> Tensor:
+        x = self.preprocessor(x)
+        return self.encoder(x)
+
+    def configure_optmizers(self):
+        pass
+
+
+class VAE(EncoderDecoder):
+    """Variational Autoencoder Model."""
+
+    def forward(self, x: Tensor) -> Tensor:
+        mu, log_var = self.encoder(x)
+        z = self.reparametrize(mu, log_var)
+        x = self.decoder(z)
+        return x  # noqa: RET504
+
+    def reparametrize(self, mu: Tensor, log_var: Tensor) -> Tensor:
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return mu + eps * std

src/auto_cast/models/encoder_processor_decoder.py

@@ -0,0 +1,70 @@
+from typing import Any
+
+import lightning as L
+import torch
+from torch import nn
+
+from auto_cast.models.encoder_decoder import EncoderDecoder
+from auto_cast.processors.base import Processor
+from auto_cast.types import Batch, RolloutOutput, Tensor
+
+
+class EncoderProcessorDecoder(L.LightningModule):
+    """Encoder-Processor-Decoder Model."""
+
+    encoder_decoder: EncoderDecoder
+    processor: Processor
+    teacher_forcing_ratio: float
+    stride: int
+    max_rollout_steps: int
+    loss_func: nn.Module
+
+    def __init__(self): ...
+
+    def from_encoder_processor_decoder(
+        self, encoder_decoder: EncoderDecoder, processor: Processor
+    ) -> None:
+        self.encoder_decoder = encoder_decoder
+        self.processor = processor
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
+        return self.encoder_decoder.decoder(
+            self.processor(self.encoder_decoder.encoder(*args, **kwargs))
+        )
+
+    def training_step(self, batch: Batch, batch_idx: int) -> Tensor:  # noqa: ARG002
+        output = self(batch)
+        loss = self.processor.loss_func(output, batch.output_fields)
+        return loss  # noqa: RET504
+
+    def configure_optimizers(self): ...
+
+    def rollout(self, batch: Batch) -> RolloutOutput:
+        """Rollout over multiple time steps."""
+        pred_outs, gt_outs = [], []
+        for _ in range(0, self.max_rollout_steps, self.stride):
+            x = self.encoder_decoder.encoder(batch)
+            pred_outs.append(self.processor.map(x))
+            # TODO: combining teacher forcing logic
+            gt_outs.append(batch.output_fields)  # This assumes we have output fields
+        return torch.stack(pred_outs), torch.stack(gt_outs)
+
+
+# TODO: consider if separate rollout class would be better
+class Rollout:
+    max_rollout_steps: int
+    stride: int
+
+    def rollout(
+        self,
+        batch: Batch,
+        model: Processor | EncoderProcessorDecoder,
+    ) -> RolloutOutput:
+        """Rollout over multiple time steps."""
+        pred_outs, gt_outs = [], []
+        for _ in range(0, self.max_rollout_steps, self.stride):
+            output = model(batch)
+            pred_outs.append(output)
+            # TODO: logic for moving window with teacher forcing that assigns
+            gt_outs.append(batch.output_fields)  # This assumes we have output fields
+        return torch.stack(pred_outs), torch.stack(gt_outs)

src/auto_cast/processors/base.py

@@ -0,0 +1,67 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+import lightning as L
+import torch
+from torch import nn
+
+from auto_cast.types import Batch, EncodedBatch, RolloutOutput, Tensor
+
+
+class Processor(L.LightningModule):
+    """Processor Base Class."""
+
+    teacher_forcing_ratio: float
+    stride: int
+    max_rollout_steps: int
+    loss_func: nn.Module
+
+    def forward(self, *args, **kwargs: Any) -> Any:
+        """Forward pass through the Processor."""
+        msg = "To implement."
+        raise NotImplementedError(msg)
+
+    def training_step(self, batch: EncodedBatch, batch_idx: int) -> Tensor:  # noqa: ARG002
+        output = self.map(batch.encoded_inputs)
+        loss = self.loss_func(output, batch.encoded_output_fields)
+        return loss  # noqa: RET504
+
+    @abstractmethod
+    def map(self, x: Tensor) -> Tensor:
+        """Map input window of states/times to output window."""
+
+    def configure_optimizers(self): ...
+
+    def rollout(self, batch: EncodedBatch) -> RolloutOutput:
+        """Rollout over multiple time steps."""
+        pred_outs, gt_outs = [], []
+        for _ in range(0, self.max_rollout_steps, self.stride):
+            pred_outs.append(self.map(batch.encoded_inputs))
+            # TODO: combining teacher forcing logic
+            gt_outs.append(
+                batch.encoded_output_fields
+            )  # This assumes we have output fields
+        return torch.stack(pred_outs), torch.stack(gt_outs)
+
+
+class DiscreteProcessor(Processor, ABC):
+    """DiscreteProcessor."""
+
+    @abstractmethod
+    def map(self, x: Tensor) -> Tensor:
+        ...
+        # Map input window of states/times to output window
+
+    def rollout(self, batch: EncodedBatch) -> RolloutOutput:
+        ...
+
+        # Use self.map to generate trajectory
+
+
+class FlowBasedGenerativeProcessor(DiscreteProcessor):
+    """Flow-based generative processor."""
+
+    def map(self, x: Tensor) -> Tensor:
+        ...
+        # Sample generative model    def loss(self, ...):...
+        # Flow matc

src/auto_cast/types/__init__.py

@@ -0,0 +1,38 @@
+from dataclasses import dataclass
+
+import torch
+from torch.utils.data import DataLoader
+
+Tensor = torch.Tensor
+Input = Tensor | DataLoader
+RolloutOutput = tuple[Tensor, None] | tuple[Tensor, Tensor]
+
+# Batch = dict[str, Tensor]
+# EncodedBatch = dict[str, Tensor]
+
+
+# TODO: Could be a dataclass if we want more structure
+@dataclass
+class Batch:
+    input_fields: Tensor
+    output_fields: Tensor
+    constant_scalars: Tensor
+    constant_fields: Tensor
+
+
+@dataclass
+class EncodedBatch:
+    encoded_inputs: Tensor
+    encoded_output_fields: Tensor
+    encoded_info: dict[str, Tensor]
+
+
+class EncoderForBatch:
+    """EncoderForBatch."""
+
+    def __call__(self, batch: Batch) -> EncodedBatch:
+        return EncodedBatch(
+            encoded_inputs=batch.input_fields,
+            encoded_output_fields=batch.output_fields,
+            encoded_info={},
+        )