Liutex.py

import numpy as np
import pyvista as pv
from scipy.spatial.transform import Rotation as R

# Set global theme to allow empty meshes
pv.global_theme.allow_empty_mesh = True

def load_mesh(file_path):
    try:
        mesh = pv.read(file_path)
        return mesh
    except Exception as e:
        print(f"Failed to load mesh from {file_path}: {e}")
        return None

# Compute gradients of the vector field from available vector fields
def compute_gradients(mesh):
    if "Velocity" in mesh.point_data:
        print("Computing derivative for vector field: Velocity")
        result = mesh.compute_derivative(scalars="Velocity", gradient=True)
        if result.n_points == 0:
            print("Gradient computation resulted in an empty mesh.")
        return result
    else:
        raise KeyError("Mesh does not contain 'Velocity' for gradient computation.")

def gradients_to_dict(arr):
    keys = np.array(
        ["du/dx", "du/dy", "du/dz", "dv/dx", "dv/dy", "dv/dz", "dw/dx", "dw/dy", "dw/dz"]
    ).reshape(3, 3).ravel()
    return dict(zip(keys, arr.T))

def update_mesh_with_gradients(mesh, gradients_dict):
    print("Updating mesh with gradients")
    mesh.point_data.update(gradients_dict)

def create_nablav_matrices(gradients_dict):
    print("Creating NablaV matrices")
    du_dx = gradients_dict.get("du/dx")
    du_dy = gradients_dict.get("du/dy")
    du_dz = gradients_dict.get("du/dz")
    dv_dx = gradients_dict.get("dv/dx")
    dv_dy = gradients_dict.get("dv/dy")
    dv_dz = gradients_dict.get("dv/dz")
    dw_dx = gradients_dict.get("dw/dx")
    dw_dy = gradients_dict.get("dw/dy")
    dw_dz = gradients_dict.get("dw/dz")

    dim1, dim2 = 3, 3
    dim3 = len(dv_dy)
    nablav = np.zeros((dim1, dim2, dim3))

    for i in range(dim3):
        nablav[:, :, i] = np.array([
            [du_dx[i], du_dy[i], du_dz[i]],
            [dv_dx[i], dv_dy[i], dv_dz[i]],
            [dw_dx[i], dw_dy[i], dw_dz[i]]
        ])

    return nablav

def determine_rotation_strength(A):
    eigenvalues, eigenvectors = np.linalg.eig(A)
    real_indices = np.isreal(eigenvalues)
    real_eigenvalues = eigenvalues[real_indices]
    real_eigenvectors = eigenvectors[:, real_indices]

    lambda_r = real_eigenvalues[0].real
    r_vector = real_eigenvectors[:, 0].real

    r_vector /= np.linalg.norm(r_vector)
    z_axis = np.array([0, 0, 1])
    axis = np.cross(r_vector, z_axis)
    angle = np.arccos(np.dot(r_vector, z_axis))

    if np.linalg.norm(axis) != 0:
        axis /= np.linalg.norm(axis)
        rot = R.from_rotvec(axis * angle)
        q_rotation = rot.as_matrix()
    else:
        q_rotation = np.eye(3)

    return lambda_r, q_rotation

def calculate_rortex(R, r):
    return R * r

def update_mesh_with_omega_liutex(mesh, a, b):
    epsilon = 0.001 * np.maximum(np.max(b**2 - a**2), 1)

    # Calculate R_Omega
    R_Omega = b**2 / (a**2 + b**2 + epsilon)
    
    mesh.point_data["R_Omega"] = R_Omega

    return mesh

