🌊 Reservoir PINN Simulator 🚀

💧 This 3D Reservoir PINN Simulator Streamlit application leverages Physics-Informed Neural Networks (PINNs) to model pressure dynamics within a 3D reservoir. It tackles a complex 3D partial differential equation (PDE) for single-phase fluid flow, seamlessly integrating sparse data, boundary conditions, and initial conditions as part of the loss function of the neural network architecture.

What you can do with the simulator:

• ⚙️ Input your own reservoir parameters.
• 🧠 Train new PINN models or load existing ones.
• 📊 Visualize pressure fields with static slices and dynamic animations.
• 🗣️ Query pressure values using cutting-edge LLMs (Deepseek and OpenAI)
• 📈 Track experiments and manage models with MLflow integration.
• 📖 Explore a detailed Jupyter Notebook explaining the PINN implementation.

---

Demo Video

Youtube Demo Link

✨ Features

• 📝 Input Parameters: Configure reservoir settings like porosity, permeability, viscosity, compressibility, injection rate, well position, base pressure, and pressure scale.
• 💾 Data Upload: Easily upload CSV files (x, y, z, t, p columns) or use the provided data/reservoir_data.csv. Visualize your data with a 3D scatter plot in the sidebar!
• 🏋️ Model Training: Train new PINN models with real-time loss updates displayed in the sidebar integrated with MLflow.
• 📂 Model Loading: Load previously trained models directly from MLflow runs for quick visualizations and querying.
• 🖼️ Visualization: Generate static 2D slices of the pressure field at user-defined times and z-levels.
• 🎬 Animation: Create captivating animated GIFs showcasing the evolution of the pressure field over time.
• 🤖 LLM Queries: Interact with your simulation by querying pressure at specific 3D coordinates (x, y, z, t) using DeepSeek or OpenAI LLMs.
• 📊 MLflow Tracking: Keep track of all your experiments! Parameters, metrics (losses), and artifacts (models, loss plots) are logged for comprehensive experiment management.
• 🐳 Docker Support: Run the entire application in a GPU-compatible Docker container for easy deployment and reproducibility.

---

📁 File Structure

reservoir_pinn_simulator/
├── .env.example                   # 🔑 Template for environment variables
├── .gitignore                    # 🙈 Git ignore file
├── README.md                     # 📖 This awesome project documentation
├── requirements.txt              # 📜 Python dependencies, including mlflow
├── Dockerfile                    # 🐳 Docker configuration for GPU support
├── data/
│   └── reservoir_data.csv        # 📉 Default dataset with 50 pressure points
├── notebooks/
│   └── reservoir_pinn_notebook.ipynb  # 📓 Jupyter Notebook explaining PINN
├── src/
│   ├── __init__.py               # 📦 Makes src a Python package
│   ├── api/
│   │   ├── __init__.py
│   │   └── server.py             # 🌐 FastAPI server for LLM queries
│   ├── models/
│   │   ├── __init__.py
│   │   └── pinn.py               # 🧠 PINN model definition
│   ├── utils/
│   │   ├── __init__.py
│   │   ├── plotting.py           # 📈 Plotting functions for 3D visualization
│   │   └── training.py           # 💪 PINN training logic
│   └── app.py                    # 🎈 Streamlit app with training and loading
├── assets/
│   └── reservoir_demo.gif        # 🏗️ Demo video of simulator
└── mlruns/                       #  mlflow Auto-generated by MLflow for experiment tracking

🛠️ Prerequisites

• 🐍 Python 3.12+: Required for local setup.
• 🐳 Docker: For containerized deployment (with NVIDIA Container Toolkit for GPU support).
• 💻 GPU (Optional): NVIDIA GPU with CUDA 11.8 for accelerated training.
• 🧠 Ollama: For local DeepSeek R1 1.5B LLM.
• 🖼️ Pillow: For GIF animation rendering (included in requirements.txt).
• 📊 MLflow: For experiment tracking and model persistence.
• 🐍Anaconda: For managing python environment.

---

⚙️ Local Setup

1. Clone the Repository:

   git clone https://github.com/Danny024/reservoir_pinn_simulator.git
   cd reservoir_pinn_simulator

