ParticlePath.py

import sys
import os
import pandas as pd

def isCSVandExist(file, genName):
    splitFileName = file.split('.')
    for i in splitFileName:
        print(i)
    if(len(splitFileName) < 2):
        return False

    lastName = splitFileName[len(splitFileName) - 1]
    lastName = lastName.lower()

    if lastName != str('csv'):
        return False

    genName.append(splitFileName[0] + '.vtk')
    return True

def doGenerate(filename):
    genName = []
    if isCSVandExist(filename, genName):
        print('deal with file : ', filename, ',generating track file : ' , genName[0])
    else:
        print('please input a exist IB file (csv)')
        return
    
    data = pd.read_csv(filename)
    X_location = data['X']
    Y_location = data['Y']
    Z_location = data['Z']

    U = data['Vx']
    V = data['Vy']
    W = data['Vz']

    Fx = data['Fx']
    Fy = data['Fy']
    Fz = data['Fz']

    with open(genName[0], 'w') as out:
        out.write('# vtk DataFile Version 4.0\n'
                  'Generated by LXZ-pycode\n'
                  'ASCII\n'
                  'DATASET POLYDATA\n'
                  'POINTS ' + str(X_location.size) + ' double\n')
        # write location
        for i in range(0, len(X_location), 1):
            out.write(str(X_location[i]) + ' ' + str(Y_location[i]) + ' ' + str(Z_location[i]) + '\n')
        # write VERTICES
        out.write('VERTICES ' + str(X_location.size) + ' ' + str(X_location.size * 2) + '\n')
        for i in range(0, len(X_location), 1):
            out.write('1 ' + str(i) + '\n')
        # write velocity data
        out.write('\nPOINT_DATA '+ str(X_location.size) + '\nFIELD FieldData 2\n')
        out.write('v 3 ' + str(X_location.size) + ' double\n')
        for i in range(0, len(X_location), 1):
            out.write(str(U[i]) + ' ' + str(V[i]) + ' ' + str(W[i]) + '\n')
        # write ib force
        out.write('ib_f 3 ' + str(X_location.size) + ' double\n')
        for i in range(0, len(X_location), 1):
            out.write(str(Fx[i]) + ' ' + str(Fy[i]) + ' ' + str(Fz[i]) + '\n')


if __name__ == '__main__':
    arg1 = sys.argv[1]
    if(arg1 == 'help'):
        print('run with Particle IB file name , sunch as :\n\n'
              '/{your python env}/python ParticlePath.py Particle.csv\n\n'
              'will be generate VTK data file , you can open it(*.vtk) with paraview\n'
              '*** but run this file, need pandas to deal data ***\n'
              '*** written by LXZ , 2024/5 ***\n')
        exit()
    doGenerate(sys.argv[1])