Update main.py

main.py

@@ -10,8 +10,9 @@ from pydantic import BaseModel
 from contextlib import asynccontextmanager
 
 # ==========================================
-# 1. SHARED MODEL ARCHITECTURES
+# 1. UNIQUE ARCHITECTURES (PER AUDIT)
 # ==========================================
+
 class Mish(nn.Module):
     def forward(self, x): return x * torch.tanh(nn.functional.softplus(x))
 
@@ -23,11 +24,45 @@ class FourierFeatureMapping(nn.Module):
         proj = 2 * np.pi * (x @ self.B)
         return torch.cat([torch.sin(proj), torch.cos(proj)], dim=-1)
 
-# --- A. Voltage Model (Grid) ---
+# --- Solar: 128-neuron State-Space PINN ---
+class SolarPINN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.backbone = nn.Sequential(
+            nn.Linear(4, 128), Mish(),
+            nn.Linear(128, 128), Mish(),
+            nn.Linear(128, 1)
+        )
+        # Physics Tensors from Audit
+        self.log_thermal_mass = nn.Parameter(torch.tensor(7.1546))
+        self.log_h_conv = nn.Parameter(torch.tensor(1.8767))
+    def forward(self, x): return self.backbone(x)
+
+# --- Load: Fourier Residual Architecture ---
+class LoadForecastPINN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fourier = FourierFeatureMapping(9, 32)
+        self.input_layer = nn.Linear(64, 128)
+        self.res_blocks = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(128, 128), 
+                nn.BatchNorm1d(128), 
+                Mish(), 
+                nn.Linear(128, 128)
+            ) for _ in range(3)
+        ])
+        self.output_layer = nn.Linear(128, 1)
+    def forward(self, x):
+        x = self.input_layer(self.fourier(x))
+        for block in self.res_blocks: x = x + block(x)
+        return self.output_layer(x)
+
+# --- Voltage: 256-dim Multi-Layer PINN ---
 class VoltagePINN(nn.Module):
     def __init__(self):
         super().__init__()
-        self.fourier = FourierFeatureMapping(input_dim=7, mapping_size=32)
+        self.fourier = FourierFeatureMapping(7, 32)
         self.network = nn.Sequential(
             nn.Linear(64, 256), nn.LayerNorm(256), Mish(),
             nn.Linear(256, 128), nn.LayerNorm(128), Mish(),
@@ -36,11 +71,11 @@ class VoltagePINN(nn.Module):
         )
     def forward(self, x): return self.network(self.fourier(x))
 
-# --- B. Battery Model (Storage) ---
+# --- Battery: 24-dim Linear PINN ---
 class BatteryPINN(nn.Module):
     def __init__(self):
         super().__init__()
-        self.fourier = FourierFeatureMapping(input_dim=5, mapping_size=12)
+        self.fourier = FourierFeatureMapping(5, 12)
         self.network = nn.Sequential(
             nn.Linear(24, 64), Mish(),
             nn.Linear(64, 64), Mish(),
@@ -48,100 +83,45 @@ class BatteryPINN(nn.Module):
         )
     def forward(self, x): return self.network(self.fourier(x))
 
-# --- C. Frequency Model (Stability) ---
+# --- Frequency: Stability Twin ---
 class FrequencyPINN(nn.Module):
     def __init__(self):
         super().__init__()
-        self.fourier = FourierFeatureMapping(input_dim=4, mapping_size=32)
-        self.network = nn.Sequential(
+        self.fourier = FourierFeatureMapping(4, 32)
+        self.net = 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))
-
-# --- 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)
+            nn.Linear(128, 2)
         )
