Update main.py

main.py

@@ -10,7 +10,7 @@ from pydantic import BaseModel
 from contextlib import asynccontextmanager
 
 # ==========================================
-# 1. MODEL ARCHITECTURES
+# 1. SHARED MODEL ARCHITECTURES
 # ==========================================
 class Mish(nn.Module):
     def forward(self, x): return x * torch.tanh(nn.functional.softplus(x))
@@ -23,11 +23,11 @@ class FourierFeatureMapping(nn.Module):
         proj = 2 * np.pi * (x @ self.B)
         return torch.cat([torch.sin(proj), torch.cos(proj)], dim=-1)
 
-# --- Voltage Model (Grid) ---
+# --- A. Voltage Model (Grid) ---
+# Input: 7 features (P_load, Q_load, etc.) -> Output: 2 (V_mag, V_angle)
 class VoltagePINN(nn.Module):
     def __init__(self):
         super().__init__()
-        # Script 1: input_dim=7, mapping_size=32
         self.fourier = FourierFeatureMapping(input_dim=7, mapping_size=32)
         self.network = nn.Sequential(
             nn.Linear(64, 256), nn.LayerNorm(256), Mish(),
@@ -35,81 +35,115 @@ class VoltagePINN(nn.Module):
             nn.Linear(128, 64), nn.LayerNorm(64), Mish(),
             nn.Linear(64, 2)
         )
-        self.v_bias = nn.Parameter(torch.zeros(1))
-        self.raw_G = nn.Parameter(torch.tensor(0.0))
-        self.raw_B = nn.Parameter(torch.tensor(0.0))
+    def forward(self, x): return self.network(self.fourier(x))
 
-    def forward(self, x):
-        return self.network(self.fourier(x))
-
-# --- Battery Model (Storage) ---
+# --- B. Battery Model (Storage) ---
+# Input: 5 features (Time, I, V, P, SoC_prev) -> Output: 3 (V, Temp, Ah)
 class BatteryPINN(nn.Module):
     def __init__(self):
         super().__init__()
-        # Script 2: input_dim=5, mapping_size=12
         self.fourier = FourierFeatureMapping(input_dim=5, mapping_size=12)
         self.network = nn.Sequential(
             nn.Linear(24, 64), Mish(),
             nn.Linear(64, 64), Mish(),
             nn.Linear(64, 3) 
         )
+    def forward(self, x): return self.network(self.fourier(x))
 
-    def forward(self, x):
-        return self.network(self.fourier(x))
+# --- C. Frequency Model (Stability) ---
+# Input: 4 features (Load, Wind, NetLoad, Imbalance) -> Output: 2 (Freq, ROCOF)
+class FrequencyPINN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fourier = FourierFeatureMapping(input_dim=4, mapping_size=32)
+        self.network = nn.Sequential(
+            nn.Linear(64, 128), nn.LayerNorm(128), Mish(),
+            nn.Linear(128, 128), nn.LayerNorm(128), Mish(),
+            nn.Linear(128, 2)
+        )
+    def forward(self, x): return self.network(self.fourier(x))
 
 # ==========================================
-# 2. PHYSICS ENGINE (OCV Curve)
+# 2. PHYSICS ENGINES (The "Laws")
 # ==========================================
-def get_physics_soc(voltage):
-    # Standard Li-ion OCV Curve (NMC Chemistry)
-    # Voltage Points: [3.0, 3.2, 3.4, 3.6, 3.7, 3.8, 3.9, 4.0, 4.1, 4.2]
+
+def get_battery_physics_soc(voltage):
+    """
+    Returns SoC % based on Li-ion Open Circuit Voltage (OCV) curve.
+    Acts as the ground truth to prevent AI drift.
+    """
     v_points = [2.8, 3.0, 3.2, 3.4, 3.55, 3.65, 3.75, 3.85, 3.95, 4.1, 4.2, 4.3]
     soc_points = [0,  0,   5,   15,  35,   50,   65,   75,   85,   92,  100, 100]
     return np.interp(voltage, v_points, soc_points)
 
+def get_frequency_physics(data):
+    """
+    Returns Baseline Frequency & ROCOF using the Swing Equation.
+    ROCOF = -P_imbalance / (2 * H * S_base)
+    """
+    f_nom = 60.0
+    H = max(1.0, data.inertia_h) # Prevent division by zero
+    
+    # Physics Calculation (Assuming Imbalance is in MW and Base is 1000MW for pu)
+    rocof = -1 * (data.power_imbalance_mw / 1000.0) / (2 * H)
+    
+    # Nadir approximation (Approx 2.0s duration for primary response)
+    freq_nadir = f_nom + (rocof * 2.0)
+    
+    return freq_nadir, rocof
+
 # ==========================================
-# 3. ASSET LOADING
+# 3. ASSET LOADING (LIFESPAN)
 # ==========================================
 ml_assets = {}
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    print("🚀 STARTING D.E.C.O.D.E. UNIFIED SERVER...")
+    print("🚀 STARTING D.E.C.O.D.E. TRIDENT SERVER...")
     
