##----------------------------------------------------------------------------
## Name:     RotarySpeaker.py
##
##    Compute the angular pressure (beampattern) of the rotary speaker.
##    The wall vibrations can be included with -v & -W
##
##    Coordinates are: x - horizontal (x=0 along wall/baffle)
##                     y - vertical   (y=0 along wall/baffle)
##                     z - fan axis 
##                     The beampattern is taken in the x-z plane
##                     Origin (0,0) is the fan center
##
##    Depends on Python SciPy and NumPy
##             
## Author:       J Park
##
## Created:      4/9/08
##
## Modified:     
##----------------------------------------------------------------------------

import os, sys, string, getopt, time

Version = sys.argv[0]+" : "+str( time.asctime( ( time.gmtime(time.time()) ) ) )
DEBUG = False

# Python array module and numpy/scipy are NOT compatible as standalone modules.
# numpy/scipy requires that the array constructor must be passed an object
# that is either an array, or list: a = array([]) will create an empty array.
#
# With the python array module, one can define an array object as:
# a = array('d'), and then call a.append(1.2) to build the array as
# needed. In numpy/scipy, you have to first create an array object a,
# then call b = append(a, 1.2) to return an array b that appends the
# value 1.2 to the existing array object a. 
# 
try:
    from numpy import *  # www.numpy.org numpy.scipy.org
except ImportError: 
    print "Failed to import numpy."
    
try:
    import pylab as mp  # matplotlib.sourceforge.net
except ImportError: 
    print "Failed to import pylab."

try:
    from scipy import special as sp # www.scipy.org
except ImportError: 
    print "Failed to import scipy."

# Global variables
# Default Parameters
PlotNum      = 0               # -p
WriteOutput  = False           # -w
Monopole     = False           # -m compute monopole field only
FlexBaffle   = False           # -v include baffle
AxisOnly     = False           # -z 
OutputFile   = "piston.dat"    # -o
Frequency    = 16.             # -f blade actuation freq (Hz)
R_i          = 4.  * 0.0254    # -a piston inner radius (in) to (m)
R_o          = 10. * 0.0254    # -b piston outer radius (in) to (m)
dL           = 1.0 * 0.0254    # -l piston grid resolution (in) to (m)
F_r          = 10.             # -R Field range (m)
range_inc    = 0.              # -r recurring range calculation increment (m)
B_theta_max  = 35. * pi/180.   # -x blade max actuation angle (deg) to (rad)
                               #    this is set by Amp Gain and Input Voltage (RMS)
#Frequency_fan= 11.02          # -F fan rotational speed (Hz) Cfan=23, 661 RPM
Frequency_fan= 16.8            # -F fan rotational speed (Hz)  Cfan=35, 1006 RPM
delta_wall   = 0.003           # -W wall displacement (m)

c_0          = 343.            # sound speed (m/s)
rho          = 1.2             # air density (kg/m^3)
Omega        = 2.*pi*Frequency # (rad/s)
Omega_fan    = 2.*pi*Frequency_fan # fan rotational speed (rad/s)
K            = Omega / c_0     # wavenumber
T            = 0.              # time
A_total      = pi*(pow(R_o,2) - pow(R_i,2)) # area of 'piston' (m^2)

V_r          = Omega_fan * (R_i + (2.*(R_o - R_i)/3.)) # blade eff forward velocity (m/s)
V_r_y        = V_r * sin( B_theta_max ) # forward velocity normal to blade
B_w          = 2.  * 0.0254    # blade width from actuator axis (m)
V_b          = Omega * 2.* B_w / 3.# blade eff actuation velocity normal to blade (m/s)
V_b_avg      = V_r_y + V_b     # effective particle velocity normal to blade (m/s)
C_fan        = ( rho * c_0 * V_b_avg * K ) / ( 2. * pi )
V_wall       = 2. * delta_wall * Frequency # wall velocity (m/s)
C_wall       = ( rho * c_0 * V_wall * K ) / ( 2. * pi )

F_theta      = arange( 0., pi/2., pi/180. ) # Field angles (rad)

