program

Signed-off-by: 葛峻恺 <202115006@mail.sdu.edu.cn>

utils/layers_generator.py

@@ -0,0 +1,93 @@
+import numpy as np
+import random
+
+def generate_layers_1d(grid_size, num_layers, min_size, transition_size, first_layer_minsize, first_layer_maxsize): 
+    """Generate a 1D stratified model with transition zones."""
+    usable_size = grid_size - (num_layers - 1) * transition_size
+    grid = np.zeros(grid_size, dtype=int)
+    remaining_size = usable_size
+    layer_sizes = []
+
+    # Determine layer sizes
+    for i in range(num_layers - 1):
+        if i == 0:
+            size = random.randint(first_layer_minsize, first_layer_maxsize)
+        else:
+            size = random.randint(min_size, remaining_size - (num_layers - i - 1) * min_size)
+        layer_sizes.append(size)
+        remaining_size -= size
+    layer_sizes.append(remaining_size)
+
+    # Assign layers and transitions to grid
+    current_position = 0
+    for i, size in enumerate(layer_sizes):
+        grid[current_position:current_position + size] = i + 1
+        current_position += size
+        if i < num_layers - 1:
+            grid[current_position:current_position + transition_size] = 0  # 0 indicates transition zone
+            current_position += transition_size
+
+    return grid
+
+def assign_permittivity(grid, num_layers, permittivity_range=(1, 81), transition_size=5):
+    """Assign permittivity values to layers and interpolate values across transition zones."""
+    permittivity_grid = np.zeros_like(grid, dtype=float)
+    layer_permittivities = {}
+
+    # Assign permittivity to each layer
+    for layer in range(1, num_layers + 1):
+        value = random.uniform(*permittivity_range)
+        layer_permittivities[layer] = value
+        permittivity_grid[grid == layer] = value
+
+    # Smooth transitions between layers using linear interpolation
+    for i in range(1, num_layers):
+        start_val = layer_permittivities[i]
+        end_val = layer_permittivities[i + 1]
+        transition_start = np.where(grid == 0)[0][(i - 1) * transition_size]
+        for j in range(transition_size):
+            t = j / (transition_size - 1)
+            permittivity_grid[transition_start + j] = (1 - t) * start_val + t * end_val
+
+    return permittivity_grid
+
+def assign_permittivity_with_smooth_transition(grid, num_layers, first_layer_eps_range, permittivity_range=(1, 81), transition_size=5):
+    """Assign permittivity values using smooth cosine transition between layers."""
+    permittivity_grid = np.zeros_like(grid, dtype=float)
+    layer_permittivities = {}
+
+    for layer in range(1, num_layers + 1):
+        if layer == 1:
+            value = random.randint(*first_layer_eps_range)
+        else:
+            value = random.uniform(*permittivity_range)
+        layer_permittivities[layer] = value
+        permittivity_grid[grid == layer] = value
+
+    # Cosine-smooth transitions
+    for i in range(1, num_layers):
+        start_val = layer_permittivities[i]
+        end_val = layer_permittivities[i + 1]
+        transition_start = np.where(grid == 0)[0][(i - 1) * transition_size]
+        for j in range(transition_size):
+            t = j / (transition_size - 1)
+            smooth = 0.5 * (1 - np.cos(np.pi * t))
+            permittivity_grid[transition_start + j] = start_val + (end_val - start_val) * smooth
+
+    return permittivity_grid
+
+def assign_integer_values(permittivity_grid):
+    """Convert permittivity values to discrete integers representing unique material zones."""
+    layer_id = np.zeros_like(permittivity_grid, dtype=float)
+    unique_permittivities = {}
+    integer_value = 0
+
+    for idx, val in enumerate(permittivity_grid):
+        val = round(val, 5)
+        if idx == 0 or val != round(permittivity_grid[idx - 1], 5):
+            unique_permittivities[val] = integer_value
+            integer_value += 1
+        layer_id[idx] = unique_permittivities[val]
+
+    return layer_id
+

utils/plot.py

