From 26af105195613da8465d5e9b2c097d2f02eb4603 Mon Sep 17 00:00:00 2001 From: Mayeul Aubin Date: Thu, 5 Jun 2025 17:30:26 +0200 Subject: [PATCH 1/2] ChatGPT models --- sCOCA_ML/models/UNet_models.py | 76 ++++++++++++++++++++++++++ sCOCA_ML/models/__init__.py | 0 sCOCA_ML/models/base_class_models.py | 38 +++++++++++++ sCOCA_ML/models/e3nn_models.py | 79 ++++++++++++++++++++++++++++ 4 files changed, 193 insertions(+) create mode 100644 sCOCA_ML/models/UNet_models.py create mode 100644 sCOCA_ML/models/__init__.py create mode 100644 sCOCA_ML/models/base_class_models.py create mode 100644 sCOCA_ML/models/e3nn_models.py diff --git a/sCOCA_ML/models/UNet_models.py b/sCOCA_ML/models/UNet_models.py new file mode 100644 index 0000000..56ee21e --- /dev/null +++ b/sCOCA_ML/models/UNet_models.py @@ -0,0 +1,76 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from .base_class_models import BaseModel + +class FiLM(nn.Module): + def __init__(self, num_features, style_dim): + super(FiLM, self).__init__() + self.gamma = nn.Linear(style_dim, num_features) + self.beta = nn.Linear(style_dim, num_features) + + def forward(self, x, style): + gamma = self.gamma(style).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) + beta = self.beta(style).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) + return gamma * x + beta + +class UNetBlock(nn.Module): + def __init__(self, in_channels, out_channels, style_dim=None): + super(UNetBlock, self).__init__() + self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1) + self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1) + self.norm = nn.BatchNorm3d(out_channels) + self.relu = nn.ReLU(inplace=True) + self.film = FiLM(out_channels, style_dim) if style_dim else None + + def forward(self, x, style=None): + x = self.relu(self.norm(self.conv1(x))) + x = self.relu(self.norm(self.conv2(x))) + if self.film: + x = self.film(x, style) + return x + +class UNet3D(BaseModel): + def __init__(self, N: int = 128, + in_channels: int = 2, + out_channels: int = 1, + style_dim: int = 2, + device: torch.device = torch.device('cpu')): + """ + 3D U-Net model with optional FiLM layers for style conditioning. + Parameters: + - N: Size of the input data: data will have the shape (B, C, N, N, N) with B the batch size, C the number of channels. + - in_channels: Number of input channels (default is 2). + - out_channels: Number of output channels (default is 1). + - style_dim: Dimension of the style vector (default is 2). + - device: Device to load the model onto (default is CPU). + This model implements a 3D U-Net architecture with downsampling and upsampling blocks. + The model uses convolutional layers with ReLU activations and batch normalization. + The FiLM layers are used to condition the feature maps on style parameters. + """ + + super().init(N=N, + in_channels=in_channels, + out_channels=out_channels, + style_parameters=style_dim, + device=device) + + self.enc1 = UNetBlock(in_channels, 32, style_dim) + self.pool1 = nn.MaxPool3d(2) + self.enc2 = UNetBlock(32, 64, style_dim) + self.pool2 = nn.MaxPool3d(2) + self.bottleneck = UNetBlock(64, 128, style_dim) + + self.up2 = nn.ConvTranspose3d(128, 64, kernel_size=2, stride=2) + self.dec2 = UNetBlock(128, 64) + self.up1 = nn.ConvTranspose3d(64, 32, kernel_size=2, stride=2) + self.dec1 = UNetBlock(64, 32) + self.final = nn.Conv3d(32, out_channels, kernel_size=1) + + def forward(self, x, style): + e1 = self.enc1(x, style) + e2 = self.enc2(self.pool1(e1), style) + b = self.bottleneck(self.pool2(e2), style) + d2 = self.dec2(torch.cat([self.up2(b), e2], dim=1)) + d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1)) + return self.final(d1) diff --git a/sCOCA_ML/models/__init__.py b/sCOCA_ML/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sCOCA_ML/models/base_class_models.py b/sCOCA_ML/models/base_class_models.py new file mode 100644 index 0000000..3ce66e6 --- /dev/null +++ b/sCOCA_ML/models/base_class_models.py @@ -0,0 +1,38 @@ +import torch + + +class BaseModel(torch.nn.Module): + def __init__(self, + N:int=128, + in_channels:int=2, + out_channels:int=1, + style_parameters:int=2, + device: torch.device = torch.device('cpu')): + """ + Base class for all models. + Parameters: + - N: Size of the input data: data will have the shape (B, C, N,N,N) with B the batch size, C the number of channels. + - in_channels: Number of input channels (default is 2). + - out_channels: Number of output channels (default is 1). + - style_parameters: Number of style parameters (default is 2). + - device: Device to load the model onto (default is CPU). + """ + super().__init__() + self.N = N + self.in_channels = in_channels + self.out_channels = out_channels + self.style_parameters = style_parameters + self.device = device + self.to(self.device) + + def forward(self, x, style): + """ + Forward pass of the model. + Should be implemented in subclasses. + Parameters: + - x: Input tensor of shape (B, C, N, N, N) where B is the batch size, C is the number of channels. + - style: Style parameters tensor of shape (B, S) where S is the number of style parameters. + Returns: + - Output tensor of shape (B, C_out, N, N, N) where C_out is the number of output channels. + """ + raise NotImplementedError("Forward method must be implemented in subclasses.") diff --git a/sCOCA_ML/models/e3nn_models.py b/sCOCA_ML/models/e3nn_models.py new file mode 100644 index 0000000..a35dc86 --- /dev/null +++ b/sCOCA_ML/models/e3nn_models.py @@ -0,0 +1,79 @@ +import torch +from .base_class_models import BaseModel +from torch import nn +from e3nn import o3 +from e3nn.o3 import Irreps +from e3nn.nn import Gate +from e3nn.o3 import FullyConnectedTensorProduct +from e3nn.nn.models.v2106.gate_points_networks import SimpleNetwork + +# Scalar field = trivial irreducible representation + +class E3nnNet(BaseModel): + + def __init__(self, + N: int = 128, + in_channels: int = 2, + out_channels: int = 1, + style_parameters: int = 2, + hidden_channels: int = 16, + num_layers: int = 4, + radius: float = 2.5, + num_neighbors: int = 12, + device: torch.device = torch.device('cpu')): + """ + E3nn-based model. + Parameters: + - N: Size of the input data: data will have the shape (B, C, N, N, N) with B the batch size, C the number of channels. + - in_channels: Number of input channels (default is 2). + - out_channels: Number of output channels (default is 1). + - style_parameters: Number of style parameters (default is 2). + - hidden_channels: Number of hidden channels (default is 16). + - num_layers: Number of hidden layers (default is 4). + - radius: Radius for the neighborhood search (default is 2.5). + - num_neighbors: Number of neighbors to consider (default is 12). + - device: Device to load the model onto (default is CPU). + This model uses e3nn to handle the spherical harmonics and irreducible representations. + The input is expected to be a scalar field, which is represented as a trivial irreducible representation (0e). + The model consists of a simple network with fully connected layers and a gate mechanism. + The input is reshaped to a 2D tensor where each voxel's position is concatenated with the style parameters. + The output is reshaped back to the original 3D shape with a single output channel. + """ + + super().__init__(N=N, + in_channels=in_channels, + out_channels=out_channels, + style_parameters=style_parameters, + device=device) + + irreps_input = Irreps(f"{in_channels+style_parameters}x0e") # input channels + style parameters + irreps_hidden = Irreps(f"{hidden_channels}x0e") # hidden layers + irreps_output = Irreps(f"{out_channels}x0e") # scalar output + + self.model = SimpleNetwork( + irreps_in=irreps_input, + irreps_out=irreps_output, + layers=[irreps_hidden] * num_layers, + radius=radius, + num_neighbors=num_neighbors, + ) + + + def forward(self, x, style): + # Reshape x: (B, C, N, N, N) -> (B*N^3, C) + B, C, N, _, _ = x.shape + x = x.permute(0, 2, 3, 4, 1).reshape(-1, C) + pos = torch.stack(torch.meshgrid( + torch.linspace(-1, 1, N), + torch.linspace(-1, 1, N), + torch.linspace(-1, 1, N), + indexing='ij' + ), dim=-1).reshape(-1, 3).repeat(B, 1, 1).reshape(-1, 3).to(x.device) + + # Expand style to each voxel + style = style.unsqueeze(1).expand(-1, N**3, -1).reshape(-1, style.shape[-1]) + x = torch.cat([x, style], dim=-1) # Simple concat of style params + + out = self.model(pos, x) + out = out.reshape(B, N, N, N, 1).permute(0, 4, 1, 2, 3) + return out From 24c2d546dbe24931c0921ee263ae28f6a90cdd8c Mon Sep 17 00:00:00 2001 From: Mayeul Aubin Date: Thu, 5 Jun 2025 17:30:38 +0200 Subject: [PATCH 2/2] dataset improvement --- sCOCA_ML/dataset/gravpot_dataset.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/sCOCA_ML/dataset/gravpot_dataset.py b/sCOCA_ML/dataset/gravpot_dataset.py index 37b75ae..d369801 100644 --- a/sCOCA_ML/dataset/gravpot_dataset.py +++ b/sCOCA_ML/dataset/gravpot_dataset.py @@ -7,6 +7,7 @@ from glob import glob import re + def read_cosmo_and_time_file(cosmo_and_time_file): with open(cosmo_and_time_file, 'r') as f: lines = f.readlines() @@ -154,6 +155,10 @@ class GravPotDataset(Dataset): def __getitem__(self, idx): from pysbmy.field import read_field_chunk_3D_periodic + from io import BytesIO + from sbmy_control.low_level import stdout_redirector, stderr_redirector + f = BytesIO() + ID, t, ox, oy, oz = self.samples[idx] # Filepaths @@ -165,11 +170,12 @@ class GravPotDataset(Dataset): style_path = os.path.join(self.root_dir, self.STYLE_DIR, f'{self.style_files}_{ID}_nforce{t}.txt') # Read 3D chunks - input_arrays = [ - read_field_chunk_3D_periodic(file, self.N,self.N,self.N, ox,oy,oz, name=varname).array - for file, varname in zip(input_paths, self.initial_conditions_variables) - ] - target_array = read_field_chunk_3D_periodic(target_path, self.N, self.N, self.N, ox, oy, oz, name=self.target_variable).array + with stdout_redirector(f): + input_arrays = [ + read_field_chunk_3D_periodic(file, self.N,self.N,self.N, ox,oy,oz, name=varname).array + for file, varname in zip(input_paths, self.initial_conditions_variables) + ] + target_array = read_field_chunk_3D_periodic(target_path, self.N, self.N, self.N, ox, oy, oz, name=self.target_variable).array # Stack the input arrays input_tensor = np.stack(input_arrays, axis=0)