Improve README documentation

Graph_Representation_Learning_Rushil_Singha/README.md

@@ -1,55 +1,104 @@
 # JetNet Graph Diffusion Model
 
-A PyTorch/PyTorch-Geometric implementation of a **graph-based diffusion model** for generating realistic jets from the [JetNet dataset](https://huggingface.co/datasets/jetnet).  
+## Overview
+This project implements a graph-based diffusion model using PyTorch and PyTorch Geometric to generate realistic particle jets from the JetNet dataset.
 
-This project builds **k-nearest neighbor (kNN) jet graphs**, learns **Chebyshev GCN (ChebNet) embeddings**, trains a **diffusion model in latent space**, and decodes back into particle-level jets.
+The model constructs k-Nearest Neighbor (kNN) graphs from particle data, encodes them using a Graph Neural Network (Chebyshev GCN), applies a diffusion process in latent space, and reconstructs the jets using a decoder network.
 
 ---
 
-## 🚀 Features
-- kNN graph construction from jet particle clouds  
-- Graph encoder using **Chebyshev GCN** (`SimpleChebNet`)  
-- Latent **diffusion process** with denoising MLP  
-- Jet particle **decoder** network  
-- Evaluation with **KL divergence** & **Wasserstein distance**  
-- Visualization utilities for jet properties  
+## Features
+- kNN graph construction from particle data  
+- Graph encoding using Chebyshev Graph Convolution (ChebNet)  
+- Diffusion model in latent space  
+- Decoder for reconstructing particle-level jets  
+- Evaluation using KL Divergence and Wasserstein Distance  
+- Visualization of results  
 
 ---
 
-## ⚙️ Installation
+## Project Structure
+Graph_Representation_Learning_Rushil_Singha/
+├── code.py
+├── requirements.txt
+├── results/
+└── README.md
 
-Clone the repo and install dependencies:
+---
+
+## Installation
 
-```bash
-git clone https://github.com/your-username/jetnet-graph-diffusion.git
-cd jetnet-graph-diffusion
+### Step 1: Clone the repository
+git clone https://github.com/ML4SCI/GENIE.git
 
+cd GENIE/Graph_Representation_Learning_Rushil_Singha
+
+### Step 2: Install dependencies
 pip install -r requirements.txt
 
-requirements.txt
 
-numpy==1.24.3
-torch==2.0.0
-torch-geometric
-torch-scatter
-torch-sparse
-torch-cluster
-networkx
-scikit-learn
-jetnet
-```
-# This script:
+---
 
-->Encodes jets into latent space
+## How to Run
 
-->Runs diffusion training
+python code.py
 
-->Decodes jets back into particle space
 
-->Logs evaluation metrics
+---
+
+## Workflow
 
-->Saves visualizations to results/
+The script performs the following steps:
 
+1. Loads the JetNet dataset  
+2. Converts particle data into kNN graphs  
+3. Encodes graphs into latent representations  
+4. Applies diffusion process  
+5. Trains a denoising model  
+6. Decodes latent vectors into particle-level jets  
+7. Evaluates results using statistical metrics  
+8. Saves output visualizations  
 
+---
+
+## Output
+
+After running the code, the following outputs are generated:
+
+- Generated jet samples  
+- Evaluation metrics (KL Divergence, Wasserstein Distance)  
+- Visualization plots  
+
+All outputs are saved in the results/ directory.
+
+---
 
+## Requirements
+- Python 3.x  
+- PyTorch  
+- PyTorch Geometric  
+- NumPy  
+- NetworkX  
+- Scikit-learn  
+- JetNet  
+
+Install all dependencies using:
+
+pip install -r requirements.txt
+
+
+---
+
+## Notes
+- GPU is recommended for faster training  
+- Make sure PyTorch Geometric dependencies are installed correctly  
+
+---
+
+## Contributing
+Contributions are welcome. You can improve documentation, optimize code, or add new features.
+
+---
 
+## Acknowledgement
+This project is part of the ML4SCI GENIE repository and focuses on applying machine learning techniques in particle physics.

Graph_Representation_Learning_Rushil_Singha/code.py

@@ -1,4 +1,25 @@
-
+"""
+JetNet Graph Diffusion Model Pipeline
+
+This script implements:
+- Graph construction using kNN
+- ChebNet-based GNN encoding
+- Diffusion model in latent space
+- Particle reconstruction using decoder
+
+Workflow:
+1. Load dataset
+2. Build graphs
+3. Train autoencoder
+4. Train diffusion model
+5. Generate new jets
+6. Evaluate and visualize results
+
+Contribution:
+- Added documentation
+- Improved readability
+- Structured code sections
+"""
 import warnings
 warnings.filterwarnings('ignore', category=UserWarning)