def main():
    file_path = r"C:\Users\arman\OneDrive\Dokumente\Master\Studienarbeit\VortexIdentification_Armand\T106C_Re80K.vtu"
    mesh = load_mesh(file_path)
    if mesh is None:
        return  # Exit if the mesh couldn't be loaded

    print(f"Mesh loaded with {mesh.n_points} points and {mesh.n_cells} cells.")
    print("Available point data arrays in mesh:", mesh.point_data.keys())

    mesh_g = compute_gradients(mesh)
    gradients = mesh_g["gradient"]
    gradients_dict = gradients_to_dict(gradients)

    update_mesh_with_gradients(mesh_g, gradients_dict)

    nablav = create_nablav_matrices(gradients_dict)

    rotation_list = []
    q_rot_list = []
    r = np.zeros((3, 1, nablav.shape[2]))
    vector = np.array([[0], [0], [1]])
    alpha = np.zeros((1, nablav.shape[2]))
    beta = np.zeros_like(alpha)
    Rortex = np.zeros((3, 1, nablav.shape[2]))
    n = nablav.shape[2]
    Rortex_magnitude = np.zeros(n)
    NABLAV_Array = np.zeros((nablav.shape[0], nablav.shape[1], n))
    a = np.zeros(n)
    b = np.zeros(n)
    Q_criterion = np.zeros(n)  # Initialize Q-criterion array
    Lambda2 = np.zeros(n)  # Initialize Lambda2 array

    for i in range(n):
        A = nablav[:, :, i]
        
        # Calculate symmetric (C) and antisymmetric (S) parts
        C = 0.5 * (nablav[:, :, i] + nablav[:, :, i].T)
        S = 0.5 * (nablav[:, :, i] - nablav[:, :, i].T)
        
        a[i] = np.linalg.norm(C, 'fro')
        b[i] = np.linalg.norm(S, 'fro')

        rotation_strength, q_rotation = determine_rotation_strength(A)
        rotation_list.append(rotation_strength)
        q_rot_list.append(q_rotation)

        NABLAV_i = q_rotation @ A @ q_rotation.T
        NABLAV_i[np.abs(NABLAV_i) < 1e-10] = 0
        NABLAV_Array[:, :, i] = NABLAV_i

        if i in [0, 1, 10, 50, 67, 100, 200, 300, 345, 890, 10000]:
            print(f'The {i}th 2D array of NABLAV:\n', NABLAV_i)

        r[:, :, i] = q_rotation.T @ vector

        alpha[0, i] = 0.5 * np.sqrt((NABLAV_i[0, 0] - NABLAV_i[1, 1]) ** 2 + (NABLAV_i[1, 0] + NABLAV_i[0, 1]) ** 2)
        beta[0, i] = 0.5 * (NABLAV_i[1, 0] - NABLAV_i[0, 1])

        if beta[0, i] < 0:
            beta[0, i] = -beta[0, i]
            r[:, :, i] = -r[:, :, i]

        if alpha[0, i] ** 2 - beta[0, i] ** 2 < 0:
            if beta[0, i] > 0:
                Rortex[:, :, i] = calculate_rortex(beta[0, i] - alpha[0, i], r[:, :, i])
            else:
                Rortex[:, :, i] = calculate_rortex(beta[0, i] + alpha[0, i], r[:, :, i])
        else:
            Rortex[:, :, i] = 0
            
        # Calculate the magnitude of Rortex at index i
        Rortex_magnitude[i] = np.linalg.norm(Rortex[:, :, i])
        
        # Calculate Q-criterion for the velocity gradient tensor at index i
        Q_criterion[i] = 0.5 * (np.sum(S * S) - np.sum(C * C))
        
        # Calculate Lambda2 for the velocity gradient tensor at index i
        tensor = S @ S + S @ S  # Calculate S^2 + Ω^2 (note Ω = -S, so Ω^2 = S^2)
        eigenvalues = np.linalg.eigvals(tensor)
        Lambda2[i] = np.sort(eigenvalues)[1]  # Second largest eigenvalue

    # Now, let's add the RortexMagnitude to the mesh point data
    mesh.point_data["RortexMagnitude"] = Rortex_magnitude
    mesh.point_data["RortexVector"] = Rortex.T.reshape((-1, 3))
    mesh.point_data["Q_Criterion"] = Q_criterion  # Store Q-criterion in the mesh
    mesh.point_data["Lambda2"] = Lambda2  # Store Lambda2 in the mesh

    # Compute Vorticity using PyVista
    Vorticity = mesh.compute_derivative(scalars='Velocity', vorticity=True)
    vorticity_array = Vorticity["vorticity"]  # Extract vorticity values from the result
    mesh.point_data["Vorticity"] = vorticity_array  # Assign vorticity values to the mesh
    vorticity_magnitude = np.linalg.norm(vorticity_array, axis=1)  # Ensure correct array name
    #mesh.point_data["VorticityMagnitude"] = vorticity_magnitude

    # Calculate Omega in the main function
    epsilon = 0.001 * np.maximum(np.max(b**2 - a**2), 1)
    F_a = a**2
    F_b = b**2
    Omega = F_b / (F_a + F_b + epsilon)
    mesh.point_data["Omega"] = Omega

    # Call the function to update mesh with R_Omega
    mesh = update_mesh_with_omega_liutex(mesh, a, b)

    # Save the mesh with vorticity and omega data to a VTK file
    mesh.save("Liu_Rortex_T106C_Re80K_19.09.vtu")
    print("Dataset saved as Liu_Rortex_T106C_Re80K_19.09.vtu")

    # Selecting seed points in high vorticity regions
    vorticity_threshold = np.percentile(vorticity_magnitude, 90)  # Top 10% vorticity magnitude
    seed_points = mesh.points[vorticity_magnitude > vorticity_threshold]
    seed_polydata = pv.PolyData(seed_points)

    plotter = pv.Plotter(window_size=[1600, 1200])
    plotter.add_mesh(mesh, scalars='RortexMagnitude', cmap='viridis')
    
    # Adding iso-surface contour for RortexMagnitude
    contours = mesh.contour(isosurfaces=10, scalars='RortexMagnitude')
    plotter.add_mesh(contours, opacity=0.4, cmap='viridis', label='Rortex Iso-surface')

    # Adding streamlines for RortexVector
    streamline = mesh.streamlines_from_source(seed_polydata, vectors='RortexVector')
    if streamline.n_points > 0:
        plotter.add_mesh(streamline, line_width=2, label='Streamlines', opacity=0.6, cmap='cool')
    else:
        print("Warning: No streamlines generated. Check seed points and flow vector field.")

    # Save vorticity lines
    vorticity_mesh = mesh.streamlines_from_source(seed_polydata, vectors='Velocity')
    vorticity_mesh.save("Vorticity_Lines_Blade_19.09.vtu")
    print("Vorticity lines saved as Vorticity_Lines_Blade_19.09.vtu")

    plotter.add_axes()
    plotter.view_isometric()
    plotter.show(screenshot='rortex_visualization.png')
    print("Plot saved as rortex_visualization.png")


if __name__ == "__main__":
    main()