How to create custom dataset. (MGN)

[yuki51192]

Hello, I have a question regarding the datasets used in Learning to Simulate, such as water-3d.

My goal is to build datasets for other types of liquids and train a MeshGraphNet (MGN) model that can predict the dynamics of those materials. I have tried different simulation engines (e.g., PhysX and MPM-based solvers), but the ground-truth data I generated is not ideal, and the trained models do not function properly.

One issue I noticed is that the value ranges of my dataset are very different from the original datasets.
For example:

The original datasets (e.g., water-3d) use values within the range 0.125–0.875

My simulated data has values ranging from 0 to 80

As a result, the predicted outputs of my model differ significantly from the ground truth.

If there are specific requirements or considerations when generating new datasets—such as:

whether the dataset needs to be normalized to a fixed value range,

recommended preprocessing steps,

or how the original ground-truth data (e.g., water-3d) was generated,

I would greatly appreciate any guidance or clarification.

Best regards.

[NiharJani2002]

Solution: Creating Custom Datasets for MeshGraphNet (MGN)

Hello @yuki51192,

I understand your challenge with creating custom datasets for MGN. The value range discrepancy you're experiencing (0.125–0.875 vs 0–80) is indeed a critical issue. Based on research into DeepMind's implementation and best practices, here's a comprehensive solution:

🔍 Understanding the Problem

The original DeepMind datasets use normalized values for numerical stability during neural network training. Data normalization to zero mean and unit variance is applied to stabilize training, with statistics computed during the training process. Your raw simulation data needs similar preprocessing.

✅ Complete Solution

1. Data Normalization Strategy

DeepMind's approach uses two normalization methods:

For Positions: Min-max scaling to the range [0.125, 0.875]

First normalize to [0, 1] based on simulation bounds
Then scale to the target range
For Velocities/Accelerations: Z-score normalization (zero mean, unit variance)

Standardize using global mean and standard deviation
This handles the high variability in velocity data

2. Python Implementation

Here's a production-ready normalization pipeline:

import numpy as np
import json
from pathlib import Path

class MGNDatasetNormalizer:
"""
Normalizes simulation data to match DeepMind's Learning to Simulate format.
Handles particle positions, velocities, and other features.
"""

def __init__(self, target_min=0.125, target_max=0.875):
    """
    Args:
        target_min: Minimum value for normalized range (default: 0.125)
        target_max: Maximum value for normalized range (default: 0.875)
    """
    self.target_min = target_min
    self.target_max = target_max
    self.target_range = target_max - target_min
    
    # Store normalization parameters
    self.position_bounds = None
    self.velocity_stats = None
    
def fit(self, trajectories):
    """
    Compute normalization parameters from training data.
    
    Args:
        trajectories: List of dicts with keys 'positions', 'velocities'
                     Each should be shape (n_timesteps, n_particles, n_dims)
    """
    all_positions = []
    all_velocities = []
    
    for traj in trajectories:
        all_positions.append(traj['positions'])
        all_velocities.append(traj['velocities'])
    
    all_positions = np.concatenate(all_positions, axis=0)
    all_velocities = np.concatenate(all_velocities, axis=0)
    
    # Get global position bounds (per dimension)
    self.position_bounds = {
        'min': np.min(all_positions, axis=(0, 1)),  # Shape: (n_dims,)
        'max': np.max(all_positions, axis=(0, 1))
    }
    
    # Get velocity statistics for standardization
    self.velocity_stats = {
        'mean': np.mean(all_velocities, axis=(0, 1)),
        'std': np.std(all_velocities, axis=(0, 1)) + 1e-8  # Avoid div by zero
    }
    
    print(f"Position bounds: {self.position_bounds}")
    print(f"Velocity stats: {self.velocity_stats}")
    