wall_dL       = 0.25
wallHeightMin = -0.3
wallHeightMax = 3.2
wallWidthMin  = -5
wallWidthMax  = 3.8
wallHeight    = wallHeightMax - wallHeightMin
wallWidth     = wallWidthMax  - wallWidthMin
bldg_f0       = c_0/30. # c/lambda_0 = 11.43 Hz, lambda_0 = 2 L
bldg_omega0   = 2. * pi * bldg_f0
bldg_k0       = bldg_omega0 / c_0

# List of field points, list preserves order
Field = []
# array of field point pressures for output
P_n_amp  = array([], dtype=float32)
P_n_real = array([], dtype=float32)
P_n_imag = array([], dtype=float32)

fanAxisPres = complex(0., 0.)

# Dictionary of piston grid points
# key is piston point (x,y) tuple, value is class object
PistonGrid = {}
WallGrid   = {}

# Dictionary of rays from piston grid points to field points
# key is tuple of source/field (PistonPoint, FieldPoint), value is class object
Rays = {}

##----------------------------------------------------------------------------
class PistonGridPoint:
    def __init__( self, x, y ):
        self.x = x
        self.y = y
        self.z = 0.
        self.A = 0.
        self.C = C_fan
        PistonGrid[(x,y)] = self
        
##----------------------------------------------------------------------------
class WallGridPoint:
    def __init__( self, x, y ):
        self.x = x
        self.y = y
        self.z = 0.
        self.A = 0.
        #self.C = C_wall
        # sin distribution of velocities, nodes at bottom/top
        self.C = C_wall * sin( pi * (self.y - wallHeightMin) / wallHeight ) * \
                          sin( pi * (self.x - wallWidthMin)  / wallWidth)
        WallGrid[(x,y)] = self
        
##----------------------------------------------------------------------------
class FieldPoint:
    def __init__( self, r, theta ):
        self.r            = r
        self.theta        = theta
        #self.y            = 0.
        self.y            = -0.34 # the fan and microphone are not planar
        self.x            = r * sin( theta )
        self.z            = r * cos( theta )
        self.pressure     = complex(0., 0.)
        self.rays         = []
        # insert object into Field container
        Field.append(self)

    def SumPressure( self ):
        global fanAxisPres
        # sum the amplitude of this field point from all rays
        # this is complex addition of ray amp/phase
        # exp[j(kr+wt)] = exp[j phi] = cos(phi) + j sin(phi)
        for ray in self.rays:
            # Scale for element area, piston/wall velocity and range decay
            C = ( ray.gridPoint.A * ray.gridPoint.C / self.r )
            
            self.pressure = self.pressure + C * ( \
                            complex( cos(ray.phi), sin(ray.phi) ) )
            
##             # This adds a plane wave from the wall
##             self.pressure = self.pressure + 0.5 * C * ( \
##                             complex( cos(ray.phi), sin(ray.phi) ) +
##                             complex( cos(K*self.z), sin(K*self.z) ) )

            if FlexBaffle:
                # compute fan pressure alone on the axis
                if self.theta == 0.:
                    if isinstance(ray.gridPoint, PistonGridPoint):
                        fanAxisPres = fanAxisPres + C * ( \
                                      complex( cos(ray.phi), sin(ray.phi) ) )
        if self.theta == 0.:
            if not FlexBaffle:
                fanAxisPres = self.pressure
            print "fan axis pressure", 20.*log10(abs(fanAxisPres))


##----------------------------------------------------------------------------
class Ray:
    def __init__( self, gp, fp ):
        self.gridPoint   = gp
        self.fieldPoint  = fp
        self.r           = self.distance(gp, fp)
        self.amp         = 1./(pow(self.r,2))
        self.loss        = 20 * log10(self.r)
        if isinstance(self.gridPoint, PistonGridPoint):
            # propagation from the fan
            self.phi = Omega * T - K * self.r
        else:
            # wall points are out of phase with piston
            self.phi = bldg_omega0 * T - K * self.r - bldg_k0 * self.r # + \
                       #sqrt(pow(K * gp.x, 2) + pow(K * gp.y, 2))  # + 180 * pi/180.

        # save reference to this ray in the field point
        fp.rays.append( self )
        # add ray to main Rays container
        Rays[(gp,fp)] = self

    def distance( self, gp, fp ):
        dx = gp.x - fp.x
        dy = gp.y - fp.y
        dz = gp.z - fp.z
        r  = sqrt( pow(dx,2) + pow(dy,2) + pow(dz,2) )
        return r
    
