def subtract_wave(density):
averagedDensity = np.average(density, axis = 1)
peak_rad, peak_density = find_peak(averagedDensity)
end_radius = min(peak_rad + 7 * scale_height, 2.5)
used_rad, wave_locations = util.getWaveLocation(density, rad, theta, end_radius = end_radius)
# Interpolate in Wave Region
delta_theta = 30.0 * (np.pi / 180.0)
for i, (r_i, wave_location) in enumerate(zip(used_rad, wave_locations)):
rad_index = np.searchsorted(rad, r_i)
lower_theta = wave_location - delta_theta
upper_theta = wave_location + delta_theta
if lower_theta < 0.0:
lower_theta_i = np.searchsorted(theta, lower_theta + 2 * np.pi)
upper_theta_i = np.searchsorted(theta, upper_theta)
interpolated_thetas = np.concatenate((theta[lower_theta_i :] - 2 * np.pi, theta[ : upper_theta_i]))
thetas = np.array([lower_theta, upper_theta])
values = np.array([density[rad_index, lower_theta_i], density[rad_index, upper_theta_i]])
interpolated_values = np.interp(interpolated_thetas, thetas, values)
density[rad_index, lower_theta_i :] = interpolated_values[: len(theta[lower_theta_i :])]
density[rad_index, : upper_theta_i] = interpolated_values[len(theta[lower_theta_i :]) :]
elif upper_theta > 2 * np.pi:
lower_theta_i = np.searchsorted(theta, lower_theta)
upper_theta_i = np.searchsorted(theta, upper_theta - 2 * np.pi)
interpolated_thetas = np.concatenate((theta[lower_theta_i :], theta[ : upper_theta_i] + 2 * np.pi))
thetas = np.array([lower_theta, upper_theta])
values = np.array([density[rad_index, lower_theta_i], density[rad_index, upper_theta_i]])
interpolated_values = np.interp(interpolated_thetas, thetas, values)
density[rad_index, lower_theta_i :] = interpolated_values[: len(theta[lower_theta_i :])]
density[rad_index, : upper_theta_i] = interpolated_values[len(theta[lower_theta_i :]) :]
else:
lower_theta_i = np.searchsorted(theta, lower_theta)
upper_theta_i = np.searchsorted(theta, upper_theta)
interpolated_thetas = theta[lower_theta_i : upper_theta_i]
thetas = np.array([lower_theta_i, upper_theta_i])
values = np.array([density[rad_index, lower_theta_i], density[rad_index, upper_theta_i]])
density[rad_index, lower_theta_i : upper_theta_i] = np.interp(interpolated_thetas, thetas, values)
return density
if show:
plot.show()
plot.close(fig) # Close Figure (to avoid too many figures)
##### Plot One File or All Files #####
if len(sys.argv) > 1:
frame_number = int(sys.argv[1])
if frame_number == -1:
# Plot Sample
max_frame = util.find_max_frame()
sample = np.linspace(50, max_frame, 10) # 10 evenly spaced frames
for i in sample:
density = (fromfile("gasdens%d.dat" % i).reshape(num_rad, num_theta)) / surface_density_zero
radii, wave_locations = util.getWaveLocation(density, rad)
make_plot(i, density, radii, wave_locations)
else:
# Plot Single
density = (fromfile("gasdens%d.dat" % frame).reshape(num_rad, num_theta)) / surface_density_zero
radii, wave_locations = util.getWaveLocation(density, rad)
make_plot(frame_number, density, radii, wave_locations, show = True)
else:
# Search for maximum frame
density_files = glob.glob("gasdens*.dat")
max_frame = find_max_frame()
num_frames = max_frame + 1
#for i in range(num_frames):
# make_plot(i)