[mesh motion] generalized RBF network, Gaussian converges

Author: tmaric

src/displacementSmartSimMotionSolver/pytorchApproximationModels/rbf_network.py

@@ -1,61 +1,78 @@
 import torch
 import torch.nn as nn
-import math
 
-class WendlandLinearNetwork(nn.Module):
-    def __init__(self, centers, r_max, smoothness):
+# 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
+
+def wendland_rbf(r):
+    """Compactly supported Wendland C^4 RBF."""
+    mask = (r < 1).float()
+    rm = (1 - r).clamp(min=0.0)
+    return mask * (1 + 6*r + (35/3)*r**2) * rm**6
+
+def multiquadric_rbf(r):
+    """Multiquadric RBF with compact support."""
+    mask = (r < 1).float()
+    return mask * torch.sqrt(1 + r**2)
+
+def inverse_multiquadric_rbf(r):
+    """Inverse multiquadric RBF with compact support."""
+    mask = (r < 1).float()
+    return mask / torch.sqrt(1 + r**2)
+
+# Create an RBF function dictionary
+rbf_dict = {
+    "gaussian": gaussian_rbf,
+    "wendland": wendland_rbf,
+    "multiquadric": multiquadric_rbf,
+    "inverse_multiquadric": inverse_multiquadric_rbf
+}
+
+class RadialBasisFunctionNetwork(nn.Module):
+    def __init__(self, centers, r_max, rbf_dict, rbf_type):
         """
-        Initialize the Wendland RBF network with linear polynomial terms.
+        Generalized RBF network with user-selectable RBF functions.
 
         Parameters:
-        centers (torch.Tensor): shape (num_centers, dimension) RBF centers.
-        r_max (float): radius of compact support.
-        smoothness (int): even integer specifying desired smoothness (C^smoothness).
+        centers (torch.Tensor): shape (num_centers, dimension), RBF centers.
+        r_max (float): radius of compact support (applies to all RBFs).
+        rbf_dict (dict): Dictionary mapping RBF type names to function implementations.
+        rbf_type (str): Type of RBF function to use (must be in rbf_dict).
         """
         super().__init__()
 
-        if smoothness != 4:
-            raise NotImplementedError("Only smoothness=4 (C⁴) is currently implemented explicitly.")
-
-        self.centers = centers.clone().detach()
+        self.centers = centers.clone().detach()  # Fixed RBF centers
         self.r_max = r_max
-        self.smoothness = smoothness
-        self.k = smoothness // 2
         self.num_centers, self.dimension = centers.shape
+        self.rbf_type = rbf_type.lower()  # Store selected RBF type as an attribute
 
-        # Trainable parameters (explicitly initialized!)
-        self.weights = nn.Parameter(torch.zeros(self.num_centers))
-        self.a0 = nn.Parameter(torch.tensor(0.0))
-        self.a = nn.Parameter(torch.zeros(self.dimension))
+        # Ensure rbf_type is valid
+        if self.rbf_type not in rbf_dict:
+            raise ValueError(f"Invalid RBF type '{self.rbf_type}'. Available options: {list(rbf_dict.keys())}")
 
-    #def rbf(self, x):
-    #    r = torch.cdist(x, self.centers) / self.r_max
-    #    mask = (r < 1).float()
-    #    rm = (1 - r).clamp(min=0.0)
+        self.rbf_function = rbf_dict[self.rbf_type]  # Store selected RBF function
 
-    #    phi = (1 + 6*r + (35/3)*r**2) * rm**6  # Only compute once
-    #    return phi * mask  # Ensure compact support
+        # Trainable parameters (weights for RBFs) - match Wendland's model!
+        self.weights = nn.Parameter(torch.zeros(self.num_centers))  # Initialize to zeros
+        self.a0 = nn.Parameter(torch.tensor(0.0))  # Bias term initialized as 0 (like Wendland)
 
     def rbf(self, x):
         """
-        Compute Gaussian RBF instead of Wendland.
+        Compute the RBF values for input x using the selected RBF function.
         """
-        r = torch.cdist(x, self.centers) / self.r_max
-        return torch.exp(-r**2)  # Infinitely smooth Gaussian RBF
-
-    #def rbf(self, x):
-    #    """
-    #    Compute Gaussian RBF with enforced compact support.
-    #    """
-    #    r = torch.cdist(x, self.centers) / self.r_max
-    #    mask = (r < 1).float()  # Ensure compact support
-    #    return mask * torch.exp(-r**2)  # Compactly supported Gaussian
+        r = torch.cdist(x, self.centers) / self.r_max  # Compute normalized distance
+        return self.rbf_function(r)  # Apply the selected RBF function
 
     def forward(self, x):
         """
-        Forward pass: Wendland RBF with explicit linear polynomial extension.
+        Forward pass: Compute RBF output.
         """
         rbf_output = self.rbf(x)
         rbf_term = rbf_output @ self.weights
