Merge pull request #18 from PolymathicAI/gaia

Adding codecs for Gaia

Author: EiffL

.gitignore

@@ -172,3 +172,5 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+notebooks
+data

aion/codecs/__init__.py

@@ -1,6 +1,12 @@
 from .image import ImageCodec
-from .scalar import ScalarCodec, LogScalarCodec
+from .scalar import ScalarCodec, LogScalarCodec, MultiScalarCodec
 from .spectrum import SpectrumCodec
 
 
-__all__ = ["ImageCodec", "ScalarCodec", "LogScalarCodec", "SpectrumCodec"]
+__all__ = [
+    "ImageCodec",
+    "ScalarCodec",
+    "LogScalarCodec",
+    "MultiScalarCodec",
+    "SpectrumCodec",
+]

aion/codecs/quantizers/scalar.py

@@ -3,6 +3,7 @@
 
 import scipy.interpolate
 import torch
+import torch.nn as nn
 
 from aion.codecs.quantizers import Quantizer
 
@@ -11,7 +12,7 @@ class ScalarReservoirQuantizer(Quantizer):
     """
     Scalar quantizer module.
 
-    The sclar quantizer module takes a batch of scalars and quantizes them using a CDF codebook.
+    The scalar quantizer module takes a batch of scalars and quantizes them using a CDF codebook.
     The CDF estimate is updated using reservoir sampling, allowing you to stream through data.
 
     Args:
@@ -182,7 +183,7 @@ class ScalarLogReservoirQuantizer(ScalarReservoirQuantizer):
     """
     Scalar quantizer module.
 
-    The sclar quantizer module takes a batch of scalars and quantizes them using a CDF codebook.
+    The scalar quantizer module takes a batch of scalars and quantizes them using a CDF codebook.
     The CDF estimate is updated using reservoir sampling, allowing you to stream through data.
 
     Args:
@@ -271,3 +272,235 @@ def decode(self, codes: torch.Tensor) -> torch.Tensor:
                 Decoded sample.
         """
         return torch.exp(super().decode(codes))
+
+
+class ScalarCompressedReservoirQuantizer(ScalarReservoirQuantizer):
+    """
+    Scalar quantizer module with compression/decompression functions.
+
+    The scalar quantizer module takes a batch of scalars, applies compression functions,
+    and quantizes them using a CDF codebook. The CDF estimate is updated using reservoir
+    sampling, allowing you to stream through data.
+
+    Args:
+        compression_fns: list[str]
+            List of torch function names to apply for compression (e.g., ['arcsinh']).
+        decompression_fns: list[str]
+            List of torch function names to apply for decompression (e.g., ['sinh']).
+        codebook_size: int
+            The number of codes in the codebook.
+        reservoir_size: int
+            The size of the reservoir to keep in memory.
+        reservoir_default: float
+            Optional default value of reservoir samples. Only relevant if there
+            are fewer samples in your dataset than the size of your codebook.
+    """
+
+    def __init__(
+        self,
+        compression_fns: list[str],
+        decompression_fns: list[str],
+        codebook_size: int,
+        reservoir_size: int,
+        reservoir_default: Optional[float] = 0.0,
+    ):
+        super().__init__(codebook_size, reservoir_size, reservoir_default)
+        assert len(compression_fns) == len(decompression_fns), (
+            "Mismatched compression/decompression functions"
+        )
+        self.compression_fns = compression_fns
+        self.decompression_fns = decompression_fns
+
+        assert self._check_identity(torch.tensor([1.0])), (
+            "Identity check failed, compression/decompression functions are not inverses."
+        )
+
+    def compress(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply compression functions to input tensor.
+
+        Args:
+            x: torch.Tensor
+                Input tensor to compress.
+
+        Returns:
+            torch.Tensor
+                Compressed tensor.
+        """
+        for c in self.compression_fns:
+            x = getattr(torch, c)(x)
+        return x
+
+    def decompress(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply decompression functions to input tensor.
+
+        Args:
+            x: torch.Tensor
+                Input tensor to decompress.
+
+        Returns:
+            torch.Tensor
+                Decompressed tensor.
+        """
+        for c in self.decompression_fns[::-1]:
+            x = getattr(torch, c)(x)
+        return x
+
+    def _check_identity(self, x: torch.Tensor) -> bool:
+        """Check if compression and decompression are inverses.
+
+        Args:
+            x: torch.Tensor
+                Test tensor.
+
+        Returns:
+            bool
+                True if compress(decompress(x)) ≈ x.
+        """
+        return torch.allclose(self.decompress(self.compress(x)), x)
+
+    def _update_reservoirs(self, z_e: torch.Tensor):
+        z_e = self.compress(z_e)
+        super()._update_reservoirs(z_e)
+
+    def encode(self, z: torch.Tensor) -> torch.Tensor:
+        z = self.compress(z)
+        return super().encode(z)
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        return self.decompress(super().decode(codes))
+
+
+class MultiScalarCompressedReservoirQuantizer(Quantizer):
+    """
+    Multi-channel scalar quantizer with compression.
+
+    Wraps multiple ScalarCompressedReservoirQuantizers to quantize multi-channel tensors.
+    Each channel is quantized independently with its own reservoir.
+
+    Args:
+        compression_fns: list[str]
+            List of torch function names to apply for compression (e.g., ['arcsinh']).
+        decompression_fns: list[str]
+            List of torch function names to apply for decompression (e.g., ['sinh']).
+        codebook_size: int
+            The number of codes in the codebook.
+        reservoir_size: int
+            The size of the reservoir to keep in memory.
+        reservoir_default: float
+            Optional default value of reservoir samples.
+        num_quantizers: int
+            Number of channels/quantizers to create.
+    """
+
+    def __init__(
+        self,
+        compression_fns: list[str],
+        decompression_fns: list[str],
+        codebook_size: int,
+        reservoir_size: int,
+        reservoir_default: Optional[float] = 0.0,
+        num_quantizers: int = 1,
+    ):
+        super().__init__()
+        self.quantizers = nn.ModuleList(
+            [
+                ScalarCompressedReservoirQuantizer(
+                    compression_fns,
+                    decompression_fns,
+                    codebook_size,
+                    reservoir_size,
+                    reservoir_default,
+                )
+                for _ in range(num_quantizers)
+            ]
+        )
+        self.num_quantizers = num_quantizers
+
+    def encode(self, z: torch.Tensor) -> torch.Tensor:
+        """Encodes the input tensor z, returns the corresponding
+        codebook index.
+
+        Args:
+            z: torch.Tensor (B, C)
+                The input tensor to be encoded.
+
+        Returns:
+            codes: torch.Tensor (B, C)
+                Encoded tensor.
+        """
+        return torch.stack(
+            [q.encode(z[:, i]) for i, q in enumerate(self.quantizers)],
+            dim=1,
+        )
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decodes the input code index into corresponding codebook entry of
+        dimension (embedding_dim).
+
+        Args:
+            codes: torch.Tensor (B, C)
+                Codes to be decoded.
+
+        Returns:
+            z: torch.Tensor (B, C)
+                Decoded sample.
+        """
+        return torch.stack(
+            [q.decode(codes[:, i]) for i, q in enumerate(self.quantizers)],
+            dim=1,
+        )
+
+    def quantize(self, z: torch.Tensor) -> torch.Tensor:
+        """Quantize the input tensor z, returns corresponding
+        codebook entry.
+
+        Args:
+            z: torch.Tensor (B, C)
+                The input tensor to be quantized.
+
+        Returns:
+            z: torch.Tensor (B, C)
+                Quantized tensor.
+        """
+        return self.decode(self.encode(z))
+
+    def _update_reservoirs(self, z_e: torch.Tensor):
+        for i, q in enumerate(self.quantizers):
+            q._update_reservoirs(z_e[:, i])
+
+    def forward(
+        self, z_e: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Performs a forward pass through the vector quantizer.
+        Args:
+            z_e: torch.Tensor (B, C, ...)
+                The input tensor to be quantized.
+        Returns:
+            z_q: torch.Tensor
+                The quantized tensor.
+            loss: torch.Tensor
+                The embedding loss for the quantization.
+            codebook_usage: torch.Tensor
+                The fraction of codes used in the codebook.
+        """
+        self._update_reservoirs(z_e)
+        indices = self.encode(z_e)
+        z_q = self.decode(indices)
+        num_unique = sum([len(torch.unique(c)) for c in indices.T])
+        codebook_usage = num_unique / (self.codebook_size * self.num_quantizers)
+        return z_q, torch.nn.functional.mse_loss(z_q, z_e), torch.tensor(codebook_usage)
+
+    @property
+    def codebook_size(self) -> int:
+        """Returns the size of the codebook."""
+        return self.quantizers[0].codebook_size
+
+    @property
+    def codebook(self) -> torch.Tensor:
+        """Returns the codebook."""
+        return self.quantizers[0].codebook
+
+    @property
+    def embedding_dim(self) -> int:
+        """Returns the dimension of the codebook entries."""
+        return 1

aion/codecs/scalar.py

@@ -8,6 +8,7 @@
 from aion.codecs.quantizers.scalar import (
     ScalarLogReservoirQuantizer,
     ScalarReservoirQuantizer,
+    MultiScalarCompressedReservoirQuantizer,
 )
 from aion.codecs.base import Codec
 from aion.modalities import ScalarModality, ScalarModalities
@@ -85,3 +86,70 @@ def __init__(
             reservoir_size=reservoir_size,
             min_log_value=min_log_value,
         )
+
+
+class MultiScalarCodec(BaseScalarIdentityCodec):
+    """Codec for multi-channel scalar quantities with compression.
+
+    A codec that handles multi-channel scalar modalities using compression
+    and decompression functions before quantization. This is particularly useful
+    for spectral coefficients or other multi-dimensional scalar data that
+    benefits from preprocessing transformations.
+
+    Each channel is quantized independently using a compressed reservoir quantizer,
+    allowing for different statistical distributions across channels while
+    maintaining the ability to handle streaming data.
+
+    Args:
+        modality: str
+            The name of the modality this codec is designed for. Must match
+            a modality name defined in the ScalarModalities registry.
+        compression_fns: list[str]
+            List of PyTorch function names to apply for compression (e.g., ['arcsinh']).
+            These functions are applied in order to transform the data before quantization.
+        decompression_fns: list[str]
+            List of PyTorch function names to apply for decompression (e.g., ['sinh']).
+            These functions are applied in reverse order during decoding to restore
+            the original data range.
+        codebook_size: int
+            The number of codes in each quantizer's codebook.
+        reservoir_size: int
+            The size of the reservoir to keep in memory for each channel's quantizer.
+        num_quantizers: int
+            Number of channels/quantizers to create, corresponding to the number
+            of dimensions in the multi-channel scalar data.
+
+    Note:
+        The compression and decompression functions must be mathematical inverses
+        of each other. The codec will verify this during initialization and raise
+        an assertion error if the functions are not properly inverse.
+
+    Example:
+        >>> codec = MultiScalarCodec(
+        ...     modality="bp_coefficients",
+        ...     compression_fns=["arcsinh"],
+        ...     decompression_fns=["sinh"],
+        ...     codebook_size=1024,
+        ...     reservoir_size=10000,
+        ...     num_quantizers=55
+        ... )
+    """
+
+    def __init__(
+        self,
+        modality: str,
+        compression_fns: list[str],
+        decompression_fns: list[str],
+        codebook_size: int,
+        reservoir_size: int,
+        num_quantizers: int,
+    ):
+        super().__init__()
+        self._modality_class = next(m for m in ScalarModalities if m.name == modality)
+        self._quantizer = MultiScalarCompressedReservoirQuantizer(
+            compression_fns=compression_fns,
+            decompression_fns=decompression_fns,
+            codebook_size=codebook_size,
+            reservoir_size=reservoir_size,
+            num_quantizers=num_quantizers,
+        )