Update main.py

main.py

@@ -24,7 +24,6 @@ class FourierFeatureMapping(nn.Module):
         return torch.cat([torch.sin(proj), torch.cos(proj)], dim=-1)
 
 # --- 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__()
@@ -38,7 +37,6 @@ class VoltagePINN(nn.Module):
     def forward(self, x): return self.network(self.fourier(x))
 
 # --- 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__()
@@ -46,12 +44,11 @@ class BatteryPINN(nn.Module):
         self.network = nn.Sequential(
             nn.Linear(24, 64), Mish(),
             nn.Linear(64, 64), Mish(),
-            nn.Linear(64, 3) 
+            nn.Linear(64, 3)
         )
     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__()
@@ -63,33 +60,33 @@ class FrequencyPINN(nn.Module):
         )
     def forward(self, x): return self.network(self.fourier(x))
 
+# --- D. Load Model (Forecast) ---
+class LoadPINN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fourier = FourierFeatureMapping(input_dim=9, 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, 64), nn.LayerNorm(64), Mish(),
+            nn.Linear(64, 1)
+        )
+    def forward(self, x): return self.network(self.fourier(x))
+
 # ==========================================
-# 2. PHYSICS ENGINES (The "Laws")
+# 2. PHYSICS ENGINES
 # ==========================================
 
 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)
+    H = max(1.0, data.inertia_h)
     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
 
 # ==========================================
@@ -99,51 +96,51 @@ ml_assets = {}
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    print("🚀 STARTING D.E.C.O.D.E. TRIDENT SERVER...")
+    print("🚀 STARTING D.E.C.O.D.E. UNIFIED SERVER...")
     
-    # --- 1. Load Voltage Assets ---
+    # 1. Load Voltage
     try:
-        if os.path.exists("scaling_stats_v3.joblib"):
-            ml_assets["v_scaler"] = joblib.load("scaling_stats_v3.joblib")
-            
+        if os.path.exists("voltage_model_v3.pt"):
             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")
-    except Exception as e: print(f"⚠️ Grid Module Error: {e}")
+    except Exception as e: print(f"⚠️ Voltage Error: {e}")
 
-    # --- 2. Load Battery Assets ---
+    # 2. Load Battery
     try:
         if os.path.exists("battery_model.joblib"):
             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')
+            ml_assets["b_stats"] = raw['stats'] if 'stats' in raw else raw
             model = BatteryPINN()
-            model.load_state_dict(ckpt if isinstance(ckpt, dict) else ckpt.state_dict(), strict=False)
+            model.load_state_dict(torch.load("battery_model.pt", map_location='cpu'), strict=False)
             model.eval()
             ml_assets["b_model"] = model
-            print("✅ Battery Module: Loaded")
-    except Exception as e: print(f"⚠️ Battery Module Error: {e}")
+    except Exception as e: print(f"⚠️ Battery Error: {e}")
 
-    # --- 3. Load Frequency Assets ---
+    # 3. Load Frequency
     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.load_state_dict(ckpt['model_state_dict'] if isinstance(ckpt, dict) else 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}")
+    except Exception as e: print(f"⚠️ Frequency Error: {e}")
+
+    # 4. Load Forecast
+    try:
+        if os.path.exists("load_model.pt"):
+            model = LoadPINN()
+            model.load_state_dict(torch.load("load_model.pt", map_location='cpu'), strict=False)
+            model.eval()
+            ml_assets["l_model"] = model
+            stats = joblib.load("Load_stats.joblib")
+            ml_assets["l_stats"] = stats
+    except Exception as e: print(f"⚠️ Load Error: {e}")
 
     yield
     ml_assets.clear()
@@ -152,129 +149,80 @@ app = FastAPI(title="D.E.C.O.D.E. Unified API", lifespan=lifespan)
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
 # ==========================================
-# 4. ENDPOINTS
+# 4. SCHEMAS & ENDPOINTS
 # ==========================================
 
-# --- Input Schemas ---
 class GridData(BaseModel):
-    p_load: float
-    q_load: float
-    wind_gen: float
-    solar_gen: float
-    hour: int
+    p_load: float; q_load: float; wind_gen: float; solar_gen: float; hour: int
 
 class BatteryData(BaseModel):
-    time_sec: float
-    current: float
-    voltage: float
-    temperature: float 
-    soc_prev: float
+    time_sec: float; current: float; voltage: float; temperature: float; soc_prev: float
 
 class FreqData(BaseModel):
-    load_mw: float
-    wind_mw: float
-    inertia_h: float
-    power_imbalance_mw: float
+    load_mw: float; wind_mw: float; inertia_h: float; power_imbalance_mw: float
+
+class LoadData(BaseModel):
+    temperature_c: float; hour: int; month: int; wind_mw: float = 0.0; solar_mw: float = 0.0
 
 @app.get("/")
