#%%%

# dir_in = '/Users/throop/

# dir_in = '/Users/throop/data/ORT4/throop/ort4_bc3_10cbr2_dph/'

stretch_percent = 99

stretch = astropy.visualization.PercentileInterval(stretch_percent)

# dir_data = os.path.expanduser('~/Data/')

# dir_in git

do_compress = False # Do we use .gzip compression on the Track-4 input grids?

# If we used compression on the track4_calibrate routine, we must use it here too.

# dir_track4 = os.path.join(dir_data, name_ort, 'throop', 'track4')

if do_compress:

files = glob.glob(os.path.join(dir_in, '*.grid4d.gz'))

else:

files = glob.glob(os.path.join(dir_in, '*.grid4d'))

files = glob.glob(os.path.join(dir_in, '*.dust.pkl'))

# Alphabetize file list

files = sorted(files)

plt.set_cmap('plasma')

utc_ca = '2019 1 Jan 05:33:00'

dt_before = 1*u.hour

dt_after = 1*u.hour

# area_sc = (1*u.m)**2

frame = '2014_MU69_SUNFLOWER_ROT'

name_target = 'MU69'

origin = 'lower' # Required plotting order for imshow

name_observer = 'New Horizons'

hbt.figsize((8,6))

hbt.set_fontsize(12)

dt = 1*u.s # Sampling time through the flyby. Astropy units.

# Create an output table, Astropy format

t = Table(names = ['trajectory', 'speed', 'q_dust', 'albedo', 'rho',

'tau_max', 'tau_typical', 'iof_max', 'iof_typical'],

dtype = ['U30', float, float, float, float,

float, float, float, float] )

# Start up SPICE if needed. Unload old kernels just as a safety precaution.

sp.unload('kernels_kem_prime.tm')

sp.unload('kernels_kem_alternate.tm')

sp.furnsh(f'kernels_kem_{name_trajectory}.tm')

do_short = False

if do_short:

files = files[0:4]

i=3

file = files[i]

num_files = len(files)

name_run = dir_out.split('/')[-2]

#%%%

for i,file in enumerate(files):

#%%%

print(f'Starting file {i}/{len(files)}')

grid = nh_ort_track4_grid(file) # Load the grid from disk. Uses gzip, so it is quite slow (10 sec/file)

print(f'Loading file {file}')

# Load the trajectory parameters

et_ca = int( sp.utc2et(utc_ca) ) # Force this to be an integer, just to make output cleaner.

et_start = et_ca - dt_before.to('s').value

et_end = et_ca + dt_after.to('s').value

grid.frame = frame

grid.name_target = name_target

grid.name_trajectory = name_trajectory

# And call the method to fly through it!

# The returned density values etc are available within the instance variables, not returned explicitly.

grid.fly_trajectory(name_observer, et_start, et_end, dt)

# If the first time thru loop, make plot of our path through the system

do_plots_geometry = True

if (do_plots_geometry and (i==0)):

grid.plot_trajectory_geometry()

# Make slice plots thru the grid

do_plot_slices_xyz = False

if do_plot_slices_xyz:

hbt.fontsize(8)

hbt.figsize((20,5))

grid.plot(axis_sum=0)

grid.plot(axis_sum=1)

grid.plot(axis_sum=2)

hbt.fontsize(10)

# Make a plot of optical depth

do_plot_tau = True

if do_plot_tau:

grid.plot_tau()

# =============================================================================

# Make some plots of count rate vs. time!

# =============================================================================

# Make a plot of the instantaneous count rate

hbt.figsize((10,15))

# Make a plot of the actual density that we give to Doug Mehoke

# Define a list of colors. This is so we can use colors= argument to set

# a marker to show grain size, rather than let plot() auto-assign.

# colors = plt.rcParams['axes.prop_cycle'].by_key()['color'] # This is the default color iterator.

colors = ['antiquewhite',

'tomato',

'blueviolet',

'skyblue',

'gold',

'darkcyan',

'thistle',

'olive',

'red',

'sienna',

'deepskyblue',

'lightsalmon',

'pink',

]

# 'aqua']

# 'antiquewhite4', 'aqua', 'aquamarine4', 'black', 'blue', 'blueviolet',

# 'brown1', 'chartreuse1', 'darkgreen', 'darkorange1', 'dodgerblue1', 'lightpink', 'magenta']

# Make a plot of dust number density. This is straight out of the grid, and for comparison with MRS.

vals_fiducial = [1e-10, 1e-8, 1e-6,1e-4, 1e-2, 1e-0]

plt.subplot(3,1,1)

for j,s in enumerate(grid.s):

plt.plot(grid.delta_et_t, grid.number_t[j],

label = 's={:.2f} mm'.format(s),

color=colors[j])

plt.legend()

plt.title('Dust number density'.format(grid.area_sc))

plt.xlabel('ET from C/A')

plt.yscale('log')

plt.ylim((1e-10, 1e2)) # Match to MRS.

plt.axvline(0, color='black', alpha=0.05)