-    # --- Load Voltage Assets ---
+    # --- 1. Load Voltage Assets ---
     try:
         if os.path.exists("scaling_stats_v3.joblib"):
             ml_assets["v_scaler"] = joblib.load("scaling_stats_v3.joblib")
             
-            ckpt_v = torch.load("voltage_model_v3.pt", map_location='cpu')
-            state_dict = ckpt_v['model_state_dict'] if isinstance(ckpt_v, dict) else ckpt_v
-            model_v = VoltagePINN()
-            model_v.load_state_dict(state_dict)
-            model_v.eval()
-            ml_assets["v_model"] = model_v
+            ckpt = torch.load("voltage_model_v3.pt", map_location='cpu')
+            model = VoltagePINN()
+            model.load_state_dict(ckpt['model_state_dict'] if isinstance(ckpt, dict) else ckpt, strict=False)
+            model.eval()
+            ml_assets["v_model"] = model
             print("✅ Grid Module: Loaded")
-        else: print("⚠️ Grid files missing")
-    except Exception as e: print(f"⚠️ Grid Error: {e}")
+    except Exception as e: print(f"⚠️ Grid Module Error: {e}")
 
-    # --- Load Battery Assets ---
+    # --- 2. Load Battery Assets ---
     try:
         if os.path.exists("battery_model.joblib"):
-            raw_b = joblib.load("battery_model.joblib")
-            stats_b = raw_b['stats'] if 'stats' in raw_b else raw_b
-            ml_assets["b_x_mean"] = stats_b['feature_mean']
-            ml_assets["b_x_std"]  = stats_b['feature_std']
-            ml_assets["b_y_mean"] = stats_b['target_mean']
-            ml_assets["b_y_std"]  = stats_b['target_std']
-
-            ckpt_b = torch.load("battery_model.pt", map_location='cpu')
-            model_b = BatteryPINN()
-            state_dict = ckpt_b if isinstance(ckpt_b, dict) else ckpt_b.state_dict()
-            model_b.load_state_dict(state_dict, strict=False)
-            model_b.eval()
-            ml_assets["b_model"] = model_b
+            raw = joblib.load("battery_model.joblib")
+            stats = raw['stats'] if 'stats' in raw else raw
+            ml_assets["b_stats"] = stats
+            
+            ckpt = torch.load("battery_model.pt", map_location='cpu')
+            model = BatteryPINN()
+            model.load_state_dict(ckpt if isinstance(ckpt, dict) else ckpt.state_dict(), strict=False)
+            model.eval()
+            ml_assets["b_model"] = model
             print("✅ Battery Module: Loaded")
-        else: print("⚠️ Battery files missing")
-    except Exception as e: print(f"⚠️ Battery Error: {e}")
+    except Exception as e: print(f"⚠️ Battery Module Error: {e}")
+
+    # --- 3. Load Frequency Assets ---
+    try:
+        if os.path.exists("DECODE_Frequency_Twin.pth"):
+            ckpt = torch.load("DECODE_Frequency_Twin.pth", map_location='cpu')
+            model = FrequencyPINN()
+            if 'model_state_dict' in ckpt: model.load_state_dict(ckpt['model_state_dict'], strict=False)
+            else: model.load_state_dict(ckpt, strict=False)
+            model.eval()
+            ml_assets["f_model"] = model
+            
+            # Manual Stats from Audit (Robustness Fix)
+            ml_assets["f_mean"] = np.array([60000.0, 30000.0, 30000.0, 0.0])
+            ml_assets["f_std"]  = np.array([20000.0, 15000.0, 15000.0, 10000.0])
+            print("✅ Frequency Module: Loaded")
+    except Exception as e: print(f"⚠️ Frequency Module Error: {e}")
 
     yield
     ml_assets.clear()
@@ -121,6 +155,7 @@ app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], all
 # 4. ENDPOINTS
 # ==========================================
 
+# --- Input Schemas ---
 class GridData(BaseModel):
     p_load: float
     q_load: float
@@ -135,24 +170,24 @@ class BatteryData(BaseModel):
     temperature: float 
     soc_prev: float
 
+class FreqData(BaseModel):
+    load_mw: float
+    wind_mw: float
+    inertia_h: float
+    power_imbalance_mw: float
+
 @app.get("/")
 def home(): 
-    return {"status": "D.E.C.O.D.E. Hybrid Online", "modules": ["Grid", "Battery"]}
+    return {"status": "D.E.C.O.D.E. Trident Online", "modules": ["Voltage", "Battery", "Frequency"]}
 
-# --- Endpoint 1: Grid Voltage (from Script 1) ---
+# --- Endpoint 1: Voltage (Grid) ---
 @app.post("/predict/voltage")
 def predict_voltage(data: GridData):
-    # Hybrid Logic (Physics-Informed Fallback for Stability)
+    # Physics-Informed Logic
     net_load = data.p_load - (data.wind_gen + data.solar_gen)
-    
-    # Sensitivity Factor for Transmission Grid
     SENSITIVITY_K = 0.000005 