-    def forward(self, x): return self.network(self.fourier(x))
-
-# ==========================================
-# 2. PHYSICS ENGINES
-# ==========================================
-
-def get_battery_physics_soc(voltage):
-    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):
-    f_nom = 60.0
-    H = max(1.0, data.inertia_h)
-    rocof = -1 * (data.power_imbalance_mw / 1000.0) / (2 * H)
-    freq_nadir = f_nom + (rocof * 2.0)
-    return freq_nadir, rocof
+    def forward(self, x): return self.net(self.fourier(x))
 
 # ==========================================
-# 3. ASSET LOADING (LIFESPAN)
+# 2. ASSET LOADING (LIFESPAN)
 # ==========================================
 ml_assets = {}
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    print("🚀 STARTING D.E.C.O.D.E. UNIFIED SERVER...")
-    
-    # 1. Load Voltage
-    try:
-        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
-    except Exception as e: print(f"⚠️ Voltage Error: {e}")
-
-    # 2. Load Battery
-    try:
-        if os.path.exists("battery_model.joblib"):
-            raw = joblib.load("battery_model.joblib")
-            ml_assets["b_stats"] = raw['stats'] if 'stats' in raw else raw
-            model = BatteryPINN()
-            model.load_state_dict(torch.load("battery_model.pt", map_location='cpu'), strict=False)
-            model.eval()
-            ml_assets["b_model"] = model
-    except Exception as e: print(f"⚠️ Battery Error: {e}")
-
-    # 3. Load Frequency
-    try:
-        if os.path.exists("DECODE_Frequency_Twin.pth"):
-            ckpt = torch.load("DECODE_Frequency_Twin.pth", map_location='cpu')
-            model = FrequencyPINN()
-            model.load_state_dict(ckpt['model_state_dict'] if isinstance(ckpt, dict) else ckpt, strict=False)
-            model.eval()
-            ml_assets["f_model"] = model
-            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])
-    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}")
-
+    # Load Models based on Audit Shapes
+    loaders = {
+        "solar": ("solar_model.pt", SolarPINN()),
+        "load": ("load_model.pt", LoadForecastPINN()),
+        "voltage": ("voltage_model_v3.pt", VoltagePINN()),
+        "battery": ("battery_model.pt", BatteryPINN()),
+        "freq": ("DECODE_Frequency_Twin.pth", FrequencyPINN())
+    }
+    for key, (path, model) in loaders.items():
+        if os.path.exists(path):
+            ckpt = torch.load(path, map_location='cpu')
+            sd = ckpt['model_state_dict'] if isinstance(ckpt, dict) and 'model_state_dict' in ckpt else ckpt
+            model.load_state_dict(sd, strict=False)
+            ml_assets[key] = model.eval()
+
+    # Load All Stats
+    if os.path.exists("Load_stats.joblib"): ml_assets["l_stats"] = joblib.load("Load_stats.joblib")
+    if os.path.exists("battery_model.joblib"): ml_assets["b_stats"] = joblib.load("battery_model.joblib")
+    if os.path.exists("scaling_stats_v3.joblib"): ml_assets["v_stats"] = joblib.load("scaling_stats_v3.joblib")
     yield
     ml_assets.clear()
 
@@ -149,80 +129,60 @@ app = FastAPI(title="D.E.C.O.D.E. Unified API", lifespan=lifespan)
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
 # ==========================================
-# 4. SCHEMAS & ENDPOINTS
+# 3. SCHEMAS & PHYSICS
 # ==========================================
+def get_ocv_soc(v): return np.interp(v, [2.8, 3.4, 3.7, 4.2], [0, 15, 65, 100])
 
-class GridData(BaseModel):
-    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
-
-class FreqData(BaseModel):
-    load_mw: float; wind_mw: float; inertia_h: float; power_imbalance_mw: float
+class SolarData(BaseModel): irradiance_stream: list[float]; ambient_temp_stream: list[float]; wind_speed_stream: list[float]
+class LoadData(BaseModel): temperature_c: float; hour: int; month: int; wind_mw: float = 0; solar_mw: float = 0
+class BatteryData(BaseModel): 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
+class GridData(BaseModel): p_load: float; q_load: float; wind_gen: float; solar_gen: float; hour: int
 
