sphereSolver/plot.py at main · SpeedyTurtle599/sphereSolver · 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
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm, gridspec

def read_velocity_field(filename):
    # Read the header to get dimensions and sphere info
    with open(filename, 'r') as f:
        _ = f.readline()  # Skip title
        _ = f.readline()  # Skip timestep
        radius_line = f.readline()
        sphere_radius = float(radius_line.split(':')[1])
        dim_line = f.readline()
        dims = {k: int(v) for k, v in
                [item.split('=') for item in dim_line.strip('# \n').split()]}
        # Skip the column headers line
        f.readline()

        # Read the data
        data = np.loadtxt(f)

    # Reshape the data
    x = data[:, 0].reshape(dims['nz'], dims['ny'], dims['nx'])
    y = data[:, 1].reshape(dims['nz'], dims['ny'], dims['nx'])
    z = data[:, 2].reshape(dims['nz'], dims['ny'], dims['nx'])
    u = data[:, 3].reshape(dims['nz'], dims['ny'], dims['nx'])
    v = data[:, 4].reshape(dims['nz'], dims['ny'], dims['nx'])
    w = data[:, 5].reshape(dims['nz'], dims['ny'], dims['nx'])
    vel_mag = data[:, 6].reshape(dims['nz'], dims['ny'], dims['nx'])

    return x, y, z, u, v, w, vel_mag, sphere_radius

def plot_flow_field(filename):
    # Update unpacking to include sphere_radius
    x, y, z, u, v, w, vel_mag, sphere_radius = read_velocity_field(filename)

    # Setup figure with GridSpec
    fig = plt.figure(figsize=(15, 10))
    gs = gridspec.GridSpec(1, 2, width_ratios=[20, 1])

    # Create 3D axis
    ax = fig.add_subplot(gs[0], projection='3d')

    # Take slices at the middle
    mid_z = x.shape[0] // 2

    # Downsample for quiver
    skip = 4 # Skip every 4th point
    x_sub = x[::skip, ::skip, ::skip]
    y_sub = y[::skip, ::skip, ::skip]
    z_sub = z[::skip, ::skip, ::skip]
    u_sub = u[::skip, ::skip, ::skip]
    v_sub = v[::skip, ::skip, ::skip]
    w_sub = w[::skip, ::skip, ::skip]
    vel_mag_sub = vel_mag[::skip, ::skip, ::skip]

    # Set up colormap
    norm = plt.Normalize(vel_mag_sub.min(), vel_mag_sub.max())
    colors = cm.viridis(norm(vel_mag_sub))

    # Plot quiver
    quiv = ax.quiver(x_sub, y_sub, z_sub,
                    u_sub, v_sub, w_sub,
                    length=0.5,
                    normalize=True,
                    colors=colors.reshape(-1, 4))

    # Add slice
    slice_x = ax.contourf(x[mid_z,:,:], y[mid_z,:,:], vel_mag[mid_z,:,:],
                         zdir='z', offset=z[mid_z,0,0],
                         levels=20, cmap='viridis', alpha=0.5)

    # Use actual sphere radius from file header
    sphere_center = np.mean([x.min(), x.max()])

    phi = np.linspace(0, 2 * np.pi, 20)
    theta = np.linspace(0, np.pi, 20)
    phi, theta = np.meshgrid(phi, theta)

    sphere_x = sphere_center + sphere_radius * np.cos(phi) * np.sin(theta)
    sphere_y = sphere_center + sphere_radius * np.sin(phi) * np.sin(theta)
    sphere_z = sphere_center + sphere_radius * np.cos(theta)
    ax.plot_surface(sphere_x, sphere_y, sphere_z, color='gray', alpha=0.3)

    # Labels
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_title('3D Flow Field Visualization')

    # Add colorbar with proper axes
    cax = fig.add_subplot(gs[1])
    sm = plt.cm.ScalarMappable(cmap='viridis', norm=norm)
    sm.set_array([])
    plt.colorbar(sm, cax=cax, label='Velocity Magnitude')

    # Cubic aspect ratio
    ax.set_box_aspect([1, 1, 1])

    plt.tight_layout()
    plt.show()

# Run visualization
filename = "velocity_field_000000.dat"
plot_flow_field(filename)