def normalize_positions(self, positions):
    """
    Normalize positions to [target_min, target_max] range.
    
    Args:
        positions: Array of shape (n_timesteps, n_particles, n_dims)
    
    Returns:
        Normalized positions in target range
    """
    if self.position_bounds is None:
        raise ValueError("Must call fit() before normalize_positions()")
    
    pos_min = self.position_bounds['min']
    pos_max = self.position_bounds['max']
    pos_range = pos_max - pos_min
    
    # Normalize to [0, 1]
    normalized = (positions - pos_min) / (pos_range + 1e-8)
    
    # Scale to [target_min, target_max]
    scaled = normalized * self.target_range + self.target_min
    
    return scaled

def normalize_velocities(self, velocities):
    """
    Standardize velocities (zero mean, unit variance).
    
    Args:
        velocities: Array of shape (n_timesteps, n_particles, n_dims)
    
    Returns:
        Standardized velocities
    """
    if self.velocity_stats is None:
        raise ValueError("Must call fit() before normalize_velocities()")
    
    mean = self.velocity_stats['mean']
    std = self.velocity_stats['std']
    
    return (velocities - mean) / std

def denormalize_positions(self, normalized_positions):
    """Inverse transform for positions."""
    if self.position_bounds is None:
        raise ValueError("Must call fit() before denormalize_positions()")
    
    pos_min = self.position_bounds['min']
    pos_max = self.position_bounds['max']
    pos_range = pos_max - pos_min
    
    # Scale from [target_min, target_max] to [0, 1]
    normalized = (normalized_positions - self.target_min) / self.target_range
    
    # Scale to original range
    return normalized * pos_range + pos_min

def denormalize_velocities(self, normalized_velocities):
    """Inverse transform for velocities."""
    if self.velocity_stats is None:
        raise ValueError("Must call fit() before denormalize_velocities()")
    
    mean = self.velocity_stats['mean']
    std = self.velocity_stats['std']
    
    return normalized_velocities * std + mean

def save_normalization_params(self, filepath):
    """Save normalization parameters to JSON file."""
    params = {
        'position_bounds': {
            'min': self.position_bounds['min'].tolist(),
            'max': self.position_bounds['max'].tolist()
        },
        'velocity_stats': {
            'mean': self.velocity_stats['mean'].tolist(),
            'std': self.velocity_stats['std'].tolist()
        },
        'target_min': self.target_min,
        'target_max': self.target_max
    }
    
    with open(filepath, 'w') as f:
        json.dump(params, f, indent=2)
    
    print(f"Saved normalization params to {filepath}")

def load_normalization_params(self, filepath):
    """Load normalization parameters from JSON file."""
    with open(filepath, 'r') as f:
        params = json.load(f)
    
    self.position_bounds = {
        'min': np.array(params['position_bounds']['min']),
        'max': np.array(params['position_bounds']['max'])
    }
    self.velocity_stats = {
        'mean': np.array(params['velocity_stats']['mean']),
        'std': np.array(params['velocity_stats']['std'])
    }
    self.target_min = params['target_min']
    self.target_max = params['target_max']
    self.target_range = self.target_max - self.target_min
    
    print(f"Loaded normalization params from {filepath}")

Example usage

def process_custom_dataset():
"""Example of how to use the normalizer with your PhysX/MPM simulation data."""

# Load your simulation data
trajectories = []
for i in range(num_trajectories):
    # Replace with your actual data loading from PhysX/MPM
    positions = load_positions_from_simulation(i)  # Shape: (timesteps, particles, 3)
    velocities = load_velocities_from_simulation(i)
    
    trajectories.append({
        'positions': positions,
        'velocities': velocities
    })

# Initialize normalizer
normalizer = MGNDatasetNormalizer(target_min=0.125, target_max=0.875)

# Fit on training data to compute normalization parameters
print("Computing normalization parameters...")
normalizer.fit(trajectories)

# Save normalization parameters (CRITICAL: use same params during inference)
normalizer.save_normalization_params('normalization_params.json')

# Normalize all trajectories
normalized_trajectories = []
for traj in trajectories:
    normalized_pos = normalizer.normalize_positions(traj['positions'])
    normalized_vel = normalizer.normalize_velocities(traj['velocities'])
    
    normalized_trajectories.append({
        'positions': normalized_pos,
        'velocities': normalized_vel
    })

return normalized_trajectories, normalizer

3. Dataset Generation Best Practices