##----------------------------------------------------------------------------
## Main module
def main():
    print  Version
    global fdOutFile, P_n_amp, P_n_real, P_n_imag, F_r, Field, Rays, fanAxisPres
    
    ParseCmdLine()

    if DEBUG:
        print "K R_o=%.2f" % ( K*R_o ),
        print "V_r=%.2f(m/s)  V_r_y=%.2f(m/s)  V_b=%.2f(m/s)  V_b_avg=%.2f(m/s)" % \
              (V_r, V_r_y, V_b, V_b_avg)
        print "A_total=%.3f" % (A_total),
        print "C_fan=%.2f" % (C_fan),
        print "C_wall=%.2f" % (C_wall)
    
    # Create a grid of points for the piston
    if Monopole:
        PistonGridPoint( 0., 0. )
    else:
        for x in arange( -2*R_o, 2*R_o, dL ):
            for y in arange( -2*R_o, 2*R_o, dL ):
                xsqr = pow(x,2)
                ysqr = pow(y,2)
                if xsqr == 0. and ysqr == 0.:
                    continue
                if (xsqr + ysqr < R_o) and (xsqr + ysqr > R_i):
                    PistonGridPoint( x, y )

    # area of each grid point
    numPistonGridPoints = len(PistonGrid)
    for pg in PistonGrid.values():
        pg.A = A_total / numPistonGridPoints

    print "%d piston grid points" % (numPistonGridPoints),

    # Create a grid for the wall
    numWallGridPoints = 0
    if FlexBaffle:
        for x in arange( wallWidthMin, wallWidthMax, wall_dL ):
            for y in arange( wallHeightMin, wallHeightMax, wall_dL ):
                xsqr = pow(x,2)
                ysqr = pow(y,2)
                if (xsqr + ysqr < R_o) :
                    continue
                WallGridPoint( x, y )
    
        # area of each grid point
        numWallGridPoints = len(WallGrid)
        for wg in WallGrid.values():
            wg.A = wall_dL * wall_dL

        print "%d wall grid points" % (numWallGridPoints)

    if range_inc:
        stop = 3
    else:
        stop = 1

    #++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    for range_num in range(stop):

        # reset the containers if multiple ranges
        if range_inc:
            Field    = []
            Rays     = {}
            P_n_amp  = array([], dtype=float32)
            P_n_real = array([], dtype=float32)
            P_n_imag = array([], dtype=float32)
            fanAxisPres = complex(0., 0.)
            
        # Create the field points
        if DEBUG:
            print 'Range', F_r,
        for angle in F_theta:
            FieldPoint( F_r, angle )
            if AxisOnly:
                if angle == 0:
                    break

        numRays = len(Field) * (numPistonGridPoints + numWallGridPoints)
        print "%d field points," % (len(Field)),
        if numRays > 1.2E6:
            print "Too many rays:", numRays
            sys.exit(1)

        # From each piston grid point, compute rays to the field points
        for pg in PistonGrid.values():
            for fp in Field:
                Ray( pg, fp )

        print "%d piston rays" % (len(Rays)),

        # From each wall grid point, compute rays to the field points
        for wg in WallGrid.values():
            for fp in Field:
                Ray( wg, fp )

        print "%d total rays." % (len(Rays))

        # At each field point sum the ray pressures
        for fp in Field:
            fp.SumPressure()

        # Extract field results into arrays for plotting
        for fp in Field:
            #P_n_amp  = append(P_n_amp,  abs(fp.pressure))
            # normalize to the fan axis pressure
            P_n_amp  = append(P_n_amp,  abs(fp.pressure) * 
                                        (abs(fanAxisPres) / abs(Field[0].pressure)))
            P_n_real = append(P_n_real, fp.pressure.real)
            P_n_imag = append(P_n_imag, fp.pressure.imag)

