ComputePSD.py

import numpy as np
import vtk
from glob import glob
import argparse
from utilities import *
from scipy.fftpack import fftfreq,fft,ifft
import scipy.signal as signal
import os
from matplotlib import pyplot as plt
class ComputePSD():
	def __init__(self,args):
		self.Args=args
		if self.Args.OutputFolder is None:
			self.Args.OutputFolder = self.Args.InputFolder

	def Main(self):
		#Read the velocity file
		InputFileName=glob(self.Args.InputFolder+"/VelocityData.npy")
		InputFileName+=glob(self.Args.InputFolder+"/VelocityData.npz")		
		print ("The Input FileName is: %s"%InputFilename

		VelocityData=np.load(InputFileName)
		
		#Get all the array sizes
		Nprobes=len(VelocityData.item().keys())
		Nts=len(VelocityData.item().get(0)[0,:])
		print ("--- The number of time steps: %d"%Nts)
		print ("--- The number of probe points: %d"%Nprobes)


		#Compute frequencies
		fs=int(Nts/self.Args.Period)
		time=np.linspace(0,self.Args.Period,Nts)

		#Loop over all of the probe points
		print ("\n")
		for i in range(Nprobes):
			print ("------ Looping over Probe Point %d"%i)
			#Create the file to store the data
			outfile=open("%s/PSD_Probe%d.dat"%(self.Args.OutputFolder,i),'w')
			outfile.write("Frequencies PSD\n")			

			#Create an array to store the frequencies
			PSD_=0

			#Loop over all of the probe locations
			for j in range(0,len(VelocityData.item().get(i)[:,0])):
				f_, Pxx_den_ = signal.welch(VelocityData.item().get(0)[j,:], fs, nperseg=Nts/4.)
				PSD_+=Pxx_den_

			#Now Average the frequency content
			PSD_=PSD_/j

			print("------ Writing Sphere-averaged PSD to: %s/PSD_Probe%d.dat"%(self.Args.OutputFolder,i))
			for j in range(0,len(f_)):
				outfile.write("%.05f %.5e\n"%(f_[j],PSD_[j]))
			outfile.close()	
		

if __name__=="__main__":
        #Description
	parser = argparse.ArgumentParser(description="This script will compute power spectral density from velocity probe data")

        #Provide a path to the Magnitude Images
	parser.add_argument('-InputFolder', '--InputFolder', type=str, required=True, dest="InputFolder",help="This folder containing the velocity.npy file generated by ExtractProbePoint.py file.")
        
	parser.add_argument('-Period', '--Period', type=float, required=True, dest="Period",help="The period of the cardiac cycle")
        
	parser.add_argument('-OutputFolder', '--OutputFolder', type=int, required=False, dest="OutputFolder",help="The output folder to store the data")

        #Put all the arguments together
	args=parser.parse_args()

        #Call your Class
	ComputePSD(args).Main()