-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. Unified Digital Twin Online"}
+# ==========================================
+# 4. CALIBRATED ENDPOINTS
+# ==========================================
 
-@app.post("/predict/voltage")
-def predict_voltage(data: GridData):
-    net_load = data.p_load - (data.wind_gen + data.solar_gen)
-    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"}
-
-@app.post("/predict/battery")
-def predict_battery(data: BatteryData):
-    soc_physics = get_battery_physics_soc(data.voltage)
-    temp_est = 25.0
-    if "b_model" in ml_assets:
-        try:
-            stats = ml_assets["b_stats"]
-            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]
-                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"}
-
-@app.post("/predict/frequency")
-def predict_frequency(data: FreqData):
-    f_phys, r_phys = get_frequency_physics(data)
-    f_ai = 60.0
-    if "f_model" in ml_assets:
-        try:
-            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]
-                f_ai = 60.0 + (preds[0] * 0.5)
-        except: pass
-    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/solar")
+def predict_solar(data: SolarData):
+    # Constraint: Recursive Simulation @ 900s dt
+    curr_temp = data.ambient_temp_stream[0] + 5.0
+    sim = []
+    with torch.no_grad():
+        for i in range(len(data.irradiance_stream)):
+            x = torch.tensor([[(data.irradiance_stream[i]-450)/250, (data.ambient_temp_stream[i]-25)/10, 
+                               data.wind_speed_stream[i]/10.0, (curr_temp-35)/15]], dtype=torch.float32)
+            # Physical Clamping
+            next_t = max(10.0, min(75.0, ml_assets["solar"](x).item()))
+            eff = 0.20 * (1 - 0.004 * (next_t - 25.0))
+            sim.append({"temp": round(next_t, 2), "mw": round((5000 * data.irradiance_stream[i] * max(0, eff)) / 1e6, 4)})
+            curr_temp = next_t 
+    return {"simulation": sim}
 
 @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)))
-    
-    x_norm = np.array([t_norm, max(0, data.temperature_c-18)/10, max(0, 18-data.temperature_c)/10, 
+    # Constraint: Hard Z-Score Clamping at +/-3 to prevent Inverted Load Paradox
+    t_norm = max(-3.0, min(3.0, (data.temperature_c - 15.38) / 4.12))
+    x = torch.tensor([[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)
-    
-    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
-
-    # 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"}
+                       data.wind_mw/10000, data.solar_mw/10000]], dtype=torch.float32)
+    load_mw = 35000.0
+    if "load" in ml_assets:
+        with torch.no_grad(): load_mw = (ml_assets["load"](x).item() * 9773.8) + 35000.0
+    # Physical Safety Correction
+    if data.temperature_c > 32: load_mw = max(load_mw, 45000 + (data.temperature_c - 32) * 1200)
+    elif data.temperature_c < 5: load_mw = max(load_mw, 42000 + (5 - data.temperature_c) * 900)
+    return {"mw": round(load_mw, 2)}
+
+@app.post("/predict/battery")
+def predict_battery(data: BatteryData):
+    # Constraint: Feature Engineering (Power Product V*I)
+    p_prod = data.voltage * data.current
+    stats = ml_assets["b_stats"].get('stats', ml_assets["b_stats"])
+    raw = np.array([data.time_sec, data.current, data.voltage, p_prod, data.soc_prev])
+    x_scaled = (raw - stats['feature_mean']) / (stats['feature_std'] + 1e-6)
+    with torch.no_grad():
+        preds = ml_assets["battery"](torch.tensor([x_scaled], dtype=torch.float32)).numpy()[0]
+        temp = preds[1] * stats['target_std'][1] + stats['target_mean'][1]
+    return {"soc": round(get_ocv_soc(data.voltage), 2), "temp": round(temp, 2)}