Based on the original water-3d dataset which used SPH simulation with 14k particles over 800 steps, here are key considerations:

Simulation Setup:

Use a consistent domain size (e.g., unit cube [1.0, 1.0, 1.0])
Maintain uniform particle density across trajectories
Use consistent time step resolution (typically dt = 0.001-0.01s)
Apply proper boundary conditions (walls, periodic boundaries)
Material Properties:

Keep physical parameters within reasonable ranges
For liquids: viscosity ~0.001-0.1 Pa·s
Density should be consistent across similar materials
Ensure your solver is numerically stable

4. Critical Implementation Steps

Step 1: Generate raw simulation data with PhysX/MPM

Ensure your simulation outputs:

- Particle positions: (n_timesteps, n_particles, 3)

- Particle velocities: (n_timesteps, n_particles, 3)

- Particle types: (n_particles,) [e.g., 0=fluid, 1=boundary]

Step 2: Fit normalization parameters on training set ONLY

normalizer = MGNDatasetNormalizer()
normalizer.fit(training_trajectories) # Only use training data!
normalizer.save_normalization_params('norm_params.json')
Step 3: Apply same normalization to train/val/test sets

Apply to all datasets using the SAME normalizer

train_normalized = [normalize_trajectory(t, normalizer) for t in train_data]
val_normalized = [normalize_trajectory(t, normalizer) for t in val_data]
test_normalized = [normalize_trajectory(t, normalizer) for t in test_data]
Step 4: During inference, load the same normalization parameters

At inference time

normalizer = MGNDatasetNormalizer()
normalizer.load_normalization_params('norm_params.json')

Normalize input

normalized_input = normalizer.normalize_positions(raw_positions)

Get model prediction

prediction = model(normalized_input)

Denormalize output to physical units

physical_prediction = normalizer.denormalize_positions(prediction)

5. Debugging Checklist
✓ Verify position range after normalization: should be ≈[0.125, 0.875]
✓ Verify velocity statistics: mean ≈ 0, std ≈ 1
✓ Check denormalization: denormalize(normalize(x)) ≈ x
✓ Compare data statistics with original datasets (check metadata.json)
✓ Visualize trajectories for physical plausibility
✓ Start with small dataset (2-3 trajectories) for rapid iteration

6. Understanding Original Data Generation
The original DeepMind datasets used SPH (Smoothed Particle Hydrodynamics) for water simulations and MPM (Material Point Method) for sand and goop. Key characteristics:

Normalized spatial coordinates within defined bounds
Particle-based representation (not grid-based)
Multiple material types encoded as particle attributes
Carefully tuned simulation parameters for stability

📊 Expected Results
After proper normalization, you should see:

Original position range: [0.000, 80.000]
Normalized position range: [0.125, 0.875] ✓

Original velocity range: [-50.000, 50.000]
Normalized velocity mean: 0.001
Normalized velocity std: 0.998 ✓

Denormalization error: < 1e-6 ✓
🔗 Additional Resources

DeepMind Research Repository: https://github.com/google-deepmind/deepmind-research/tree/master/learning_to_simulate
Paper: "Learning to Simulate Complex Physics with Graph Networks" (ICML 2020)
Dataset metadata: Each dataset includes a metadata.json file with normalization statistics

⚠️ Important Notes
Save normalization parameters with your model - inference MUST use the same normalization
Never fit normalization on test data - compute statistics only from training set
Validate physical plausibility - normalized data should still show realistic dynamics
Monitor training loss - large initial losses indicate scaling issues

Disclaimer: I am not a maintainer of this repository. This solution is based on my analysis of the DeepMind Learning to Simulate codebase, related research papers, and best practices for physics-based neural network training. I hope this helps resolve your dataset generation challenges!

If you have specific questions about implementing this solution with your PhysX/MPM simulator, feel free to ask. Good luck with your research!

Best regards

NIHAR MAHESH JANI

niharmaheshjani@gmail.com

[yuki51192]

Thanks. I'll try it.
I originally extracted only the particle positions from my simulator and used the following repository to generate the TFRecords:
https://github.com/nsrawat0333/deepmind-research/tree/solve-tfrecord-generation-204

