from pygasflow.shockwave import PressureDeflectionLocus
import matplotlib.pyplot as plt

gamma = 1.4
M1 = 3
theta = 25
locus1 = PressureDeflectionLocus(M=M1, gamma=gamma)
locus2 = locus1.new_locus_from_shockwave(theta)

theta1, pr1 = locus1.pressure_deflection(include_mirror=True)
theta2, pr2 = locus2.pressure_deflection(include_mirror=True)
theta_segment, pr_segment = locus1.pressure_deflection_segment(theta)

fig, ax = plt.subplots()
ax.plot(theta1, pr1, ":", label="M1")
ax.plot(theta2, pr2, ":", label="M2")
ax.plot(theta_segment, pr_segment, label="segment")
ax.set_xlabel(r"Deflection Angle $  heta$ [deg]")
ax.set_ylabel("Pressure Ratio to Freestream")
ax.legend()
ax.minorticks_on()
ax.grid(which='major', linestyle='-', alpha=0.7)
ax.grid(which='minor', linestyle=':', alpha=0.5)
plt.show()