plt.ylabel(r'Dust,, # km$^{-3}$')

for val in vals_fiducial:

plt.axhline(val, color='black', alpha=0.05)

# Make a plot of impact rate. This assumes a s/c area.

plt.subplot(3,1,2)

for j,s in enumerate(grid.s): # 's' is dust size

plt.plot(grid.delta_et_t, grid.number_sc_t[j],

label = f's={s:.2f} mm',

color=colors[j])

plt.title('Impact rate, A={}'.format(grid.area_sc))

plt.yscale('log')

plt.xlabel('ET from C/A')

plt.legend()

plt.ylabel(r'Dust, # Impacts sec$^{{-1}}$')

# Make a plot of the cumulative count rate. Mark grain sizes here too.

plt.subplot(3,1,3)

for j,s in enumerate(grid.s): # Loop over size

plt.plot(grid.delta_et_t, grid.number_sc_cum_t[j], # Main plot line

label = 's={:.2f} mm'.format(s), color=colors[j])

plt.plot([grid.delta_et_t[-1]], [grid.number_sc_cum_t[j,-1].value], # Circle to indicate grain size

markersize=(7-j)*2, marker = 'o', # Use same color as prev line!

color=colors[j])

hbt.figsize(5,5)

plt.legend()

plt.title('Number of impacts (cumulative), A={}'.format(grid.area_sc))

plt.xlabel('ET from C/A')

plt.yscale('log')

plt.ylabel('# of Impacts')

plt.axhline(y = 1, linestyle = '--', alpha = 0.1)

plt.tight_layout()

plt.show()

# Make a plot of size distibution.

# Make two curves: one for n(r) for the entire grid, and one for n(r) that hits s/c

# Now add an entry to the table. This is a table that lists all of the results --

# e.g., max_tau, count rate etc

# One line per grid.

t.add_row(vals=[grid.name_trajectory, grid.speed, grid.q, grid.albedo, grid.rho,

grid.tau_max, grid.tau_typ,

grid.iof_max, grid.iof_typ])

# Get size dist along path

number_path = grid.number_sc_cum_t[:,-1].value

# Take the full particle grid, and sum along all spatial axes, leaving just the size axis left.

number_grid = np.sum(np.sum(np.sum(grid.density, axis=1), axis=1), axis=1)

# Normalize the size dists both

number_grid = hbt.normalize(number_grid)

number_path = hbt.normalize(number_path)

plt.plot(grid.s, number_path, label = 'Along s/c path')

plt.plot(grid.s, number_grid, label = 'In grid, total')

plt.yscale('log')

plt.xscale('log')

plt.ylim( (hbt.minval(np.array([number_grid, number_path]))/2, 1) )

plt.xlabel('Radius [mm]')

plt.ylabel('Particle number [arbitrary]')

plt.legend(loc = 'lower right')

plt.show()

# Output the dust population for this run to a file. This is the file that Doug Mehoke will read.

grid.output_trajectory(name_run=name_run, do_positions=False, dir_out=dir_out)

print('---')

#%%%

# Print the table

t.pprint(max_width=-1)

# Save the table as output

file_out = os.path.join(dir_out, f'nh_{name_trajectory}_track4_table.pkl')

lun = open(file_out, 'wb')

pickle.dump(t,lun)

lun.close()

print(f'Wrote: {file_out}')

#%%%

# Now that all files have been created, compress results into an archive (.tar.gz) for Doug Mehoke

inits_track4 = 'hbt'

# if 'hamilton' in files[0]:

# inits_track3 = 'dph'

# if 'kauf' in files[0]:

# inits_track3 = 'dk'

file_out = f'{name_trajectory}_{name_run}_{inits_track4}_n{num_files}.tgz'

str = f'cd {dir_out}; tar -czf {file_out} *{name_trajectory}*.dust'

_ = subprocess.Popen(str, shell=True)

print(str)

file_in = '/Users/throop/Data/ORT2/throop/track4/nh_ort_track4_table.pkl'

lun = open(file_in, 'rb')

t = ge.load(lun)

lun.close()

hbt.figsize((10,8))

xvals = ['albedo', 'speed', 'rho', 'q_dust']

i = 0

for xval in xvals:

plt.subplot(2,2,i+1)

# plt.plot(t[xval], t['iof_max'], marker = 'o', linestyle='none', label = 'I/F Max')

plt.plot(t[xval], t['iof_typical'], marker = '+', linestyle='none', label = 'I/F Typical')

plt.xlabel(xval)

plt.yscale('log')

plt.legend()

i +=1

plt.tight_layout()

plt.show()

t.sort(['speed', 'q_dust', 'albedo', 'rho'])

"""

This is a one-off utility function for MRS.

In it, I just do a flyby of MU69, and output the X Y Z grid indices (as well as positions and timestamps).

I don't output density at all -- just the s/c positions.

I do this for both prime and alternate trajectories.

This is just because he hasn't integrated SPICE into his grid code.

This function is stand-alone. It doesn't rely on the grid class.

It is included in this file because it directly relates to the grids.

"""