@@ -0,0 +1,135 @@
+import os
+import sys
+import torch
+from torchsummary import summary
+from torchvision.utils import make_grid, save_image
+import numpy as np
+import matplotlib.pyplot as plt
+import scipy.ndimage
+# Add parent directory to path for config import
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from config import Network_train_Config as cfg
+
+
+def plot_BSCAN_data(data, path, line_length=100, time_length=200, ratio=1):
+    """
+    Plot the inverted permittivity constant map and adjust the colormap range based on the ratio parameter.
+    If ratio < 1, values exceeding ratio * max(abs(data)) will be saturated at the colormap maximum.
+    
+    :param data: 2D NumPy array (N, M), where N is time/depth and M is survey line direction
+    :param path: Path to save the output image
+    :param line_length: Survey line length in meters (default: 400m)
+    :param time_length: Time range in nanoseconds (default: 200ns)
+    :param ratio: Scaling factor for colormap range (default: 1)
+    """
+    num_points, num_lines = data.shape
+
+    # Compute the maximum absolute value for normalization
+    max_abs = np.max(np.abs(data))
+    vmin = -ratio * max_abs
+    vmax = ratio * max_abs
+
+    # Set font style
+    plt.rcParams.update({'font.family': 'Times New Roman', 'font.size': 20})
+
+    # Plot the permittivity image
+    plt.figure(figsize=(10, 4))
+    im = plt.imshow(data, aspect='auto', cmap='gray',
+                    extent=[0, line_length, time_length, 0],
+                    vmin=vmin, vmax=vmax)
+    
+    # Configure axis labels and ticks
+    plt.xlabel('Distance (m)', fontsize=20)
+    plt.xticks([0, 20, 40, 60, 80, 100, line_length])
+    plt.ylabel('Time (ns)', fontsize=20)
+    plt.yticks([0, 50, 100, 150, 200, time_length])
+
+    # Customize tick and border appearance
+    plt.tick_params(axis='both', direction='in', width=1)
+    for spine in plt.gca().spines.values():
+        spine.set_linewidth(1)
+    
+    # Remove grid lines, adjust layout, and save the figure
+    plt.grid(False)
+    plt.tight_layout()
+    plt.savefig(path, dpi=300)
+    plt.show()
+
+
+
+def plot_permittivity_constant(data, path, line_length=100, time_length=200):
+    """
+    Plot the inverted permittivity constant.
+    
+    :param data: 2D NumPy array, shape (N, M), where N is depth (or time) and M is distance along the survey line.
+    :param path: Path to save the output image.
+    :param line_length: Survey line length (meters), default is 100m.
+    :param time_length: Time range (nanoseconds), default is 200ns.
+    """
+    num_points, num_lines = data.shape
+    
+    # Configure font settings
+    plt.rcParams.update({'font.family': 'Times New Roman', 'font.size': 20})
+    
+    # Plot the permittivity map
+    plt.figure(figsize=(12, 4))
+    im = plt.imshow(data, aspect='auto', cmap='rainbow_r', extent=[0, 100, 200, 0], vmin=9, vmax=26)
+    
+    # Set axis labels and ticks
+    plt.xlabel('Distance (m)', fontsize=20)
+    plt.xticks([0, 20, 40, 60, 80, 100])
+    plt.ylabel('Time (ns)', fontsize=20)
+    plt.yticks([0, 50, 100, 150, 200])
+    
+    # Add colorbar
+    cbar = plt.colorbar(im)
+    cbar.set_label('Permittivity', fontsize=20)
+    
+    # Adjust axis formatting
+    plt.tick_params(axis='both', direction='in', width=1)
+    for spine in plt.gca().spines.values():
+        spine.set_linewidth(1)
+    
+    plt.grid(False)
+    plt.tight_layout()
+    plt.savefig(path, dpi=300)
+    plt.show()
+
+
+def plot_depth_permittivity_constant(data, path, line_length=100, time_length=200):
+    """
+    Plot the inverted permittivity constant as a 2D colormap.
+    
+    :param data: 2D NumPy array (depth/time, distance), representing permittivity values.
+    :param path: File path to save the output image.
+    :param line_length: Length of the survey line in meters (default: 100m).
+    :param time_length: Time range in nanoseconds (default: 200ns).
+    """
+    num_points, num_lines = data.shape
+    
+    # Set font properties
+    plt.rcParams.update({'font.family': 'Times New Roman', 'font.size': 20})
+    
+    # Create the plot
+    plt.figure(figsize=(12, 4))
+    im = plt.imshow(data, aspect='auto', cmap='rainbow_r', extent=[0, line_length, time_length, 0], vmin=8, vmax=30)
+    
+    # Label axes and set ticks
+    plt.xlabel('Distance (m)', fontsize=20)
+    plt.xticks([0, 20, 40, 60, 80, 100])
+    plt.ylabel('Time (ns)', fontsize=20)
+    plt.yticks([0, 2, 4, 6, 8])
+    
+    # Add colorbar
+    cbar = plt.colorbar(im)
+    cbar.set_label('Permittivity', fontsize=20)
+    
+    # Format axis appearance
+    plt.tick_params(axis='both', direction='in', width=1)
+    for spine in plt.gca().spines.values():
+        spine.set_linewidth(1)
+    
+    plt.grid(False)
+    plt.tight_layout()
+    plt.savefig(path, dpi=300)
+    plt.show()

utils/train_val_lr.py

@@ -0,0 +1,63 @@
+import os
+import sys
+import torch
+from torchsummary import summary
+from torchvision.utils import make_grid, save_image
+
+# Add parent directory to path for config import
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from config import Network_train_Config as cfg
+
+# Training function
+def train(train_loader, model, loss_func, optimizer, epoch):
+    model.train()
+    total_loss = 0
+    batch_count = 0
+
+    for data, label in train_loader:
+        data = data.cuda()
+        label = label.cuda()
+
+        output = model(data.type(torch.cuda.FloatTensor))
+        loss = loss_func(output.float(), label.float())
+
+        total_loss += loss.item()
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        batch_count += 1
+
+    avg_loss = total_loss / batch_count
+    print(f"Epoch {epoch}: Training Loss = {avg_loss:.6f}")
+    return avg_loss
+
+# Validation function
+def validate(val_loader, model, loss_func):
+    model.eval()
+    total_loss = 0
+    batch_count = 0
+
+    with torch.no_grad():
+        for data, label in val_loader:
+            data = data.cuda()
+            label = label.cuda()
+
+            output = model(data.type(torch.cuda.FloatTensor))
+            loss = loss_func(output.float(), label.float())
+
+            total_loss += loss.item()
+            batch_count += 1
+
+    avg_loss = total_loss / batch_count
+    return avg_loss
+
+# Learning rate scheduler
+def adjust_learning_rate(optimizer, epoch, start_lr):
+    """Exponentially decays learning rate based on epoch and configured decay rate."""
+    lr = start_lr * (cfg.lr_decrease_rate ** epoch)
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+    return lr
+
+
+