Canopus51/Cosmic-Large-Scale-Structure-Generator

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import numpy as np
5
6class UNet(nn.Module):
7    """
8    A Light-weight UNet with Time Embedding for Diffusion-based 
9    Large Scale Structure (LSS) generation.
10    """
11    def __init__(self, time_dim=128):
12        super().__init__()
13        
14        self.time_dim = time_dim
15
16        # --- Time Embedding MLP ---
17        self.time_mlp = nn.Sequential(
18            nn.Linear(time_dim, time_dim * 2),
19            nn.SiLU(),
20            nn.Linear(time_dim * 2, time_dim),
21        )
22
23        # --- Encoder (Downsampling) ---
24        # Input: (B, 1, 128, 128)
25        self.enc1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)             
26        self.enc2 = nn.Conv2d(32, 64, kernel_size=3, padding=1, stride=2)  
27        self.enc3 = nn.Conv2d(64, 128, kernel_size=3, padding=1, stride=2) 
28
29        # --- Middle ---
30        self.mid = nn.Conv2d(128, 128, kernel_size=3, padding=1)           
31
32        # --- Decoder (Upsampling) ---
33        self.dec1 = nn.Conv2d(128 + 64, 64, kernel_size=3, padding=1)      
34        self.dec2 = nn.Conv2d(64 + 32, 32, kernel_size=3, padding=1)       
35        self.dec3 = nn.Conv2d(32, 1, kernel_size=3, padding=1)             
36
37        self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
38
39        # --- Time Modulation Layers ---
40        self.t_enc1 = nn.Linear(time_dim, 32)
41        self.t_enc2 = nn.Linear(time_dim, 64)
42        self.t_mid  = nn.Linear(time_dim, 128)
43
44        self.act = nn.SiLU()
45
46    def sinusoidal_embedding(self, t, dim):
47        """
48        Standard sinusoidal embedding as used in DDPM.
49        """
50        device = t.device
51        half_dim = dim // 2
52        emb = torch.arange(half_dim, device=device).float()
53        emb = torch.exp(-np.log(10000) * emb / (half_dim - 1))
54        emb = t.float().unsqueeze(1) * emb.unsqueeze(0)
55        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
56        if dim % 2 == 1:
57            emb = F.pad(emb, (0, 1))
58        return emb
59
60    def forward(self, x, t):
61        # 1. Time Embedding
62        t_emb = self.sinusoidal_embedding(t, self.time_dim)
63        t_emb = self.time_mlp(t_emb)
64
65        # 2. Encoder
66        # Add time info via broadcasting: (B, C, 1, 1)
67        x1 = self.act(self.enc1(x) + self.t_enc1(t_emb)[:, :, None, None])
68        x2 = self.act(self.enc2(x1) + self.t_enc2(t_emb)[:, :, None, None])
69        x3 = self.act(self.enc3(x2))
70
71        # 3. Middle
72        x_mid = self.act(self.mid(x3) + self.t_mid(t_emb)[:, :, None, None])
73
74        # 4. Decoder with Skip Connections
75        # 32x32 -> 64x64
76        x_up1 = self.upsample(x_mid)
77        x_cat1 = torch.cat([x_up1, x2], dim=1)
78        x_d1 = self.act(self.dec1(x_cat1))
79
80        # 64x64 -> 128x128
81        x_up2 = self.upsample(x_d1)
82        x_cat2 = torch.cat([x_up2, x1], dim=1)
83        x_d2 = self.act(self.dec2(x_cat2))
84
85        # Final Projection
86        out = self.dec3(x_d2)
87        return out
88
89if __name__ == "__main__":
90    # Quick architecture test
91    model = UNet(time_dim=128)
92    test_input = torch.randn(1, 1, 128, 128)
93    test_time = torch.tensor([10])
94    output = model(test_input, test_time)
95    print(f"Input shape: {test_input.shape}")
96    print(f"Output shape: {output.shape}")
97