257 lines
9.8 KiB
Python
257 lines
9.8 KiB
Python
import pickle
|
|
import os
|
|
import matplotlib.pyplot as plt
|
|
from pprint import pprint
|
|
import numpy as np
|
|
from matplotlib.ticker import ScalarFormatter
|
|
import common
|
|
from datetime import date
|
|
|
|
"""
|
|
These functions plot stuff. One can but does not have to use them.
|
|
"""
|
|
def plot_state_densities_and_force_field_with_trajectories(Isotope, file_name, date_str = date.today().isoformat()):
|
|
folder_name = "Data\\"+date_str
|
|
sols_folder = folder_name+'\Solutions'
|
|
|
|
other_data_folder = folder_name+"\OtherData\\"
|
|
|
|
other_data_list = []
|
|
for filename in os.listdir(other_data_folder):
|
|
if Isotope in filename and file_name in filename and filename.endswith(".npz"):
|
|
other_data_list.append(np.load(other_data_folder+filename))
|
|
print(other_data_folder+filename)
|
|
|
|
#other_data_list.append(np.load("Z:\Users\Lenny\Lenny\ (Master)\Simulations\ and\ other\ Code\AtomLightSimulations\CleanedUpSimulation\Data\\2025-04-07\OtherData\data_dir_161into2DMOT.npz"))
|
|
if len(other_data_list) > 1:
|
|
print("There were multiple 'Other Data' Files in the directory,"
|
|
"that belong to you Isotope. The list contains:\n", other_data_list)
|
|
elif len(other_data_list) == 0:
|
|
print("Couldn't find a forcefield to plot. List seems to be empty.")
|
|
|
|
other_data_dir = other_data_list[0]
|
|
|
|
Sols = {}
|
|
|
|
for filename in os.listdir(sols_folder):
|
|
if Isotope in filename and file_name in filename and filename.endswith(".pkl"):
|
|
key = os.path.splitext(filename)[0] # Remove .pkl from filename
|
|
full_path = os.path.join(sols_folder, filename)
|
|
with open(full_path, 'rb') as f:
|
|
Sols[key] = pickle.load(f)
|
|
|
|
|
|
|
|
basis_g = other_data_dir['basis_g']
|
|
basis_e = other_data_dir['basis_e']
|
|
x0 = other_data_dir['x0']/100
|
|
t0 = other_data_dir['t0']
|
|
|
|
|
|
|
|
for names in Sols:
|
|
sol = Sols[names]
|
|
data = sol.N
|
|
num_states = data.shape[0]
|
|
num_timesteps = data.shape[1]
|
|
|
|
v0 = sol.v[0,0]
|
|
|
|
# Normalize (optional)
|
|
#data = data / data.max()
|
|
# Assume your data is in `data` with shape (timesteps, num_states)
|
|
data_g = np.transpose(data[:len(basis_g), :]) # First g states
|
|
data_e = np.transpose(data[len(basis_g):, :]) # Last e states
|
|
fig, axes = plt.subplots(1, 2, figsize=(14, 6), sharey=True)
|
|
|
|
# Plot System 1
|
|
im1 = axes[0].imshow(data_g, aspect='auto', origin='lower',
|
|
cmap='viridis', interpolation='none')
|
|
axes[0].set_title('J = 8')
|
|
axes[0].set_xlabel('State Index')
|
|
axes[0].set_ylabel('position x[m]')
|
|
|
|
|
|
state_labels_str = [f"{int(s[0])},{int(s[1])}" for s in basis_g[::]]
|
|
axes[0].set_xticks(np.arange(len(basis_g)))
|
|
axes[0].set_xticklabels(state_labels_str, rotation=90, fontsize=8)
|
|
|
|
|
|
indices = np.linspace(0, len(sol.t) - 1, 40, dtype=int) #write 40 labels
|
|
position_label = sol.r[0]*x0
|
|
pretty_position_label = [f"{x:.3f}" for x in position_label[indices]]
|
|
|
|
axes[0].set_yticks(np.arange(len(sol.t))[indices])
|
|
axes[0].set_yticklabels(pretty_position_label, rotation=0, fontsize=8)
|
|
fig.colorbar(im1, ax=axes[0], label='State Density')
|
|
|
|
# Plot System 2
|
|
im2 = axes[1].imshow(data_e, aspect='auto', origin='lower', cmap='viridis', interpolation='none')
|
|
axes[1].set_title('J = 9')
|
|
axes[1].set_xlabel('State Index')
|
|
state_labels_str = [f"{int(s[0])},{int(s[1])}" for s in basis_e[::-1]]
|
|
axes[1].set_xticks(np.arange(len(basis_e)))
|
|
axes[1].set_xticklabels(state_labels_str, rotation=90, fontsize=8)
|
|
fig.colorbar(im2, ax=axes[1], label='State Density')
|
|
fig.suptitle("Density of mF states at v_0="+str(v0*(x0/t0))+"mps")
|
|
plt.tight_layout(pad=2.0)
|
|
|
|
plotname = "StateDensityDy"+Isotope+"at"+str(round(v0*(x0/t0)))+"mps"
|
|
filedir = common.get_data_filepath("Plots", plotname, ".png")
|
|
plt.show()
|
|
#plt.savefig(filedir, dpi=200)
|
|
|
|
plt.close()
|
|
|
|
|
|
|
|
Fx_RE = other_data_dir['acceleration_profile']
|
|
x = other_data_dir['x']
|
|
v = other_data_dir['v']
|
|
|
|
X,V = np.meshgrid(x, v)
|
|
formatter = ScalarFormatter(useMathText=True)
|
|
formatter.set_powerlimits((-2, 2)) # Forces scientific notation outside this range
|
|
|
|
fig, ax0 = plt.subplots(nrows=1, ncols=1)
|
|
|
|
im0 = ax0.pcolormesh(X*x0, V*x0/t0, np.nan_to_num(Fx_RE))#, cmap="hot")
|
|
cb0 = plt.colorbar(im0)
|
|
cb0.ax.yaxis.set_major_formatter(formatter)
|
|
cb0.set_label('$a [m/s^2]$')
|
|
ax0.set_xlabel('$x[m]$')
|
|
ax0.set_ylabel('$v[m/s]$')
|
|
ax0.set_title("Acceleration from rate equaitons")
|
|
|
|
for names in Sols:
|
|
sol = Sols[names]
|
|
ax0.plot(sol.r[0]*x0, sol.v[0]*x0/t0, "w")
|
|
ax0.set_xlim(x[0]*x0, x[-1]*x0)
|
|
plt.tight_layout()
|
|
|
|
plotname = "PhaseSpaceDy"+Isotope
|
|
filedir = common.get_data_filepath("Plots", plotname, ".png")
|
|
plt.show()
|
|
#plt.savefig(filedir, dpi=200)
|
|
plt.close()
|
|
|
|
def plot_force_field_and_trajectories(Isotope, file_name, date_str = date.today().isoformat()):
|
|
folder_name = "Data\\"+date_str
|
|
sols_folder = folder_name+'\Solutions'
|
|
|
|
other_data_folder = folder_name+"\OtherData\\"
|
|
|
|
other_data_list = []
|
|
for filename in os.listdir(other_data_folder):
|
|
if Isotope in filename and file_name in filename and filename.endswith(".npz"):
|
|
other_data_list.append(np.load(other_data_folder+filename))
|
|
print(other_data_folder+filename)
|
|
|
|
#other_data_list.append(np.load("Z:\Users\Lenny\Lenny\ (Master)\Simulations\ and\ other\ Code\AtomLightSimulations\CleanedUpSimulation\Data\\2025-04-07\OtherData\data_dir_161into2DMOT.npz"))
|
|
if len(other_data_list) > 1:
|
|
print("There were multiple 'Other Data' Files in the directory,"
|
|
"that belong to you Isotope. The list contains:\n", other_data_list)
|
|
elif len(other_data_list) == 0:
|
|
print("Couldn't find a forcefield to plot. List seems to be empty.")
|
|
|
|
other_data_dir = other_data_list[0]
|
|
#other_data_dir = np.load("Data\\2025-04-14\OtherData\data_dir_Dy161RijIntoPushBeam0Gauss0DetPol-1.npz")
|
|
Sols = {}
|
|
|
|
for filename in os.listdir(sols_folder):
|
|
if Isotope in filename and file_name in filename and filename.endswith(".pkl"):
|
|
key = os.path.splitext(filename)[0] # Remove .pkl from filename
|
|
full_path = os.path.join(sols_folder, filename)
|
|
with open(full_path, 'rb') as f:
|
|
Sols[key] = pickle.load(f)
|
|
|
|
|
|
|
|
basis_g = other_data_dir['basis_g']
|
|
basis_e = other_data_dir['basis_e']
|
|
x0 = other_data_dir['x0']/100
|
|
t0 = other_data_dir['t0']
|
|
|
|
|
|
Fx_RE = other_data_dir['acceleration_profile']
|
|
x = other_data_dir['x']
|
|
v = other_data_dir['v']
|
|
|
|
X,V = np.meshgrid(x, v)
|
|
formatter = ScalarFormatter(useMathText=True)
|
|
formatter.set_powerlimits((-2, 2)) # Forces scientific notation outside this range
|
|
|
|
fig, ax0 = plt.subplots(nrows=1, ncols=1)
|
|
|
|
im0 = ax0.pcolormesh(X*x0, V*x0/t0, np.nan_to_num(Fx_RE))
|
|
cb0 = plt.colorbar(im0)
|
|
cb0.ax.yaxis.set_major_formatter(formatter)
|
|
cb0.set_label('$a [m/s^2]$')
|
|
ax0.set_xlabel('$x[m]$')
|
|
ax0.set_ylabel('$v[m/s]$')
|
|
ax0.set_title("Acceleration from rate equaitons")
|
|
|
|
for names in Sols:
|
|
sol = Sols[names]
|
|
ax0.plot(sol.r[0]*x0, sol.v[0]*x0/t0, "w")
|
|
ax0.set_xlim(x[0]*x0, x[-1]*x0)
|
|
plt.tight_layout()
|
|
|
|
plotname = "PhaseSpaceDy"+Isotope
|
|
filedir = common.get_data_filepath("Plots", plotname, ".png")
|
|
plt.show()
|
|
#plt.savefig(filedir, dpi=200)
|
|
plt.close()
|
|
|
|
def plot_force_field(Isotope, file_name, date_str = date.today().isoformat()):
|
|
folder_name = "Data\\"+date_str
|
|
sols_folder = folder_name+'\Solutions'
|
|
|
|
other_data_folder = folder_name+"\OtherData\\"
|
|
|
|
other_data_list = []
|
|
for filename in os.listdir(other_data_folder):
|
|
if file_name in filename and Isotope in filename and filename.endswith(".npz"):
|
|
other_data_list.append(np.load(other_data_folder+filename))
|
|
print(other_data_folder+filename)
|
|
|
|
#other_data_list.append(np.load("Z:\Users\Lenny\Lenny\ (Master)\Simulations\ and\ other\ Code\AtomLightSimulations\CleanedUpSimulation\Data\\2025-04-07\OtherData\data_dir_161into2DMOT.npz"))
|
|
if len(other_data_list) > 1:
|
|
print("There were multiple 'Other Data' Files in the directory,"
|
|
"that belong to you Isotope. The first entry will be chosen. The list contains:\n", other_data_list)
|
|
elif len(other_data_list) == 0:
|
|
print("Couldnt read a file in 'Other Data'. Check that you chose the correct date, Isotope and filename.")
|
|
|
|
other_data_dir = other_data_list[0]
|
|
|
|
basis_g = other_data_dir['basis_g']
|
|
basis_e = other_data_dir['basis_e']
|
|
x0 = other_data_dir['x0']/100
|
|
t0 = other_data_dir['t0']
|
|
|
|
|
|
Fx_RE = other_data_dir['acceleration_profile']
|
|
x = other_data_dir['x']
|
|
v = other_data_dir['v']
|
|
|
|
X,V = np.meshgrid(x, v)
|
|
formatter = ScalarFormatter(useMathText=True)
|
|
formatter.set_powerlimits((-2, 2)) # Forces scientific notation outside this range
|
|
|
|
fig, ax0 = plt.subplots(nrows=1, ncols=1)
|
|
|
|
im0 = ax0.pcolormesh(X*x0, V*x0/t0, np.nan_to_num(Fx_RE))
|
|
cb0 = plt.colorbar(im0)
|
|
cb0.ax.yaxis.set_major_formatter(formatter)
|
|
cb0.set_label('$a [m/s^2]$')
|
|
ax0.set_xlabel('$x[m]$')
|
|
ax0.set_ylabel('$v[m/s]$')
|
|
ax0.set_title("Acceleration from rate equaitons")
|
|
plt.tight_layout()
|
|
plotname = "PhaseSpaceDy"+Isotope
|
|
filedir = common.get_data_filepath("Plots", plotname, ".png")
|
|
plt.show()
|
|
#plt.savefig(filedir, dpi=200)
|
|
plt.close()
|
|
|