Deep Learning for Molecular Cloud Collapse Forecasting

A spatiotemporal deep learning framework designed to forecast the structural evolution of molecular clouds in Magnetohydrodynamic (MHD) simulations. This project leverages 3D Attention U-Nets combined with ConvGRUs to learn the complex, non-linear dynamics of star-forming regions from multi-terabyte volumetric datasets.

xy projection of density channel of model input and output

Key Features

• Hybrid Architecture: Combines 3D CNNs for spatial feature extraction with ConvGRUs for temporal memory.
• HPC Optimized: Built for distributed training on Slurm clusters using PyTorch DistributedDataParallel (DDP).
• Custom ETL Pipeline: Implements a chunked HDF5 data loader to efficiently stream TB-scale datasets into GPU memory with minimal I/O overhead.
• Physics-Aware: Validated against physical constraints like mass conservation and Jeans instability criteria.

Installation

# Clone the repository  
git clone [https://github.com/AstroDnerd/StarFormationForecasting.git](https://github.com/AstroDnerd/StarFormationForecasting.git)  
cd StarFormationForecasting

# Install dependencies and the package in editable mode  
pip install -r requirements.txt  
pip install -e .

Usage

Training

To train the model on a SLURM cluster or local machine:

python scripts/train.py \
    --config configs/default_config.yaml \
   --data_path /path/to/dataset.hdf5

Inference

To run inference on a trained checkpoint:

python scripts/model_inference.py \  
    --checkpoint outputs/best_model.pth \ 
    --data_path /path/to/test_data.hdf5

Project Structure

├── configs/          # YAML configuration files for hyperparameters  
├── src/              # Core source code  
│   ├── data\_loader.py    # Custom HDF5 streaming logic  
│   ├── model.py          # 3D U-Net + ConvGRU Architecture  
│   ├── trainer.py        # DDP Training Loop  
│   └── utils.py          # Metrics and helper functions  
├── scripts/          # Execution scripts (train, inference)  
└── notebooks/        # Exploratory Data Analysis (EDA)

Architecture Details

The model utilizes an encoder-decoder structure:

1. Encoder: 3D Convolutional blocks with Attention Gates to focus on high-density regions (star-forming cores).
2. Bottleneck: A Bi-directional ConvGRU captures temporal evolution across simulation snapshots.
3. Decoder: Transposed convolutions reconstruct the predicted volumetric density field.

Author: Nikhil P. S. Bisht