##     # Normalize amplitude
##     P_min = min(P_n_amp)
##     P_max = max(P_n_amp)
##     if P_max > P_min:
##         P_n_amp = (P_n_amp - P_min)/(P_max - P_min)
##         for fp in Field:
##             fp.amp = (fp.amp - P_min)/(P_max - P_min)

##     # compute ideal directivity for piston in rigid infinite baffle
##     P_J1 = array([], dtype=float32)
##     for fp in Field:
##         arg_j = K * R_o * sin( fp.theta )
##         if fp.theta != 0.:
##             j1 = 2. * sp.j1( arg_j ) / arg_j
##         else:
##             j1 = 1.
##         P_J1 = append( P_J1, j1 )
            
##     mp.figure(1)
##     mp.subplot(2,1,2)
##     mp.plot ( F_theta*(180./pi), P_J1, linewidth=2 )
    
        
        #==========================================================================
        # Plotting & Output
        #
        freqLabel  = "Frequency %.0f Hz" % (Frequency)
        R_oLabel   = "R_o=%.2f in"  % (R_o / 0.0254)
        R_iLabel   = "R_i=%.2f in"  % (R_i / 0.0254)
        rangeLabel = "Range=%.1f m" % (F_r) # for legend() w/ multiple ranges
        
        #=========================================================
        if PlotNum > 5: # plot the field points
            x = array([], dtype=float32)
            z = array([], dtype=float32)
            for fp in Field:
                #print fp.x, fp.z
                x = append(x, fp.x)
                z = append(z, fp.z)
            
            mp.figure(6)
            mp.plot ( x, z, '.' )

        #=========================================================
        # On-axis pressure for a piston in baffle Kinsler & Frey pg 176
        if PlotNum > 4:
            r_axis = array([], dtype=float32)
            for r in arange( 1E-3, 2 * F_r, 2 * F_r / 100. ):
                r_axis = append( r_axis, r )
            P_axis = abs( sin( (K*r_axis/2) * sqrt( 1 + pow((R_o/r_axis),2)) - 1 ) ) 
            mp.figure(5)
            mp.title("On-axis pressure for a piston in baffle")
            mp.xlabel("Range (m)", size='16' )
            mp.ylabel("Relative Pressure",   size='16')
            mp.figtext(0.7, 0.25, freqLabel, size='16')
            mp.figtext(0.7, 0.2,  R_oLabel,  size='16')
            mp.plot( r_axis, P_axis, linewidth=2 )

        #=========================================================
        if PlotNum > 3: # real and imaginary pressure vs. angle
            mp.figure(4)
            mp.subplot(2,1,1)
            mp.title("Pressure vs. angle")
            mp.ylabel("Real Pressure re Pa",  size='16')
            mp.figtext(0.6, 0.75, freqLabel,  size='16')
            mp.figtext(0.6, 0.7,  R_oLabel,   size='16')
            mp.figtext(0.6, 0.65, R_iLabel ,  size='16')
            mp.figtext(0.6, 0.6,  rangeLabel, size='16')
            mp.plot ( F_theta*(180./pi), P_n_real, linewidth=2 )
            
            mp.subplot(2,1,2)
            mp.xlabel("Angle (deg)", size='16' )
            mp.ylabel("Imag Pressure re Pa", size='16' )
            mp.plot ( F_theta*(180./pi), P_n_imag, linewidth=2 )

        #=========================================================
        if PlotNum > 2: # plot the piston grid
            xp = array([], dtype=float32)
            yp = array([], dtype=float32)
            xw = array([], dtype=float32)
            yw = array([], dtype=float32)
            for pg in PistonGrid.values():
                #print pg.x, pg.y 
                xp = append(xp, pg.x)
                yp = append(yp, pg.y)
            for wg in WallGrid.values():
                #print pg.x, pg.y 
                xw = append(xw, wg.x)
                yw = append(yw, wg.y)
            
            mp.figure(3)
            mp.plot ( xp, yp, 'r.' )
            mp.hold(True)
            mp.plot ( xw, yw, 'b.' )

        #=========================================================
        # polar plot field pressure
        if PlotNum > 1:
            mp.figure(2)
            # This is cool, in numpy.array [::-1] is an 'extended' slice that reverses
            polarAngles = concatenate( (-F_theta[::-1], F_theta[1:]) )
            polarAmp    = concatenate( (P_n_amp[::-1],  P_n_amp[1:]) )
            mp.polar ( polarAngles, polarAmp, linewidth=2, label=rangeLabel )
            #mp.polar ( F_theta, P_n_amp, -F_theta, P_n_amp, linewidth=2, label=rangeLabel )
            mp.title("Normalized Beampattern")
            mp.figtext(0.28, 0.7,  freqLabel,  size='16')
            mp.figtext(0.28, 0.65, R_oLabel,   size='16')
            mp.figtext(0.28, 0.6,  R_iLabel,   size='16')
            if not range_inc:
                mp.figtext(0.28, 0.55,  rangeLabel, size='16')

        #=========================================================
        # line plot field pressure vs theta
        if PlotNum > 0:
            mp.figure(1)
            #mp.subplot(2,1,1)
            mp.title("Pressure Amplitude vs. angle")
            mp.xlabel("Angle (deg)", size='16' )
            mp.ylabel("Pressure dB re Pa",   size='16' )
            mp.figtext(0.6, 0.45, freqLabel, size='16' )
            mp.figtext(0.6, 0.4,  R_oLabel,  size='16' )
            mp.figtext(0.6, 0.35, R_iLabel,  size='16' )
            if not range_inc:
                mp.figtext(0.6, 0.3, rangeLabel, size='16' )
            
            mp.plot ( F_theta*(180./pi), 20.*log10(P_n_amp), linewidth=2, label=rangeLabel )

        # increment the field range for the next run
        F_r = F_r + range_inc

    #++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
    
    if WriteOutput:
        fd = OpenFile( OutputFile, 'w' )
        fd.write( "Theta(rad)\tAmp\n" )
        for fp in Field:
            fd.write( "%.3f\t%.3f\n" % ( fp.theta, fp.pressure ) )
        fd.close()
        print "Wrote ", len(Field), " points to ", OutputFile

    if PlotNum > 0:
        # Write out graphic of the plot
        plotFileName = OutputFile + ".png"
        mp.savefig( plotFileName )
        print "Wrote png image to ", plotFileName
        
        if range_inc:
            mp.figure(1)
            mp.legend( loc='lower left' )
            if PlotNum > 1:
                mp.figure(2)
                mp.legend( loc='lower left' )
            
        # Matplotlib show for all figures (must be only one call)
        mp.show()

    print "Normal Exit"
