chore: remove legacy JointCVAE

tools/models.py

@@ -12,7 +12,6 @@ import math
 
 import torch
 import torch.nn as nn
-from torch.distributions import Normal
 
 
 # ---------------------------------------------------------------------------
@@ -157,80 +156,3 @@ class TrajectoryFlowModel(nn.Module):
 
         # Output projection
         return self.output_proj(h)  # (B, T, 3)
-
-
-# ---------------------------------------------------------------------------
-# Legacy JointCVAE — kept for backward compatibility with generator.py
-# ---------------------------------------------------------------------------
-
-
-class JointCVAE(nn.Module):
-    """Joint Conditional VAE for mouse trajectory generation (legacy).
-
-    Kept for backward compatibility with the existing generator.
-    See TrajectoryFlowModel for the new approach.
-    """
-
-    def __init__(
-        self,
-        seq_len: int = 64,
-        latent_dim: int = 32,
-        hidden: int = 128,
-        cond_dim: int = 3,
-    ):
-        super().__init__()
-        self.seq_len = seq_len
-        self.latent_dim = latent_dim
-        self.hidden = hidden
-        self.cond_dim = cond_dim
-        self.feat_dim = 3
-
-        self.enc_gru = nn.GRU(
-            input_size=self.feat_dim + cond_dim,
-            hidden_size=hidden,
-            num_layers=2,
-            batch_first=True,
-            bidirectional=True,
-        )
-        self.enc_mu = nn.Linear(hidden * 2, latent_dim)
-        self.enc_logvar = nn.Linear(hidden * 2, latent_dim)
-
-        self.dec_h0 = nn.Linear(latent_dim + cond_dim, hidden * 2)
-        self.dec_gru = nn.GRU(
-            input_size=latent_dim + cond_dim,
-            hidden_size=hidden,
-            num_layers=2,
-            batch_first=True,
-        )
-        self.dec_out = nn.Linear(hidden, self.feat_dim)
-
-    def encode(self, seq: torch.Tensor, cond: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        B, T, _ = seq.shape
-        c_exp = cond.unsqueeze(1).expand(B, T, self.cond_dim)
-        x_in = torch.cat([seq, c_exp], dim=-1)
-        _, h_n = self.enc_gru(x_in)
-        h_fwd = h_n[-2]
-        h_bwd = h_n[-1]
-        h_cat = torch.cat([h_fwd, h_bwd], dim=-1)
-        return self.enc_mu(h_cat), self.enc_logvar(h_cat)
-
-    def decode(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
-        B = z.shape[0]
-        zc = torch.cat([z, cond], dim=-1)
-        h0_flat = self.dec_h0(zc)
-        h0 = h0_flat.view(B, 2, self.hidden).permute(1, 0, 2).contiguous()
-        inp = zc.unsqueeze(1).expand(B, self.seq_len, -1)
-        out, _ = self.dec_gru(inp, h0)
-        return self.dec_out(out)
-
-    def reparameterise(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
-        std = torch.exp(0.5 * logvar)
-        return Normal(mu, std).rsample()
-
-    def forward(
-        self, seq: torch.Tensor, cond: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        mu, logvar = self.encode(seq, cond)
-        z = self.reparameterise(mu, logvar)
-        recon = self.decode(z, cond)
-        return recon, mu, logvar