upload: model.py

0ef36e7 verified 6 months ago

14.7 kB

"""
model.py — Specialised Hα encoder (Stage 1 Hα) with MAE.
Compact version of SpectralEncoder adapted for spectra cropped
to 128 bins around Hα (6512.8–6612.8 Å).
Architecture :
  Spectre Hα [128] → Patches [31, 8] → Projection [31, 128]
  → + Wavelength PE → Masquage 60% → [CLS] + Visibles [~13, 128]
  → Transformer 4L → CLS embedding z_halpha [128]
  → MAE Decoder [reconstruction]
Key differences from the full encoder (stage1/model.py):
  - Input : 128 bins (vs 4096)
  - Patches : 8 px, overlap 4, step 4 → 31 patches (vs 511)
  - d_model : 128 (vs 256)
  - n_layers : 4 (vs 6)
  - No GRL/discriminator (negative result confirmed)
  - ~300K params (vs ~5.3M)
"""

import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from config import ModelConfig


# ══════════════════════════════════════════════════════════════════════════════
# MASQUAGE CONTIGU
# ══════════════════════════════════════════════════════════════════════════════

def contiguous_masking(n_patches: int, mask_ratio: float = 0.60,
                       n_blocks: int = 3) -> np.ndarray:
    """
    Generate a contiguous-block mask for MAE.
    Adapted for 31 patches (vs 511 for the full encoder):
      - 60% masking → ~19 masked patches, ~12 visible
      - 3 contiguous blocks (vs 4) because the sequence is short
    Avec 12 patches visibles + 1 CLS = 13 tokens pour l'encodeur.
    C'est suffisant pour un Transformer 4 couches.
    """
    n_masked = int(n_patches * mask_ratio)

    if n_masked == 0 or n_patches < n_blocks:
        return np.zeros(n_patches, dtype=bool)

    mask = np.zeros(n_patches, dtype=bool)
    block_size = max(1, n_masked // n_blocks)

    possible_starts = np.arange(0, max(1, n_patches - block_size))

    if len(possible_starts) < n_blocks:
        starts = possible_starts
    else:
        starts = np.sort(
            np.random.choice(possible_starts, size=n_blocks, replace=False)
        )

    for s in starts:
        end = min(s + block_size, n_patches)
        mask[s:end] = True

    # Fill if necessary
    current = mask.sum()
    if current < n_masked:
        unmasked = np.where(~mask)[0]
        extra = min(n_masked - current, len(unmasked))
        if extra > 0:
            chosen = np.random.choice(unmasked, size=extra, replace=False)
            mask[chosen] = True

    return mask


# ══════════════════════════════════════════════════════════════════════════════
# ENCODAGE POSITIONNEL PAR LONGUEUR D'ONDE
# ══════════════════════════════════════════════════════════════════════════════

class WavelengthPE(nn.Module):
    """
    Sinusoidal positional encoding based on wavelength (Å).
    Identique au full encoder : PE(λ) = sin/cos(λ / 10000 × div_term).
    Even over a 100 Å window, the relative position is physically
    significative : les ailes bleue (λ < 6562.8) et rouge (λ > 6562.8)
    of Hα have different physical meanings (V/R ratio, asymmetries).
    """

    def __init__(self, d_model: int = 128):
        super().__init__()
        self.d_model = d_model
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        self.register_buffer("div_term", div_term)

    def forward(self, lambda_means: torch.Tensor) -> torch.Tensor:
        """
        lambda_means: [B, N_patches] — λ moyen de chaque patch (en Å).
        Retourne: [B, N_patches, d_model].
        """
        pos = lambda_means.unsqueeze(-1) / 10000.0
        pe = torch.zeros(*lambda_means.shape, self.d_model,
                         device=lambda_means.device)
        pe[..., 0::2] = torch.sin(pos * self.div_term)
        pe[..., 1::2] = torch.cos(pos * self.div_term)
        return pe


# ══════════════════════════════════════════════════════════════════════════════
# DÉCODEUR MAE
# ══════════════════════════════════════════════════════════════════════════════

class MAEDecoder(nn.Module):
    """
    Lightweight MAE decoder for reconstructing masked Hα patches.
    Smaller than the full encoder decoder:
      - d_decoder = 64 (vs 128)
      - n_layers = 2
      - head projette vers patch_size = 8 (vs 16)
    """

    def __init__(self, d_encoder: int, d_decoder: int, n_layers: int,
                 n_heads: int, patch_size: int):
        super().__init__()
        self.d_decoder = d_decoder

        # Projection encoder → decoder
        self.proj = nn.Linear(d_encoder, d_decoder)

        # Mask token appris
        self.mask_token = nn.Parameter(torch.randn(1, 1, d_decoder) * 0.02)

        # Decoder mini-Transformer
        decoder_layer = nn.TransformerEncoderLayer(
            d_model=d_decoder, nhead=n_heads,
            dim_feedforward=d_decoder * 4,
            activation="gelu", batch_first=True, dropout=0.1,
        )
        self.decoder = nn.TransformerEncoder(decoder_layer, num_layers=n_layers)

        # Prediction head: d_decoder → patch_size
        self.head = nn.Linear(d_decoder, patch_size)

    def forward(self, encoded_visible, visible_pe, full_pe, mask):
        """
        Reconstruct masked patches.
        Inputs:
          encoded_visible : [B, N_vis, d_encoder]
          visible_pe      : [B, N_vis, d_decoder]
          full_pe         : [B, N_all, d_decoder]
          mask            : [B, N_all] — True = masked
        Output:
          [B, N_all, patch_size]
        """
        B, N_all = mask.shape

        vis = self.proj(encoded_visible)

        full_seq = self.mask_token.expand(B, N_all, -1).clone()
        vis_positions = (~mask)
        full_seq[vis_positions] = vis.reshape(-1, self.d_decoder)

        full_seq = full_seq + full_pe

        decoded = self.decoder(full_seq)
        return self.head(decoded)


# ══════════════════════════════════════════════════════════════════════════════
# ENCODEUR SPECTRAL Hα
# ══════════════════════════════════════════════════════════════════════════════

class SpectralEncoderHalpha(nn.Module):
    """
    MAE Transformer encoder for cropped Hα spectra (128 bins).
    Pipeline interne :
      1. Patchify    : split into 31 patches of 8 pixels (step=4)
      2. Projection  : chaque patch [8] → token [128]
      3. PE          : positional encoding based on λ
      4. Masquage    : ne garde que ~40% des patches (12 visibles)
      5. [CLS]       : global summary token
      6. Transformer : 4 couches d'auto-attention
      7. Sortie      : z_halpha = LayerNorm(CLS) ∈ ℝ^128
    Avec 128 bins au lieu de 4096 :
      - 31 patches au lieu de 511
      - 12 visibles au lieu de 153
      - ~10× faster to run
    """

    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.cfg = cfg
        self.patch_size = cfg.patch_size
        self.patch_overlap = cfg.patch_overlap
        self.d_model = cfg.d_model
        self.step = cfg.patch_size - cfg.patch_overlap  # = 4

        # Projection patch → token
        self.patch_proj = nn.Linear(cfg.patch_size, cfg.d_model)

        # Encodage positionnel
        self.wave_pe = WavelengthPE(cfg.d_model)

        # Token CLS
        self.cls_token = nn.Parameter(torch.randn(1, 1, cfg.d_model) * 0.02)

        # Transformer
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=cfg.d_model,
            nhead=cfg.n_heads,
            dim_feedforward=cfg.d_ff,
            dropout=cfg.dropout,
            activation="gelu",
            batch_first=True,
        )
        self.encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=cfg.n_layers
        )

        # Normalisation finale
        self.norm = nn.LayerNorm(cfg.d_model)

    def patchify(self, flux, wavelengths, validity):
        """
        Split the 128-bin spectrum into 31 patches of 8 pixels.
        Avec step=4 : N = (128 - 8) / 4 + 1 = 31 patches
        """
        patches = flux.unfold(-1, self.patch_size, self.step)
        lam_patches = wavelengths.unfold(-1, self.patch_size, self.step)
        val_patches = validity.unfold(-1, self.patch_size, self.step)

        lambda_means = lam_patches.mean(-1)
        patch_valid = val_patches.mean(-1)

        return patches, lambda_means, patch_valid

    def forward(self, flux, wavelengths, validity, mask=None):
        """
        Forward pass.
        Inputs:
          flux        : [B, 128]
          wavelengths : [B, 128]
          validity    : [B, 128]
          mask        : [B, 31] optional
        Outputs:
          z            : [B, d_model]       — CLS embedding (z_halpha)
          encoded      : [B, N_vis, d_model] — encoded visible tokens
          patches      : [B, 31, 8]          — all patches
          lambda_means : [B, 31]             — mean λ per patch
          mask         : [B, 31]             — mask used
          wpe          : [B, 31, d_model]    — positional encoding
        """
        B = flux.shape[0]

        # Patchify
        patches, lambda_means, patch_valid = self.patchify(
            flux, wavelengths, validity
        )
        N = patches.shape[1]

        # Projection + PE
        tokens = self.patch_proj(patches)
        wpe = self.wave_pe(lambda_means)
        tokens = tokens + wpe

        # Masquage MAE
        if mask is None:
            mask = torch.zeros(B, N, dtype=torch.bool, device=flux.device)

        visible_mask = ~mask
        n_visible = visible_mask[0].sum().item()

        visible_tokens = tokens[visible_mask].view(B, n_visible, self.d_model)
        visible_pe = wpe[visible_mask].view(B, n_visible, self.d_model)

        # CLS + tokens visibles
        cls = self.cls_token.expand(B, -1, -1)
        input_tokens = torch.cat([cls, visible_tokens], dim=1)

        # Masque d'attention
        visible_valid = patch_valid[visible_mask].view(B, n_visible)
        attn_pad = torch.cat([
            torch.ones(B, 1, device=flux.device),
            (visible_valid > 0.1).float()
        ], dim=1)
        src_key_padding_mask = (attn_pad == 0)

        # Transformer
        encoded = self.encoder(
            input_tokens, src_key_padding_mask=src_key_padding_mask
        )

        # Embedding CLS
        z = self.norm(encoded[:, 0])

        return z, encoded[:, 1:], patches, lambda_means, mask, wpe


