Adding code for catalog

Author: EiffL

aion/codecs/__init__.py

@@ -1,12 +1,15 @@
 from .image import ImageCodec
 from .scalar import ScalarCodec, LogScalarCodec, MultiScalarCodec
 from .spectrum import SpectrumCodec
-
+from .catalog import CatalogCodec
+from .base import Codec
 
 __all__ = [
     "ImageCodec",
     "ScalarCodec",
     "LogScalarCodec",
     "MultiScalarCodec",
     "SpectrumCodec",
+    "CatalogCodec",
+    "Codec",
 ]

aion/codecs/catalog.py

@@ -8,47 +8,38 @@
 
 from aion.codecs.base import Codec
 from aion.codecs.quantizers import Quantizer
-from aion.codecs.quantizers.scalar import ComposedScalarQuantizer
+from aion.codecs.quantizers.scalar import ComposedScalarQuantizer, IdentityQuantizer, ScalarReservoirQuantizer
 from aion.modalities import Catalog
 
-__all__ = ["CatalogIdentityCodec"]
 
-
-class CatalogIdentityCodec(Codec, PyTorchModelHubMixin):
+class CatalogCodec(Codec, PyTorchModelHubMixin):
     """Codec for catalog quantities.
 
     A codec that embeds catalog quantities through an identity mapping. A
     quantizer is applied if specified.
-
-    Args:
-        catalog_keys: List[str]
-            List of catalog keys to encode.
-        quantizers: Optional[List[Quantizer]]
-            Optional list of quantizers for each catalog key.
-        mask_value: int
-            Value used to indicate masked/missing data.
     """
 
     def __init__(
         self,
-        catalog_keys: List[str],
-        quantizers: Optional[List[Quantizer]] = None,
-        mask_value: int = 9999,
+        mask_value: int =9999,
     ):
         super().__init__()
         self._modality = Catalog
-        self._catalog_keys = catalog_keys
+        catalog_keys = ['X', 'Y', 'SHAPE_E1', 'SHAPE_E2', 'SHAPE_R']
+        quantizers = [IdentityQuantizer(96),
+                      IdentityQuantizer(96),
+                      ScalarReservoirQuantizer(1024, 100000),
+                      ScalarReservoirQuantizer(1024, 100000),
+                      ScalarReservoirQuantizer(1024, 100000)]
         self.mask_value = mask_value
-        if quantizers:
-            assert len(catalog_keys) == len(quantizers), (
-                "Number of catalog keys and quantizers must match"
-            )
-            _quantizer = OrderedDict()
-            for key, quantizer in zip(catalog_keys, quantizers):
-                _quantizer[key] = quantizer
-            self._quantizer = ComposedScalarQuantizer(_quantizer)
-        else:
-            self._quantizer = None
+        self._catalog_keys = catalog_keys
+        assert len(catalog_keys) == len(quantizers), (
+            "Number of catalog keys and quantizers must match"
+        )
+        _quantizer = OrderedDict()
+        for key, quantizer in zip(catalog_keys, quantizers):
+            _quantizer[key] = quantizer
+        self._quantizer = ComposedScalarQuantizer(_quantizer)
 
     @property
     def modality(self) -> Type[Catalog]:
@@ -61,7 +52,7 @@ def quantizer(self) -> Optional[Quantizer]:
     def _encode(self, x: Catalog) -> Dict[str, Tensor]:
         encoded = OrderedDict()
         for key in self._catalog_keys:
-            catalog_value = x[self.modality][key]
+            catalog_value = getattr(x, key)
             mask = catalog_value != self.mask_value
             catalog_value = catalog_value[mask]
             encoded[key] = catalog_value
@@ -87,7 +78,7 @@ def encode(self, x: Catalog) -> Float[Tensor, "b c1 *code_shape"]:
         return encoded
 
     def _decode(self, z: Dict[str, Tensor]) -> Catalog:
-        return Catalog(data=z)
+        return Catalog(**z)
 
     def decode(self, z: Float[Tensor, "b c1 *code_shape"]) -> Catalog:
         B, LC = z.shape

aion/codecs/quantizers/scalar.py

@@ -623,3 +623,68 @@ def codebook_size(self) -> int:
     def embedding_dim(self) -> int:
         """Returns the dimension of the codebook entries."""
         return 1