For velocities, I followed the same approach as the original DeepMind datasets:
velocity = (pos[t+1] – pos[t]) / dt, and acceleration was computed in the same way.

However, regarding normalization, in my current experiments I only compute the mean and standard deviation and write them into metadata.json, but I do not apply normalization to the actual simulation data before writing it into the TFRecords.
Could this be the reason why my training results are very poor?
(Even though my position values fall roughly within the range 0.125–0.875, the performance is still far from acceptable.)
By the way, the training data I generated does not include the velocity feature.

Another question I have is about particle uniformity.
In the official DeepMind datasets, the particles appear to be evenly spaced in every example.
But in my case, because many of my datasets originate from mesh-to-particle sampling (using particle samplers), the resulting particles do not have uniform spacing.
Do you think this will significantly affect training or inference?
Does MeshGraphNet require uniformly spaced particles to work properly?
If so, it would greatly limit its practicality in many real-world applications.

Best regards.

[NiharJani2002]

CRITICAL ISSUE #1: You Must Normalize Data BEFORE Writing TFRecords

The Problem

@yuki51192 You stated:

"In my current experiments I only compute the mean and standard deviation and write them into metadata.json, but I do not apply normalization to the actual simulation data before writing it into the TFRecords."

This is causing your training failure. The DeepMind implementation expects pre-normalized data in the TFRecords, NOT raw simulation data.

How DeepMind's System Actually Works

Based on analysis of learned_simulator.py and the codebase:

1. During Dataset Generation:

     1.1 Compute normalization statistics (mean, std) from training data
     1.2 Apply normalization to ALL data (positions, velocities)
     1.3 Write normalized data to TFRecords
     1.4 Save statistics to metadata.json
   

2. During Training:

   2.1 Read already-normalized data from TFRecords
   2.2 Apply additional online normalization for stability (zero mean, unit variance)

3. Train the model

   3.1 During Inference:

        3.1.1 Normalize inputs using saved statistics
        3.1.2 Get predictions (normalized)
        3.1.3 Denormalize outputs back to physical units
   

Why Your Current Approach Fails

Your position values "roughly fall within 0.125–0.875" by coincidence, but:

1. Velocities are not normalized: This causes massive gradient instabilities
2. Accelerations are not normalized: Training targets are in wrong scale
3. No consistent scaling: The model sees different value ranges than it expects

The paper explicitly states:

"We normalize all input and target vectors elementwise to zero mean and unit variance, using statistics computed online during training."

This means:

Data should be pre-normalized in TFRecords (z-score: zero mean, unit variance)
Additional online normalization is applied during training for further stability

CORRECTED IMPLEMENTATION

Step 1: Generate Raw Simulation Data

import numpy as np
import tensorflow as tf

# Your PhysX/MPM simulator outputs
positions = []  # Shape: (n_trajectories, n_timesteps, n_particles, 3)
particle_types = []  # Shape: (n_trajectories, n_particles)

# Extract from your simulator
for traj_idx in range(n_trajectories):
    pos = extract_particle_positions(traj_idx)  # (timesteps, particles, 3)
    types = extract_particle_types(traj_idx)    # (particles,)
    
    positions.append(pos)
    particle_types.append(types)
Step 2: Compute Velocities from Positions

def compute_velocities(positions, dt=0.01):
    """
    Compute velocities from position sequence.
    
    Args:
        positions: (n_timesteps, n_particles, n_dims)
        dt: time step size
    
    Returns:
        velocities: (n_timesteps, n_particles, n_dims)
    """
    # v[t] = (p[t+1] - p[t]) / dt
    velocities = np.zeros_like(positions)
    velocities[:-1] = (positions[1:] - positions[:-1]) / dt
    # For last timestep, use previous velocity
    velocities[-1] = velocities[-2]
    
    return velocities

# Compute for all trajectories
velocities = [compute_velocities(pos) for pos in positions]
Step 3: Compute Normalization Statistics (TRAINING SET ONLY)