2. Set up PYTHONPATH (Important!):

   export PYTHONPATH=$PYTHONPATH:/path/to/your/reservoir_pinn_simulator

3. Create a Conda Environment:

   conda create -n reservoir_simulator python=3.12
   conda activate reservoir_simulator

4. Install Dependencies:

   pip install -r requirements.txt

5. Configure Environment:

   • Copy .env.example to .env:

     cp .env.example .env

   • Edit .env to add your OpenAI API key (optional):

     OPENAI_API_KEY=your_openai_api_key_here

     Leave blank if you plan to use DeepSeek:

     OPENAI_API_KEY=

6. Install and Configure Ollama:

   • Install Ollama: https://ollama.ai
   • Pull DeepSeek model:

     ollama pull deepseek-r1:1.5b

   • Start Ollama server (usually runs in the background after installation, but you can manually start if needed):

     ollama serve

7. Prepare Assets:

   • Ensure assets/pinn_architecture.png exists. You can generate it using notebooks/reservoir_pinn_notebook.ipynb or provide a placeholder image.

8. Run MLflow Server:

   #mlflow server --host 0.0.0.0 --port 5000 --backend-store-uri file://$(pwd)/mlruns --default-artifact-root file://$(pwd)/artifacts
   mlfow ui

9. Run FastAPI Server (for LLM queries):

   uvicorn src.api.server:app --host 0.0.0.0 --port 8000

10. Run Streamlit App:

    streamlit run src/app.py

---

🐳 Docker Setup

