An advanced AI system combining Graph Neural Networks and Reinforcement Learning for predictive maintenance, anomaly detection, and intelligent flow optimization in wastewater networks.
- Overview
- Key Features
- Architecture
- Project Structure
- Installation
- Usage
- Results
- Technologies
- Documentation
- Contributing
- License
This project tackles the critical challenge of wastewater network management using state-of-the-art AI techniques. By leveraging Graph Neural Networks (GNN) for anomaly detection and Proximal Policy Optimization (PPO) for dynamic flow control, we provide a comprehensive solution for:
- Anomaly Detection: Identifying potential bottlenecks and infrastructure issues
- Flow Optimization: Minimizing node overload while maintaining system efficiency
- Predictive Maintenance: Early warning system for network components
- Real-time Monitoring: Continuous analysis of 17,706 time-step sequences
Municipal wastewater networks face challenges including:
- Unexpected capacity overloads
- Inefficient flow distribution
- Reactive rather than predictive maintenance
- Limited real-time optimization capabilities
We developed a dual-AI approach:
- Unsupervised Learning (Graph Autoencoder) detects anomalous behavior patterns
- Reinforcement Learning (PPO) learns optimal flow regulation policies
- Input: 23-node network with 2 features per node (flow rate + depth)
- Architecture: GCN-based encoder-decoder with 8D latent space
- Output: Anomaly scores for each network node
- Training: 200 epochs with edge reconstruction loss
- Algorithm: Proximal Policy Optimization (PPO)
- Environment: Custom OpenAI Gym environment
- State Space: Flow rate and depth for all nodes
- Action Space: Continuous flow regulation coefficients [0, 1]
- Objective: Minimize capacity violations while maintaining throughput
- Network topology with anomaly heatmaps
- Flow comparison plots (before/after optimization)
- Node overload frequency analysis
- Stress reduction metrics
┌─────────────────────────────────────────────────────────────┐
│ DATA COLLECTION │
│ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │
│ │ Flow Rate │ │ Flow Depth │ │ Network │ │
│ │ Data │ │ Data │ │ Topology │ │
│ │ (17,706 ×23) │ │ (17,706 ×23) │ │ (23N, 22E) │ │
│ └──────┬───────┘ └──────┬───────┘ └──────┬───────┘ │
└─────────┼──────────────────┼──────────────────┼─────────────┘
│ │ │
└──────────────────┴──────────────────┘
│
┌─────────────────▼─────────────────┐
│ DATA PREPROCESSING │
│ • Normalization (MinMaxScaler) │
│ • NaN handling │
│ • Feature stacking │
│ • Graph construction │
└─────────────────┬─────────────────┘
│
┌─────────────────┴─────────────────┐
│ │
┌─────▼─────┐ ┌───────▼────────┐
│ GNN │ │ RL │
│ PIPELINE │ │ PIPELINE │
└─────┬─────┘ └───────┬────────┘
│ │
┌─────▼──────────┐ ┌─────────▼────────┐
│ Graph │ │ Environment │
│ Autoencoder │ │ (FlowEnv) │
│ │ │ │
│ • GCN Encoder │ │ State: [46] │
│ • Latent: 8D │ │ Action: [23] │
│ • Decoder │ │ Reward: -penalty │
└─────┬──────────┘ └─────────┬────────┘
│ │
┌─────▼──────────┐ ┌─────────▼────────┐
│ Anomaly │ │ PPO Agent │
│ Detection │ │ │
│ │ │ • Policy Net │
│ • Score nodes │ │ • Value Net │
│ • Rank by │ │ • 100K steps │
│ deviation │ │ │
└─────┬──────────┘ └─────────┬────────┘
│ │
└──────────────┬─────────────────────┘
│
┌──────────────▼──────────────┐
│ VISUALIZATION │
│ │
│ • Anomaly heatmaps │
│ • Flow comparisons │
│ • Optimization metrics │
│ • Network topology │
└─────────────────────────────┘
Input (23 nodes × 2 features)
↓
GCN Layer 1 (2 → 16)
↓
ReLU
↓
GCN Layer 2 (16 → 8)
↓
Latent Space (8D)
↓
Edge Decoder
↓
Reconstruction Loss
State (46 features)
↓
┌────┴────┐
↓ ↓
Policy Net Value Net
(MLP) (MLP)
↓ ↓
Actions State Value
(23) (1)
wastewater-predictive-maintenance/
│
├── wastewater-predictive-maintenance.ipynb # Main implementation notebook
│
├── datasets/ # Data files
│ ├── Flow_depth_v3.csv # Depth measurements (4.99 MB)
│ ├── Flow_rate_v3.csv # Flow rate data (4.94 MB)
│ ├── WW01_edge.csv # Network edges (22 pipes)
│ └── WW01_node.csv # Network nodes (23 junctions)
│
├── figures/ # Generated visualizations
│ ├── autoencoder_node_anomaly_scores.png
│ ├── flow_comparison_plot.png
│ ├── node_overload_frequency.png
│ └── nodestress_reduction_after_PPO.png
│
├── .gitignore # Git ignore file
├── LICENSE # MIT License
├── README.md # This file
└── requirements.txt # Python dependencies
- Python 3.11+
- CUDA 12.4+ (for GPU acceleration)
- 8GB+ RAM recommended
- Clone the repository
git clone https://github.com/yourusername/wastewater-predictive-maintenance.git
cd wastewater-predictive-maintenance- Create virtual environment
python -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate- Install dependencies
pip install -r requirements.txt- Install PyTorch Geometric
pip install torch-scatter torch-sparse torch-geometric torch-cluster \
-f https://data.pyg.org/whl/torch-2.6.0+cu124.html- Run the notebook
jupyter notebook wastewater-predictive-maintenance.ipynb# Load and preprocess data
depth_df = pd.read_csv('datasets/Flow_depth_v3.csv', header=[0, 1, 2])
rate_df = pd.read_csv('datasets/Flow_rate_v3.csv', header=[0, 1, 2])
# Build graph
data = Data(x=node_features, edge_index=edge_index)
# Train model
model = GAE(GCNEncoder(in_channels=2, out_channels=8))
for epoch in range(200):
loss = train()
# Get anomaly scores
z = model.encode(data.x, data.edge_index)
anomaly_scores = compute_anomaly_scores(z)# Create environment
env = FlowEnv(node_ids=valid_nodes, flow_tensor=node_features)
# Train PPO agent
model = PPO("MlpPolicy", env, learning_rate=0.0003)
model.learn(total_timesteps=100000)
# Evaluate policy
obs, _ = env.reset()
for _ in range(1000):
action, _ = model.predict(obs)
obs, reward, done, truncated, info = env.step(action)| Node ID | Anomaly Score | Interpretation |
|---|---|---|
| 92090090 | 2.297 | Critical - Highest anomaly |
| 92090042 | 1.795 | High - Requires attention |
| 92090100 | 1.684 | High - Monitor closely |
| OF-1 | 1.663 | High - Outlet irregularity |
| 92100230 | 1.103 | Moderate - Watch for trends |
Figure 1: Network topology with anomaly scores. Red intensity indicates anomaly severity.
Figure 2: Flow patterns before and after PPO optimization.
Graph Autoencoder:
- Training Loss: 1.35 → 1.13 (200 epochs)
- Convergence: ~180 epochs
- Inference Time: <100ms per evaluation
PPO Optimization:
- Training Steps: 100,000
- Episode Length: 17,706 steps
- Capacity Violations: Reduced by ~40%
Figure 3: Node stress metrics before and after optimization.
Figure 4: Frequency distribution of node overload events.
- PyTorch (2.6.0) - Deep learning framework
- PyTorch Geometric - Graph neural networks
- Stable-Baselines3 (2.6.0) - RL algorithms
- Gymnasium (1.1.1) - RL environments
- Pandas (2.2.2) - Data manipulation
- NumPy (2.0.2) - Numerical computing
- Scikit-learn - Data preprocessing
- Matplotlib (3.10.0) - Plotting and visualization
- CUDA 12.4 - GPU acceleration
- cuDNN 9.1.0 - Deep learning optimization
This project implements concepts from:
Detailed project documentation including proposal, final report, and data description were part of the academic submission.
- GPU: NVIDIA GPU with CUDA 12.4 support
- Memory: 16GB+ recommended
- Storage: 2GB for datasets and models
Graph Autoencoder:
learning_rate: 0.01
hidden_dim: 16
latent_dim: 8
epochs: 200
optimizer: AdamPPO:
learning_rate: 0.0003
n_steps: 2048
batch_size: 64
n_epochs: 10
gamma: 0.99
gae_lambda: 0.95
clip_range: 0.2This project was developed as part of an AI course at Texas A&M University - Corpus Christi (TAMUCC). It demonstrates the practical application of machine learning techniques to civil infrastructure challenges.
Team: Group 1
Course: AI Applications
Institution: TAMUCC
Contributions are welcome! Please feel free to submit a Pull Request. For major changes:
- Fork the repository
- Create your feature branch (
git checkout -b feature/AmazingFeature) - Commit your changes (
git commit -m 'Add some AmazingFeature') - Push to the branch (
git push origin feature/AmazingFeature) - Open a Pull Request
This project is licensed under the MIT License - see the LICENSE file for details.
- Texas A&M University - Corpus Christi
- Course instructors and advisors
- Open-source community for excellent ML frameworks
For questions or collaboration opportunities, please open an issue or reach out through GitHub.
Built with ❤️ using PyTorch and Reinforcement Learning
⭐ Star this repo if you find it helpful!