+
+
+class IdentityQuantizer(Quantizer):
+    """
+    Identity quantizer module.
+
+    The identity quantizer module takes a batch of tensors and returns the same tensor.
+
+    Args:
+        codebook_size: int
+            The number of labels to be used as signature for the codebook.
+    """
+
+    def __init__(self, codebook_size: int):
+        super().__init__()
+        self.register_buffer("_codebook_size", torch.tensor(codebook_size))
+
+    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.
+        """
+        codebook_usage = z_e.unique().numel() / self._codebook_size.item()
+        return z_e, torch.tensor(0), torch.tensor(codebook_usage)
+
+    def quantize(self, z: torch.Tensor) -> torch.Tensor:
+        """Quantize the input tensor z, returns corresponding
+        codebook entry.
+        """
+        return z
+
+    def encode(self, z: torch.Tensor) -> torch.Tensor:
+        """Encodes the input tensor z, returns the corresponding codebook index."""
+        return z
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decodes the input code index into corresponding codebook entry of
+        dimension (embedding_dim).
+        """
+        return codes
+
+    @property
+    def codebook_size(self) -> int:
+        """Returns the size of the codebook."""
+        return int(self._codebook_size.item())
+
+    @property
+    def codebook(self) -> torch.Tensor:
+        """Returns the codebook."""
+        return torch.arange(self._codebook_size.item())
+
+    @property
+    def embedding_dim(self) -> int:
+        """Returns the dimension of the codebook entries."""
+        return 1

aion/modalities.py

@@ -2,7 +2,7 @@
 
 from typing import List, Union, ClassVar
 from pydantic import BaseModel, Field, ConfigDict
-from jaxtyping import Float, Bool, Dict
+from jaxtyping import Float, Bool, Int
 from torch import Tensor
 
 
@@ -54,6 +54,31 @@ def __repr__(self) -> str:
         return repr_str
 
 
+# Catalog modality
+class Catalog(Modality):
+    """Catalog modality data.
+
+    Represents a catalog of scalar values from the 
+    Legacy Survey.
+    """
+
+    X: Int[Tensor, "batch n"] = Field(
+        description="X position of the object in the image."
+    )
+    Y: Int[Tensor, "batch n"] = Field(
+        description="Y position of the object in the image."
+    )
+    SHAPE_E1: Float[Tensor, "batch n"] = Field(
+        description="First ellipticity component of the object."
+    )
+    SHAPE_E2: Float[Tensor, "batch n"] = Field(
+        description="Second ellipticity component of the object."
+    )
+    SHAPE_R: Float[Tensor, "batch n"] = Field(
+        description="Size of the object."
+    )
+
+
 class ScalarModality(Modality):
     """Base class for scalar modality data.
 
@@ -70,18 +95,6 @@ def __repr__(self) -> str:
         return f"{self.__class__.__name__}(shape={list(self.value.shape)})"
 
 
-# Catalog modality
-class Catalog(Modality):
-    """Catalog modality data.
-
-    Represents a catalog of scalar values.
-    """
-
-    data: Dict[str, Dict[str, Float[Tensor, "b t"]]] = Field(
-        description="Dictionary of dictionaries of scalar values."
-    )
-
-
 # Flux measurements in different bands
 class FluxG(ScalarModality):
     """G-band flux measurement."""

tests/test_data/catalog_codec_decoded_batch.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:436f2a56496712c65fa571c91d85e15b58268673ac0ab77facebe510219466d5
+size 575936

tests/test_data/catalog_codec_encoded_batch.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6091832e50b9060014065e3ab85130250fe02c2960dd16a10b78561ca20abbb3
+size 820544

tests/test_data/catalog_codec_input_batch.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:627dc9ab69c3d26f3293029f9106fba20868e5aa4b245d00f1723f947f6e6c35
+size 330158

tests/tokenizers/test_catalog_tokenizer.py

@@ -0,0 +1,30 @@
+import torch
+
+from aion.codecs import CatalogCodec
+from aion.modalities import Catalog
+
+def test_catalog_tokenizer(data_dir):
+    codec = CatalogCodec.from_pretrained(
+        f"polymathic-ai/aion-catalog-codec"
+    )
+    codec.eval()
+    input_batch = torch.load(
+        data_dir / f"catalog_codec_input_batch.pt", weights_only=False
+    )
+    reference_encoded_batch = torch.load(
+        data_dir / f"catalog_codec_encoded_batch.pt", weights_only=False
+    )
+    reference_decoded_batch = torch.load(
+        data_dir / f"catalog_codec_decoded_batch.pt", weights_only=False
+    )
+
+    with torch.no_grad():
+        output = codec.encode(Catalog(**input_batch))
+        decoded_output = codec.decode(output)
+
+    assert torch.allclose(output, reference_encoded_batch)
+    assert torch.allclose(decoded_output.X, reference_decoded_batch['X'], atol=1e-5)
+    assert torch.allclose(decoded_output.Y, reference_decoded_batch['Y'], atol=1e-5)
+    assert torch.allclose(decoded_output.SHAPE_E1, reference_decoded_batch['SHAPE_E1'], atol=1e-5)
+    assert torch.allclose(decoded_output.SHAPE_E2, reference_decoded_batch['SHAPE_E2'], atol=1e-5)
+    assert torch.allclose(decoded_output.SHAPE_R, reference_decoded_batch['SHAPE_R'], atol=1e-5)