1. Install Docker and NVIDIA Container Toolkit:

   • Docker (Ubuntu):

     sudo apt-get update
     sudo apt-get install -y docker.io
     sudo systemctl start docker
     sudo systemctl enable docker

   • NVIDIA Container Toolkit (for GPU support):

     curl -fsSL [https://nvidia.github.io/libnvidia-container/gpgkey](https://nvidia.github.io/libnvidia-container/gpgkey) | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
     && curl -s -L [https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list](https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list) | \
     sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
     sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
     sudo apt-get update
     sudo apt-get install -y nvidia-container-toolkit
     sudo nvidia-ctk runtime configure --runtime=docker
     sudo systemctl restart docker

2. Build Docker Image:

   docker build -t reservoir-pinn-simulator .

3. Run Docker Container:

   • With GPU:

     docker run --gpus all -p 8501:8501 -p 8000:8000 -p 5000:5000 -p 11434:11434 reservoir-pinn-simulator

   • CPU-only:

     docker run -p 8501:8501 -p 8000:8000 -p 5000:5000 -p 11434:11434 reservoir-pinn-simulator

   (The internal Ollama server within Docker typically runs on port 11434)

4. Persist MLflow Data (Optional): You can manually create directories on your host machine if they don't exist:

   mkdir -p mlruns artifacts

   Then run the container with volume mounts:

   docker run --gpus all -p 8501:8501 -p 8000:8000 -p 5000:5000 -p 11434:11434 \
   -v $(pwd)/mlruns:/app/mlruns \
   -v $(pwd)/artifacts:/app/artifacts \
   reservoir-pinn-simulator

---

🚀 Usage

1. Access the Application:

   • Streamlit GUI: http://localhost:8501
   • MLflow UI: http://localhost:5000
   • (Ollama usually runs on http://localhost:11434 if run locally, or is accessed internally within Docker)

2. ⚙️ Configure Parameters:

   • In the "Reservoir Parameters" section of the Streamlit app, set values for $\phi$ (porosity), $k$ (permeability), $\mu$ (viscosity), $c_t$ (total compressibility), $q_0$ (injection rate), well position, $p_{base}$, and $p_{scale}$.

3. 💾 Upload Data:

   • Upload a CSV file containing columns: x, y, z, t, p.
   • Alternatively, use the default data/reservoir_data.csv.
   • A 3D scatter plot of the uploaded data points will appear in the sidebar.

4. 🧠 Train or Load Model:

   • Train New Model: Navigate to "Training Parameters", specify the number of epochs, and click "Train Model". Monitor training losses and the MLflow run ID in the sidebar.
   • Load Trained Model: Select an existing MLflow run ID from the dropdown and click "Load Trained Model" to use a pre-trained model.

5. 📊 Visualize Pressure:

   • Go to the "Visualize 3D Pressure Field Slice" section.
   • Select the desired time ($t$) and z-slice level.
   • Click "Visualize Pressure Field" to see the 2D slice.

6. 🎬 Generate Animation:

   • In the "Generate 3D Pressure Field Animation" section, choose a z-slice.
   • Click "Generate Pressure Animation" to create and view an animated GIF of the pressure field evolving over time.

7. 🗣️ Query Pressure with LLM:

   • Head to the "Query Pressure with LLM" section.
   • Select either "DeepSeek" or "OpenAI" as your LLM provider.
   • Enter your query in natural language (e.g., "What is the pressure at x=0.5, y=0.5, z=0.5, t=0.5?").
   • Click "Submit Query" to get the pressure prediction.

---

🔗 MLflow Integration

• 🔍 Access Runs: Open your MLflow UI (default: http://localhost:5000) to view the 3D_Reservoir_PINN experiment. Here you can find:
  • Parameters: $p_{base}$, $p_{scale}$, learning rate, epochs, etc.
  • Metrics: Training losses (e.g., total loss, PDE loss, data loss).
  • Artifacts: Saved model files, loss plots.
• 🔄 Load Models Programmatically: Use the run ID (visible in the Streamlit sidebar or MLflow UI) to load models in your Python scripts or notebooks:

  import mlflow.pytorch
  
  run_id = "<your_run_id>" # Replace with actual run ID
  model_uri = f"runs:/{run_id}/model"
  model = mlflow.pytorch.load_model(model_uri)

• ⚖️ Compare Runs: Utilize the MLflow UI to compare different training runs based on their loss metrics and configured parameters. This is invaluable for hyperparameter tuning and model selection!

---

🤔 Troubleshooting

• ModuleNotFoundError:

  • Ensure all dependencies from requirements.txt are installed in your active environment: pip install -r requirements.txt.
  • Verify PYTHONPATH is set correctly if running locally.

• MLflow Issues:

  • Confirm the MLflow server is running and accessible (default: http://localhost:5000).
  • Check the mlruns/ directory for experiment data. Ensure it's correctly mounted if using Docker with persistent storage.

• Docker Build Fails:

  • Make sure all necessary files exist (e.g., data/reservoir_data.csv, assets/pinn_architecture.png).
  • Verify you have internet access for downloading packages and base images during the build.

• Docker Container Fails to Start/Run:

  • Check container logs for errors: docker logs <container_id_or_name>.
  • Port Conflicts: If ports 8501, 8000, 5000, or 11434 are already in use on your host, map to different host ports. For example:

    docker run --gpus all -p 8502:8501 -p 8001:8000 -p 5001:5000 -p 11435:11434 reservoir-pinn-simulator

    Then access Streamlit at http://localhost:8502, FastAPI at http://localhost:8001, etc.

• GPU Issues (Docker):

  • Ensure NVIDIA drivers are correctly installed on your host machine.
  • Verify the NVIDIA Container Toolkit is properly installed and configured. Test with:

    docker run --gpus all nvidia/cuda:11.8.0-base-ubuntu22.04 nvidia-smi

    This command should output GPU information if everything is set up correctly.

• Ollama Fails / LLM Issues:

  • Local: Ensure the Ollama server is running (ollama serve). Confirm you've pulled the model (ollama pull deepseek-coder:1.3b or similar).
  • Docker: Check container logs for messages related to Ollama starting or model pulling. The Dockerfile attempts to pull the model during the build.
  • If using OpenAI, ensure your OPENAI_API_KEY in the .env file is correct and has credit.

• GIF Animation Issues:

  • Make sure the pillow library is installed: pip show pillow. It should be included in requirements.txt.

---

📌 TO DO

1. Make use of real reservoir dataset
2. Design PDE for multiphase flow
3. Write Code base in C++ and C# for high performance and efficiency

---

📜 License

This project is licensed under the MIT License. See the LICENSE