modnn 2.0.0__tar.gz → 2.0.2__tar.gz

{modnn-2.0.0 → modnn-2.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modnn
-Version: 2.0.0
+Version: 2.0.2
 Summary: Physics-Informed Modularized Neural Network for Building Energy Modeling
 Author-email: Zixin Jiang <zjiang19@syr.edu>
 License-Expression: MIT

{modnn-2.0.0 → modnn-2.0.2}/modnn/Config.py

@@ -66,7 +66,7 @@ def _args(**kwargs):
         "envelop": _envelops(**envelop_overrides),
         # Paths and device
         "datapath": "/home/zjiang19/Documents/GitHub/Eplus_ModNN_Compare/dataset/Eplus/EPlus_train_AC_off_2month.csv", #"../Dataset/EPlus.csv",
-        "device": "cuda:1",
+        "device": "cuda:0",
         "save_name": "Mid",
 
         # Data settings
@@ -77,6 +77,8 @@ def _args(**kwargs):
         "trainday": 180, # Training data selection
         "testday": 1, # Testing data selection
         "training_batch": 1024*1,
+        "envelop_mdl": "physics", # We provide "physics" and "datadriven"
+                                  # "physics" rely on heatbalance equation, "datadriven" is a blackbox
         "ext_mdl": "LSTM", # We provide LSTM and RNN module, for RNN, we can apply positive constraint easily
                            # But LSTM has Hadamard product, making this constraint hard to integrate
                            # However, disturbance variables always have similiar distribution, in other word, is this constraint really necessary?

modnn-2.0.2/modnn/Models/ModNN_phy.py

@@ -0,0 +1,381 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# --- Solar Radiation Module ---
+class SolarRadiation(nn.Module):
+    def __init__(self, envelop, name="Solar"):
+        super().__init__()
+        # Initialization
+        self.envelop = envelop
+
+        # Load transmission and absorption coefficient
+        self.direct_gain = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
+        self.abs_wall = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
+        self.abs_roof = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
+        self.name = name
+
+    def forward(self, radiation, n_wall, n_roof):
+        """
+        Calculate solar heat gains [transmission and absorption] in a vectorized way.
+
+        Args:
+            radiation: Solar radiation input [batch, time, 1] or [batch, time, directions]
+                      (e.g., [batch, time, 4] for E,S,W,N directions)
+            n_wall: Number of wall components
+            n_roof: Number of roof components
+
+        Returns:
+            direct_gain: Heat directly entering zone
+            wall_gains: List of absorbed heat for each wall [n_wall elements]
+            roof_gains: List of absorbed heat for each roof [n_roof elements]
+        """
+        # Apply softplus to ensure positive coefficients
+        direct_coef = F.softplus(self.direct_gain) * self.envelop["direct_coef"]
+        abs_wall_coef = F.softplus(self.abs_wall) * self.envelop["abs_wall_coef"]
+        abs_roof_coef = F.softplus(self.abs_roof) * self.envelop["abs_roof_coef"]
+
+        # Transmission solar gain
+        direct_gain = direct_coef * radiation
+
+        # Absorption solar gain
+        # Wall
+        batch_size, time_steps = radiation.shape[0], radiation.shape[1]
+        if radiation.shape[-1] > 1:  # TODO: add solar angle model to distribute solar radiations
+            # Handle multi-directional radiation
+            wall_gains = []
+            for i in range(n_wall):
+                dir_idx = min(i, radiation.shape[-1] - 1)
+                wall_gain = abs_wall_coef * radiation[:, :, dir_idx:dir_idx + 1]
+                wall_gains.append(wall_gain)
+        else:
+            wall_gain = abs_wall_coef * radiation
+            wall_gains = [wall_gain] * n_wall
+
+        # Roof
+        roof_gain = abs_roof_coef * radiation
+        roof_gains = [roof_gain] * n_roof
+
+        return direct_gain, wall_gains, roof_gains
+
+
+# --- Envelope Module [RC, R] ---
+class EnvelopeRC(nn.Module):
+    def __init__(self, envelop, n_wall, n_roof, n_window):
+        super().__init__()
+        self.n_wall = n_wall
+        self.n_roof = n_roof
+        self.n_window = n_window
+        self.n_rc = n_wall + n_roof
+        self.envelop = envelop
+
+        # Initialization
+        wall_r_init = torch.tensor([1.0 / (10.0 + i * 1.0) for i in range(n_wall)], dtype=torch.float32)
+        roof_r_init = torch.tensor([1.0 / (20.0 + i * 2.0) for i in range(n_roof)], dtype=torch.float32)
+        self.rc_R_inv = nn.Parameter(torch.cat([wall_r_init, roof_r_init]))
+
+        wall_c_init = torch.tensor([1.0 / (80.0 + i * 10.0) for i in range(n_wall)], dtype=torch.float32)
+        roof_c_init = torch.tensor([1.0 / (60.0 + i * 10.0) for i in range(n_roof)], dtype=torch.float32)
+        self.rc_C_inv = nn.Parameter(torch.cat([wall_c_init, roof_c_init]))
+
+        if n_window > 0:
+            window_r_init = torch.tensor([1.0 / (1.0 + i * 0.1) for i in range(n_window)], dtype=torch.float32)
+            self.window_R_inv = nn.Parameter(window_r_init)
+
+    def forward(self, T_amb, Tmid_list, Tzone, solar_gains_wall, solar_gains_roof):
+        """
+        Envelope heat transfer calculation
+
+        Args:
+            T_amb: Ambient temperature [batch, time, 1]
+            Tmid_list: List of envelope temperatures [n_rc elements each [batch, time, 1]]
+            Tzone: Zone temperature [batch, time, 1]
+            solar_gains_wall: List of wall solar gains [n_wall elements each [batch, time, 1]]
+            solar_gains_roof: List of roof solar gains [n_roof elements each [batch, time, 1]]
+
+        Returns:
+            dTmid_dt_list: List of RC temperature derivatives
+            q_env_total: Total heat flow to zone
+        """
+        batch_size = T_amb.shape[0]
+
+        # Stack all envelope temperatures into one tensor [batch, time, n_rc]
+        Tmid_stacked = torch.cat(Tmid_list, dim=2)
+
+        # Expand ambient and zone temperatures for broadcasting
+        T_amb_expanded = T_amb.expand(-1, -1, self.n_rc)
+        Tzone_expanded = Tzone.expand(-1, -1, self.n_rc)
+
+        # Apply softplus to parameters and expand for broadcasting
+        r_inv_bounded = F.softplus(self.rc_R_inv) * (1.0 / self.envelop["r_opaque_coef"])
+        c_inv_bounded = F.softplus(self.rc_C_inv) * (1.0 / self.envelop["c_opaque_coef"])
+        r_inv_expanded = r_inv_bounded.view(1, 1, -1).expand(batch_size, T_amb.shape[1], -1)
+        c_inv_expanded = c_inv_bounded.view(1, 1, -1).expand(batch_size, T_amb.shape[1], -1)
+
+        # Energy balance constraints
+        q_in = (T_amb_expanded - Tmid_stacked) * r_inv_expanded
+        q_out = (Tmid_stacked - Tzone_expanded) * r_inv_expanded
+
+        # Combine solar gains into single tensor with same shape as Tmid_stacked
+        # Add solar radiation to the first capacitance node
+        solar_gains = []
+        for i in range(self.n_wall):
+            solar_gains.append(solar_gains_wall[i])
+        for i in range(self.n_roof):
+            solar_gains.append(solar_gains_roof[i])
+
+        solar_stacked = torch.cat(solar_gains, dim=2) if solar_gains else torch.zeros_like(Tmid_stacked)
+
+        # Calculate temperature derivatives
+        dTmid_dt_stacked = (q_in - q_out + solar_stacked) * c_inv_expanded
+
+        # Window heat transfer (Assume no storage capacity)
+        q_window_total = 0
+        if self.n_window > 0:
+            window_r_inv_bounded = F.softplus(self.window_R_inv) * (1.0 / self.envelop["r_transparent_coef"])
+            window_r_inv_expanded = window_r_inv_bounded.view(1, 1, -1).expand(batch_size, T_amb.shape[1], -1)
+            T_amb_window = T_amb.expand(-1, -1, self.n_window)
+            Tzone_window = Tzone.expand(-1, -1, self.n_window)
+            q_window = (T_amb_window - Tzone_window) * window_r_inv_expanded
+            q_window_total = torch.sum(q_window, dim=2, keepdim=True)
+
+        # Total flow calculation
+        q_rc_total = torch.sum(q_out, dim=2, keepdim=True)
+        q_env_total = q_rc_total + q_window_total
+
+        # Split to each state node
+        dTmid_dt_list = torch.split(dTmid_dt_stacked, 1, dim=2)
+
+        return dTmid_dt_list, q_env_total
+
+
+class ModNN(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        para = args["para"]
+        envelop = args["envelop"]
+        self.enc_len = args["enLen"]
+        self.device = args["device"]
+        self.delta_t = 900  # 15 minutes in seconds
+
+        # Model configuration
+        self.n_wall = envelop.get("n_wall", 0)
+        self.n_roof = envelop.get("n_roof", 0)
+        self.n_window = envelop.get("n_window", 0)
+
+        print(f"Initializing Vectorized ModNN with "
+              f"{self.n_wall} walls, "
+              f"{self.n_roof} roofs, "
+              f"{self.n_window} windows")
+
+        # Modules Assemble
+        self.Int = Internal(para["Int_in"], para["Int_h"], para["Int_out"])
+        self.HVAC = HVAC()
+        self.solar = SolarRadiation(envelop=envelop)
+        self.envelopes = EnvelopeRC(envelop, self.n_wall, self.n_roof, self.n_window)
+        self.zone_update = ZoneUpdate(envelop=envelop)
+
+    def forward(self, input_X):
+        """
+        input_X order: [T_zone, T_ambient, solar, day_sin, day_cos, occ, phvac]
+        Shape: [batch_size, time_steps, features]
+        """
+        batch_size = input_X.shape[0]
+        time_steps = input_X.shape[1]
+
+        # Extract input data
+        T_amb = input_X[:, :, [1]]  # Ambient air temperature
+        Solar_X = input_X[:, :, [2]]  # Solar radiation
+        Int_X = input_X[:, :, [3, 4, 5]]  # Time_sin/cos, Occ
+        HVAC_X = input_X[:, :, [6]]  # pHVAC
+        Tzone_gt = input_X[:, :, [0]]  # Zone air temperature groundtruth
+
+        # Pre-allocate results arrays
+        Tzone_pred = torch.zeros(batch_size, time_steps, 1, device=input_X.device)
+        q_int_pred = torch.zeros(batch_size, time_steps, 1, device=input_X.device)
+        q_hvac_pred = torch.zeros(batch_size, time_steps, 1, device=input_X.device)
+        q_env_pred = torch.zeros(batch_size, time_steps, 1, device=input_X.device)
+        q_solar_direct_pred = torch.zeros(batch_size, time_steps, 1, device=input_X.device)
+
+        # Pre-compute internal and HVAC
+        q_int_pred = self.Int(Int_X)
+        q_hvac_pred = self.HVAC(HVAC_X)
+
+        # Initialize zone temperature
+        Tzone = Tzone_gt[:, 0:1, :]
+
+        # Initialize envelope temperatures
+        Tmid = []
+        for _ in range(self.n_wall + self.n_roof):
+            tmid_init = 0.7 * Tzone + 0.3 * T_amb[:, 0:1, :]
+            Tmid.append(tmid_init)
+
+        # Pre-allocate envelope temperature histories
+        Tmid_hist = [torch.zeros(batch_size, time_steps, 1, device=input_X.device)
+                     for _ in range(self.n_wall + self.n_roof)]
+
+        # --- Encoder phase ---
+        for i in range(self.enc_len):
+            # The purpose is to stablize the initial envelop temperatures so we use ground truth temperature
+            Tzone_pred[:, i:i + 1, :] = Tzone_gt[:, i:i + 1, :]
+            Tzone_i = Tzone_gt[:, i:i + 1, :]
+
+            # Get inputs for this timestep
+            T_amb_i = T_amb[:, i:i + 1, :]
+            Solar_i = Solar_X[:, i:i + 1, :]
+
+            # Calculate solar distribution
+            q_solar_direct, q_solar_wall, q_solar_roof = self.solar(
+                Solar_i, self.n_wall, self.n_roof
+            )
+            q_solar_direct_pred[:, i:i + 1, :] = q_solar_direct
+
+            # Envelope heat transfer calculation
+            dTmid_dt_list, q_env_total = self.envelopes(
+                T_amb_i, Tmid, Tzone_i, q_solar_wall, q_solar_roof
+            )
+            q_env_pred[:, i:i + 1, :] = q_env_total
+
+            # Envelope temperatures update
+            # Use RK4 if want higher accuracy
+            for idx, dTmid_dt in enumerate(dTmid_dt_list):
+                # Clamp derivatives for stability
+                dTmid_dt_clamped = torch.clamp(dTmid_dt, -3.0 / self.delta_t, 3.0 / self.delta_t)
+                delta = torch.clamp(self.delta_t * dTmid_dt_clamped, -3.0, 3.0)
+                Tmid[idx] = Tmid[idx] + delta
+                Tmid_hist[idx][:, i:i + 1, :] = Tmid[idx]
+
+        # --- Decoder phase (this is the actual prediction stage) ---
+        Tzone = Tzone_gt[:, self.enc_len - 1:self.enc_len, :]
+
+        for i in range(time_steps - self.enc_len):
+            idx = self.enc_len + i
+
+            # Initialize
+            T_amb_i = T_amb[:, idx:idx + 1, :]
+            Solar_i = Solar_X[:, idx:idx + 1, :]
+            q_int_i = q_int_pred[:, idx:idx + 1, :]
+            q_hvac_i = q_hvac_pred[:, idx:idx + 1, :]
+
+            # Calculate solar distribution
+            q_solar_direct, q_solar_wall, q_solar_roof = self.solar(
+                Solar_i, self.n_wall, self.n_roof
+            )
+            q_solar_direct_pred[:, idx:idx + 1, :] = q_solar_direct
+
+            # Envelope heat transfer calculation
+            dTmid_dt_list, q_env_total = self.envelopes(
+                T_amb_i, Tmid, Tzone, q_solar_wall, q_solar_roof
+            )
+            q_env_pred[:, idx:idx + 1, :] = q_env_total
+
+            # Envelope temperatures update
+            for j, dTmid_dt in enumerate(dTmid_dt_list):
+                # Clamp derivatives for stability
+                dTmid_dt_clamped = torch.clamp(dTmid_dt, -3.0 / self.delta_t, 3.0 / self.delta_t)
+                delta = torch.clamp(self.delta_t * dTmid_dt_clamped, -3.0, 3.0)
+                Tmid[j] = Tmid[j] + delta
+                Tmid_hist[j][:, idx:idx + 1, :] = Tmid[j]
+
+            # Calculate zone temperature derivative
+            dTzone_dt = self.zone_update(q_env_total, q_int_i, q_hvac_i, q_solar_direct)
+            # TODO: consider to integrated with Neural ODE
+
+            # Update zone temperature with stability limits
+            dTzone_dt_clamped = torch.clamp(dTzone_dt, -2.0 / self.delta_t, 2.0 / self.delta_t)
+            delta = torch.clamp(self.delta_t * dTzone_dt_clamped, -2.0, 2.0)
+            Tzone = Tzone + delta
+            Tzone_pred[:, idx:idx + 1, :] = Tzone
+
+        return Tzone_pred, Tzone_gt, (q_int_pred, q_hvac_pred, q_env_pred, q_solar_direct_pred)
+
+
+# --- Runge-Kutta Integration ---
+def rk4_step(f, x, dt, max_delta=3.0):
+    k1 = f(x)
+    k1 = torch.clamp(k1, -max_delta / dt, max_delta / dt)
+
+    k2 = f(x + 0.5 * dt * k1)
+    k2 = torch.clamp(k2, -max_delta / dt, max_delta / dt)
+
+    k3 = f(x + 0.5 * dt * k2)
+    k3 = torch.clamp(k3, -max_delta / dt, max_delta / dt)
+
+    k4 = f(x + dt * k3)
+    k4 = torch.clamp(k4, -max_delta / dt, max_delta / dt)
+
+    dx = (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
+    return x + torch.clamp(dx, -max_delta, max_delta)
+
+
+# --- Internal Heat Gain Module ---
+class Internal(nn.Module):
+    """
+    A simple MLP is used to predict internal heat gain (convection ONLY)
+    #TODO: add a RNN/MLP with lookback window to consider internal radiation
+
+    In general, the internal heat gain comes from Lighting, Occupant, Appliance, etc,.
+    But they can be represented by a factor (alpha) multiply with a "schedule" (sch), for example:
+    q_light = alpha_light * sch_light
+    q_human = alpha_human * sch_human
+    q_cooking = alpha_cooking * sch_cooking
+
+    If detailed information is available, we can replace it by Physics Equations
+    For example, 80~100W * Number_of_people
+    Otherwise, we learn form periodic features
+    """
+
+    def __init__(self, input_size, hidden_size, output_size):
+        super().__init__()
+        self.FC1 = nn.Linear(input_size - 1, hidden_size)
+        self.FC2 = nn.Linear(hidden_size, output_size)
+        self.gain = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
+
+        # Initialize weights
+        nn.init.xavier_uniform_(self.FC1.weight, gain=0.1)
+        nn.init.zeros_(self.FC1.bias)
+        nn.init.xavier_uniform_(self.FC2.weight, gain=0.1)
+        nn.init.zeros_(self.FC2.bias)
+
+    def forward(self, x):
+        x1 = F.relu(self.FC1(x[:, :, :-1]))
+        sch = torch.sigmoid(self.FC2(x1)) + x[:, :, -1:]
+        return F.softplus(self.gain) * 0.1 * sch
+
+
+# --- HVAC Module ---
+class HVAC(nn.Module):
+    """
+    A simplified linear module is used here for AIR SIDE SYSTEM ONLY
+    Change it to any type of RNN or add look back window for RADIATION SYSTEM
+    The input is pre-calculated pHVAC (thermal load)
+
+    But if the raw data is HVAC energy, or supply air flow/temperature, fan speed, coil load, etc
+    #TODO: Add More Types of HVAC Modules
+
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.gain = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
+
+    def forward(self, x):
+        return (F.softplus(self.gain) * 0.1 * x)
+
+
+# --- Zone Heat Balance Module ---
+class ZoneUpdate(nn.Module):
+    def __init__(self, envelop, c_init=10000.0):
+        super().__init__()
+        self.C_inv = nn.Parameter(torch.tensor(1.0 / c_init, dtype=torch.float32))
+        self.envelop = envelop
+
+    def forward(self, q_env_total, q_int, q_hvac, q_solar_direct=None):
+        c_inv_bounded = F.softplus(self.C_inv) * (1.0 / self.envelop["c_zone"])
+
+        # Add direct solar gain if provided
+        if q_solar_direct is not None:
+            return (q_env_total + q_int + q_hvac + q_solar_direct) * c_inv_bounded
+        else:
+            return (q_env_total + q_int + q_hvac) * c_inv_bounded

{modnn-2.0.0 → modnn-2.0.2}/modnn/Play.py

@@ -73,7 +73,7 @@ class EarlyStopping:
 
 def train_model(model, train_loader, valid_loader, test_loader, lr, epochs,
                 patience, tempscal, fluxscal, enlen, delen,
-                rawdf, plott, modeltype, scale, device, ext_mdl):
+                rawdf, plott, modeltype, scale, device, ext_mdl, envelop_mdl):
     num_update = 0
     total_time = 0
     time_every_update = {}
@@ -97,6 +97,8 @@ def train_model(model, train_loader, valid_loader, test_loader, lr, epochs,
         "PI-modnn|L": (0, 1),
         "PI-modnn|LC": (0, 0)
     }.get(modeltype, (None, None))
+    if envelop_mdl=="physics":
+        Phy_loss, Phy_cons = 0, 0
     if Phy_loss is None:
         raise ValueError("Unknown modeltype")
 
@@ -213,6 +215,7 @@ def train_model(model, train_loader, valid_loader, test_loader, lr, epochs,
                     hvac_out_denorm.append(
                         fluxscal.inverse_transform(hvac_out[[idx], enlen:, :].reshape(-1, 1)))
 
+                    #TODO: collect different heat fluxes, need match with different model, leave it  later
                     # inter_out_denorm.append(
                     #     fluxscal.inverse_transform(inter_out[[idx], enlen:, :].reshape(-1, 1)))
                     # external_out_denorm.append(

{modnn-2.0.0 → modnn-2.0.2}/modnn/utils.py

@@ -8,7 +8,7 @@ import numpy as np
 import torch
 import matplotlib.dates as dates
 import matplotlib.pyplot as plt
-from modnn.Models import ModNN, BaseNN
+from modnn.Models import ModNN_phy, BaseNN, ModNN_data
 from modnn.Play import train_model, test_model, check_model, dynamic_check_model, grad_model
 import matplotlib
 import seaborn as sns
@@ -37,7 +37,10 @@ class Mod:
 
     def train(self):
         if "modnn" in self.args["modeltype"]:
-            model = ModNN.ModNN(self.args).to(self.device)
+            if self.args["envelop_mdl"] == "physics":
+                model = ModNN_phy.ModNN(self.args).to(self.device)
+            else:
+                model = ModNN_data.ModNN(self.args).to(self.device)
         if self.args["modeltype"] == "LSTM":
             model = BaseNN.Baseline(self.args).to(self.device)
         start_time = time.time()
@@ -58,7 +61,8 @@ class Mod:
                                                 modeltype = self.args["modeltype"],
                                                 scale = self.args["scale"],
                                                 device=self.device,
-                                                ext_mdl = self.args["ext_mdl"])
+                                                ext_mdl = self.args["ext_mdl"],
+                                                envelop_mdl = self.args["envelop_mdl"],)
         print("--- %s seconds ---" % (time.time() - start_time))
 
         folder_name = ("../Saved/{}/Trained_mdl".format(self.args['save_name']) +
@@ -138,7 +142,10 @@ class Mod:
         start_time = time.time()
         print("Loading model")
         if "modnn" in self.args["modeltype"]:
-            model = ModNN.ModNN(self.args).to(self.device)
+            if self.args["envelop_mdl"] == "physics":
+                model = ModNN_phy.ModNN(self.args).to(self.device)
+            else:
+                model = ModNN_data.ModNN(self.args).to(self.device)
         if self.args["modeltype"] == "LSTM":
             model = BaseNN.Baseline(self.args).to(self.device)
 

{modnn-2.0.0 → modnn-2.0.2}/modnn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modnn
-Version: 2.0.0
+Version: 2.0.2
 Summary: Physics-Informed Modularized Neural Network for Building Energy Modeling
 Author-email: Zixin Jiang <zjiang19@syr.edu>
 License-Expression: MIT

{modnn-2.0.0 → modnn-2.0.2}/modnn.egg-info/SOURCES.txt

@@ -13,6 +13,6 @@ modnn.egg-info/dependency_links.txt
 modnn.egg-info/requires.txt
 modnn.egg-info/top_level.txt
 modnn/Models/BaseNN.py
-modnn/Models/ModNN.py
-modnn/Models/Modback.py
+modnn/Models/ModNN_data.py
+modnn/Models/ModNN_phy.py
 modnn/Models/__init__.py

{modnn-2.0.0 → modnn-2.0.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "modnn"
-version = "2.0.0"
+version = "2.0.2"
 description = "Physics-Informed Modularized Neural Network for Building Energy Modeling"
 readme = "README.md"
 requires-python = ">=3.7"

modnn-2.0.0/modnn/Models/ModNN.py

@@ -1,402 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-# --- Solar Radiation Module ---
-class SolarRadiation(nn.Module):
-    def __init__(self, envelop, name="Solar"):
-        super().__init__()
-        # Initialization
-        self.envelop = envelop
-
-        # Solar transmission coefficient
-        self.direct_gain = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
-
-        # Solar absorption coefficient
-        self.abs_wall = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
-        self.abs_roof = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
-        self.name = name
-
-    def forward(self, radiation, n_wall, n_roof):
-        """
-        Calculate solar heat gains [transmission and absorption]
-
-        Args:
-            radiation: Solar radiation input [batch, time, 1] or [batch, time, directions]
-                      (e.g., [batch, time, 4] for E,S,W,N directions)
-            n_wall: Number of wall components
-            n_roof: Number of roof components
-
-        Returns:
-            direct_gain: Heat directly entering zone
-            wall_gains: List of absorbed heat for each wall [n_wall elements]
-            roof_gains: List of absorbed heat for each roof [n_roof elements]
-        """
-        # Apply softplus to ensure positive coefficients
-        direct_coef   = F.softplus(self.direct_gain) * self.envelop["direct_coef"]
-        abs_wall_coef = F.softplus(self.abs_wall) * self.envelop["abs_wall_coef"]
-        abs_roof_coef = F.softplus(self.abs_roof) * self.envelop["abs_roof_coef"]
-
-        # Calculate direct solar gain to zone
-        direct_gain = direct_coef * radiation
-        # TODO: Distribute transmission solar radiation to each surfaces
-
-        # Calculate absorption by walls
-        wall_gains = []
-        for i in range(n_wall):
-            # Different walls can have different orientations
-            # TODO: ASHRAE Clear-Sky Model
-            if radiation.shape[-1] > 1 and i < radiation.shape[-1]:
-                # Use specific direction for this wall
-                wall_gain = abs_wall_coef * radiation[:, :, i % radiation.shape[-1]]
-                wall_gains.append(wall_gain.unsqueeze(-1))
-            else:
-                # Use total radiation
-                wall_gain = abs_wall_coef * radiation
-                wall_gains.append(wall_gain)
-
-        # Calculate absorption by roof (assume horizontal)
-        roof_gains = []
-        for i in range(n_roof):
-            roof_gain = abs_roof_coef * radiation
-            roof_gains.append(roof_gain)
-
-        return direct_gain, wall_gains, roof_gains
-
-
-# --- Envelope Module [RC] ---
-class EnvelopeRC(nn.Module):
-    def __init__(self, envelop, name="RC", r_init=10.0, c_init=100.0):
-        super().__init__()
-        self.R_inv = nn.Parameter(torch.tensor(1.0 / r_init, dtype=torch.float32))
-        self.C_inv = nn.Parameter(torch.tensor(1.0 / c_init, dtype=torch.float32))
-        self.name = name
-        self.envelop = envelop
-
-    def forward(self, Tamb, Tmid, Tzone, solar_gain=None):
-        # Bounded parameters for stability
-        r_inv_bounded = F.softplus(self.R_inv) * (1.0 / self.envelop["r_opaque_coef"])
-        c_inv_bounded = F.softplus(self.C_inv) * (1.0 / self.envelop["c_opaque_coef"])
-
-        # Calculate heat flows
-        q_in = (Tamb - Tmid) * r_inv_bounded
-        q_out = (Tmid - Tzone) * r_inv_bounded
-
-        # Add solar gain if provided
-        if solar_gain is not None:
-            # Add solar radiation to the first capacitance node
-            dTmid_dt = (q_in - q_out + solar_gain) * c_inv_bounded
-        else:
-            dTmid_dt = (q_in - q_out) * c_inv_bounded
-
-        return dTmid_dt, q_out
-
-
-# --- Envelope Module [R] ---
-class EnvelopeROnly(nn.Module):
-    def __init__(self, envelop, name="R-only", r_init=0.1):
-        super().__init__()
-        self.R_inv = nn.Parameter(torch.tensor(1/r_init, dtype=torch.float32))
-        self.name = name
-        self.envelop = envelop
-
-    def forward(self, Tamb, Tzone):
-        r_inv_bounded = F.softplus(self.R_inv) * (1.0 / self.envelop["r_transparent_coef"])
-
-        return (Tamb - Tzone) * r_inv_bounded
-
-
-# --- Internal Heat Gain Module ---
-class Internal(nn.Module):
-    """
-    A simple MLP is used to predict internal heat gain (convection ONLY)
-    #TODO: add a RNN/MLP with lookback window to consider internal radiation
-
-    In general, the internal heat gain comes from Lighting, Occupant, Appliance, etc,.
-    But they can be represented by a factor (alpha) multiply with a "schedule" (sch), for example:
-    q_light = alpha_light * sch_light
-    q_human = alpha_human * sch_human
-    q_cooking = alpha_cooking * sch_cooking
-
-    If detailed information is available, we can replace it by Physics Equations
-    For example, 80~100W * Number_of_people
-    Otherwise, we learn form periodic features
-    """
-    def __init__(self, input_size, hidden_size, output_size):
-        super().__init__()
-        self.FC1 = nn.Linear(input_size - 1, hidden_size)
-        self.FC2 = nn.Linear(hidden_size, output_size)
-        self.gain = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
-
-        # Initialize weights
-        nn.init.xavier_uniform_(self.FC1.weight, gain=0.1)
-        nn.init.zeros_(self.FC1.bias)
-        nn.init.xavier_uniform_(self.FC2.weight, gain=0.1)
-        nn.init.zeros_(self.FC2.bias)
-
-    def forward(self, x):
-        x1 = F.relu(self.FC1(x[:, :, :-1]))
-        sch = torch.sigmoid(self.FC2(x1)) + x[:, :, -1:]
-        return F.softplus(self.gain) * 0.1 * sch
-
-class HVAC(nn.Module):
-    """
-        A simplified linear module is used here for AIR SIDE SYSTEM ONLY
-        Change it to any type of RNN or add look back window for RADIATION SYSTEM
-        The input is pre-calculated pHVAC (thermal load)
-
-        But if the raw data is HVAC energy, or supply air flow/temperature, fan speed, coil load, etc
-        #TODO: Add More Types of HVAC Modules
-
-    """
-    def __init__(self):
-        super().__init__()
-        self.gain = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
-
-    def forward(self, x):
-        return (F.softplus(self.gain) * 0.1 * x)
-
-
-class ZoneUpdate(nn.Module):
-    def __init__(self, envelop, c_init=10000.0):
-        super().__init__()
-        self.C_inv = nn.Parameter(torch.tensor(1.0 / c_init, dtype=torch.float32))
-        self.envelop = envelop
-
-    def forward(self, q_env_total, q_int, q_hvac, q_solar_direct=None):
-        # Bounded parameter
-        c_inv_bounded = F.softplus(self.C_inv) * (1.0 / self.envelop["c_zone"])
-
-        # Add direct solar gain if provided
-        if q_solar_direct is not None:
-            return (q_env_total + q_int + q_hvac + q_solar_direct) * c_inv_bounded
-        else:
-            return (q_env_total + q_int + q_hvac) * c_inv_bounded
-
-
-# --- Runge-Kutta Integration ---
-def rk4_step(f, x, dt, max_delta=3.0):
-    k1 = f(x)
-    k1 = torch.clamp(k1, -max_delta / dt, max_delta / dt)
-
-    k2 = f(x + 0.5 * dt * k1)
-    k2 = torch.clamp(k2, -max_delta / dt, max_delta / dt)
-
-    k3 = f(x + 0.5 * dt * k2)
-    k3 = torch.clamp(k3, -max_delta / dt, max_delta / dt)
-
-    k4 = f(x + dt * k3)
-    k4 = torch.clamp(k4, -max_delta / dt, max_delta / dt)
-
-    dx = (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
-    return x + torch.clamp(dx, -max_delta, max_delta)
-
-
-class ModNN(nn.Module):
-    def __init__(self, args):
-        super().__init__()
-        para = args["para"]
-        envelop = args["envelop"]
-        self.enc_len = args["enLen"]
-        self.device = args["device"]
-        self.delta_t = 900  # 15 minutes in seconds
-
-        # Model configuration
-        self.n_wall = envelop.get("n_wall", 0)
-        self.n_roof = envelop.get("n_roof", 0)
-        self.n_window = envelop.get("n_window", 0)
-
-        print(f"Initializing ModNN with {self.n_wall} walls, {self.n_roof} roofs, {self.n_window} windows")
-
-        # Assemble modules
-        # Internal heat gain module
-        self.Int = Internal(para["Int_in"], para["Int_h"], para["Int_out"])
-        # HVAC module
-        self.HVAC = HVAC()
-        # Envelope module
-        self.solar = SolarRadiation(envelop=envelop)
-        self.rc_envelopes = nn.ModuleList([
-                                              EnvelopeRC(envelop=envelop,
-                                                         name=f"wall_{i}",
-                                                         r_init=10.0 + i * 1.0,
-                                                         c_init=80.0 + i * 10.0)
-                                              for i in range(self.n_wall)
-                                          ] + [
-                                              EnvelopeRC(envelop=envelop,
-                                                         name=f"roof_{i}",
-                                                         r_init=20.0 + i * 2.0,
-                                                         c_init=60.0 + i * 10.0)
-                                              for i in range(self.n_roof)
-                                          ])
-
-        self.r_only_envelopes = nn.ModuleList([
-            EnvelopeROnly(envelop=envelop,
-                          name=f"window_{i}",
-                          r_init=1 + i * 0.1)
-            for i in range(self.n_window)
-        ])
-
-        self.zone_update = ZoneUpdate(envelop=envelop)
-
-    def forward(self, input_X):
-        """
-        input_X order: [T_zone, T_ambient, solar, day_sin, day_cos, occ, phvac]
-        Shape: [batch_size, time_steps, features]
-        """
-        batch_size = input_X.shape[0]
-        time_steps = input_X.shape[1]
-
-        # Extract input data
-        T_amb = input_X[:, :, [1]]  # Ambient air temperature
-        Solar_X = input_X[:, :, [2]]  # Solar radiation
-        Int_X = input_X[:, :, [3, 4, 5]]  # Time_sin/cos, Occ
-        HVAC_X = input_X[:, :, [6]]  # pHVAC
-        Tzone_gt = input_X[:, :, [0]]  # Zone air temperature groundtruth
-
-        # Pre-allocate results arrays - more efficient than appending
-        Tzone_pred = torch.zeros(batch_size, time_steps, 1, device=input_X.device)
-        q_int_pred = torch.zeros(batch_size, time_steps, 1, device=input_X.device)
-        q_hvac_pred = torch.zeros(batch_size, time_steps, 1, device=input_X.device)
-        q_env_pred = torch.zeros(batch_size, time_steps, 1, device=input_X.device)
-        q_solar_direct_pred = torch.zeros(batch_size, time_steps, 1, device=input_X.device)
-
-        # Pre-allocate envelope temperature histories
-        Tmid_hist = [torch.zeros(batch_size, time_steps, 1, device=input_X.device)
-                     for _ in range(len(self.rc_envelopes))]
-
-        # Initialize zone temperature for first step
-        Tzone = Tzone_gt[:, 0:1, :]
-
-        # Initialize envelope temperatures (only once)
-        Tmid = []
-        for _ in range(len(self.rc_envelopes)):
-            tmid_init = 0.7 * Tzone + 0.3 * T_amb[:, 0:1, :]
-            Tmid.append(tmid_init)
-
-        # --- Encoder phase (optimization: only calculate what's needed) ---
-        for i in range(self.enc_len):
-            # Simply record the ground truth temperature
-            Tzone_pred[:, i:i + 1, :] = Tzone_gt[:, i:i + 1, :]
-
-            # For envelope temperatures, we need to calculate and track these
-            # even during encoder phase as they're not directly observed
-            T_amb_i = T_amb[:, i:i + 1, :]
-            Solar_i = Solar_X[:, i:i + 1, :]
-            Tzone_i = Tzone_gt[:, i:i + 1, :]
-
-            # Calculate solar distribution (needed for envelope temperatures)
-            q_solar_direct, q_solar_wall, q_solar_roof = self.solar(
-                Solar_i, self.n_wall, self.n_roof
-            )
-            q_solar_direct_pred[:, i:i + 1, :] = q_solar_direct
-
-            # Update envelope temperatures
-            for idx, env in enumerate(self.rc_envelopes):
-                # Assign solar gains
-                if idx < self.n_wall and q_solar_wall:
-                    solar_gain = q_solar_wall[idx]
-                elif idx >= self.n_wall and q_solar_roof:
-                    solar_gain = q_solar_roof[idx - self.n_wall]
-                else:
-                    solar_gain = None
-
-                # Evolution function for rk4
-                def dTmid(Tmid_local):
-                    dTdt, _ = env(T_amb_i, Tmid_local, Tzone_i, solar_gain)
-                    return dTdt
-
-                # Update envelope temperature using RK4
-                Tmid[idx] = rk4_step(dTmid, Tmid[idx], self.delta_t)
-                Tmid_hist[idx][:, i:i + 1, :] = Tmid[idx]
-
-            # Only calculate these if needed for analysis, otherwise leave as zeros
-            # in encoder phase since they don't affect the state evolution
-            Int_i = Int_X[:, i:i + 1, :]
-            HVAC_i = HVAC_X[:, i:i + 1, :]
-            q_int_pred[:, i:i + 1, :] = self.Int(Int_i)
-            q_hvac_pred[:, i:i + 1, :] = self.HVAC(HVAC_i)
-
-            # Optional: Calculate envelope heat transfer contributions
-            # if needed for analysis. If not, can skip this.
-            q_env_total = torch.zeros_like(Tzone)
-            for idx, env in enumerate(self.rc_envelopes):
-                if idx < self.n_wall and q_solar_wall:
-                    solar_gain = q_solar_wall[idx]
-                elif idx >= self.n_wall and q_solar_roof:
-                    solar_gain = q_solar_roof[idx - self.n_wall]
-                else:
-                    solar_gain = None
-                _, q_env = env(T_amb_i, Tmid[idx], Tzone_i, solar_gain)
-                q_env_total += q_env
-
-            for env in self.r_only_envelopes:
-                q_env_r = env(T_amb_i, Tzone_i)
-                q_env_total += q_env_r
-
-            q_env_pred[:, i:i + 1, :] = q_env_total
-
-        # --- Decoder phase (predicting future temps) ---
-        # Starting from the last encoder step
-        Tzone = Tzone_gt[:, self.enc_len - 1:self.enc_len, :]
-
-        for i in range(time_steps - self.enc_len):
-            idx = self.enc_len + i
-            T_amb_i = T_amb[:, idx:idx + 1, :]
-            Int_i = Int_X[:, idx:idx + 1, :]
-            HVAC_i = HVAC_X[:, idx:idx + 1, :]
-            Solar_i = Solar_X[:, idx:idx + 1, :]
-
-            # Process internal gains and HVAC - only calculate once
-            q_int = self.Int(Int_i)
-            q_hvac = self.HVAC(HVAC_i)
-            q_int_pred[:, idx:idx + 1, :] = q_int
-            q_hvac_pred[:, idx:idx + 1, :] = q_hvac
-
-            # Distribute solar radiation
-            q_solar_direct, q_solar_wall, q_solar_roof = self.solar(
-                Solar_i, self.n_wall, self.n_roof
-            )
-            q_solar_direct_pred[:, idx:idx + 1, :] = q_solar_direct
-
-            # Prepare for zone temperature update
-            q_env_total = torch.zeros_like(Tzone)
-
-            # Update envelope temperatures and calculate heat flow
-            for j, env in enumerate(self.rc_envelopes):
-                # Assign solar gains
-                if j < self.n_wall and q_solar_wall:
-                    solar_gain = q_solar_wall[j]
-                elif j >= self.n_wall and q_solar_roof:
-                    solar_gain = q_solar_roof[j - self.n_wall]
-                else:
-                    solar_gain = None
-
-                def dTmid(Tmid_local):
-                    dTdt, _ = env(T_amb_i, Tmid_local, Tzone, solar_gain)
-                    return dTdt
-
-                # Update envelope temperature
-                Tmid[j] = rk4_step(dTmid, Tmid[j], self.delta_t)
-                Tmid_hist[j][:, idx:idx + 1, :] = Tmid[j]
-
-                # Calculate heat flow contribution (calculate once)
-                _, q_env = env(T_amb_i, Tmid[j], Tzone, solar_gain)
-                q_env_total += q_env
-
-            # Add window heat flow
-            for env in self.r_only_envelopes:
-                q_env_r = env(T_amb_i, Tzone)
-                q_env_total += q_env_r
-
-            q_env_pred[:, idx:idx + 1, :] = q_env_total
-
-            # Update zone temperature using RK4
-            def dTzone(Tzone_local):
-                return self.zone_update(q_env_total, q_int, q_hvac, q_solar_direct)
-
-            # Update zone temperature with stability limits
-            Tzone = rk4_step(dTzone, Tzone, self.delta_t, max_delta=2.0)
-            Tzone_pred[:, idx:idx + 1, :] = Tzone
-
-        return Tzone_pred, Tzone_gt, (q_int_pred, q_hvac_pred, q_env_pred, q_solar_direct_pred)