def compute_normalization_stats(train_positions, train_velocities):
    """
    Compute statistics from TRAINING data only.
    
    Returns:
        stats: dict with 'mean' and 'std' for positions and velocities
    """
    # Concatenate all training trajectories
    all_positions = np.concatenate(train_positions, axis=0)  # (total_steps, particles, dims)
    all_velocities = np.concatenate(train_velocities, axis=0)
    
    # Flatten to compute global statistics
    pos_flat = all_positions.reshape(-1, all_positions.shape[-1])  # (N, dims)
    vel_flat = all_velocities.reshape(-1, all_velocities.shape[-1])
    
    stats = {
        'position': {
            'mean': np.mean(pos_flat, axis=0),  # Per-dimension mean
            'std': np.std(pos_flat, axis=0) + 1e-8  # Avoid division by zero
        },
        'velocity': {
            'mean': np.mean(vel_flat, axis=0),
            'std': np.std(vel_flat, axis=0) + 1e-8
        }
    }
    
    return stats

# Compute from training set ONLY
norm_stats = compute_normalization_stats(train_positions, train_velocities)
Step 4: Apply Normalization to ALL Data

def normalize_data(positions, velocities, stats):
    """
    Apply z-score normalization.
    
    Returns:
        Normalized positions and velocities (zero mean, unit variance)
    """
    pos_mean = stats['position']['mean']
    pos_std = stats['position']['std']
    vel_mean = stats['velocity']['mean']
    vel_std = stats['velocity']['std']
    
    # Apply normalization
    positions_norm = (positions - pos_mean) / pos_std
    velocities_norm = (velocities - vel_mean) / vel_std
    
    return positions_norm, velocities_norm

# Normalize ALL splits using SAME statistics
train_pos_norm = [normalize_data(p, v, norm_stats)[0] for p, v in zip(train_positions, train_velocities)]
train_vel_norm = [normalize_data(p, v, norm_stats)[1] for p, v in zip(train_positions, train_velocities)]

valid_pos_norm = [normalize_data(p, v, norm_stats)[0] for p, v in zip(valid_positions, valid_velocities)]
valid_vel_norm = [normalize_data(p, v, norm_stats)[1] for p, v in zip(valid_positions, valid_velocities)]

test_pos_norm = [normalize_data(p, v, norm_stats)[0] for p, v in zip(test_positions, test_velocities)]
test_vel_norm = [normalize_data(p, v, norm_stats)[1] for p, v in zip(test_positions, test_velocities)]
Step 5: Write Normalized Data to TFRecords

def write_tfrecord(filename, positions, velocities, particle_types):
    """
    Write NORMALIZED data to TFRecord.
    
    Args:
        positions: list of (n_timesteps, n_particles, n_dims) arrays - NORMALIZED
        velocities: list of (n_timesteps, n_particles, n_dims) arrays - NORMALIZED
        particle_types: list of (n_particles,) arrays
    """
    with tf.io.TFRecordWriter(filename) as writer:
        for traj_idx, (pos, vel, ptypes) in enumerate(zip(positions, velocities, particle_types)):
            
            # Create SequenceExample
            context = tf.train.Features(feature={
                'particle_type': tf.train.Feature(
                    bytes_list=tf.train.BytesList(value=[ptypes.astype(np.int32).tobytes()])
                ),
                'key': tf.train.Feature(
                    int64_list=tf.train.Int64List(value=[traj_idx])
                )
            })
            
            # Position sequence
            position_features = []
            for t in range(len(pos)):
                position_features.append(
                    tf.train.Feature(
                        bytes_list=tf.train.BytesList(
                            value=[pos[t].astype(np.float32).tobytes()]
                        )
                    )
                )
            
            feature_lists = tf.train.FeatureLists(feature_list={
                'position': tf.train.FeatureList(feature=position_features)
            })
            
            example = tf.train.SequenceExample(
                context=context,
                feature_lists=feature_lists
            )
            
            writer.write(example.SerializeToString())

