Merge pull request #64 from PolymathicAI/fix/spectrum_crop

automatically crop spectra

Author: al-jshen

aion/codecs/preprocessing/spectrum.py

@@ -0,0 +1,26 @@
+import torch.nn.functional as F
+
+from aion.modalities import Spectrum
+
+
+def pad_spectrum(x: Spectrum) -> Spectrum:
+    """Pad a Spectrum object to its specified pad_length.
+
+    Note: Each of the sequence attributes (flux, ivar,
+          mask, wavelength) should be 2D tensors.
+    """
+    padding_values = {"lambda": 99999, "mask": True, "ivar": 0}
+
+    for k in ["flux", "ivar", "mask", "wavelength"]:
+        setattr(
+            x,
+            k,
+            F.pad(
+                getattr(x, k),
+                (0, x.pad_length - getattr(x, k).shape[-1]),
+                mode="constant",
+                value=padding_values[k] if k in padding_values else 0,
+            ),
+        )
+
+    return x

aion/codecs/spectrum.py

@@ -6,6 +6,7 @@
 from aion.codecs.base import Codec
 from aion.codecs.modules.convnext import ConvNextDecoder1d, ConvNextEncoder1d
 from aion.codecs.modules.spectrum import LatentSpectralGrid
+from aion.codecs.preprocessing.spectrum import pad_spectrum
 from aion.codecs.quantizers import LucidrainsLFQ, Quantizer, ScalarLinearQuantizer
 from aion.codecs.utils import CodecPytorchHubMixin
 from aion.modalities import Spectrum
@@ -54,6 +55,9 @@ def quantizer(self) -> Quantizer:
         return self._quantizer
 
     def _encode(self, x: Spectrum) -> Float[torch.Tensor, "b c t"]:
+        if hasattr(x, "pad_length"):
+            x = pad_spectrum(x)
+
         # Extract fields from Spectrum instance
         flux = x.flux
         ivar = x.ivar

aion/modalities.py

@@ -97,6 +97,7 @@ class Spectrum(Modality):
     ivar: Float[Tensor, " batch length"]
     mask: Bool[Tensor, " batch length"]
     wavelength: Float[Tensor, " batch length"]
+    pad_length: ClassVar[int]
 
     def __repr__(self) -> str:
         repr_str = (
@@ -112,13 +113,15 @@ class DESISpectrum(Spectrum):
 
     token_key: ClassVar[str] = "tok_spectrum_desi"
     num_tokens: ClassVar[int] = 273
+    pad_length: ClassVar[int] = 7808
 
 
 class SDSSSpectrum(Spectrum):
     """SDSS spectrum modality data."""
 
     token_key: ClassVar[str] = "tok_spectrum_sdss"
     num_tokens: ClassVar[int] = 273
+    pad_length: ClassVar[int] = 4800
 
 
 # Catalog modality