-        linear_term = x @ self.a
-        return self.a0 + rbf_term #+ linear_term
+        return self.a0 + rbf_term  # No polynomial correction term
+
+    def get_rbf_type(self):
+        """Return the currently selected RBF type."""
+        return self.rbf_type

src/displacementSmartSimMotionSolver/pytorchApproximationModels/test_rbf_network_stream_function.py

@@ -7,7 +7,7 @@
 import pandas as pd
 import os
 
-from rbf_network import WendlandLinearNetwork
+from rbf_network import rbf_dict, RadialBasisFunctionNetwork 
 
 def psi(x, y):
     """
@@ -28,7 +28,7 @@ def visualize_psi(x, y, psi_values, centers, title="Stream Function"):
     plt.figure(figsize=(6, 6))
     plt.contourf(x, y, psi_values, levels=20, cmap='viridis')
     plt.colorbar(label='ψ')
-    plt.title(title)
+    plt.title(title + f" num_centers {len(centers)}")
     plt.xlabel('x')
     plt.ylabel('y')
     plt.grid()
@@ -64,7 +64,7 @@ def generate_centers(num_centers):
 def estimate_convergence_order(csv_filename):
     """
     Opens a CSV file containing numerical convergence results and estimates the
-    convergence order for each row using log-log error reduction.
+    convergence order for each error column using log-log error reduction.
 
     The last row's convergence order is set equal to the second-to-last row.
     
@@ -75,33 +75,36 @@ def estimate_convergence_order(csv_filename):
     df = pd.read_csv(csv_filename)
 
     # Ensure required columns exist
-    required_columns = {"point_dist", "err_validation"}
+    required_columns = {"point_dist", "err_mean", "err_max"}
     if not required_columns.issubset(df.columns):
         raise ValueError(f"Missing required columns in CSV: {required_columns - set(df.columns)}")
 
-    # Compute convergence order using log-log slope formula
-    convergence_orders = []
+    # List of error columns to process
+    error_columns = ["err_mean", "err_max"]
+
+    for error_col in error_columns:
+        convergence_orders = []  # Store convergence orders for this error type
 
-    for i in range(len(df) - 1):  # Iterate up to the second-to-last row
-        h_coarse, h_fine = df.iloc[i]["point_dist"], df.iloc[i + 1]["point_dist"]
-        err_coarse, err_fine = df.iloc[i]["err_validation"], df.iloc[i + 1]["err_validation"]
+        for i in range(len(df) - 1):  # Iterate up to the second-to-last row
+            h_coarse, h_fine = df.iloc[i]["point_dist"], df.iloc[i + 1]["point_dist"]
+            err_coarse, err_fine = df.iloc[i][error_col], df.iloc[i + 1][error_col]
 
-        if err_coarse > 0 and err_fine > 0:  # Avoid log errors due to zero or negative values
-            p = np.log(err_coarse / err_fine) / np.log(h_coarse / h_fine)
-            convergence_orders.append(p)
-        else:
-            convergence_orders.append(np.nan)
+            if err_coarse > 0 and err_fine > 0:  # Avoid log errors due to zero or negative values
+                p = np.log(err_coarse / err_fine) / np.log(h_coarse / h_fine)
+                convergence_orders.append(p)
+            else:
+                convergence_orders.append(np.nan)
 
-    # Ensure last row gets the same convergence order as the previous row
-    convergence_orders.append(convergence_orders[-1] if len(convergence_orders) > 0 else np.nan)
+        # Ensure last row gets the same convergence order as the previous row
+        convergence_orders.append(convergence_orders[-1] if len(convergence_orders) > 0 else np.nan)
 
-    # Add convergence order column
-    df["error_convergence_order"] = convergence_orders
+        # Add convergence order column to DataFrame
+        df[f"{error_col}_convergence_order"] = convergence_orders
 
     # Save the updated CSV file
     df.to_csv(csv_filename, index=False)
 
-    print(f"Updated {csv_filename} with convergence orders.")
+    print(f"Updated {csv_filename} with convergence orders for {error_columns}.")
 
 def main(num_points):
     # Generate training data
@@ -123,7 +126,7 @@ def main(num_points):
     smoothness = 4  # C^4 smoothness
 
     # Initialize model
-    model = WendlandLinearNetwork(centers, r_max, smoothness)
+    model = RadialBasisFunctionNetwork(centers, r_max, rbf_dict, rbf_type="gaussian")
 
     # Optimizer and loss
     optimizer = optim.Adam(model.parameters(), lr=0.05)
@@ -175,8 +178,8 @@ def main(num_points):
     csv_filename = "stream_function_validation.csv"
 
     # Define the header and the values to be appended
-    header = ["num_points", "point_dist", "r_max", "err_validation"]
-    data = [num_points, 1.0 / num_points, r_max, np.mean(err_val)]
+    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)]
 
     # Check if file exists
     file_exists = os.path.isfile(csv_filename)