all of the modifications are done to update the old PR one_step

graphphysics/dataset/dataset.py

@@ -26,6 +26,7 @@ def __init__(
         add_edge_features: bool = True,
         use_previous_data: bool = False,
         world_pos_parameters: Optional[dict] = None,
+        target_same_frame: bool = False
     ):
         with open(meta_path, "r") as fp:
             meta = json.load(fp)
@@ -50,7 +51,7 @@ def __init__(
         self.new_edges_ratio = new_edges_ratio
         self.add_edge_features = add_edge_features
         self.use_previous_data = use_previous_data
-
+        self.target_same_frame = target_same_frame
         self.world_pos_index_start = None
         self.world_pos_index_end = None
         if world_pos_parameters is not None:
@@ -77,12 +78,20 @@ def get_traj_frame(self, index: int) -> Tuple[int, int]:
         Returns:
             Tuple[int, int]: A tuple containing the trajectory number and the frame number within that trajectory.
         """
-        traj = index // (self.trajectory_length - 1)
-        frame = index % (self.trajectory_length - 1) + int(self.use_previous_data)
+        if (self.target_same_frame):
+            traj = index // self.trajectory_length
+            frame = index % self.trajectory_length
+        else:
+            traj = index // (self.trajectory_length - 1)
+            frame = index % (self.trajectory_length - 1) + int(self.use_previous_data)
         return traj, frame
 
     def __len__(self) -> int:
-        return self.size_dataset * (self.trajectory_length - 1)
+        self.target_same_frame=True # OJO !!!!!!!!!!!!!!!!!!!!!!!!
+        if (self.target_same_frame):
+            return self.size_dataset*self.trajectory_length
+        else:
+            return self.size_dataset * (self.trajectory_length - 1)
 
     @abstractmethod
     def __getitem__(self, index: int) -> Data:

graphphysics/dataset/xdmf_dataset.py

@@ -24,6 +24,7 @@ def __init__(
         add_edge_features: bool = True,
         use_previous_data: bool = False,
         switch_to_val: bool = False,
+        target_same_frame: bool = False,
         random_prev: int = 1,  # If we use previous data, we will fetch one previous frame between [-1, -random_prev]
         random_next: int = 1,  # The target will be the frame : t + [1, random_next]
     ):
@@ -35,6 +36,7 @@ def __init__(
             new_edges_ratio=new_edges_ratio,
             add_edge_features=add_edge_features,
             use_previous_data=use_previous_data,
+            target_same_frame=target_same_frame,
         )
 
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -56,12 +58,13 @@ def __init__(
 
         self.xdmf_folder = xdmf_folder
         self.meta_path = meta_path
+        self.target_same_frame = target_same_frame
 
         # Get list of XDMF files in the folder
         self.file_paths: List[str] = [
             os.path.join(xdmf_folder, f)
             for f in os.listdir(xdmf_folder)
-            if os.path.isfile(os.path.join(xdmf_folder, f)) and f.endswith(".xdmf")
+            if os.path.isfile(os.path.join(xdmf_folder, f)) and (f.endswith(".xdmf") or f.endswith(".xmf"))
         ]
         self._size_dataset: int = len(self.file_paths)
 
@@ -94,32 +97,57 @@ def __getitem__(self, index: int) -> Union[Data, Tuple[Data, torch.Tensor]]:
         _previous_data_index = random.randint(1, self.random_prev)
 
         # Read XDMF file
-        with meshio.xdmf.TimeSeriesReader(xdmf_file) as reader:
-            num_steps = reader.num_steps
-
-            if frame - _previous_data_index < 0:
-                _previous_data_index = 1
-            if frame + _target_data_index > num_steps - 1:
-                _target_data_index = 1
-
-            if frame >= num_steps - 1:
-                raise IndexError(
-                    f"Frame index {frame} out of bounds for trajectory {traj_index} with {num_steps} frames."
-                )
-
-            points, cells = reader.read_points_cells()
-            time, point_data, _ = reader.read_data(frame)
-            _, target_point_data, _ = reader.read_data(frame + _target_data_index)
-
-            if self.use_previous_data:
-                _, previous_data, _ = reader.read_data(frame - _previous_data_index)
-
+        # --- ADD ---
+        # Option without time serie
+        if self.target_same_frame:
+            mesh = meshio.read(xdmf_file)
+            points, cells = mesh.points, mesh.cells
+            point_data = dict(mesh.point_data)
+            # --- AJOUT : inclure mesh_pos et wall_mask ---
+            point_data["mesh_pos"] = mesh.points.astype(
+                self.meta["features"]["mesh_pos"]["dtype"]
+            )
+            if "wall_mask" in mesh.point_data.keys():
+                point_data["wall_mask"] = np.array(
+                    mesh.point_data["wall_mask"]
+                ).astype(self.meta["features"]["wall_mask"]["dtype"])
+            # --- FIN AJOUT ---
+            target_point_data = point_data  # même frame = même cible
+            previous_data = None
+            num_steps = 1
+            target_frame = 1
+
+        # --- END ADD ---
+        else:
+            with meshio.xdmf.TimeSeriesReader(xdmf_file) as reader:
+                num_steps = reader.num_steps
+
+                if frame - _previous_data_index < 0:
+                    _previous_data_index = 1
+                if frame + _target_data_index > num_steps - 1:
+                    _target_data_index = 1
+
+                if frame >= num_steps - 1 and (not self.target_same_frame):
+                    raise IndexError(
+                        f"Frame index {frame} out of bounds for trajectory {traj_index} with {num_steps} frames.")
+
+                points, cells = reader.read_points_cells()
+                time, point_data, _ = reader.read_data(frame)
+                target_frame = frame + 1
+                if self.target_same_frame: target_frame = frame
+                _, target_point_data, _ = reader.read_data(target_frame)
+
+                if self.use_previous_data:
+                    _, previous_data, _ = reader.read_data(frame - _previous_data_index)
+
         # Prepare the mesh data
         mesh = meshio.Mesh(points, cells, point_data=point_data)
 
         # Get faces or cells
         if "triangle" in mesh.cells_dict:
             cells = mesh.cells_dict["triangle"]
+        elif "line" in mesh.cells_dict:
+            cells = mesh.cells_dict["line"]    
         elif "tetra" in mesh.cells_dict:
             cells = torch.tensor(mesh.cells_dict["tetra"], dtype=torch.long)
         else:
@@ -128,18 +156,46 @@ def __getitem__(self, index: int) -> Union[Data, Tuple[Data, torch.Tensor]]:
             )
 
         # Process point data and target data
-        point_data = {
-            k: np.array(mesh.point_data[k]).astype(self.meta["features"][k]["dtype"])
-            for k in self.meta["features"]
-            if k in mesh.point_data.keys()
-        }
-
-        target_data = {
-            k: np.array(target_point_data[k]).astype(self.meta["features"][k]["dtype"])
-            for k in self.meta["features"]
-            if k in target_point_data.keys()
-            and self.meta["features"][k]["type"] == "dynamic"
-        }
+        if self.target_same_frame:
+            selected_features = ["mesh_pos", "wall_mask"]
+            point_data = {
+                k: np.array(mesh.point_data[k]).astype(self.meta["features"][k]["dtype"])
+                for k in self.meta["features"]
+                if k in mesh.point_data.keys() and k in selected_features
+            }
+
+            target_data = {
+                k: np.array(target_point_data[k]).astype(self.meta["features"][k]["dtype"])
+                for k in self.meta["features"]
+                if k in target_point_data.keys()
+                and k == "Velocity"
+            }
+        else:
+            point_data = {
+                k: np.array(mesh.point_data[k]).astype(self.meta["features"][k]["dtype"])
+                for k in self.meta["features"]
+                if k in mesh.point_data.keys()
+            }
+
+            target_data = {
+                k: np.array(target_point_data[k]).astype(self.meta["features"][k]["dtype"])
+                for k in self.meta["features"]
+                if k in target_point_data.keys()
+                and self.meta["features"][k]["type"] == "dynamic"
+            }    
+        # --- DEBUG PRINT ESSENTIEL ---
+        #print('-----DEBUG--------')
+        #def summarize_data(name, data_dict):
+        #    print(f"\n{name} summary:")
+        #    for k, v in data_dict.items():
+        #        print(f"  {k:<20} shape={v.shape}, dtype={v.dtype}")
+        #    print(f"Total {len(data_dict)} features.\n")
+
+        #summarize_data("point_data", point_data)
+        #summarize_data("target_data", target_data)
+        #print('-----------------')
+        # --- FIN DEBUG ---
+        #print(f"[DEBUG] GPU Mem Alloc: {torch.cuda.memory_allocated()/1024**2:.1f} MB, Reserved: {torch.cuda.memory_reserved()/1024**2:.1f} MB")
 
         def _reshape_array(a: dict):
             for k, v in a.items():
@@ -150,17 +206,28 @@ def _reshape_array(a: dict):
         _reshape_array(target_data)
 
         # Create graph from mesh data
-        graph = meshdata_to_graph(
-            points=points.astype(np.float32),
-            cells=cells,
-            point_data=point_data,
-            time=time,
-            target=target_data,
-            id=mesh_id,
-        )
+        if self.target_same_frame:
+            graph = meshdata_to_graph(
+                points=points.astype(np.float32),
+                cells=cells,
+                point_data=point_data,
+                time=0,
+                target=target_data,
+                id=mesh_id,
+            )
+        else:
+            graph = meshdata_to_graph(
+                points=points.astype(np.float32),
+                cells=cells,
+                point_data=point_data,
+                time=time,
+                target=target_data,
+                id=mesh_id,
+            )
+
         # TODO: add target_dt and previous_dt as features per node.
         graph.target_dt = _target_data_index * self.dt
-
+        self.use_previous_data = False #### I force use_previous_data to be false
         if self.use_previous_data:
             previous = {
                 k: np.array(previous_data[k]).astype(self.meta["features"][k]["dtype"])
@@ -174,7 +241,7 @@ def _reshape_array(a: dict):
 
         graph = graph.to(self.device)
 
-        graph = self._apply_preprocessing(graph)
+        #graph = self._apply_preprocessing(graph) #### I remove preprocess
         graph = self._apply_k_hop(graph, traj_index)
         graph = self._may_remove_edges_attr(graph)
         graph = self._add_random_edges(graph)

graphphysics/models/layers.py

@@ -5,6 +5,59 @@
 import torch.nn as nn
 from torch_geometric.nn import MessagePassing
 
+# --- Helper SpMM: force l'op en FP32 pour éviter mismatch sous BF16 ---
+import torch
+from contextlib import contextmanager
+
+# Activation checkpoint 
+from torch.utils.checkpoint import checkpoint
+
+@contextmanager
+def _no_autocast_cuda():
+    # évite que autocast BF16 ré-intercepte SpMM
+    if torch.is_autocast_enabled():
+        with torch.cuda.amp.autocast(enabled=False):
+            yield
+    else:
+        yield
+
+def spmm_fp32(adj, x):
+    """
+    Effectue Y = adj @ x en FP32 (hors autocast), puis cast Y vers x.dtype.
+    Pourquoi: DGL SpMM exige même dtype entre valeurs(A) et X; BF16 n'est pas toujours supporté.
+    """
+    x32 = x.float()
+    # Certaines versions DGL ont .astype ; sinon, le backend convertira côté C.
+    try:
+        adj32 = adj.astype(torch.float32)  # DGL SparseMatrix
+    except AttributeError:
+        adj32 = adj  # fallback: beaucoup de builds acceptent A en fp32 par défaut
+
+    with _no_autocast_cuda():
+        y32 = adj32 @ x32  # déclenche torch.ops.dgl_sparse.spmm(...)
+    return y32.to(x.dtype)
+
+# --- Helpers DGL sparse en FP32 (BF16-safe) ---
+def bsddmm_fp32(mask, q, kT):
+    q32, kT32 = q.float(), kT.float()
+    try:
+        mask32 = mask.astype(torch.float32)
+    except AttributeError:
+        mask32 = mask
+    with _no_autocast_cuda():
+        out = dglsp.bsddmm(mask32, q32, kT32)
+    return out  # SparseMatrix (valeurs fp32)
+
+def bspmm_fp32(attn, v, out_dtype):
+    v32 = v.float()
+    try:
+        attn32 = attn.astype(torch.float32)
+    except AttributeError:
+        attn32 = attn
+    with _no_autocast_cuda():
+        y32 = dglsp.bspmm(attn32, v32)
+    return y32.to(out_dtype)
+
 try:
     import dgl.sparse as dglsp
     from dgl.sparse import SparseMatrix
@@ -343,7 +396,8 @@ def scaled_query_key_softmax(
     q = q / scaling_factor
 
     if att_mask is not None and HAS_DGL_SPARSE:
-        attn = dglsp.bsddmm(att_mask, q, k.transpose(1, 0))
+        #attn = dglsp.bsddmm(att_mask, q, k.transpose(1, 0))
+        attn = bsddmm_fp32(att_mask, q, k.transpose(1, 0)).softmax()
         attn = attn.softmax()
     else:
         attn = q @ k.transpose(-2, -1)
@@ -378,7 +432,8 @@ def scaled_dot_product_attention(
 
     # Compute the output
     if att_mask is not None and HAS_DGL_SPARSE:
-        y = dglsp.bspmm(attn, v)
+        #y = dglsp.bspmm(attn, v)
+        y = bspmm_fp32(attn, v, v.dtype)
     else:
         y = attn @ v
 
@@ -520,12 +575,14 @@ def __init__(
         self.activation = activation_layer()
         self.norm1, self.norm2 = RMSNorm(output_dim), RMSNorm(output_dim)
         self.gated_mlp = build_gated_mlp(
-            in_size=output_dim, hidden_size=output_dim, out_size=output_dim
+            in_size=output_dim, hidden_size=output_dim, out_size=output_dim,
+            expansion_factor=2   # <-- change clé : 3 -> 2
         )
 
         self.use_adjacency = HAS_DGL_SPARSE
+        self.use_activation_checkpointing = False  # togglé depuis l'extérieur
 
-    def forward(
+    '''def forward(
         self, x: torch.Tensor, adj, return_attention: bool = False
     ) -> torch.Tensor:
         """
@@ -556,7 +613,38 @@ def forward(
             return x, attn
         else:
             return x
+    '''
+    def forward(self, x: torch.Tensor, adj, return_attention: bool = False) -> torch.Tensor:
+        if not self.use_adjacency:
+            adj = None
 
+        if return_attention:
+            # En validation/diag, on évite le checkpoint pour récupérer les poids d’attention
+            x_, attn = self.attention(self.norm1(x), adj, return_attention=True)
+            x = x + x_
+            x = x + self.gated_mlp(self.norm2(x))
+            return x, attn
+
+        # --- En TRAIN: checkpoint fin pour réduire les activations gardées en mémoire ---
+
+        # 1) Attention
+        def _attn_only(tensor_x):
+            # adj est capturé par la closure (non Tensor) => OK pour checkpoint
+            return self.attention(self.norm1(tensor_x), adj, return_attention=False)
+
+        # 2) MLP
+        def _mlp_only(tensor_x):
+            return self.gated_mlp(self.norm2(tensor_x))
+
+        if self.training:
+            # use_reentrant=False consomme moins de mémoire avec PyTorch ≥1.12/2.0
+            x = x + checkpoint(_attn_only, x, use_reentrant=False)
+            x = x + checkpoint(_mlp_only,  x, use_reentrant=False)
+        else:
+            x = x + self.attention(self.norm1(x), adj, return_attention=False)
+            x = x + self.gated_mlp(self.norm2(x))
+
+        return x
 
 class GraphNetBlock(MessagePassing):
     """