comsol2ca/plot_field.py at main · rawrex/comsol2ca · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
import os
import argparse

def read_domain_info(filepath):
    """Read domain information from the domain_info.txt file."""
    domain_info = {}
    with open(filepath, 'r') as f:
        for line in f:
            key, value = line.strip().split()
            try:
                domain_info[key] = float(value)
            except ValueError:
                domain_info[key] = value  # Keep non-numeric values as strings
    return domain_info

def read_boundary_conditions(filepath):
    """Read boundary conditions from the boundary_conditions.txt file."""
    boundary_conditions = {'inlets': [], 'outlets': [], 'walls': []}
    current_type = None

    with open(filepath, 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue

            if line in ['inlets', 'outlets', 'walls']:
                current_type = line
            elif line.startswith('coords'):
                coords = list(map(float, line.split()[1:]))
                boundary_conditions[current_type].append({
                    'coords': [(coords[0], coords[1]), (coords[2], coords[3])]
                })

    return boundary_conditions

def plot_simulation_field(data):
    """Plot the simulation field with velocity vectors and pressure contours."""
    # Read domain info
    domain_info = read_domain_info(os.path.join(data, 'domain_info.txt'))
    nx, ny = int(domain_info['nx']), int(domain_info['ny'])

    # Read data
    u = np.loadtxt(os.path.join(data, 'u.txt')).reshape(ny, nx)
    v = np.loadtxt(os.path.join(data, 'v.txt')).reshape(ny, nx)
    pressure = np.loadtxt(os.path.join(data, 'pressure.txt')).reshape(ny, nx)
    velocity_magnitude = np.loadtxt(os.path.join(data, 'velocity_magnitude.txt')).reshape(ny, nx)

    # Create coordinate grids
    x = np.linspace(domain_info['x_min'], domain_info['x_max'], nx)
    y = np.linspace(domain_info['y_min'], domain_info['y_max'], ny)
    X, Y = np.meshgrid(x, y)

    # Read boundary conditions
    boundary_conditions = read_boundary_conditions(os.path.join(data, 'boundary_conditions.txt'))

    # Plot 1: Velocity field with streamlines
    plt.figure(figsize=(10, 8))
    skip = 5  # Skip points for better visualization
    plt.streamplot(X[::skip, ::skip], Y[::skip, ::skip],
                  u[::skip, ::skip], v[::skip, ::skip],
                  color=velocity_magnitude[::skip, ::skip],
                  cmap='viridis', density=1.5)
    plt.title('Velocity Field')
    plt.xlabel('x (mm)')
    plt.ylabel('y (mm)')

    # Plot boundary conditions on velocity field
    for bc_type, bc_list in boundary_conditions.items():
        for bc in bc_list:
            coords = bc['coords']
            x_coords = [coords[0][0], coords[1][0]]
            y_coords = [coords[0][1], coords[1][1]]
            if bc_type == 'inlets':
                plt.plot(x_coords, y_coords, 'g-', linewidth=2, label='Inlet' if bc == bc_list[0] else "")
            elif bc_type == 'outlets':
                plt.plot(x_coords, y_coords, 'r-', linewidth=2, label='Outlet' if bc == bc_list[0] else "")
            elif bc_type == 'walls':
                plt.plot(x_coords, y_coords, 'k-', linewidth=2, label='Wall' if bc == bc_list[0] else "")

    # Add legend for velocity field
    handles, labels = plt.gca().get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    plt.legend(by_label.values(), by_label.keys(), loc='upper right')
    plt.tight_layout()
    plt.savefig(os.path.join(data, 'velocity_field.png'), dpi=300, bbox_inches='tight')

    # Plot 2: Pressure distribution
    plt.figure(figsize=(10, 8))
    im = plt.contourf(X, Y, pressure, levels=20, cmap='RdBu_r')
    plt.title('Pressure Distribution')
    plt.xlabel('x (mm)')
    plt.ylabel('y (mm)')
    plt.colorbar(im, label='Pressure (Pa)')

    # Plot boundary conditions on pressure field
    for bc_type, bc_list in boundary_conditions.items():
        for bc in bc_list:
            coords = bc['coords']
            x_coords = [coords[0][0], coords[1][0]]
            y_coords = [coords[0][1], coords[1][1]]
            if bc_type == 'inlets':
                plt.plot(x_coords, y_coords, 'g-', linewidth=2, label='Inlet' if bc == bc_list[0] else "")
            elif bc_type == 'outlets':
                plt.plot(x_coords, y_coords, 'r-', linewidth=2, label='Outlet' if bc == bc_list[0] else "")
            elif bc_type == 'walls':
                plt.plot(x_coords, y_coords, 'k-', linewidth=2, label='Wall' if bc == bc_list[0] else "")

    # Add legend for pressure field
    handles, labels = plt.gca().get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    plt.legend(by_label.values(), by_label.keys(), loc='upper right')
    plt.tight_layout()
    plt.savefig(os.path.join(data, 'pressure_field.png'), dpi=300, bbox_inches='tight')

    plt.show()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Plot simulation field from data directory')
    parser.add_argument('--data', type=str, help='Directory containing the exported data')
    args = parser.parse_args()
    plot_simulation_field(args.data)