## Main module
##----------------------------------------------------------------------------


##----------------------------------------------------------------------------
def usage():
    print "Usage: ", sys.argv[0]
    print "\t -h (help)"
    print "\t -d (debug)"
    print "\t -p # plots"
    print "\t -w (WriteOutput)"
    print "\t -o OutputFile"
    print "\t -v FlexBaffle"
    print "\t -z AxisOnly"
    print "\t -f Frequency (Hz)"
    print "\t -F Fan Rotation Frequency (Hz)"
    print "\t -a Inner Radius (in)"
    print "\t -b Outer Radius (in)"
    print "\t -l Grid Resolution (in)"
    print "\t -R Field range (m)"
    print "\t -r range increment (m)"
    print "\t -W wall displacement (m)"
    print "\t -x blade actuation angle (deg)"
    print "\t -m (monopole)"

##----------------------------------------------------------------------------
def ParseCmdLine():
    ##
    global DEBUG
    global delta_wall, C_wall, V_wall, B_theta_max, V_r, V_r_y, V_b, V_b_avg
    global Frequency_fan, Omega_fan, C_fan
    global PlotNum, WriteOutput, OutputFile, Monopole, FlexBaffle, AxisOnly
    global Frequency, R_i, R_o, dL, F_r, Omega, K, range_inc

    try:
	## Parse the command line arguments in sys.argv via getopt()
	opts, args = getopt.getopt(sys.argv[1:], "hdp:wo:f:a:b:l:R:mr:vzW:x:F:")

    except getopt.GetoptError, error:
	print "ParseCmdLine(): GetoptError: ",error
	return

    if len(opts) == 0:
        usage()
        return

    argCount = 0
    for optElement in opts:
        argCount = argCount + 1
        if optElement[0] == "-h":
            usage()
            sys.exit(1)
        if optElement[0] == "-d":
            DEBUG = True
        if DEBUG and 0:
            print "opts element %d:" % (argCount), optElement
        if optElement[0] == "-p":
            PlotNum = int( optElement[1] )
        if optElement[0] == "-w":
            WriteOutput = True
        if optElement[0] == "-o":
            OutputFile = optElement[1]
            WriteOutput = True
        if optElement[0] == "-f":
            Frequency = float( optElement[1] )
            Omega   = 2.*pi*Frequency
            K       = Omega / c_0
            V_wall  = 2. * delta_wall * Frequency
            V_b     = Omega * B_w / 2.
            V_b_avg = V_r_y + V_b
            C_fan   = ( rho * c_0 * V_b_avg * K ) / ( 2. * pi )
            C_wall  = ( rho * c_0 * V_wall * K ) / ( 2. * pi )
        if optElement[0] == "-F":
            Frequency_fan = float( optElement[1] )
            Omega_fan     = 2.*pi*Frequency_fan
            V_r           = Omega_fan * (R_o + (R_i - R_o/2.))
            V_r_y         = V_r * sin( B_theta_max )
            V_b_avg       = V_r_y + V_b
            C_fan         = ( rho * c_0 * V_b_avg * K ) / ( 2. * pi )
        if optElement[0] == "-a":
            R_i = float( optElement[1] ) * 0.0254 
        if optElement[0] == "-b":
            R_o = float( optElement[1] ) * 0.0254 
        if optElement[0] == "-l":
            dL = float( optElement[1] ) * 0.0254 
        if optElement[0] == "-R":
            F_r = float( optElement[1] )
        if optElement[0] == "-r":
            range_inc = float( optElement[1] )
        if optElement[0] == "-m":
            Monopole = True
        if optElement[0] == "-v":
            FlexBaffle = True
        if optElement[0] == "-W":
            delta_wall = float( optElement[1] )
            V_wall = 2. * delta_wall * Frequency
            C_wall = ( rho * c_0 * V_wall * K ) / ( 2. * pi )
        if optElement[0] == "-x":
            B_theta_max = float( optElement[1] ) * pi/180.
            V_r_y       = V_r * sin( B_theta_max )
            V_b_avg     = V_r_y + V_b
            C_fan       = ( rho * c_0 * V_b_avg * K ) / ( 2. * pi )
        if optElement[0] == "-z":
            AxisOnly = True
            PlotNum  = 0

    if DEBUG:
        print "-p", PlotNum
        print "-o", OutputFile
        print "-f", Frequency
        print "-F", Frequency_fan
        print "-a", R_i
        print "-b", R_o
        print "-l", dL
        print "-R", F_r
        print "-r", range_inc
        print "-m", Monopole
        print "-v", FlexBaffle
        print "-z", AxisOnly
        print "-W", delta_wall
        print "-x", B_theta_max

    return

##----------------------------------------------------------------------------
def OpenFile(file, mode):
    try:
        # open() will raise an IOError exception if it fails
        print "Opening file:", file
        fd = open(file,mode)

    except IOError, error:
        print "OpenFile(): I/O error:", error
        sys.exit("OpenFile()")

    except OSError, error:
        print "OpenFile(): OS error:", error
        sys.exit("OpenFile()")
        
    except:
        print "OpenFile(): Unexpected error:", sys.exc_info()[0]
        raise
        sys.exit("OpenFile()")

    else: # executed after the try if no exceptions
        return fd
    
##----------------------------------------------------------------------------
## Provide for cmd line invocation
if __name__ == "__main__":
    main()