benchmark_controllers.py

import numpy as np
import matplotlib.pyplot as plt
from system_id import WellSimulator
from mpc_controller import ChokeMPCController

class PIDController:
    """Standard Industrial PID Controller with Anti-Windup"""
    def __init__(self, Kp=0.8, Ki=0.15, Kd=0.05, max_ramp=5.0):
        self.Kp = Kp
        self.Ki = Ki
        self.Kd = Kd
        self.max_ramp = max_ramp
        self.integral = 0.0
        self.prev_error = 0.0

    def compute(self, current_Q, target_Q, u_prev):
        error = target_Q - current_Q
        self.integral += error
        self.integral = np.clip(self.integral, -50.0, 50.0)
        derivative = error - self.prev_error
        
        du = self.Kp * error + self.Ki * self.integral + self.Kd * derivative
        du = np.clip(du, -self.max_ramp, self.max_ramp)
        self.prev_error = error
        
        u_next = np.clip(u_prev + du, 0.0, 100.0)
        return float(u_next)

class RuleBasedController:
    """Standard Rule-Based Heuristic Controller"""
    def __init__(self, step_size=2.0, max_ramp=5.0):
        self.step_size = step_size
        self.max_ramp = max_ramp

    def compute(self, current_Q, target_Q, u_prev, current_WHP, whp_min=215.0):
        if current_WHP <= whp_min + 2.0:
            # Emergency close choke if near pressure limit
            du = -self.max_ramp
        elif current_Q < target_Q - 2.0:
            du = self.step_size
        elif current_Q > target_Q + 2.0:
            du = -self.step_size
        else:
            du = 0.0
            
        du = np.clip(du, -self.max_ramp, self.max_ramp)
        return float(np.clip(u_prev + du, 0.0, 100.0))

def run_benchmark_comparison():
    sim_mpc = WellSimulator()
    sim_pid = WellSimulator()
    sim_rule = WellSimulator()

    mpc = ChokeMPCController(simulator=sim_mpc)
    pid = PIDController()
    rule = RuleBasedController()

    target_profile = lambda t: 100.0 if t < 15 else 145.0
    total_hours = 40

    hist_mpc, hist_pid, hist_rule = [], [], []
    u_mpc, u_pid, u_rule = 30.0, 30.0, 30.0

    state_mpc = {'Q': 93.0, 'WHP': 268.0, 'FLP': 188.0, 'BHP': 3130.0, 'WHT': 110.0, 'AP': 500.0}
    state_pid = state_mpc.copy()
    state_rule = state_mpc.copy()

    for t in range(total_hours):
        target_q = target_profile(t)

        # MPC Action
        u_mpc = mpc.compute_choke_action(state_mpc, u_mpc, target_q, current_time=t)
        sim_mpc.current_state = state_mpc.copy()
        sim_mpc.current_time = t
        q_m, whp_m, flp_m, bhp_m = sim_mpc.step(u_mpc)
        state_mpc = sim_mpc.current_state.copy()
        hist_mpc.append({'Q': q_m, 'WHP': whp_m, 'u': u_mpc})

        # PID Action
        u_pid = pid.compute(state_pid['Q'], target_q, u_pid)
        sim_pid.current_state = state_pid.copy()
        sim_pid.current_time = t
        q_p, whp_p, flp_p, bhp_p = sim_pid.step(u_pid)
        state_pid = sim_pid.current_state.copy()
        hist_pid.append({'Q': q_p, 'WHP': whp_p, 'u': u_pid})

        # Rule-Based Action
        u_rule = rule.compute(state_rule['Q'], target_q, u_rule, state_rule['WHP'])
        sim_rule.current_state = state_rule.copy()
        sim_rule.current_time = t
        q_r, whp_r, flp_r, bhp_r = sim_rule.step(u_rule)
        state_rule = sim_rule.current_state.copy()
        hist_rule.append({'Q': q_r, 'WHP': whp_r, 'u': u_rule})

    # Plot Comparison
    fig, axes = plt.subplots(3, 1, figsize=(10, 9), sharex=True)
    fig.suptitle('Benchmark Comparison: PID vs Rule-Based vs Autonomous MPC', fontsize=15, fontweight='bold')

    time_vec = np.arange(total_hours)
    target_vec = [target_profile(t) for t in time_vec]

    # Subplot 1: Flow Rate Tracking
    axes[0].plot(time_vec, target_vec, 'k--', linewidth=2, label='Target Rate')
    axes[0].plot(time_vec, [h['Q'] for h in hist_mpc], 'b-', linewidth=2.5, label='Autonomous MPC (Proposed)')
    axes[0].plot(time_vec, [h['Q'] for h in hist_pid], 'r:', linewidth=2, label='Classical PID')
    axes[0].plot(time_vec, [h['Q'] for h in hist_rule], 'g--', linewidth=2, label='Rule-Based Heuristic')
    axes[0].set_ylabel('Oil Rate (bbl/hr)')
    axes[0].legend(loc='lower right')
    axes[0].grid(True)

    # Subplot 2: Choke Movement (Smoothness vs Choke Hunting)
    axes[1].plot(time_vec, [h['u'] for h in hist_mpc], 'b-', linewidth=2.5, label='Autonomous MPC Choke %')
    axes[1].plot(time_vec, [h['u'] for h in hist_pid], 'r:', linewidth=2, label='PID Choke % (Oscillatory)')
    axes[1].plot(time_vec, [h['u'] for h in hist_rule], 'g--', linewidth=2, label='Rule-Based Choke %')
    axes[1].set_ylabel('Choke Position (%)')
    axes[1].legend(loc='lower right')
    axes[1].grid(True)

    # Subplot 3: Wellhead Pressure Safety Envelope
    axes[2].plot(time_vec, [h['WHP'] for h in hist_mpc], 'b-', linewidth=2.5, label='MPC WHP (Safe Enforced)')
    axes[2].plot(time_vec, [h['WHP'] for h in hist_pid], 'r:', linewidth=2, label='PID WHP')
    axes[2].plot(time_vec, [h['WHP'] for h in hist_rule], 'g--', linewidth=2, label='Rule-Based WHP')
    axes[2].axhline(y=215, color='r', linestyle='--', label='Min WHP Limit (215 psi)')
    axes[2].set_ylabel('WHP (psi)')
    axes[2].set_xlabel('Time (Hours)')
    axes[2].legend(loc='lower right')
    axes[2].grid(True)

    plt.tight_layout()
    plt.savefig('benchmark_comparison.png', dpi=300)
    plt.savefig('C:/Users/deshp/.gemini/antigravity/brain/095490e2-8d68-4a4e-9adc-311edfe18d8f/benchmark_comparison.png', dpi=300)
    plt.close()
    print("Saved benchmark comparison plot: benchmark_comparison.png")

if __name__ == "__main__":
    run_benchmark_comparison()