In [ ]:
%matplotlib inline
import os
import sys
from numpy import *
import numpy as np
import matplotlib
matplotlib.rcParams['figure.dpi'] = 150
import matplotlib.pyplot as plt

In [ ]:
sim_coneVec = array([0.0, 0.0, 1.0])

def jiggle(sim_precangle, sim_precession, sim_nutation, fname='nozzleVec.dat', ranseed=0):
    random.seed(ranseed)
    f = open(fname, 'w')
    line = '  {:10s}' + ' {:25s}'*7 + ' {:10s}\n'
    f.write(line.format('#step','time','nozzleVecX','nozzleVecY','nozzleVecZ','nozzleAngVelX','nozzleAngVelY','nozzleAngVelZ', 'randSeed'))
    sim_jetvec = array([0.0, 0.0, 1.0])
    sim_jetvecOld = array([0.0, 0.0, 1.0])
    sim_angVel =array([0.0,0.0,0.0])
    ex = array([1.0, 0.0, 0.0])
    ey = array([0.0, 1.0, 0.0])
    ez = array([0.0, 0.0, 1.0])
    
    # base vectors for precession cone
    ezprime=sim_coneVec
    
    if (ezprime[0]==ez[0] and ezprime[0]==ez[0] and ezprime[0]==ez[0]):
        exprime=ex
    else:
        exprime=cross(ezprime, -ey)
        exprime=exprime/sqrt(sum(exprime**2))
    
    eyprime=cross(sim_coneVec, exprime)
    eyprime=eyprime/sqrt(sum(eyprime**2))
    
    # precession axis in simulation coordinates
    rcart=sim_jetvec
    
    # jet axis in precession cone coordinates
    r=array([sum(rcart*exprime),\
         sum(rcart*eyprime),\
         sum(rcart*ezprime)])
    r=r/sqrt(sum(r**2))
    
    # jet angles in precession cone base
    thetajet=arccos(r[2])
    if (abs(r[1])>0 or abs(r[0])>0):
        phijet=2.0*arctan(r[1]/(r[0] + sqrt(r[0]**2 + r[1]**2)))
    else:
        phijet=pi

    
    # base vectors in precession cone coordinates
    phihat=cross(ez,r)
    if (sum(phihat**2)>0):
        phihat=phihat/sqrt(sum(phihat**2))
    else:
        phihat=ey
    
    thetahat=cross(phihat,r)
    
    vprime = sim_jetvec - sim_jetvecOld
    
    # velocity in precession cone coordinates
    v = np.array([sum(exprime*vprime), sum(eyprime*vprime), sum(ezprime*vprime)])
    if (sum(v**2)>0):
        v=v/sqrt(sum(v**2))

    sim_duration = 60*3.154E13
    ts = np.linspace(0.0, tend, tend//dt)
    dr_dt = ts[1]-ts[0]
    thetajets, phijets = [], []
    ts_plot = []
    for i, t in enumerate(ts):
        if t % (tend/200) < dr_dt: 
            sys.stdout.write('\ri = %6i / %6i (%5.1f %%)' % (i, len(ts) ,t/(tend/100))) 
        phihat=cross(ez,r)
        if (sqrt(sum(phihat**2))>0):
            phihat=phihat/sqrt(sum(phihat**2))
        else:
            phihat=ey

        thetahat=cross(phihat,r)

        # Now express velocity in new coordinate system
        vphi=sum(v*phihat)
        vtheta=sum(v*thetahat)
        
        # Random walk step
        sigma=2.0*pi*exp(-(thetajet/sim_precangle)**4)
        rn = random.random()
        psi=(rn - 0.5)*sigma - pi/2.0
        
        dv=sim_nutation*sqrt(dt/sim_duration)
        # Add to velocity
        vphi=vphi +dv*cos(psi)
        vtheta=vtheta + dv*sin(psi)
        
        # And normalize
        #dummyn=sqrt(vphi**2 + vtheta**2)
        #if (dummyn>0.0):
        #    vphi=vphi/dummyn
        #    vtheta=vtheta/dummyn
            
        # New velocity in coordinates relative to cone axis
        v=vphi*phihat + vtheta*thetahat
        
        # And normalize
        v=v/sqrt(sum(v**2))
        
        # Now move jet axis and normalize
        r=r + v*sim_precession*dt/sim_duration
        r=r/sqrt(sum(r**2))
                
        # Finally, calculate new angles
        thetajet=arccos(r[2])
        if (abs(r[0])>0 or abs(r[1])>0):
            phijet=2.0*arctan(r[1]/(r[0] + sqrt(r[0]**2 + r[1]**2)))
        else:
            phijet=pi
 

        # Now express in simulation coordinates
        jetaxis=exprime*r[0] + eyprime*r[1] + ezprime*r[2]
        
        # sim_theta and sim_phi are measured in simulations coordinates
        #sim_theta=acos(jetaxis[2]/sqrt(sum(jetaxis**2)))
        #if (jetaxis[1]>-sqrt(jetaxis[0]**2 + jetaxis[1]**2)) then
        #   sim_phi=2.0*arctan(jetaxis[1]/(jetaxis[0] + \
        #        sqrt(jetaxis[0]**2 + jetaxis[1]**2)))
        #else
        #   sim_phi=pi
        #endif
        
        # angular velocity for precession
        sim_angVel=cross(sim_jetvec,jetaxis)/dt
        
        vprime = jetaxis - sim_jetvec
        
        v = np.array([sum(exprime*vprime), sum(eyprime*vprime), sum(ezprime*vprime)])
        if (sum(v**2)>0):
            v=v/sqrt(sum(v**2)) 
        
        # move this to jiggle
        sim_jetvecOld=sim_jetvec
        sim_jetvec=jetaxis
        line = ' {:10d}' + ' {:25.18E}'*7 + ' {:10d}\n'
        #print(i+1, t, *sim_jetvec, *sim_angVel)
        f.write(line.format(i+1,t, *sim_jetvec, *sim_angVel, ranseed))

        if np.mod(t, plotdt) <= dr_dt:
            thetajets.append(thetajet)
            phijets.append(phijet)
            #thetas.append(sim_theta)
            #phis.append(sim_phi)
            #vecs.append(sim_vec)
            ts_plot.append(t)
            #print '%10.3e, %6.3f, %6.3f, %6.3f' % (sigma, cos(psi), sin(psi), thetajet)

    thetajets = np.array(thetajets)
    phijets = np.array(phijets)
    
    f.close()
    return ts_plot, thetajets, phijets

In [ ]:
sim_precangle = 0.30
sim_precession = 20
sim_nutation = 20
#tend = 1.5E14
tend = 100*3.154E13
dt = 6.4E9
plotdt = dt

i=13

fname = './nozzleVec_%02d_ang%.2f_p%i_n%i' % (i, sim_precangle, sim_precession, sim_nutation)
ts_plot, thetajets, phijets= jiggle(sim_precangle, sim_precession, sim_nutation, fname=fname+'.dat', ranseed=i)

In [ ]:
fig = plt.figure(figsize=(10,10))
ax1 = plt.subplot2grid((3,3), (0,0), rowspan=2, colspan=2, polar=True)
ax2 = plt.subplot2grid((3,3), (0,2))
ax3 = plt.subplot2grid((3,3), (1,2))
ax4 = plt.subplot2grid((3,3), (2,0))
ax5 = plt.subplot2grid((3,3), (2,1))
ax6 = plt.subplot2grid((3,12), (2,8))

theta_max = 1.1*sim_precangle/np.pi*180

for tmax in [100, 60, 20]:
    mask = np.array(ts_plot) < tmax*3.154E13
    nplot = 10
    ax1.scatter(phijets[mask][::nplot], thetajets[mask][::nplot]/pi*180.0, alpha=0.5, s=1, lw=0)
    ax1.plot(np.linspace(0, 2*np.pi, 100), np.ones(100)*sim_precangle/np.pi*180.0, color='r', linestyle='-')

    n, bins, patches = ax2.hist(thetajets[mask]/pi*180.0, bins=40, range=(0,theta_max), label='%iMyr' % tmax)
    n, bins, patches = ax3.hist(phijets[mask], bins=20)
        

ax1.set_rlim(0,theta_max)
text='precang=%.2f\nprecession=%2i  nutation=%2i\nduration=%2iMyr  nstep=%i' \
          % (sim_precangle, sim_precession, sim_nutation, tend/3.154E13, tend/dt)
ax1.annotate(text, (-0.05,1.03), xycoords='axes fraction')

ax2.set_xlabel(r'$\theta$')
ax2.legend()
ax3.set_xlabel(r'$\phi$')
ax3.set_xlim(-np.pi,np.pi)

xx = thetajets*np.cos(phijets)*180/np.pi
yy = thetajets*np.sin(phijets)*180/np.pi

hb = ax4.hist2d(xx, yy, bins=40, range=[[-theta_max, theta_max], [-theta_max, theta_max]])
ax4.set_xlim(-theta_max, theta_max)
ax4.set_ylim(-theta_max, theta_max)
ax4.set_aspect('equal')



hb = ax5.hexbin(xx, yy, gridsize=40, extent=[-theta_max, theta_max, -theta_max, theta_max])
ax5.set_aspect('equal')
ax5.set_xlim(-theta_max, theta_max)
ax5.set_ylim(-theta_max, theta_max)

cb = fig.colorbar(hb, cax=ax6)

#plt.savefig('./jiggle_patterns/jiggle_ang%.2f_p%i_n%i.png' % (sim_precangle, sim_precession, sim_nutation))
plt.show()

In [ ]:
sim_precangle = 0.30
sim_precession = 20
sim_nutation = 20
#tend = 1.5E14
tend = 100*3.154E13
dt = 6.4E9
plotdt = dt

def generate_pattern(i):
    print(i)
    fname = './jiggle_patterns/nozzleVec_%02d_ang%.2f_p%i_n%i' % (i, sim_precangle, sim_precession, sim_nutation)
    ts_plot, thetajets, phijets= jiggle(sim_precangle, sim_precession, sim_nutation, fname=fname+'.dat', ranseed=i)
    
    plt.figure(figsize=(10,10))
    ax1 = plt.subplot2grid((3,3), (0,0), rowspan=2, colspan=2, polar=True)
    ax2 = plt.subplot2grid((3,3), (0,2))
    ax3 = plt.subplot2grid((3,3), (1,2))
    ax4 = plt.subplot2grid((3,3), (2,0))
    ax5 = plt.subplot2grid((3,3), (2,1))
    ax6 = plt.subplot2grid((3,12), (2,8))

    theta_max = 1.1*sim_precangle/np.pi*180
    
    for tmax in [100, 60, 20]:
        mask = np.array(ts_plot) < tmax*3.154E13
        nplot = 10
        ax1.scatter(phijets[mask][::nplot], thetajets[mask][::nplot]/pi*180.0, alpha=0.5, s=1, lw=0)
        ax1.plot(np.linspace(0, 2*np.pi, 100), np.ones(100)*sim_precangle/np.pi*180.0, color='r', linestyle='-')

        n, bins, patches = ax2.hist(thetajets[mask]/pi*180.0, bins=40, range=(0,theta_max), label='%iMyr' % tmax)
        n, bins, patches = ax3.hist(phijets[mask], bins=20)

    
    ax1.set_rlim(0,theta_max)
    text='precang=%.2f\nprecession=%2i  nutation=%2i\nduration=%2iMyr  nstep=%i' \
              % (sim_precangle, sim_precession, sim_nutation, tend/3.154E13, tend/dt)
    ax1.annotate(text, (-0.05,1.03), xycoords='axes fraction')

    ax2.set_xlabel(r'$\theta$')
    ax2.legend()
    ax3.set_xlabel(r'$\phi$')
    ax3.set_xlim(-np.pi,np.pi)

    xx = thetajets*np.cos(phijets)*180/np.pi
    yy = thetajets*np.sin(phijets)*180/np.pi

    hb = ax4.hist2d(xx, yy, bins=40, range=[[-theta_max, theta_max], [-theta_max, theta_max]])
    ax4.set_xlim(-theta_max, theta_max)
    ax4.set_ylim(-theta_max, theta_max)
    ax4.set_aspect('equal')

    hb = ax5.hexbin(xx, yy, gridsize=40, extent=[-theta_max, theta_max, -theta_max, theta_max])
    ax5.set_aspect('equal')
    ax5.set_xlim(-theta_max, theta_max)
    ax5.set_ylim(-theta_max, theta_max)

    cb = fig.colorbar(hb, cax=ax6)
    plt.savefig(fname+'.png')

In [ ]:
from multiprocessing import Pool

p = Pool(8)
p.map(generate_pattern, range(32))
p.close()

In [ ]:
plt.figure(figsize=(12,5))
ax1 = plt.subplot2grid((1,2), (0,0))
ax2 = plt.subplot2grid((1,2), (0,1))

xx = thetajets*np.cos(phijets)*180/np.pi
yy = thetajets*np.sin(phijets)*180/np.pi

hb = ax1.hist2d(xx, yy, bins=40, range=[[-theta_max, theta_max], [-theta_max, theta_max]])
ax1.set_aspect('equal')



hb = ax2.hexbin(xx, yy, gridsize=40, extent=[-theta_max, theta_max, -theta_max, theta_max])
ax1.set_aspect('equal')
cb = fig.colorbar(hb, ax=ax2, pad=0)

In [ ]:
for sim_precangle in [0.20]:
    for sim_precession in [20, 30, 40]:
        for sim_nutation in [20, 40]:

            ts_plot, thetajets, phijets= jiggle(sim_precangle, sim_precession, sim_nutation)
            #fig = plt.figure(figsize=(10,20))
            plt.figure(figsize=(10,7))
            ax1 = plt.subplot2grid((2,3), (0,0), rowspan=2, colspan=2, polar=True)
            ax2 = plt.subplot2grid((2,3), (0,2))
            ax3 = plt.subplot2grid((2,3), (1,2))

            end = 100000
            ax1.scatter(phijets[:end], thetajets[:end]/pi*180.0, alpha=0.5, s=1, lw=0)
            ax1.set_rlim(0,15)
            text='precang=%.2f\nprecession=%2i  nutation=%2i\nduration=%2iMyr  nstep=%i' \
                  % (sim_precangle, sim_precession, sim_nutation, tend/3.154E13, tend/dt)
            ax1.annotate(text, (-0.05,1.03), xycoords='axes fraction')
            ax1.plot(np.linspace(0, 2*np.pi, 100), np.ones(100)*sim_precangle/np.pi*180.0, color='r', linestyle='-')

            n, bins, patches = ax2.hist(thetajets/pi*180.0, bins=40, range=(0,16))
            ax2.set_xlabel(r'$\theta$')

            n, bins, patches = ax3.hist(phijets, bins=40)
            ax3.set_xlabel(r'$\phi$')
            ax3.set_xlim(-np.pi,np.pi)
            plt.savefig('./jiggle_patterns/jiggle_ang%.2f_p%i_n%i.png' % (sim_precangle, sim_precession, sim_nutation))
            #plt.show()

In [ ]:
dirname='/d/d9/ychen/2018_production_runs/20180802_L438_rc10_beta07/'

simtitle = dirname.split('/')[-1]

nozzledata = np.genfromtxt(os.path.join(dirname, 'nozzleVec.dat'))

endstep=-1

ttnoz = nozzledata[:,1]
xx = nozzledata[:,2]
yy = nozzledata[:,3]
zz = nozzledata[:,4]
endtime = ttnoz[endstep]
tmyr=int(ttnoz[endstep]/3.15569E13)
print (tmyr)

thetas = np.arccos(zz[:endstep])/np.pi*180
phis = np.arctan2(yy[:endstep], xx[:endstep])

plt.figure(figsize=(10,7))
ax1 = plt.subplot2grid((2,3), (0,0), rowspan=2, colspan=2, polar=True)
ax2 = plt.subplot2grid((2,3), (0,2))
ax3 = plt.subplot2grid((2,3), (1,2))

ax1.scatter(phis[::100], thetas[::100], marker='x', s=4, lw=1, alpha=0.5)
ax1.scatter(phis[:], thetas[:], s=0.1, lw=0, alpha=0.5)
ax1.set_rlim(0,12)
text = simtitle
ax1.annotate(text, (-0.05,1.03), xycoords='axes fraction')

n, bins, patches = ax2.hist(thetas, bins=40, range=(0,16), alpha=0.5)
ax2.set_xlabel(r'$\theta$')

n, bins, patches = ax3.hist(phis, bins=40, range=(-np.pi,np.pi), alpha=0.5)
ax3.set_xlabel(r'$\phi$')
ax3.set_xlim(-np.pi,np.pi)
            

#ax.set_title(simtitle,loc='left')
#ax.plot(phis, thetas, alpha=0.5)

if True:
    #nozzledata2 = np.genfromtxt(os.path.join(dirname, '../1110_L45_M10_b1_h0_rerun/nozzleVec_input.dat'))
    nozzledata2 = np.genfromtxt(os.path.join(dirname, 'nozzleVec_input.dat'))
    ttnoz2 = nozzledata2[:,1]
    xx2 = nozzledata2[:,2]
    yy2 = nozzledata2[:,3]
    zz2 = nozzledata2[:,4]
    mask = ttnoz2 < endtime
    thetas2 = np.arccos(zz2[mask])/np.pi*180
    phis2 = np.arctan2(yy2[mask], xx2[mask])
    ax1.plot(phis2, thetas2, c='r', lw=0.5, alpha=0.5)
    
    n, bins, patches = ax2.hist(thetas2, bins=40, range=(0,16), color='r', alpha=0.5)
    n, bins, patches = ax3.hist(phis2, bins=40, range=(-np.pi,np.pi), color='r', alpha=0.5)

for i in range(1,tmyr+1):
    arg = np.argmin(abs(ttnoz-i*3.15569e13))
    ax1.scatter(phis[arg], thetas[arg], marker='x', color='r', s=4)
    ax1.annotate('%i Myr' % i, xy=(phis[arg], thetas[arg]), size=6, bbox=dict(facecolor='w', lw=0, alpha=0.6))
#plt.savefig('jiggle_h1.png')

In [ ]:
nozzleHalfL = 3.75E20
zFeatherMix = 9.424e+19
duration = 6.3E14

print ('%e' % (40*(nozzleHalfL + zFeatherMix)/duration))
print ((40*(nozzleHalfL + zFeatherMix)/duration)/(0.1*3E10))

In [ ]: