1+ # Copyright 2024-2025 Open Quantum Design
2+
3+ # Licensed under the Apache License, Version 2.0 (the "License");
4+ # you may not use this file except in compliance with the License.
5+ # You may obtain a copy of the License at
6+
7+ # http://www.apache.org/licenses/LICENSE-2.0
8+
9+ # Unless required by applicable law or agreed to in writing, software
10+ # distributed under the License is distributed on an "AS IS" BASIS,
11+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ # See the License for the specific language governing permissions and
13+ # limitations under the License.
14+
15+ import matplotlib .pyplot as plt
16+ from matplotlib .figure import Figure
17+ import numpy as np
18+ from typing import Dict , List
19+
20+ # Actual imports from oqd_core
21+ from oqd_core .interface .atomic import Ion , Level , Transition
22+
23+
24+ def get_orbital_label (l_value : float ) -> str :
25+ """Converts orbital quantum number l to its spectroscopic letter."""
26+ # The Level.orbital is NonNegativeAngularMomentumNumber (float)
27+ # For labels, we typically use integer part if it's an integer
28+ l_int = int (l_value )
29+ if l_value == l_int :
30+ labels = {0 : 's' , 1 : 'p' , 2 : 'd' , 3 : 'f' , 4 : 'g' , 5 : 'h' }
31+ return labels .get (l_int , f'l={ l_int } )
32+ return f'l={ l_value } # For half-integer or other float cases if they occur
33+
34+
35+ def generate_energy_level_plot (ion : Ion , ** plot_kwargs ) -> Figure :
36+ """
37+ Generates an energy level diagram (Grotrian-like diagram) for a given Ion instance.
38+
39+ Args:
40+ ion: An instance of the Ion class containing energy levels and transitions.
41+ **plot_kwargs: Additional keyword arguments to customize the plot.
42+ figure (dict): Arguments for plt.subplots (e.g., {'figsize': (10, 8)})
43+ level_color (str): Color for level lines.
44+ level_linewidth (float): Linewidth for level lines.
45+ level_fontsize (float): Fontsize for level labels.
46+ level_line_width_half (float): Half-width of the horizontal level lines.
47+ text_offset_x (float): Horizontal offset for level labels from line end.
48+ text_offset_y_factor (float): Factor to determine vertical offset for level labels (multiplied by energy range).
49+ show_n_quantum_number (bool): Whether to show principal quantum number in level labels.
50+ arrow_style (str): Matplotlib arrow style for transitions.
51+ transition_color (str): Color for transition arrows.
52+ arrow_shrink_a (float): Shrink factor for arrow tip.
53+ arrow_shrink_b (float): Shrink factor for arrow base.
54+ transition_linewidth (float): Linewidth for transition arrows.
55+ show_transition_labels (bool): Whether to show labels for transitions.
56+ transition_label_color (str): Color for transition labels.
57+ transition_label_fontsize (float): Fontsize for transition labels.
58+ xlabel (str): Label for the x-axis.
59+ ylabel (str): Label for the y-axis.
60+ title (str): Title for the plot.
61+
62+ Returns:
63+ A matplotlib.figure.Figure object representing the plot.
64+ """
65+ fig , ax = plt .subplots (** plot_kwargs .get ('figure' , {}))
66+
67+ ion_name = getattr (ion , 'name' , 'Ion' ) # Check if Ion has a name/label, else default
68+
69+ if not ion .levels :
70+ ax .text (0.5 , 0.5 , "No energy levels defined for this ion." ,
71+ horizontalalignment = 'center' , verticalalignment = 'center' ,
72+ transform = ax .transAxes )
73+ ax .set_title (plot_kwargs .get ('title' , f"Energy Level Diagram for { ion_name } ))
74+ return fig
75+
76+ # Create a mapping from level labels to Level objects for quick lookup if transitions use strings
77+ levels_map_by_label : Dict [str , Level ] = {level .label : level for level in ion .levels }
78+
79+ # Group levels by orbital quantum number (l)
80+ levels_by_orbital : dict [float , List [Level ]] = {}
81+ for level in ion .levels :
82+ orbital_val = float (level .orbital ) # Ensure it's float for dict key
83+ if orbital_val not in levels_by_orbital :
84+ levels_by_orbital [orbital_val ] = []
85+ levels_by_orbital [orbital_val ].append (level )
86+
87+ sorted_orbitals = sorted (levels_by_orbital .keys ())
88+ orbital_x_coords = {l_val : i for i , l_val in enumerate (sorted_orbitals )}
89+
90+ level_line_width_half = plot_kwargs .get ("level_line_width_half" , 0.35 )
91+ text_offset_x = plot_kwargs .get ("text_offset_x" , 0.05 )
92+
93+ level_coords : Dict [str , tuple [float , float ]] = {}
94+ min_energy = float ('inf' )
95+ max_energy = float ('-inf' )
96+
97+ for level in ion .levels : # Iterate once to find min/max energy for dynamic text_offset_y
98+ min_energy = min (min_energy , level .energy )
99+ max_energy = max (max_energy , level .energy )
100+
101+ energy_span = max_energy - min_energy if max_energy > min_energy else 1.0
102+ text_offset_y = plot_kwargs .get ("text_offset_y_factor" , 0.01 ) * energy_span
103+
104+
105+ for l_val , levels_in_group in levels_by_orbital .items ():
106+ x_center = orbital_x_coords [l_val ]
107+ levels_in_group .sort (key = lambda lvl : lvl .energy )
108+
109+ for level in levels_in_group :
110+ y = level .energy
111+
112+ ax .hlines (y , x_center - level_line_width_half , x_center + level_line_width_half ,
113+ colors = plot_kwargs .get ('level_color' , 'black' ),
114+ linewidth = plot_kwargs .get ('level_linewidth' , 1.5 ))
115+
116+ level_coords [level .label ] = (x_center , y )
117+
118+ label_text = f"{ level .label }
119+ if plot_kwargs .get ("show_n_quantum_number" , True ):
120+ label_text = f"n={ level .principal } { label_text }
121+
122+ ax .text (x_center + level_line_width_half + text_offset_x , y + text_offset_y , label_text ,
123+ verticalalignment = 'bottom' ,
124+ fontsize = plot_kwargs .get ('level_fontsize' , 8 ))
125+
126+ for transition in ion .transitions :
127+ # Resolve level1 and level2 if they are strings
128+ level1_obj = transition .level1 if isinstance (transition .level1 , Level ) else levels_map_by_label .get (str (transition .level1 ))
129+ level2_obj = transition .level2 if isinstance (transition .level2 , Level ) else levels_map_by_label .get (str (transition .level2 ))
130+
131+ if level1_obj and level2_obj and level1_obj .label in level_coords and level2_obj .label in level_coords :
132+ x1 , y1 = level_coords [level1_obj .label ]
133+ x2 , y2 = level_coords [level2_obj .label ]
134+
135+ ax .annotate ("" ,
136+ xy = (x2 , y2 ), xycoords = 'data' ,
137+ xytext = (x1 , y1 ), textcoords = 'data' ,
138+ arrowprops = dict (arrowstyle = plot_kwargs .get ('arrow_style' , "->" ),
139+ connectionstyle = plot_kwargs .get ('connectionstyle' , "arc3,rad=0.1" ),
140+ color = plot_kwargs .get ('transition_color' , 'gray' ),
141+ shrinkA = plot_kwargs .get ('arrow_shrink_a' , 5 ),
142+ shrinkB = plot_kwargs .get ('arrow_shrink_b' , 5 ),
143+ linewidth = plot_kwargs .get ('transition_linewidth' , 1 )),
144+ )
145+ if plot_kwargs .get ("show_transition_labels" , False ) and transition .label :
146+ mid_x = (x1 + x2 ) / 2
147+ mid_y = (y1 + y2 ) / 2
148+ # Add a small offset if x1 and x2 are the same to avoid label overlap with vertical line
149+ label_offset_x = 0.1 if x1 == x2 else 0
150+ ax .text (mid_x + label_offset_x , mid_y , transition .label ,
151+ color = plot_kwargs .get ('transition_label_color' , 'blue' ),
152+ fontsize = plot_kwargs .get ('transition_label_fontsize' , 7 ),
153+ ha = 'center' , va = 'center' ,
154+ bbox = plot_kwargs .get ('transition_label_bbox' , dict (facecolor = 'white' , alpha = 0.5 , edgecolor = 'none' , pad = 0.1 )))
155+
156+ ax .set_xlabel (plot_kwargs .get ('xlabel' , "Orbital Angular Momentum (l)" ))
157+ ax .set_ylabel (plot_kwargs .get ('ylabel' , "Energy" ))
158+ ax .set_title (plot_kwargs .get ('title' , f"Energy Level Diagram for { ion_name } ))
159+
160+ if sorted_orbitals :
161+ ax .set_xticks (list (orbital_x_coords .values ()))
162+ ax .set_xticklabels ([get_orbital_label (l ) for l in sorted_orbitals ])
163+ else :
164+ ax .set_xticks ([])
165+
166+ if min_energy != float ('inf' ) and max_energy != float ('-inf' ):
167+ padding = 0.1 * energy_span if energy_span > 0 else 1.0
168+ ax .set_ylim (min_energy - padding , max_energy + padding )
169+
170+ plt .tight_layout ()
171+ return fig
0 commit comments