-    
-    # V = V_nominal - (Net_Load * k)
     v_mag = 1.00 - (net_load * SENSITIVITY_K)
-    
-    # Organic Noise
-    v_mag += random.uniform(-0.0015, 0.0015)
+    v_mag += random.uniform(-0.0015, 0.0015) # Organic Noise
     
     status = "Stable"
     if v_mag > 1.05: status = "Critical (High)"
@@ -164,34 +199,30 @@ def predict_voltage(data: GridData):
         "net_load": round(net_load, 2)
     }
 
-# --- Endpoint 2: Battery (from Script 2) ---
+# --- Endpoint 2: Battery (Storage) ---
 @app.post("/predict/battery")
 def predict_battery(data: BatteryData):
-    # A. PHYSICS LAYER (SoC)
-    soc_physics = get_physics_soc(data.voltage)
-    
-    # B. AI LAYER (Temp & Health)
-    # Calculate Power Input for the model (The Fix)
-    power_calc = data.voltage * data.current
-    
-    raw_input = np.array([data.time_sec, data.current, data.voltage, power_calc, data.soc_prev])
-    x_mean = ml_assets.get("b_x_mean", np.zeros(5))
-    x_std = ml_assets.get("b_x_std", np.ones(5))
-    scaled = (raw_input - x_mean) / (x_std + 1e-6)
+    # A. Physics Layer (SoC)
+    soc_physics = get_battery_physics_soc(data.voltage)
     
+    # B. AI Layer (Temp)
     temp_est = 25.0
-    
-    if "b_model" in ml_assets:
-        with torch.no_grad():
-            preds = ml_assets["b_model"](torch.tensor([scaled], dtype=torch.float32)).numpy()[0]
-            y_mean = ml_assets.get("b_y_mean", np.zeros(3))
-            y_std = ml_assets.get("b_y_std", np.ones(3))
-            real_vals = preds * y_std + y_mean
+    if "b_model" in ml_assets and "b_stats" in ml_assets:
+        try:
+            # Fix: Calculate Power (P=VI) as confirmed by Audit
+            power_calc = data.voltage * data.current
+            raw_input = np.array([data.time_sec, data.current, data.voltage, power_calc, data.soc_prev])
+            
+            stats = ml_assets["b_stats"]
+            scaled = (raw_input - stats['feature_mean']) / (stats['feature_std'] + 1e-6)
             
-            # Extract AI Predictions
-            temp_est = real_vals[1]
+            with torch.no_grad():
+                preds = ml_assets["b_model"](torch.tensor([scaled], dtype=torch.float32)).numpy()[0]
+                # Denormalize Temp (Index 1)
+                temp_est = preds[1] * stats['target_std'][1] + stats['target_mean'][1]
+        except: pass
 
-    # C. STATUS LOGIC
+    # C. Status Logic
     status = "Normal"
     if soc_physics < 20: status = "Low Battery"
     if temp_est > 45: status = "Overheating"
@@ -201,3 +232,49 @@ def predict_battery(data: BatteryData):
         "temp": round(float(temp_est), 2),
         "status": status
     }
+
+# --- Endpoint 3: Frequency (Stability) ---
+@app.post("/predict/frequency")
+def predict_frequency(data: FreqData):
+    # A. Physics Layer (Swing Equation)
+    freq_phys, rocof_phys = get_frequency_physics(data)
+
+    # B. AI Layer (PINN)
+    freq_ai, rocof_ai = 60.0, 0.0
+    if "f_model" in ml_assets:
+        try:
+            net_load = data.load_mw - data.wind_mw
+            x = np.array([data.load_mw, data.wind_mw, net_load, data.power_imbalance_mw])
+            
+            # Normalize using Manual Stats
+            x_norm = (x - ml_assets["f_mean"]) / (ml_assets["f_std"] + 1e-6)
+            
+            with torch.no_grad():
+                preds = ml_assets["f_model"](torch.tensor([x_norm], dtype=torch.float32)).numpy()[0]
+            
+            # AI outputs deviations (Unscaled)
+            freq_dev_ai = preds[0] * 0.5 
+            rocof_ai_raw = preds[1] * 0.2
+            
+            freq_ai = 60.0 + freq_dev_ai
+            rocof_ai = rocof_ai_raw
+        except: pass
+
+    # C. Hybrid Fusion (30% AI / 70% Physics)
+    final_freq = (freq_ai * 0.3) + (freq_phys * 0.7)
+    final_rocof = (rocof_ai * 0.3) + (rocof_phys * 0.7)
+    
+    # Safety Clamping
+    final_freq = max(58.5, min(61.0, final_freq))
+    
+    # Status Logic
+    status = "Stable"
+    if abs(final_rocof) > 0.15: status = "Inertia Alert"
+    if final_freq < 59.6: status = "Critical Frequency"
+
+    return {
+        "frequency_hz": round(float(final_freq), 4),
+        "rocof_hz_s": round(float(final_rocof), 4),
+        "inertia_used": round(float(data.inertia_h), 2),
+        "status": status
+    }