#%%%

name_trajectory = 'alternate' # ← Set this to 'prime' or 'alternate'

# name_trajectory = 'prime' # ← Set this to 'prime' or 'alternate'

hbt.unload_kernels_all()

frame = '2014_MU69_SUNFLOWER_ROT'

name_observer = 'New Horizons'

name_target = 'MU69'

sp.furnsh(f'kernels_kem_{name_trajectory}.tm')

# Get the OD version. This might not work.

# files = hbt.list_kernels_loaded()

# for file in files:

# if

# file_in = '/Users/throop/Data/ORT2/throop/track4/ort2-ring_v2.2_q2.0_pv0.10_rho0.22.grid4d.gz'

# file_in = '/Users/throop/Data/ORT5/throop/deliveries/tuna9k/ort5_None_y3.0_q3.5_pv0.70_rho1.00.dust.txt'

file_in = '/Users/throop/Data/ORT5/kaufmann/deliveries/chr3-tunacan10k/chr3-0003' + \

'/y3.0/beta1.0e+00/subset00/model.array2'

# file_in = '/Users/throop/Data/ORT5/throop/deliveries/tuna9k/ort5_None_y3.0_q3.5_pv0.70_rho1.00.dust.pkl' #250km

# file_in = '/Users/throop/Data/ORT5/throop/deliveries/dph-sunflower10k/ort5_None_y2.2_q2.5_pv0.05_rho0.46.dust.pkl' #500km

file_in = '/Users/throop/Data/ORT5/throop/deliveries/dph-tunacan3.5kinc55/ort5_None_y2.2_q2.5_pv0.05_rho0.46.dust.pkl' #500km

grid = nh_ort_track4_grid(file_in) # Load the grid from disk. Uses gzip, so it is quite slow (10 sec/file)

resolution_km = int(grid.resolution_km[0])

utc_ca = '2019 1 Jan 05:33:00'

dt_before = 1*u.hour

dt_after = 1*u.hour

dt = 1*u.s # Sampling time through the flyby. Astropy units.

et_ca = int( sp.utc2et(utc_ca) ) # Force this to be an integer, just to make output cleaner.

et_start = et_ca - dt_before.to('s').value

et_end = et_ca + dt_after.to('s').value

grid.frame = frame

grid.name_target = name_target

grid.name_trajectory = name_trajectory

# And call the method to fly through it!

# The returned density values etc are available within the instance variables, not returned explicitly.

grid.fly_trajectory(name_observer, et_start, et_end, dt)

# Make plots

hbt.figsize((9,9))

hbt.fontsize(12)

plt.subplot(3,2,1)

plt.plot(grid.bin_x_t)

plt.ylabel('X Bin #')

plt.title(f'MU69, Trajectory = {name_trajectory}, frame = {frame}')

plt.subplot(3,2,3)

plt.plot(grid.bin_y_t)

plt.ylabel('Y Bin #')

plt.subplot(3,2,5)

plt.plot(grid.bin_z_t)

plt.ylabel('Z Bin #')

plt.xlabel('Timestep #')

t_t = grid.et_t - np.mean(grid.et_t)

bin_t = range(len(t_t))

plt.subplot(3,2,2)

plt.axhline(0, color='pink')

plt.axvline(0, color='pink')

plt.plot(t_t, grid.x_t)

plt.ylabel('X [km]')

plt.xlabel('t [sec]')

plt.subplot(3,2,4)

plt.axhline(0, color='pink')

plt.axvline(0, color='pink')

plt.plot(t_t, grid.y_t)

plt.ylabel('Y [km]')

plt.subplot(3,2,6)

plt.axhline(0, color='pink')

plt.axvline(0, color='pink')

plt.plot(t_t, grid.z_t)

plt.ylabel('Z [km]')

plt.xlabel('Time from C/A [sec]')

plt.tight_layout()

# Save the plot to a file

file_out = f'positions_trajectory_{name_trajectory}.png'

path_out = os.path.join(dir_out, file_out)

plt.savefig(path_out)

print(f'Wrote: {path_out}')

plt.show()

# Make a table

arr = {'bin' : bin_t,

'delta_et' : t_t,

'X_km' : grid.x_t,

'Y_km' : grid.y_t,

'Z_km' : grid.z_t,

'Bin_X' : grid.bin_x_t,

'Bin_Y' : grid.bin_y_t,

'Bin_Z' : grid.bin_z_t}

t = Table(arr, names=['bin', 'delta_et', 'X_km', 'Y_km', 'Z_km', 'Bin_X', 'Bin_Y', 'Bin_Z'],

dtype=['int', 'int', 'float', 'float', 'float', 'int', 'int', 'int'])

# Save the table to a file

file_out = f'positions_trajectory_{name_trajectory}_res{resolution_km}.txt'

path_out = os.path.join('/Users/throop/Data/ORT5', file_out)

t.write(path_out, format = 'ascii.csv', overwrite=True)