Add FIGConvNet to crash example

[Alexey-Kamenev]

PhysicsNeMo Pull Request

Description

• Added FIGConvNet to crash example.
• Fixed a few of the small issues.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[greptile-apps]

Greptile Overview

Greptile Summary

This PR adds FIGConvUNet model support to the crash simulation example by implementing a time-conditional rollout wrapper that predicts structural deformation at each timestep independently.

Key Changes:

• Added FIGConvUNetTimeConditionalRollout class in rollout.py:469 that wraps FIGConvUNet and implements time-conditional prediction (predicts displacement from initial position at each timestep based on normalized time)
• Created configuration file with appropriate hyperparameters for factorized implicit grids, spatial domain bounds, and KNN neighbor search
• Added model validation in train.py:66 to ensure FIGConvUNet uses point cloud datapipe
• Fixed loss tracking in train.py:230 to use .detach().item() preventing unnecessary gradient graph retention
• Cleaned up unused code: removed _normalize_thickness_tensor method and tqdm import

The implementation follows the same architectural pattern as TransolverTimeConditionalRollout and integrates cleanly with the existing crash simulation infrastructure.

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The implementation follows established patterns in the codebase (mirrors TransolverTimeConditionalRollout), includes proper validation checks, fixes minor issues (loss.detach()), and only adds new functionality without modifying existing behavior
• No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
examples/structural_mechanics/crash/conf/model/figconvunet_time_conditional.yaml	5/5	New configuration file for FIGConvUNet model with time-conditional rollout, properly configured with spatial domain and grid resolutions
examples/structural_mechanics/crash/rollout.py	5/5	Added FIGConvUNetTimeConditionalRollout class following the same time-conditional pattern as TransolverTimeConditionalRollout
examples/structural_mechanics/crash/train.py	5/5	Added validation for FIGConvUNet requiring point cloud datapipe, fixed loss tracking with .detach(), removed unused tqdm import

Sequence Diagram

sequenceDiagram
    participant Trainer
    participant Dataloader
    participant SimSample
    participant FIGConvUNetTimeConditionalRollout
    participant FIGConvUNet
    
    Trainer->>Dataloader: Load batch
    Dataloader->>SimSample: Create sample with node_features and node_target
    Note over SimSample: node_features = {coords: [N,3], features: [N,F]}
    Trainer->>FIGConvUNetTimeConditionalRollout: forward(sample, data_stats)
    
    loop For each timestep t in [0, 1]
        FIGConvUNetTimeConditionalRollout->>FIGConvUNetTimeConditionalRollout: Prepare vertices [1,N,3] from initial coords
        FIGConvUNetTimeConditionalRollout->>FIGConvUNetTimeConditionalRollout: Concat features + time → [1,N,F+1]
        FIGConvUNetTimeConditionalRollout->>FIGConvUNet: forward(vertices, features_t)
        FIGConvUNet-->>FIGConvUNetTimeConditionalRollout: Return displacement [N,3]
        FIGConvUNetTimeConditionalRollout->>FIGConvUNetTimeConditionalRollout: Add displacement to initial position
        FIGConvUNetTimeConditionalRollout->>FIGConvUNetTimeConditionalRollout: Append predicted position to outputs
    end
    
    FIGConvUNetTimeConditionalRollout-->>Trainer: Return stacked predictions [T,N,3]
    Trainer->>Trainer: Compute loss against node_target

Loading

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[mnabian]

LGTM!

[Alexey-Kamenev]

/blossom-ci

[greptile-apps]

Greptile Overview

Greptile Summary

Added FIGConvUNet support to the crash simulation example with time-conditional rollout prediction. Also included several small fixes to improve code quality.

Key Changes

• Added new FIGConvUNetTimeConditionalRollout class in rollout.py:411-469 following the same pattern as existing Transolver and MeshGraphNet implementations
• Created configuration file figconvunet_time_conditional.yaml with architecture parameters (factorized implicit grids, KNN neighbor search, position encoding)
• Added validation in train.py:66-70 to ensure FIGConvUNet models use point-cloud datapipes
• Improved loss tracking in train.py:230 by calling detach() before item() to prevent unnecessary gradient tracking
• Removed unused _normalize_thickness_tensor method from datapipe.py:353-358
• Fixed f-string formatting in error logging (vtp_reader.py:172)
• Removed unused tqdm import from train.py

Implementation Details

The FIGConvUNet implementation uses time-conditional rollout where each timestep is predicted independently based on:

• Initial coordinates (x)
• Features (thickness) concatenated with normalized time (features_t)
• Predicted displacement offsets are added to initial position

The configuration specifies in_channels: 2 (thickness + time) but the actual implementation concatenates time to features at runtime, which correctly produces F+1 input channels where F is the feature dimension.

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• All changes are well-structured and follow existing patterns. The new FIGConvUNet integration mirrors the existing Transolver and MeshGraphNet implementations. Small bug fixes (detach() call, f-string formatting, unused code removal) improve code quality. No logical errors, security issues, or breaking changes detected.
• No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
examples/structural_mechanics/crash/conf/model/figconvunet_time_conditional.yaml	5/5	New FIGConvUNet configuration for crash simulation - properly structured with all necessary parameters
examples/structural_mechanics/crash/rollout.py	5/5	Added FIGConvUNetTimeConditionalRollout class - clean implementation following existing patterns
examples/structural_mechanics/crash/train.py	5/5	Added FIGConvUNet validation, removed unused import, improved loss tracking with detach()

Sequence Diagram

sequenceDiagram
    participant Trainer
    participant FIGConvUNetTimeConditionalRollout
    participant FIGConvUNet
    participant SimSample
    
    Trainer->>FIGConvUNetTimeConditionalRollout: forward(sample, data_stats)
    FIGConvUNetTimeConditionalRollout->>SimSample: Get coords [N, 3]
    FIGConvUNetTimeConditionalRollout->>SimSample: Get features [N, F]
    
    loop For each timestep in time_seq
        FIGConvUNetTimeConditionalRollout->>FIGConvUNetTimeConditionalRollout: Create time_expanded [N, 1]
        FIGConvUNetTimeConditionalRollout->>FIGConvUNetTimeConditionalRollout: Concatenate features + time → features_t [N, F+1]
        FIGConvUNetTimeConditionalRollout->>FIGConvUNetTimeConditionalRollout: Add batch dim to vertices and features_t
        
        alt Training mode
            FIGConvUNetTimeConditionalRollout->>FIGConvUNet: checkpoint(forward(vertices, features_t))
        else Inference mode
            FIGConvUNetTimeConditionalRollout->>FIGConvUNet: forward(vertices, features_t)
        end
        
        FIGConvUNet-->>FIGConvUNetTimeConditionalRollout: displacement offset [N, 3]
        FIGConvUNetTimeConditionalRollout->>FIGConvUNetTimeConditionalRollout: y_t = x + outf
        FIGConvUNetTimeConditionalRollout->>FIGConvUNetTimeConditionalRollout: Append y_t to outputs
    end
    
    FIGConvUNetTimeConditionalRollout-->>Trainer: Stack outputs [T, N, 3]

Loading

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[Alexey-Kamenev]

/blossom-ci

[Alexey-Kamenev]

/blossom-ci

[Alexey-Kamenev]

/blossom-ci

[NickGeneva]

/blossom-ci