Rollback the add_time_patch method

modeling_reve.py

@@ -217,16 +217,36 @@ class Learnable4DPE(nn.Module):
 
 
 class FourierEmb4D(nn.Module):
-    def __init__(self, dimension: int, freqs: int, increment_time=0.1, margin: float = 0.4):
+    """
+    Fourier positional embedding for 4D positions (x, y, z, t).
+    This version allows for a reduced number of frequencies (n_freqs),
+    and ensures the output embedding has the specified dimension.
+
+    Parameters
+    ----------
+    dimension : int
+        The dimension of the output embedding. Must be an even number.
+    freqs : int
+        The number of frequencies to use for the Fourier embedding.
+    increment_time : float, optional
+        The time increment to scale the time dimension. Default is 0.1.
+    margin : float, optional
+        The margin to add to the position coordinates to avoid boundary issues. Default is 0.4.
+
+    """
+
+    def __init__(
+        self, dimension: int, freqs: int, increment_time=0.1, margin: float = 0.4
+    ):
         super().__init__()
         self.dimension = dimension
         self.freqs = freqs
         self.increment_time = increment_time
         self.margin = margin
 
-    def forward(self, positions_):
+    def forward(self, positions_: torch.Tensor) -> torch.Tensor:
         positions = positions_.clone()
-        positions[..., -1] *= self.increment_time
+        positions[:, :, -1] *= self.increment_time
         input_shape = positions.shape
         batch_dims = list(input_shape[:-1])
 
@@ -241,26 +261,68 @@ class FourierEmb4D(nn.Module):
         p_z = 2 * math.pi * freqs_z / width
         p_w = 2 * math.pi * freqs_w / width
         positions = positions[..., None, None, None, None, :]
-        loc = positions[..., 0] * p_x + positions[..., 1] * p_y + positions[..., 2] * p_z + positions[..., 3] * p_w
+        loc = (
+            positions[..., 0] * p_x
+            + positions[..., 1] * p_y
+            + positions[..., 2] * p_z
+            + positions[..., 3] * p_w
+        )
         batch_dims.append(-1)
         loc = loc.view(batch_dims)
 
         half_dim = self.dimension // 2
         current_dim = loc.shape[-1]
-
-        # FIX: Safe slicing logic
         if current_dim != half_dim:
             if current_dim > half_dim:
                 loc = loc[..., :half_dim]
             else:
                 raise ValueError(
                     f"Input dimension ({current_dim}) is too small for target "
-                    f"embedding dimension ({self.dimension}). Expected at least {half_dim}.",
+                    f"embedding dimension ({self.dimension}). Expected at least {half_dim}."
                 )
 
         emb = torch.cat([torch.cos(loc), torch.sin(loc)], dim=-1)
         return emb
 
+    @classmethod
+    def add_time_patch(cls, pos: torch.Tensor, num_patches: int) -> torch.Tensor:
+        """
+        Expand the position tensor by adding a time dimension, handling batched data.
+
+        Parameters
+        ----------
+        pos : torch.Tensor
+            Input tensor of shape (B, C, 3), where B is the batch size,
+            C is the number of channels, and 3 represents x, y, z.
+        num_patches : int
+            The number of time patches.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape (B, C * num_patches, 4), where each position is repeated with each time value.
+        """
+        batch, nchans, _ = pos.shape
+        # Repeat each position for each time step
+        pos_repeated = pos.unsqueeze(2).repeat(
+            1, 1, num_patches, 1
+        )  # Shape: (batch, nchans, num_patches, 3)
+        # Generate time values with the specified increment
+        time_values = torch.arange(
+            0, num_patches, 1, device=pos.device
+        ).float()  # Shape: (num_patches,)
+        time_values = time_values.view(1, 1, num_patches, 1).expand(
+            batch, nchans, num_patches, 1
+        )  # (batch, nchans, num_patches, 1)
+        # Concatenate the repeated positions with the time values along the last dimension
+        pos_with_time = torch.cat(
+            (pos_repeated, time_values), dim=-1
+        )  # Shape: (batch, nchans, num_patches, 4)
+        # Reshape to (batch, nchans * num_patches, 4)
+        pos_with_time = pos_with_time.view(batch, nchans * num_patches, 4)
+
+        return pos_with_time
+
 
 def patch_embedding(embed_dim, patch_size):
     to_patch_embedding = nn.Sequential(nn.Linear(patch_size, embed_dim))