# Write all splits
write_tfrecord('train.tfrecord', train_pos_norm, train_vel_norm, train_particle_types)
write_tfrecord('valid.tfrecord', valid_pos_norm, valid_vel_norm, valid_particle_types)
write_tfrecord('test.tfrecord', test_pos_norm, test_vel_norm, test_particle_types)
Step 6: Save Metadata with Normalization Statistics

import json

metadata = {
    'dim': 3,  # or 2 for 2D
    'default_connectivity_radius': 0.015,  # Adjust based on your particle spacing
    'vel_mean': norm_stats['velocity']['mean'].tolist(),
    'vel_std': norm_stats['velocity']['std'].tolist(),
    'acc_mean': [0.0] * 3,  # Accelerations typically zero-mean
    'acc_std': [1.0] * 3,   # Adjust based on your data
    'bounds': [[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]],  # Your simulation bounds
    'sequence_length': 500,  # Number of timesteps per trajectory
    'default_connectivity_radius': 0.015
}

with open('metadata.json', 'w') as f:
    json.dump(metadata, f, indent=2)

IMPORTANT: Why Velocity Features Matter

You mentioned: "the training data I generated does not include the velocity feature"

This is a major problem. The DeepMind implementation:

Computes velocities from positions: v[t] = (pos[t+1] - pos[t]) / dt
Uses velocity history (last 5 velocities) as input features
Predicts acceleration as output: a[t] = (v[t+1] - v[t]) / dt
Updates positions using Euler integration: pos[t+1] = pos[t] + v[t] * dt

Without velocity features, the model cannot:

1. Learn temporal dynamics properly
2. Distinguish between stationary and moving particles
3. Predict accelerations accurately

Solution:

Compute velocities from your position sequences as shown in Step 2 above.

ISSUE #2: Non-Uniform Particle Spacing

Your Concern:

"In my case, because many of my datasets originate from mesh-to-particle sampling (using particle samplers), the resulting particles do not have uniform spacing. Do you think this will significantly affect training or inference?"

Answer: It Will NOT Be a Dealbreaker, But Requires Consideration

Good News: GNNs are designed to handle irregular spatial structures. The architecture is permutation-invariant and uses distance-based edge connectivity, not grid-based operations.

Evidence from Research:

MeshGraphNet Paper: Uses adaptive meshes with varying resolution

"Different trajectories use different 2D triangular meshes, with edge lengths deviating throughout the mesh"

Model successfully generalizes across different mesh resolutions

X-MeshGraphNet (2024): Explicitly handles non-uniform point clouds

• Uses "uniform point cloud" but states: "Uniform distributions or geometry-aware distributions, such as those based on curvature, can be used"
• Successfully trains on automotive aerodynamics with varying point densities
• GNS Architecture: Graph connectivity is based on spatial radius
• Edges connect particles within distance R
• Number of neighbors varies per particle
• Model adapts to local density variations

Practical Considerations:

• Graph Connectivity Radius

1. Adjust default_connectivity_radius in metadata.json
2. Should be ~1.5-2.0× average particle spacing
3. Too small: particles disconnect
4. Too large: excessive edges, slow training

• Particle Density Variations If your mesh-to-particle sampling creates regions with very different densities:

1. Acceptable: 2-3× density variation (e.g., refined near boundaries)
2. Problematic: 10× density variation (e.g., sparse bulk + dense boundaries)

• Recommendations:

# Compute local particle density statistics
def analyze_particle_spacing(positions):
    from scipy.spatial import KDTree
    
    for trajectory in positions:
        for timestep in trajectory:
            tree = KDTree(timestep)
            # Find 10 nearest neighbors for each particle
            distances, _ = tree.query(timestep, k=11)  # k=11 (10 neighbors + self)
            avg_distances = distances[:, 1:].mean(axis=1)  # Exclude self
            
            print(f"Mean spacing: {avg_distances.mean():.4f}")
            print(f"Std spacing: {avg_distances.std():.4f}")
            print(f"Min spacing: {avg_distances.min():.4f}")
            print(f"Max spacing: {avg_distances.max():.4f}")
            print(f"Spacing ratio (max/min): {avg_distances.max() / avg_distances.min():.2f}")