-def home(): 
-    return {"status": "D.E.C.O.D.E. Trident Online", "modules": ["Voltage", "Battery", "Frequency"]}
+def home():
+    return {"status": "D.E.C.O.D.E. Unified Digital Twin Online"}
 
-# --- Endpoint 1: Voltage (Grid) ---
 @app.post("/predict/voltage")
 def predict_voltage(data: GridData):
-    # Physics-Informed Logic
     net_load = data.p_load - (data.wind_gen + data.solar_gen)
-    SENSITIVITY_K = 0.000005 
-    v_mag = 1.00 - (net_load * SENSITIVITY_K)
-    v_mag += random.uniform(-0.0015, 0.0015) # Organic Noise
-    
-    status = "Stable"
-    if v_mag > 1.05: status = "Critical (High)"
-    if v_mag < 0.95: status = "Critical (Low)"
+    v_mag = 1.00 - (net_load * 0.000005) + random.uniform(-0.0015, 0.0015)
+    return {"voltage_pu": round(v_mag, 4), "status": "Stable" if 0.95 < v_mag < 1.05 else "Critical"}
 
-    return {
-        "voltage_pu": round(v_mag, 4), 
-        "status": status, 
-        "net_load": round(net_load, 2)
-    }
-
-# --- Endpoint 2: Battery (Storage) ---
 @app.post("/predict/battery")
 def predict_battery(data: BatteryData):
-    # 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 and "b_stats" in ml_assets:
+    if "b_model" 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)
-            
+            scaled = (np.array([data.time_sec, data.current, data.voltage, data.voltage*data.current, data.soc_prev]) - stats['feature_mean']) / (stats['feature_std'] + 1e-6)
             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
+    return {"soc": round(float(soc_physics), 2), "temp": round(float(temp_est), 2), "status": "Normal" if temp_est < 45 else "Overheating"}
 
-    # C. Status Logic
-    status = "Normal"
-    if soc_physics < 20: status = "Low Battery"
-    if temp_est > 45: status = "Overheating"
-    
-    return {
-        "soc": round(float(soc_physics), 2),
-        "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
+    f_phys, r_phys = get_frequency_physics(data)
+    f_ai = 60.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)
-            
+            x_norm = (np.array([data.load_mw, data.wind_mw, data.load_mw-data.wind_mw, data.power_imbalance_mw]) - 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
+                f_ai = 60.0 + (preds[0] * 0.5)
         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)
+    final_f = max(58.5, min(61.0, (f_ai * 0.3) + (f_phys * 0.7)))
+    return {"frequency_hz": round(float(final_f), 4), "status": "Stable" if final_f > 59.6 else "Critical"}
+
+@app.post("/predict/load")
+def predict_load(data: LoadData):
+    stats = ml_assets.get("l_stats", {})
+    t_mean, t_std = stats.get('temp_mean', 15.38), stats.get('temp_std', 4.12)
+    t_norm = max(-3.0, min(3.0, (data.temperature_c - t_mean) / (t_std + 1e-6)))
     
-    # Safety Clamping
-    final_freq = max(58.5, min(61.0, final_freq))
+    x_norm = np.array([t_norm, max(0, data.temperature_c-18)/10, max(0, 18-data.temperature_c)/10, 
+                       np.sin(2*np.pi*data.hour/24), np.cos(2*np.pi*data.hour/24), 
+                       np.sin(2*np.pi*data.month/12), np.cos(2*np.pi*data.month/12), 
+                       data.wind_mw/10000, data.solar_mw/10000], dtype=np.float32)
     
-    # Status Logic
-    status = "Stable"
-    if abs(final_rocof) > 0.15: status = "Inertia Alert"
-    if final_freq < 59.6: status = "Critical Frequency"
+    load_mw = stats.get('load_mean', 35000.0)
+    if "l_model" in ml_assets:
+        try:
+            with torch.no_grad():
+                preds = ml_assets["l_model"](torch.tensor([x_norm], dtype=torch.float32)).numpy()[0]
+                load_mw = (preds[0] * stats.get('load_std', 9773.80)) + load_mw
+        except: pass
 
-    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
-    }
+    # Physics Overrides
+    if data.temperature_c > 32 and load_mw < 45000: load_mw = 45000 + (data.temperature_c - 32) * 1200
+    elif data.temperature_c < 5 and load_mw < 42000: load_mw = 42000 + (5 - data.temperature_c) * 900
+    
+    return {"predicted_load_mw": round(float(load_mw), 2), "status": "Normal" if load_mw < 58000 else "Peak Load Alert"}