diff --git a/tools/models.py b/tools/models.py index 2d252bb..e11deec 100644 --- a/tools/models.py +++ b/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