Upload models (2).py

pinn_electromagnetics/models (2).py

@@ -0,0 +1,117 @@
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+# --- Hash Grid Encoder (from v26) ---
+class HashGridEncoder(nn.Module):
+    def __init__(self, n_levels=16, n_features_per_level=2,
+                 log2_hashmap_size=19, base_resolution=16,
+                 per_level_scale=1.5):
+        super(HashGridEncoder, self).__init__()
+        self.n_levels = n_levels
+        self.n_features_per_level = n_features_per_level
+        self.log2_hashmap_size = log2_hashmap_size
+        self.base_resolution = base_resolution
+        self.per_level_scale = per_level_scale
+        self.hashmap_size = 2**self.log2_hashmap_size
+        self.embeddings = nn.ModuleList([
+            nn.Embedding(self.hashmap_size, self.n_features_per_level)
+            for i in range(self.n_levels)
+        ])
+        for emb in self.embeddings: nn.init.uniform_(emb.weight, -1e-4, 1e-4)
+
+    def hash_fn(self, coords_int, level):
+        primes = torch.tensor([1, 2654435761, 805459861], dtype=torch.int64, device=coords_int.device)
+        hashed_indices = (coords_int * primes).sum(dim=-1) % self.hashmap_size
+        return hashed_indices
+
+    def trilinear_interp(self, x, v000, v001, v010, v011, v100, v101, v110, v111):
+        w0 = (1 - x[..., 0:1]); w1 = x[..., 0:1]
+        c00 = v000 * w0 + v100 * w1; c01 = v001 * w0 + v101 * w1
+        c10 = v010 * w0 + v110 * w1; c11 = v011 * w0 + v111 * w1
+        c0 = c00 * (1 - x[..., 1:2]) + c10 * x[..., 1:2]
+        c1 = c01 * (1 - x[..., 1:2]) + c11 * x[..., 1:2]
+        c = c0 * (1 - x[..., 2:3]) + c1 * x[..., 2:3]
+        return c
+
+    def forward(self, x_coords): # x_coords: [N, 3] in [0, 1]
+        all_features = []
+        for l in range(self.n_levels):
+            res = int(self.base_resolution * (self.per_level_scale ** l))
+            x_scaled = x_coords * res
+            x_floor = torch.floor(x_scaled).to(torch.int64)
+            x_local = (x_scaled - x_floor)
+            corners = torch.tensor([[0,0,0],[0,0,1],[0,1,0],[0,1,1],[1,0,0],[1,0,1],[1,1,0],[1,1,1]], dtype=torch.int64, device=x_coords.device)
+            cell_corners_int = x_floor.unsqueeze(1) + corners.unsqueeze(0)
+            hashed_indices = self.hash_fn(cell_corners_int, l)
+            embedded_features = self.embeddings[l](hashed_indices)
+            interp_features = self.trilinear_interp(
+                x_local,
+                embedded_features[:, 0], embedded_features[:, 1],
+                embedded_features[:, 2], embedded_features[:, 3],
+                embedded_features[:, 4], embedded_features[:, 5],
+                embedded_features[:, 6], embedded_features[:, 7]
+            )
+            all_features.append(interp_features)
+        return torch.cat(all_features, dim=-1)
+
+# --- GeoNetHash (for geometry, from v26) ---
+class GeoNetHash(nn.Module):
+    def __init__(self, box_size_scale=1.0):
+        super(GeoNetHash, self).__init__()
+        self.box_size_scale = box_size_scale
+        self.encoder = HashGridEncoder(n_levels=16, n_features_per_level=2)
+        self.mlp = nn.Sequential(
+            nn.Linear(self.encoder.n_levels * self.encoder.n_features_per_level, 64),
+            nn.ReLU(),
+            nn.Linear(64, 1)
+        )
+        # Placeholder for min/max for normalization
+        self.min = None
+        self.max = None
+
+    def normalize_coords(self, coords):
+        # Stores min/max from the original dataset for consistent normalization
+        self.min = coords.min(dim=0, keepdim=True)[0]
+        self.max = coords.max(dim=0, keepdim=True)[0]
+        self.max[self.max == self.min] += 1e-6 # Avoid division by zero
+        return (coords - self.min) / (self.max - self.min)
+
+    def forward(self, coords_input_normalized): # input is already normalized [0, 1]
+        features = self.encoder(coords_input_normalized)
+        return self.mlp(features)
+
+# --- Positional Encoder 4D (for MaxwellPINN, from v19/v22) ---
+class PositionalEncoder4D(nn.Module):
+    def __init__(self, input_dims, num_freqs):
+        super(PositionalEncoder4D, self).__init__()
+        self.input_dims = input_dims
+        self.num_freqs = num_freqs
+        self.freq_bands = 2.**torch.linspace(0., num_freqs-1, num_freqs)
+        self.output_dims = self.input_dims * (2 * self.num_freqs + 1)
+
+    def forward(self, x): # x is [N, input_dims]
+        encoding = [x]
+        x_freq = x.unsqueeze(-1) * self.freq_bands.to(x.device)
+        encoding.append(torch.sin(x_freq).view(x.shape[0], -1))
+        encoding.append(torch.cos(x_freq).view(x.shape[0], -1))
+        return torch.cat(encoding, dim=1)
+
+# --- MaxwellPINN (for electromagnetics, from v28) ---
+class MaxwellPINN(nn.Module):
+    def __init__(self, num_freqs=8, hidden_dim=128):
+        super(MaxwellPINN, self).__init__()
+        self.encoder = PositionalEncoder4D(input_dims=4, num_freqs=num_freqs)
+        encoder_output_dims = self.encoder.output_dims
+        self.network = nn.Sequential(
+            nn.Linear(encoder_output_dims, hidden_dim), nn.Tanh(),
+            nn.Linear(hidden_dim, hidden_dim), nn.Tanh(),
+            nn.Linear(hidden_dim, hidden_dim), nn.Tanh(),
+            nn.Linear(hidden_dim, 6) # Outputs (Ex, Ey, Ez, Bx, By, Bz)
+        )
+
+    def forward(self, x, y, z, t):
+        coords = torch.cat([x, y, z, t], dim=1)
+        x_encoded = self.encoder(coords)
+        return self.network(x_encoded)