# If ratio > 5.0, consider:
# 1. Regularizing particle sampling
# 2. Using multi-scale graph architecture
# 3. Weighting loss by local density
4. If Severe Non-Uniformity: Consider using edge features that encode relative spacing:

# In your data generation, add edge features
edge_features = {
    'relative_distance': distance / connectivity_radius,
    'local_density': num_neighbors_normalized
}

Will This Limit Practicality?

• No. Real-world applications successfully use GNNs on non-uniform data:

• SPH simulations: Particles naturally cluster in high-density regions

• MPM simulations: Material points adapt to deformation

• CFD meshes: Refined near walls, coarse in bulk flow

• The key is that your non-uniformity should be:

1. Physically meaningful (not random artifacts)
2. Consistent across trajectories (same sampling strategy)
3. Within reasonable bounds (< 5× density variation)
4. Validation Checklist

After implementing these fixes, verify:

# 1. Check normalization
print("Position mean:", normalized_positions.mean(axis=(0,1,2)))  # Should be ~[0, 0, 0]
print("Position std:", normalized_positions.std(axis=(0,1,2)))    # Should be ~[1, 1, 1]
print("Velocity mean:", normalized_velocities.mean(axis=(0,1,2))) # Should be ~[0, 0, 0]
print("Velocity std:", normalized_velocities.std(axis=(0,1,2)))   # Should be ~[1, 1, 1]

# 2. Check denormalization (round-trip test)
denormalized = positions_norm * pos_std + pos_mean
print("Denormalization error:", np.abs(denormalized - original_positions).max())  # Should be < 1e-6

# 3. Check connectivity
from scipy.spatial import KDTree
tree = KDTree(positions[0][0])  # First timestep, first trajectory
neighbors = tree.query_ball_point(positions[0][0], r=connectivity_radius)
avg_neighbors = np.mean([len(n) for n in neighbors])
print(f"Average neighbors per particle: {avg_neighbors:.1f}")  # Should be 20-50

# 4. Check particle spacing uniformity
distances = [tree.query(positions[0][0], k=11)[0][:, 1:].mean() for _ in range(10)]
print(f"Spacing coefficient of variation: {np.std(distances) / np.mean(distances):.2f}")  # Should be < 0.5

If after implementing these fixes you still have issues, provide:

1. Sample of normalized data statistics
2. Training loss curve (first 10k steps)
3. Particle spacing analysis results
4. A single trajectory visualization

Good luck! This should solve your training problems.

Best regards

NIHAR MAHESH JANI

niharmaheshjani@gmail.com

Disclaimer: I am not a maintainer of this repository. This solution is based on my analysis of the DeepMind Learning to Simulate codebase, related research papers, and best practices for physics-based neural network training. I hope this helps resolve your dataset generation challenges!

[yuki51192]

Here is my situation:
In my dataset(https://drive.google.com/drive/folders/17XRJV6h_3bLmTL2mf68bmTDjvrWa_CqC?usp=drive_link), the simulation domain is already defined to be within the range 0.125–0.875, which matches the description in DeepMind’s metadata.json. Given this, I would like to confirm whether I still need to perform an additional normalization step on the data before writing it into TFRecords.

import numpy as np
data = np.load("trajectory_0000.npz")
positions = data["positions"] # shape [302, num_particles, 3]

This dataset was generated using Houdini, so the simulation algorithm is MPM.
The connectivity radius was determined using Houdini’s particle sampling parameter, with a particle separation value of 0.3.
Should this radius instead be computed using a specific formula or method? Or is it acceptable to directly use the particle separation from the sampling process?

Additionally, I noticed that in the TFRecord datasets provided by DeepMind, although the metadata.json specifies a bound of [0.125, 0.875], many particles in the actual TFRecord files exceed this boundary. Do you have any thoughts on why this happens or whether this is expected behavior?

I also tested computing the mean and standard deviation directly from the DeepMind dataset after loading the TFRecords. However, the values I obtained do not result in a mean of 0 and a standard deviation of 1.
Should the mean and standard deviation be computed separately for x, y, and z, or is there another normalization procedure that I should follow?

Best regards