Update modeling_super_linear.py

modeling_super_linear.py

@@ -293,64 +293,40 @@ class SparseNoisyMoE(nn.Module):
             self.gating_network = nn.Linear(input_dim, self.num_experts, bias=True)
 
     def get_periodogram(self, inputs, ker_len=50, con=1, n=10000):
-        n_fft   = 128
-        ker_len =12
-        if inputs.ndim == 2:                   # (B, L)
-            B, L          = inputs.shape
-            C             = 1
-            x             = inputs.unsqueeze(2)            # → (B, L, 1)
-            time_first    = True                           # time is dim-1
-        elif inputs.ndim == 3:
-            B, d1, d2     = inputs.shape
-            if d1 < d2:                                    # (B, L, C)
-                L, C  = d1, d2
-                x      = inputs
-                time_first = True
-            else:                                          # (B, C, L)
-                C, L  = d1, d2
-                x      = inputs.transpose(1, 2)            # → (B, L, C)
-                time_first = False
-        else:
-            raise ValueError("Input must be (B,L), (B,L,C) or (B,C,L)")
-    
-        # ---------- centre the signal ----------
-        x = x - x.mean(dim=1, keepdim=True)
+            if inputs.dim() == 2:
+                x_0 = inputs.unsqueeze(2)
+            else:
+                x_0 = inputs
+            x_0 = x_0 - torch.mean(x_0, dim=1, keepdim=True)
     
-        # ---------- parameter defaults ----------
-        if n_fft is None:
-            n_fft = 1 << (L - 1).bit_length()
-        if ker_len is None:
-            ker_len = min(L // 4, 50)
-        ker_half = ker_len // 2
+            v = torch.arange(0, n) / n
+            if con:
+                if ker_len is None:
+                    ker_len = n // 4
+                    ker_len = min(ker_len, 50)
     
-        # ---------- high-pass detrend ----------
-        if con and ker_len > 0:
-            x_perm = x.permute(0, 2, 1)              # (B, C, L)
-            ker    = torch.ones(1, 1, ker_len, device=x.device) / ker_len
-            trend  = F.conv1d(x_perm, ker, padding="same")
-            left   = min(ker_half, L - 1)
-            right  = min(ker_half, L - 1)
-            trend[:, :, :left]   = trend[:, :, left:left+1]
-            trend[:, :, -right:] = trend[:, :, -(right+1):-right]
-            x      = (x_perm - trend).permute(0, 2, 1)     # back to (B, L, C)
+                x_0 = x_0.permute(0, 2, 1)
+                ker = (torch.ones(1, 1, ker_len) / ker_len).to(x_0.device)
+                x_c = F.conv1d(x_0, ker, padding="same")
+                x_c[:, :, :ker_len // 2] = x_c[:, :, ker_len // 2:ker_len // 2 + 1]
+                x_c[:, :, -ker_len // 2:] = x_c[:, :, -ker_len // 2 - 1:-ker_len // 2]
+                x_0 = x_0 - x_c
+                x_0 = x_0.permute(0, 2, 1)
     
-        # ---------- FFT ----------
-        dft = torch.fft.fft(x, n=n_fft, dim=1) / np.sqrt(n_fft)
-        I   = (dft[:, : n_fft//2, :]).abs() ** 2
+            dft = torch.fft.fft(x_0, dim=1, n=n) / np.sqrt(n)
+            dft = dft[:, :n//2, :]
+            I = torch.abs(dft) ** 2
     
-        # ---------- normalise ----------
-        I_sum = I.sum(dim=1, keepdim=True)
-        I_sum[I_sum == 0] = 1
-        I /= I_sum
+            I_sum = torch.sum(I, dim=1, keepdim=True)
+            I_sum[I_sum == 0] = 1
+            I = I / I_sum
     
-        # ---------- restore original layout ----------
-        if inputs.ndim == 2:            # wanted (B, … )
-            return I.squeeze(2)
+            if torch.any(I_sum == 0):
+                print("Zeros in the sum")
+                raise ValueError
     
-        if time_first:                  # original was (B, L, C)
-            return I                    # already (B, F, C)
-        else:                           # original was (B, C, L) → (B, C, F)
-            return I.transpose(1, 2)
+            if inputs.dim() == 2:
+                I = I.squeeze(2)
 
 
     def fourier_interp_dim1(self,x, target_len: int = 512):