Update initial API

- Remove preprocessor (encoder to be used)
- Change batch types to dataclasses
- Add EncodedBatch
- Initial rollout method

Co-authored-by: Christopher Iliffe Sprague <cisprague@users.noreply.github.com>
Co-authored-by: Marjan Famili <marjanfamili@users.noreply.github.com>

Author: 3 people

src/auto_cast/decoders/base.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from torch import nn
+from torch import Tensor, nn
 
 
 class Decoder(nn.Module):

src/auto_cast/models/encoder_processor_decoder.py

@@ -1,31 +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.preprocessor.base import Preprocessor
 from auto_cast.processors.base import Processor
-from auto_cast.types import Batch, Tensor
+from auto_cast.types import Batch, RolloutOutput, Tensor
 
 
 class EncoderProcessorDecoder(L.LightningModule):
     """Encoder-Processor-Decoder Model."""
 
     encoder_decoder: EncoderDecoder
     processor: Processor
-    preprocessor: Preprocessor
+    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
-        x = self.preprocessor(batch)
-        output = self(x)
-        loss = self.processor.loss_func(output, batch["output_fields"])
+        output = self(batch)
+        loss = self.processor.loss_func(output, batch.output_fields)
         return loss  # noqa: RET504
 
-    def configure_optmizers(self): ...
+    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/preprocessor/__init__.py

@@ -5,8 +5,7 @@
 import torch
 from torch import nn
 
-from auto_cast.preprocessor import Preprocessor
-from auto_cast.types import Batch, RolloutOutput, Tensor
+from auto_cast.types import Batch, EncodedBatch, RolloutOutput, Tensor
 
 
 class Processor(L.LightningModule):
@@ -15,43 +14,45 @@ class Processor(L.LightningModule):
     teacher_forcing_ratio: float
     stride: int
     max_rollout_steps: int
-    preprocessor: Preprocessor
     loss_func: nn.Module
 
-    def forward(self, *args: Any, **kwargs: Any) -> Any:
+    def forward(self, *args, **kwargs: Any) -> Any:
         """Forward pass through the Processor."""
         msg = "To implement."
         raise NotImplementedError(msg)
 
-    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"])
+    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
 
-    def configure_optmizers(self):
-        pass
+    @abstractmethod
+    def map(self, x: Tensor) -> Tensor:
+        """Map input window of states/times to output window."""
+
+    def configure_optimizers(self): ...
 
-    def rollout(self, batch: Batch) -> RolloutOutput:
+    def rollout(self, batch: EncodedBatch) -> RolloutOutput:
         """Rollout over multiple time steps."""
-        pred_outs = []
-        gt_outs = []
-        for _time_step in range(0, self.max_rollout_steps, self.stride):
-            x = self.preprocessor(batch)
-            pred_outs.append(self(x))
-            gt_outs.append(batch["output_fields"])  # This assumes we have output fields
+        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: Batch) -> Tensor:
+    def map(self, x: Tensor) -> Tensor:
         ...
         # Map input window of states/times to output window
 
-    def rollout(self, batch: Batch) -> RolloutOutput:
+    def rollout(self, batch: EncodedBatch) -> RolloutOutput:
         ...
 
         # Use self.map to generate trajectory
@@ -60,7 +61,7 @@ def rollout(self, batch: Batch) -> RolloutOutput:
 class FlowBasedGenerativeProcessor(DiscreteProcessor):
     """Flow-based generative processor."""
 
-    def map(self, x: Batch) -> Tensor:
+    def map(self, x: Tensor) -> Tensor:
         ...
         # Sample generative model    def loss(self, ...):...
         # Flow matc

src/auto_cast/preprocessor/base.py

@@ -1,16 +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]
+# 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 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={},
+        )