[mesh motion] best gaussian approximation settingss for stream function

Author: tmaric

src/displacementSmartSimMotionSolver/pytorchApproximationModels/rbf_network.py

@@ -1,16 +1,34 @@
 import torch
 import torch.nn as nn
 
+import torch
+
 # Define various RBF functions with enforced compact support
 def gaussian_rbf(r):
-    """Infinitely smooth Gaussian RBF, matching Wendland's implementation."""
-    return torch.exp(-r**2)  # Matches WendlandLinearNetwork
+    """Infinitely smooth Gaussian RBF."""
+    return torch.exp(-r**2)
+
+def wendland_d2_c2_rbf(r):
+    """
+    Wendland's C^2 RBF for d=2.
+    Compactly supported, continuously differentiable (C^2).
+    
+    Formula: (1 - r)^4_+ (4r + 1)
+    """
+    mask = (r < 1).float()
+    rm = (1 - r).clamp(min=0.0)
+    return mask * rm**4 * (4 * r + 1)
 
-def wendland_rbf(r):
-    """Compactly supported Wendland C^4 RBF."""
+def wendland_d2_c4_rbf(r):
+    """
+    Wendland's C^4 RBF for d=2.
+    Compactly supported, twice continuously differentiable (C^4).
+    
+    Formula: (1 - r)^6_+ (35r^2 + 18r + 3)
+    """
     mask = (r < 1).float()
     rm = (1 - r).clamp(min=0.0)
-    return mask * (1 + 6*r + (35/3)*r**2) * rm**6
+    return mask * rm**6 * (35 * r**2 + 18 * r + 3)
 
 def multiquadric_rbf(r):
     """Multiquadric RBF with compact support."""
@@ -25,7 +43,8 @@ def inverse_multiquadric_rbf(r):
 # Create an RBF function dictionary
 rbf_dict = {
     "gaussian": gaussian_rbf,
-    "wendland": wendland_rbf,
+    "wendland_d2_c2": wendland_d2_c2_rbf,
+    "wendland_d2_c4": wendland_d2_c4_rbf,
     "multiquadric": multiquadric_rbf,
     "inverse_multiquadric": inverse_multiquadric_rbf
 }

src/displacementSmartSimMotionSolver/pytorchApproximationModels/test_rbf_network_stream_function.py

@@ -24,11 +24,11 @@ def compute_velocity(x, y, psi_values):
     v = -dx  # v = -∂ψ/∂x
     return u, v
 
-def visualize_psi(x, y, psi_values, centers, title="Stream Function"):
+def visualize_psi(x, y, psi_values, rbf_type, centers, title):
     plt.figure(figsize=(6, 6))
     plt.contourf(x, y, psi_values, levels=20, cmap='viridis')
     plt.colorbar(label='ψ')
-    plt.title(title + f" num_centers {len(centers)}")
+    plt.title(title + f"-rbf_type_{rbf_type}-num_centers_{len(centers)}")
     plt.xlabel('x')
     plt.ylabel('y')
     plt.grid()
@@ -39,17 +39,17 @@ def visualize_psi(x, y, psi_values, centers, title="Stream Function"):
     plt.legend()
 
     fig_name = title.replace(" ", "-")
-    plt.savefig(f"{fig_name}-num_centers-{len(centers)}.png", dpi=200)
+    plt.savefig(f"{fig_name}-rbf_type_{rbf_type}-num_centers_{len(centers)}.png", dpi=200)
 
-def visualize_velocity_field(x, y, u, v, num_centers, title="Velocity Field"):
+def visualize_velocity_field(x, y, u, v, rbf_type, num_centers, title="Velocity"):
     plt.figure(figsize=(6, 6))
     plt.quiver(x, y, u, v)
-    plt.title(title)
+    plt.title(f"{title}-{rbf_type}-n_centers_{num_centers}")
     plt.xlabel('x')
     plt.ylabel('y')
     plt.grid()
     fig_name = title.replace(" ", "-")
-    plt.savefig(f"{fig_name}-num_centers{num_centers}.png", dpi=200)
+    plt.savefig(f"{fig_name}-rbf_type_{rbf_type}-num_centers_{num_centers}.png", dpi=200)
 
 def generate_centers(num_centers):
     """
@@ -106,7 +106,7 @@ def estimate_convergence_order(csv_filename):
 
     print(f"Updated {csv_filename} with convergence orders for {error_columns}.")
 
-def main(num_points):
+def main(num_points, rbf_type):
     # Generate training data
     x = np.linspace(0, 1, num_points)
     y = np.linspace(0, 1, num_points)
@@ -119,21 +119,26 @@ def main(num_points):
     y_train = torch.tensor(psi_train, dtype=torch.float32)
 
     # Generate centers
-    # centers = generate_centers(32).clone().detach()
     centers = x_train 
     print(centers.shape)
+
+    # Gaussian 3d-order support
+    #r_max = 2.5 / num_points 
     r_max = 2.5 / num_points 
-    smoothness = 4  # C^4 smoothness
 
     # Initialize model
-    model = RadialBasisFunctionNetwork(centers, r_max, rbf_dict, rbf_type="gaussian")
+    model = RadialBasisFunctionNetwork(centers, r_max, rbf_dict, rbf_type=rbf_type)
 
     # Optimizer and loss
     optimizer = optim.Adam(model.parameters(), lr=0.05)
-    criterion = nn.MSELoss()
+    criterion = torch.nn.MSELoss()
 
     # Training loop
     epochs = 4000
+    best_loss = float("inf")  # Initialize best loss to a large value
+    best_model_state = None  # Store best model state
+    stop_loss = 1e-08
+
     for epoch in range(epochs):
         model.train()
         optimizer.zero_grad()
@@ -142,12 +147,28 @@ def main(num_points):
         loss.backward()
         optimizer.step()
 
-        if loss.item() < 1e-12:
-            print(f"Stopping early at epoch {epoch + 1} due to reaching loss {loss.item()} < 1e-05")
+        # Save the model if it has the lowest loss so far
+        if loss.item() < best_loss:
+            best_loss = loss.item()
+            best_model_state = model.state_dict().copy()  # Copy best model state
+
+        # Early stopping criterion
+        if loss.item() < stop_loss:
+            print(f"Stopping early at epoch {epoch + 1} due to reaching loss {loss.item()} < {stop_loss}")
             break
 
-        if ((epoch == 1) or ((epoch + 1) % 50 == 0)):
-            print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.14f}')
+        # Print progress every 50 epochs
+        if epoch == 1 or (epoch + 1) % 50 == 0:
+            print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.14f}, Best Loss: {best_loss:.14f}')
+
+    # Restore the best model state if training didn't reach convergence
+    if best_model_state:
+        model.load_state_dict(best_model_state)
+        print(f"Restored best model with loss: {best_loss:.14f}")
+
+    # Save the best model to file
+    torch.save(best_model_state, "best_rbf_model.pth")
+    print("Best model saved as 'best_rbf_model.pth'.")
 
     # Generate validation data
     num_points_val = 100
@@ -167,19 +188,24 @@ def main(num_points):
     psi_pred = pred.reshape(X_val.shape)
 
     # Visualize actual and predicted stream functions
-    visualize_psi(X_val, Y_val, psi_val, centers, title="Actual Stream Function")
-    visualize_psi(X_val, Y_val, psi_pred, centers, title="Predicted Stream Function")
+    visualize_psi(X_val, Y_val, psi_val, rbf_type, 
+                  centers, title="Actual Stream Function")
+    visualize_psi(X_val, Y_val, psi_pred, rbf_type,
+                  centers, title="Predicted Stream Function")
 
     err_val = np.abs(psi_pred - psi_val) / np.max(psi_val)
-    visualize_psi(X_val, Y_val, err_val, centers,
-                  title="Stream Function Relative Approximation Error")
+    visualize_psi(X_val, Y_val, err_val, rbf_type, centers,
+                  title="Stream Function Relative Error")
 
     # Define the filename
     csv_filename = "stream_function_validation.csv"
 
     # Define the header and the values to be appended
-    header = ["num_points", "point_dist", "r_max", "err_mean", "err_max"]
-    data = [num_points, 1.0 / num_points, r_max, np.mean(err_val), np.max(err_val)]
+    header = ["model_rbf_type", "num_points", "support_radius", 
+              "point_dist", "err_mean", "err_max"]
+
+    data = [model.rbf_type, num_points, r_max, 1.0 / num_points, 
+            np.mean(err_val), np.max(err_val)]
 
     # Check if file exists
     file_exists = os.path.isfile(csv_filename)
@@ -199,18 +225,19 @@ def main(num_points):
 
     # Compute and visualize actual and predicted velocity fields
     u_val, v_val = compute_velocity(X_val, Y_val, psi_val)
-    visualize_velocity_field(X_val, Y_val, u_val, v_val, num_points,
-                             title="Actual Velocity Field")
+    visualize_velocity_field(X_val, Y_val, u_val, v_val, rbf_type, num_points,
+                             title="Velocity Field")
 
     u_pred, v_pred = compute_velocity(X_val, Y_val, psi_pred)
-    visualize_velocity_field(X_val, Y_val, u_pred, v_pred, num_points,
+    visualize_velocity_field(X_val, Y_val, u_pred, v_pred, rbf_type, num_points,
                              title="Predicted Velocity Field")
 
 if __name__ == "__main__":
 
-    # Run mesh convergence study
-    for num_points in [4,8,16,32]:
-        main(num_points)
+    # Run the parameter study 
+    for rbf_type in ["gaussian"]:
+        for num_points in [4,8,16,32]:
+            main(num_points, rbf_type)
 
     # Estimate convergence order
     estimate_convergence_order("stream_function_validation.csv")