Update pinn_electromagnetics/losses.py

pinn_electromagnetics/losses.py

@@ -1,14 +1,13 @@
 
 import torch
-import numpy as np # Used for constants like pi if needed elsewhere
+import numpy as np 
 
-# --- Derivative Computation (from v28) ---
 def compute_all_derivatives(model_phys, x, y, z, t):
     outputs = model_phys(x, y, z, t)
     Ex, Ey, Ez = outputs[:, 0:1], outputs[:, 1:2], outputs[:, 2:3]
     Bx, By, Bz = outputs[:, 3:4], outputs[:, 4:5], outputs[:, 5:6]
 
-    # Use .sum() for scalar output when computing gradients over multiple components
+
     dEx_grads = torch.autograd.grad(Ex.sum(), [x, y, z, t], create_graph=True, allow_unused=True)
     dEy_grads = torch.autograd.grad(Ey.sum(), [x, y, z, t], create_graph=True, allow_unused=True)
     dEz_grads = torch.autograd.grad(Ez.sum(), [x, y, z, t], create_graph=True, allow_unused=True)
@@ -16,7 +15,7 @@ def compute_all_derivatives(model_phys, x, y, z, t):
     dBy_grads = torch.autograd.grad(By.sum(), [x, y, z, t], create_graph=True, allow_unused=True)
     dBz_grads = torch.autograd.grad(Bz.sum(), [x, y, z, t], create_graph=True, allow_unused=True)
 
-    # Ensure gradients are not None if allow_unused=True
+   
     def get_grad(grad_tuple, idx):
         return grad_tuple[idx] if grad_tuple[idx] is not None else torch.zeros_like(x)
 
@@ -30,23 +29,18 @@ def compute_all_derivatives(model_phys, x, y, z, t):
     }
     return outputs, derivs
 
-# --- Maxwell Loss (from v28) ---
 def compute_maxwell_loss(derivs):
-    # Maxwell's Equations (dimensionless units)
 
-    # Gauss's Law for E: div(E) = 0
+
     loss_gauss_E = (derivs['dEx_dx'] + derivs['dEy_dy'] + derivs['dEz_dz'])**2
 
-    # Gauss's Law for B: div(B) = 0
     loss_gauss_B = (derivs['dBx_dx'] + derivs['dBy_dy'] + derivs['dBz_dz'])**2
 
-    # Faraday's Law: curl(E) = -dB/dt
     loss_faraday_x = (derivs['dEz_dy'] - derivs['dEy_dz'] + derivs['dBx_dt'])**2
     loss_faraday_y = (derivs['dEx_dz'] - derivs['dEz_dx'] + derivs['dBy_dt'])**2
     loss_faraday_z = (derivs['dEy_dx'] - derivs['dEx_dy'] + derivs['dBz_dt'])**2
     loss_faraday = loss_faraday_x + loss_faraday_y + loss_faraday_z
-
-    # Ampere-Maxwell Law: curl(B) = dE/dt (assuming no current J)
+    
     loss_ampere_x = (derivs['dBz_dy'] - derivs['dBy_dz'] - derivs['dEx_dt'])**2
     loss_ampere_y = (derivs['dBx_dz'] - derivs['dBz_dx'] - derivs['dEy_dt'])**2
     loss_ampere_z = (derivs['dBy_dx'] - derivs['dBx_dy'] - derivs['dEz_dt'])**2
@@ -54,7 +48,6 @@ def compute_maxwell_loss(derivs):
 
     return torch.mean(loss_gauss_E + loss_gauss_B + loss_faraday + loss_ampere)
 
-# --- Data Loss (from v31.2, adapted for NPZ data) ---
 def compute_data_loss(model_phys,
                       coords_tensor, E0_tensor, B0_tensor,
                       source_pos_tensor, source_orientation_tensor,
@@ -62,7 +55,6 @@ def compute_data_loss(model_phys,
                       n_samples_ic=1024, n_samples_source=128):
     loss_data = 0.0
 
-    # 1. Initial Conditions (t=0): E=E0, B=B0
     idx_ic = torch.randperm(coords_tensor.shape[0], device=coords_tensor.device)[:n_samples_ic]
     x_ic, y_ic, z_ic = coords_tensor[idx_ic, 0:1], coords_tensor[idx_ic, 1:2], coords_tensor[idx_ic, 2:3]
     t_ic = torch.zeros_like(x_ic)
@@ -72,22 +64,16 @@ def compute_data_loss(model_phys,
 
     loss_data += torch.mean((pred_ic[:, 0:3] - E_actual_ic)**2)
     loss_data += torch.mean((pred_ic[:, 3:6] - B_actual_ic)**2)
-
-    # 2. Source Boundary (time-varying E-field, B=0)
-    # We use all time points for the source signal
     n_times_source = t_tensor.shape[0]
-    # Source position needs to be expanded to match the number of time samples
     x_s = source_pos_tensor[0:1].expand(n_times_source, -1)
     y_s = source_pos_tensor[1:2].expand(n_times_source, -1)
     z_s = source_pos_tensor[2:3].expand(n_times_source, -1)
-    t_s = t_tensor # shape (n_times_source, 1)
+    t_s = t_tensor 
 
     pred_source = model_phys(x_s, y_s, z_s, t_s)
-
-    # E_actual = signal * orientation. Expand orientation to match time samples.
     E_actual_source = source_signal_tensor * source_orientation_tensor.view(1, 3)
 
     loss_data += torch.mean((pred_source[:, 0:3] - E_actual_source)**2)
-    loss_data += torch.mean(pred_source[:, 3:6]**2) # B=0 at the source
+    loss_data += torch.mean(pred_source[:, 3:6]**2)
 
     return loss_data