# ══════════════════════════════════════════════════════════════════════════════
# MODÈLE COMPLET — STAGE 1 Hα
# ══════════════════════════════════════════════════════════════════════════════

class Stage1HalphaModel(nn.Module):
    """
    Assembly: Hα Encoder + MAE Decoder.
    No instrument discriminator (GRL removed).
    The only loss is the MAE reconstruction of masked patches.
    """

    def __init__(self, model_cfg: ModelConfig):
        super().__init__()

        self.encoder = SpectralEncoderHalpha(model_cfg)

        self.mae_decoder = MAEDecoder(
            d_encoder=model_cfg.d_model,
            d_decoder=model_cfg.d_decoder,
            n_layers=model_cfg.n_decoder_layers,
            n_heads=model_cfg.n_decoder_heads,
            patch_size=model_cfg.patch_size,
        )

        # PE projection for the decoder
        self.pe_proj = nn.Linear(model_cfg.d_model, model_cfg.d_decoder)

    def forward(self, flux, wavelengths, validity, mask):
        """
        Forward pass : encoder → reconstruction MAE.
        Inputs:
          flux        : [B, 128]
          wavelengths : [B, 128]
          validity    : [B, 128]
          mask        : [B, 31] — MAE mask
        Returns:
          z             : [B, 128]     — CLS embedding (z_halpha)
          mae_loss      : scalar       — MSE reconstruction of masked patches
          reconstructed : [B, 31, 8]   — patches reconstruits
          patches       : [B, 31, 8]   — patches originaux (cible)
          mask          : [B, 31]      — mask used
        """
        # Encode
        z, encoded_vis, patches, lambda_means, mask, wpe = self.encoder(
            flux, wavelengths, validity, mask
        )

        B, N, P = patches.shape

        # PE for the decoder
        full_pe = self.pe_proj(wpe)
        vis_mask = ~mask
        n_vis = vis_mask[0].sum().item()
        vis_pe = full_pe[vis_mask].view(B, n_vis, -1)

        # MAE decoding
        reconstructed = self.mae_decoder(encoded_vis, vis_pe, full_pe, mask)

        # MAE loss: MSE on masked patches only
        target = patches
        mae_loss = ((reconstructed - target) ** 2)
        mask_expanded = mask.unsqueeze(-1).expand_as(mae_loss)
        n_masked_total = mask_expanded.sum()

        if n_masked_total > 0:
            mae_loss = (mae_loss * mask_expanded.float()).sum() / n_masked_total
        else:
            mae_loss = mae_loss.mean()

        return {
            "z": z,
            "mae_loss": mae_loss,
            "reconstructed": reconstructed,
            "patches": patches,
            "mask": mask,
        }

    def get_embeddings(self, flux, wavelengths, validity):
        """
        Inference mode: CLS embeddings without masking.
        Retourne z_halpha ∈ ℝ^128.
        """
        z, _, _, _, _, _ = self.encoder(flux, wavelengths, validity, mask=None)
        return z

	"""
	model.py — Specialised Hα encoder (Stage 1 Hα) with MAE.

	Compact version of SpectralEncoder adapted for spectra cropped
	to 128 bins around Hα (6512.8–6612.8 Å).

	Architecture :
	Spectre Hα [128] → Patches [31, 8] → Projection [31, 128]
	→ + Wavelength PE → Masquage 60% → [CLS] + Visibles [~13, 128]
	→ Transformer 4L → CLS embedding z_halpha [128]
	→ MAE Decoder [reconstruction]

	Key differences from the full encoder (stage1/model.py):
	- Input : 128 bins (vs 4096)
	- Patches : 8 px, overlap 4, step 4 → 31 patches (vs 511)
	- d_model : 128 (vs 256)
	- n_layers : 4 (vs 6)
	- No GRL/discriminator (negative result confirmed)
	- ~300K params (vs ~5.3M)
	"""

	import math
	import numpy as np
	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	from config import ModelConfig


	# ══════════════════════════════════════════════════════════════════════════════
	# MASQUAGE CONTIGU
	# ══════════════════════════════════════════════════════════════════════════════

	def contiguous_masking(n_patches: int, mask_ratio: float = 0.60,
	n_blocks: int = 3) -> np.ndarray:
	"""
	Generate a contiguous-block mask for MAE.

	Adapted for 31 patches (vs 511 for the full encoder):
	- 60% masking → ~19 masked patches, ~12 visible
	- 3 contiguous blocks (vs 4) because the sequence is short

	Avec 12 patches visibles + 1 CLS = 13 tokens pour l'encodeur.
	C'est suffisant pour un Transformer 4 couches.
	"""
	n_masked = int(n_patches * mask_ratio)

	if n_masked == 0 or n_patches < n_blocks:
	return np.zeros(n_patches, dtype=bool)

	mask = np.zeros(n_patches, dtype=bool)
	block_size = max(1, n_masked // n_blocks)

	possible_starts = np.arange(0, max(1, n_patches - block_size))

	if len(possible_starts) < n_blocks:
	starts = possible_starts
	else:
	starts = np.sort(
	np.random.choice(possible_starts, size=n_blocks, replace=False)
	)

	for s in starts:
	end = min(s + block_size, n_patches)
	mask[s:end] = True

	# Fill if necessary
	current = mask.sum()
	if current < n_masked:
	unmasked = np.where(~mask)[0]
	extra = min(n_masked - current, len(unmasked))
	if extra > 0:
	chosen = np.random.choice(unmasked, size=extra, replace=False)
	mask[chosen] = True

	return mask


	# ══════════════════════════════════════════════════════════════════════════════
	# ENCODAGE POSITIONNEL PAR LONGUEUR D'ONDE
	# ══════════════════════════════════════════════════════════════════════════════

	class WavelengthPE(nn.Module):
	"""
	Sinusoidal positional encoding based on wavelength (Å).

	Identique au full encoder : PE(λ) = sin/cos(λ / 10000 × div_term).

	Even over a 100 Å window, the relative position is physically
	significative : les ailes bleue (λ < 6562.8) et rouge (λ > 6562.8)
	of Hα have different physical meanings (V/R ratio, asymmetries).
	"""

	def __init__(self, d_model: int = 128):
	super().__init__()
	self.d_model = d_model
	div_term = torch.exp(
	torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
	)
	self.register_buffer("div_term", div_term)

	def forward(self, lambda_means: torch.Tensor) -> torch.Tensor:
	"""
	lambda_means: [B, N_patches] — λ moyen de chaque patch (en Å).
	Retourne: [B, N_patches, d_model].
	"""
	pos = lambda_means.unsqueeze(-1) / 10000.0
	pe = torch.zeros(*lambda_means.shape, self.d_model,
	device=lambda_means.device)
	pe[..., 0::2] = torch.sin(pos * self.div_term)
	pe[..., 1::2] = torch.cos(pos * self.div_term)
	return pe


	# ══════════════════════════════════════════════════════════════════════════════
	# DÉCODEUR MAE
	# ══════════════════════════════════════════════════════════════════════════════

	class MAEDecoder(nn.Module):
	"""
	Lightweight MAE decoder for reconstructing masked Hα patches.

	Smaller than the full encoder decoder:
	- d_decoder = 64 (vs 128)
	- n_layers = 2
	- head projette vers patch_size = 8 (vs 16)
	"""

	def __init__(self, d_encoder: int, d_decoder: int, n_layers: int,
	n_heads: int, patch_size: int):
	super().__init__()
	self.d_decoder = d_decoder

	# Projection encoder → decoder
	self.proj = nn.Linear(d_encoder, d_decoder)

	# Mask token appris
	self.mask_token = nn.Parameter(torch.randn(1, 1, d_decoder) * 0.02)

	# Decoder mini-Transformer
	decoder_layer = nn.TransformerEncoderLayer(
	d_model=d_decoder, nhead=n_heads,
	dim_feedforward=d_decoder * 4,
	activation="gelu", batch_first=True, dropout=0.1,
	)
	self.decoder = nn.TransformerEncoder(decoder_layer, num_layers=n_layers)

	# Prediction head: d_decoder → patch_size
	self.head = nn.Linear(d_decoder, patch_size)

	def forward(self, encoded_visible, visible_pe, full_pe, mask):
	"""
	Reconstruct masked patches.

	Inputs:
	encoded_visible : [B, N_vis, d_encoder]
	visible_pe : [B, N_vis, d_decoder]
	full_pe : [B, N_all, d_decoder]
	mask : [B, N_all] — True = masked

	Output:
	[B, N_all, patch_size]
	"""
	B, N_all = mask.shape

	vis = self.proj(encoded_visible)

	full_seq = self.mask_token.expand(B, N_all, -1).clone()
	vis_positions = (~mask)
	full_seq[vis_positions] = vis.reshape(-1, self.d_decoder)

	full_seq = full_seq + full_pe

	decoded = self.decoder(full_seq)
	return self.head(decoded)


	# ══════════════════════════════════════════════════════════════════════════════
	# ENCODEUR SPECTRAL Hα
	# ══════════════════════════════════════════════════════════════════════════════

	class SpectralEncoderHalpha(nn.Module):
	"""
	MAE Transformer encoder for cropped Hα spectra (128 bins).

	Pipeline interne :
	1. Patchify : split into 31 patches of 8 pixels (step=4)
	2. Projection : chaque patch [8] → token [128]
	3. PE : positional encoding based on λ
	4. Masquage : ne garde que ~40% des patches (12 visibles)
	5. [CLS] : global summary token
	6. Transformer : 4 couches d'auto-attention
	7. Sortie : z_halpha = LayerNorm(CLS) ∈ ℝ^128

	Avec 128 bins au lieu de 4096 :
	- 31 patches au lieu de 511
	- 12 visibles au lieu de 153
	- ~10× faster to run
	"""

	def __init__(self, cfg: ModelConfig):
	super().__init__()
	self.cfg = cfg
	self.patch_size = cfg.patch_size
	self.patch_overlap = cfg.patch_overlap
	self.d_model = cfg.d_model
	self.step = cfg.patch_size - cfg.patch_overlap # = 4

	# Projection patch → token
	self.patch_proj = nn.Linear(cfg.patch_size, cfg.d_model)

	# Encodage positionnel
	self.wave_pe = WavelengthPE(cfg.d_model)

	# Token CLS
	self.cls_token = nn.Parameter(torch.randn(1, 1, cfg.d_model) * 0.02)

	# Transformer
	encoder_layer = nn.TransformerEncoderLayer(
	d_model=cfg.d_model,
	nhead=cfg.n_heads,
	dim_feedforward=cfg.d_ff,
	dropout=cfg.dropout,
	activation="gelu",
	batch_first=True,
	)
	self.encoder = nn.TransformerEncoder(
	encoder_layer, num_layers=cfg.n_layers
	)

	# Normalisation finale
	self.norm = nn.LayerNorm(cfg.d_model)

	def patchify(self, flux, wavelengths, validity):
	"""
	Split the 128-bin spectrum into 31 patches of 8 pixels.

	Avec step=4 : N = (128 - 8) / 4 + 1 = 31 patches
	"""
	patches = flux.unfold(-1, self.patch_size, self.step)
	lam_patches = wavelengths.unfold(-1, self.patch_size, self.step)
	val_patches = validity.unfold(-1, self.patch_size, self.step)

	lambda_means = lam_patches.mean(-1)
	patch_valid = val_patches.mean(-1)

	return patches, lambda_means, patch_valid

	def forward(self, flux, wavelengths, validity, mask=None):
	"""
	Forward pass.

	Inputs:
	flux : [B, 128]
	wavelengths : [B, 128]
	validity : [B, 128]
	mask : [B, 31] optional

	Outputs:
	z : [B, d_model] — CLS embedding (z_halpha)
	encoded : [B, N_vis, d_model] — encoded visible tokens
	patches : [B, 31, 8] — all patches
	lambda_means : [B, 31] — mean λ per patch
	mask : [B, 31] — mask used
	wpe : [B, 31, d_model] — positional encoding
	"""
	B = flux.shape[0]

	# Patchify
	patches, lambda_means, patch_valid = self.patchify(
	flux, wavelengths, validity
	)
	N = patches.shape[1]

	# Projection + PE
	tokens = self.patch_proj(patches)
	wpe = self.wave_pe(lambda_means)
	tokens = tokens + wpe

	# Masquage MAE
	if mask is None:
	mask = torch.zeros(B, N, dtype=torch.bool, device=flux.device)

	visible_mask = ~mask
	n_visible = visible_mask[0].sum().item()

	visible_tokens = tokens[visible_mask].view(B, n_visible, self.d_model)
	visible_pe = wpe[visible_mask].view(B, n_visible, self.d_model)

	# CLS + tokens visibles
	cls = self.cls_token.expand(B, -1, -1)
	input_tokens = torch.cat([cls, visible_tokens], dim=1)

	# Masque d'attention
	visible_valid = patch_valid[visible_mask].view(B, n_visible)
	attn_pad = torch.cat([
	torch.ones(B, 1, device=flux.device),
	(visible_valid > 0.1).float()
	], dim=1)
	src_key_padding_mask = (attn_pad == 0)

	# Transformer
	encoded = self.encoder(
	input_tokens, src_key_padding_mask=src_key_padding_mask
	)

	# Embedding CLS
	z = self.norm(encoded[:, 0])

	return z, encoded[:, 1:], patches, lambda_means, mask, wpe


	# ══════════════════════════════════════════════════════════════════════════════
	# MODÈLE COMPLET — STAGE 1 Hα
	# ══════════════════════════════════════════════════════════════════════════════

	class Stage1HalphaModel(nn.Module):
	"""
	Assembly: Hα Encoder + MAE Decoder.

	No instrument discriminator (GRL removed).
	The only loss is the MAE reconstruction of masked patches.
	"""

	def __init__(self, model_cfg: ModelConfig):
	super().__init__()

	self.encoder = SpectralEncoderHalpha(model_cfg)

	self.mae_decoder = MAEDecoder(
	d_encoder=model_cfg.d_model,
	d_decoder=model_cfg.d_decoder,
	n_layers=model_cfg.n_decoder_layers,
	n_heads=model_cfg.n_decoder_heads,
	patch_size=model_cfg.patch_size,
	)

	# PE projection for the decoder
	self.pe_proj = nn.Linear(model_cfg.d_model, model_cfg.d_decoder)

	def forward(self, flux, wavelengths, validity, mask):
	"""
	Forward pass : encoder → reconstruction MAE.

	Inputs:
	flux : [B, 128]
	wavelengths : [B, 128]
	validity : [B, 128]
	mask : [B, 31] — MAE mask

	Returns:
	z : [B, 128] — CLS embedding (z_halpha)
	mae_loss : scalar — MSE reconstruction of masked patches
	reconstructed : [B, 31, 8] — patches reconstruits
	patches : [B, 31, 8] — patches originaux (cible)
	mask : [B, 31] — mask used
	"""
	# Encode
	z, encoded_vis, patches, lambda_means, mask, wpe = self.encoder(
	flux, wavelengths, validity, mask
	)

	B, N, P = patches.shape

	# PE for the decoder
	full_pe = self.pe_proj(wpe)
	vis_mask = ~mask
	n_vis = vis_mask[0].sum().item()
	vis_pe = full_pe[vis_mask].view(B, n_vis, -1)

	# MAE decoding
	reconstructed = self.mae_decoder(encoded_vis, vis_pe, full_pe, mask)

	# MAE loss: MSE on masked patches only
	target = patches
	mae_loss = ((reconstructed - target) ** 2)
	mask_expanded = mask.unsqueeze(-1).expand_as(mae_loss)
	n_masked_total = mask_expanded.sum()

	if n_masked_total > 0:
	mae_loss = (mae_loss * mask_expanded.float()).sum() / n_masked_total
	else:
	mae_loss = mae_loss.mean()

	return {
	"z": z,
	"mae_loss": mae_loss,
	"reconstructed": reconstructed,
	"patches": patches,
	"mask": mask,
	}

	def get_embeddings(self, flux, wavelengths, validity):
	"""
	Inference mode: CLS embeddings without masking.
	Retourne z_halpha ∈ ℝ^128.
	"""
	z, _, _, _, _, _ = self.encoder(flux, wavelengths, validity, mask=None)
	return z