import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from scipy.stats import norm, cauchy
# FWHM = 2.355 * sigma, so sigma = FWHM / 2.355
fwhm = 1
sigma = fwhm / 2.355
width = 0.02;
# Prepare data
v = np.linspace(-1, 1, 100)
L = 1 * v
x_right = np.linspace(-1, 1, 200)
def profile_right(v0, n):
dwidth = 0.08;
pop1_raw=0.6 * norm.pdf(x_right, 0, sigma)/norm.pdf(0, 0, sigma)
pop2_raw=0.4 * norm.pdf(x_right, 0, sigma)/norm.pdf(0, 0, sigma)
depl1= 0.25*np.tanh(0.2 * (norm.pdf(x_right-v0-0.2, 0, dwidth)+norm.pdf(x_right+v0+0.2, 0, dwidth)))*pop1_raw
depl2= 0.25*np.tanh(0.2 * (norm.pdf(x_right-v0+0.2, 0, dwidth)+norm.pdf(x_right+v0-0.2, 0, dwidth)))*pop2_raw
pop1 = pop1_raw - depl1 + depl2
pop2 = pop2_raw + depl1 - depl2
exc1 = (cauchy.pdf(2*(x_right-v0-0.2)/width)+cauchy.pdf(2*(x_right+v0+0.2)/width)) / cauchy.pdf(0)
exc2 = (cauchy.pdf(2*(x_right-v0+0.2)/width)+cauchy.pdf(2*(x_right+v0-0.2)/width)) / cauchy.pdf(0)
pop1 = pop1 * ( 1 - 0.5*np.tanh(exc1/0.5) )
pop2 = pop2 * ( 1 - 0.5*np.tanh(exc2/0.5) )
if (n==1):
return pop1
elif (n==2):
return pop2
elif (n==3):
return pop1_raw
elif (n==4):
return pop2_raw
else: return 0
# Bottom plot - Gaussian
#def gaussian_b(v0):
# x_bottom = np.linspace(-1, 1, 200)
# return profile_bottom(v0, x_bottom)
# Left plot - 1 - 0.5*Gaussian
#x_top = np.linspace(-1, 1, 200)
#gaussian_top = norm.pdf(x_top, 0, sigma)/norm.pdf(0, 0, sigma)
#profile_top = 1 - 0.3 * gaussian_top
def x_top(x):
return np.linspace(-1, x, int(100*(1+x)))
def profile_top(x):
return 1 - 0.3 * norm.pdf(x, 0, sigma)/norm.pdf(0, 0, sigma)
def profile2_top(x):
return 1 - 0.5 * norm.pdf(x, 0, sigma)/norm.pdf(0, 0, sigma)
def profile_lambdip_top(x):
return 1 - 0.5 * norm.pdf(x, 0, sigma)/norm.pdf(0, 0, sigma) + (0.1 * (cauchy.pdf(2*(x-0.2)/(1.4*width))+ cauchy.pdf(2*(x+0.2)/(1.4*width)))-0.2*cauchy.pdf(2*(x)/(1.*width))) / cauchy.pdf(0)
# Create subplots
fig = make_subplots(
rows=2, cols=2,
row_heights=[0.25, 0.75],
column_widths=[0.75, 0.25],
horizontal_spacing=0.02,
vertical_spacing=0.02,
specs=[[{"type": "xy"}, None],
[{"type": "xy"}, {"type": "xy"}]]
)
# Create frames for animation/slider
v0_values = np.linspace(-1, 1, 41) # 41 steps from -1 to 1
frames = []
for v0 in v0_values:
frame_data = [
# Top plot trace
go.Scatter(x=x_top(v0), y=profile_lambdip_top(x_top(v0)), mode='lines', line=dict(color='green', width=2), showlegend=False),
# Top plot arrow
go.Scatter(x=[v0], y=[profile_lambdip_top(v0)], mode="markers", marker=dict(symbol='arrow-down', size=11, color="red")),
go.Scatter(x=[v0,v0], y=[1,profile_lambdip_top(v0)], mode='lines', line=dict( color='red', width=2), showlegend=False),
go.Scatter(x=[v0,v0], y=[0,profile_lambdip_top(v0)], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False),
# Main plot line
go.Scatter(x=v, y=L-0.2, mode='lines', line=dict(color='red', width=2), showlegend=False),
go.Scatter(x=v, y=L+0.2, mode='lines', line=dict(color='blue', width=2), showlegend=False),
go.Scatter(x=v, y=-(L-0.2), mode='lines', line=dict(color='red', width=2), showlegend=False),
go.Scatter(x=v, y=-(L+0.2), mode='lines', line=dict(color='blue', width=2), showlegend=False),
# Main plot marker
go.Scatter(x=[v0,v0], y=[-1,1], mode='lines',
line=dict(dash='dash', color='black', width=0.5), showlegend=False),
go.Scatter(x=[v0,1], y=[v0-0.2,v0-0.2], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False),
go.Scatter(x=[v0,1], y=[v0+0.2,v0+0.2], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False),
go.Scatter(x=[v0,1], y=[-(v0-0.2),-(v0-0.2)], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False),
go.Scatter(x=[v0,1], y=[-(v0+0.2),-(v0+0.2)], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False),
# Main plot dot
go.Scatter(x=[v0], y=[v0-0.2], mode='markers', marker=dict(size=10, color='red'), showlegend=False),
go.Scatter(x=[v0], y=[v0+0.2], mode='markers', marker=dict(size=10, color='red'), showlegend=False),
go.Scatter(x=[v0], y=[-(v0-0.2)], mode='markers', marker=dict(size=10, color='red'), showlegend=False),
go.Scatter(x=[v0], y=[-(v0+0.2)], mode='markers', marker=dict(size=10, color='red'), showlegend=False),
# Right plot
go.Scatter(x=profile_right(v0, 1), y=x_right, mode='lines', line=dict(color='blue', width=2), showlegend=False),
go.Scatter(x=profile_right(v0, 2), y=x_right, mode='lines', line=dict(color='red', width=2), showlegend=False),
# Right plot marker
go.Scatter(x=[0,1.1], y=[v0-0.2,v0-0.2], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False),
go.Scatter(x=[0,1.1], y=[v0+0.2,v0+0.2], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False),
go.Scatter(x=[0,1.1], y=[-(v0-0.2),-(v0-0.2)], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False),
go.Scatter(x=[0,1.1], y=[-(v0+0.2),-(v0+0.2)], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False)
]
frames.append(go.Frame(data=frame_data, name=str(v0)))
# Initial v0 = 0
v0_init = -0.95
# Add traces for initial state
# Top plot (row=1, col=1)
fig.add_trace(go.Scatter(x=x_top(v0_init), y=profile_lambdip_top(x_top(v0_init)), mode='lines', line=dict(color='green', width=2), showlegend=False), row=1, col=1)
fig.add_trace(go.Scatter(x=[v0_init], y=[profile_lambdip_top(v0_init)], mode="markers", marker=dict(symbol='arrow-down', size=11, color="red")), row=1, col=1)
fig.add_trace(go.Scatter(x=[v0_init,v0_init], y=[1,profile_lambdip_top(v0_init)], mode='lines', line=dict( color='red', width=2), showlegend=False), row=1, col=1)
fig.add_trace(go.Scatter(x=[v0_init,v0_init], y=[profile_lambdip_top(v0_init),0], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False), row=1, col=1)
# Main plot (row=1, col=2)
fig.add_trace(go.Scatter(x=v, y=L-0.2, mode='lines', line=dict(color='red', width=2), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=v, y=L+0.2, mode='lines', line=dict(color='blue', width=2), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=-v, y=L-0.2, mode='lines', line=dict(color='blue', width=2), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=-v, y=L+0.2, mode='lines', line=dict(color='red', width=2), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=[v0_init,v0_init], y=[-1,1], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=[v0_init,1], y=[v0_init-0.2,v0_init-0.2], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=[v0_init,1], y=[v0_init+0.2,v0_init+0.2], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=[v0_init,1], y=[-(v0_init-0.2),-(v0_init-0.2)], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=[v0_init,1], y=[-(v0_init+0.2),-(v0_init+0.2)], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=[v0_init], y=[v0_init-0.2], mode='markers', marker=dict(size=10, color='red'), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=[v0_init], y=[v0_init+0.2], mode='markers', marker=dict(size=10, color='red'), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=[v0_init], y=[-(v0_init-0.2)], mode='markers', marker=dict(size=10, color='red'), showlegend=False), row=2, col=1)
fig.add_trace(go.Scatter(x=[v0_init], y=[-(v0_init+0.2)], mode='markers', marker=dict(size=10, color='red'), showlegend=False), row=2, col=1)
# Right plot (row=2, col=2)
fig.add_trace(go.Scatter(x=profile_right(v0_init, 1), y=x_right, mode='lines', line=dict(color='blue', width=2), showlegend=False), row=2, col=2)
fig.add_trace(go.Scatter(x=profile_right(v0_init, 2), y=x_right, mode='lines', line=dict(color='red', width=2), showlegend=False), row=2, col=2)
fig.add_trace(go.Scatter(x=[0,1.1], y=[v0_init-0.2,v0_init-0.2], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False), row=2, col=2)
fig.add_trace(go.Scatter(x=[0,1.1], y=[v0_init+0.2,v0_init+0.2], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False), row=2, col=2)
fig.add_trace(go.Scatter(x=[0,1.1], y=[-(v0_init-0.2),-(v0_init-0.2)], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False), row=2, col=2)
fig.add_trace(go.Scatter(x=[0,1.1], y=[-(v0_init+0.2),-(v0_init+0.2)], mode='lines', line=dict(dash='dash', color='black', width=0.5), showlegend=False), row=2, col=2)
# Add frames to figure
fig.frames = frames
# Settings to draw a black border line around the plot area
BORDER_SETTINGS = dict(
showline=True,
linecolor='black',
linewidth=1,
mirror=True # This forces the line to be drawn on all four sides of the plot area
)
fig.add_hline(y=1, line_color="black", line_width=1, line_dash="solid", row=1, col=1)
fig.add_trace(go.Scatter(x=x_top(1), y=profile_top(x_top(1)), mode='lines', line=dict(color='lightgray', width=1), showlegend=False), row=1, col=1)
fig.add_trace(go.Scatter(x=x_top(1), y=profile2_top(x_top(1)), mode='lines', line=dict(color='gray', width=1), showlegend=False), row=1, col=1)
fig.add_trace(go.Scatter(x=profile_right(0, 3), y=x_right, mode='lines', line=dict(color='lightblue', width=1), showlegend=False), row=2, col=2)
fig.add_trace(go.Scatter(x=profile_right(0, 4), y=x_right, mode='lines', line=dict(color='lightcoral', width=1), showlegend=False), row=2, col=2)
# Update layout for top plot (row=1, col=1)
fig.update_xaxes( range=[-1, 1], showticklabels=False, showgrid=False, gridwidth=0.5, gridcolor='lightgray', zeroline=True, **BORDER_SETTINGS,row=1, col=1)
fig.update_yaxes( range=[0, 1.1], showgrid=True, gridwidth=0.5, gridcolor='lightgray', zeroline=True, side='left', **BORDER_SETTINGS, row=1, col=1)
# Update layout for main plot (row=1, col=2)
fig.update_xaxes( range=[-1, 1], showticklabels=True, tickvals=[0], showgrid=False, gridwidth=0.5, gridcolor='lightgray', zeroline=True, **BORDER_SETTINGS, row=2, col=1)
fig.update_yaxes( range=[-1, 1], tickvals=[0], showgrid=False, gridwidth=0.5, gridcolor='lightgray', zeroline=True, **BORDER_SETTINGS, row=2, col=1)
# Update layout for bottom plot (row=2, col=2)
fig.update_xaxes( range=[0, 1.1], showgrid=True, gridwidth=0.5, gridcolor='lightgray', zeroline=True, side='left', **BORDER_SETTINGS, row=2, col=2)
fig.update_yaxes( range=[-1, 1], showticklabels=False, showgrid=False, gridwidth=0.5, gridcolor='lightgray', zeroline=True, **BORDER_SETTINGS, row=2, col=2)
# Add slider
sliders = [dict(
active=1, # Start at v0=0 (middle of range)
yanchor="top",
y=-0.08,
xanchor="left",
x=0,
currentvalue=dict(
prefix="v0: ",
visible=False,
xanchor="right"
),
pad=dict(b=10, t=50),
len=0.75,
steps=[dict(
args=[[f.name], dict(
frame=dict(duration=0, redraw=True),
mode="immediate",
transition=dict(duration=0)
)],
label=f"{v0:.2f}",
method="animate"
) for f, v0 in zip(frames, v0_values)]
)]
fig.update_layout(
sliders=sliders,
height=650,
width=650,
showlegend=False,
plot_bgcolor='white',
margin=dict(l=80, r=50, t=75, b=75)
)
# Add annotations
#fig.add_annotation( text="Δω = k v<sub>z</sub>", xref="x2", yref="y2", x=0.6, y=0.88, showarrow=False, font=dict(size=18, color='black'), xanchor='center')
#fig.add_annotation( text="Δω = - k v<sub>z</sub>", xref="x2", yref="y2", x=-0.4, y=0.7, showarrow=False, font=dict(size=18, color='black'), xanchor='center')
# Labels using annotations to position them correctly
fig.add_annotation( text="transmitted<br>laser intensity", xref="paper", yref="paper", x=-0.075, y=0.99, showarrow=False, font=dict(size=15), textangle=-90, xanchor='center', yanchor='top')
fig.add_annotation( text="velocity v<sub>z</sub>", xref="paper", yref="paper", x=-0.05, y=0.37, showarrow=False, font=dict(size=15), textangle=-90, xanchor='center', yanchor='middle')
fig.add_annotation( text="frequency Δω", xref="paper", yref="paper", x=0.37, y=-0.025, showarrow=False, font=dict(size=15), xanchor='center', yanchor='top')
fig.add_annotation( text="ground state<br>populations N<sub>1</sub>, N<sub>2</sub>", xref="paper", yref="paper", x=0.87, y=0.74, showarrow=False, font=dict(size=15), xanchor='center', yanchor='bottom')
fig.show()