We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
better WMEL defaults


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `WrightTools/diagrams/WMEL.py`
Content:
```
1 ### import ####################################################################
2 
3 
4 import numpy as np
5 import matplotlib.pyplot as plt
6 
7 
8 ### define ####################################################################
9 
10 
11 ### subplot ###################################################################
12 
13 
14 class Subplot:
15 
16     def __init__(self, ax, energies, number_of_interactions=4, title='',
17                  title_font_size=16, state_names=None, virtual=[None],
18                  state_font_size=8, state_text_buffer=0.5, label_side='left'):
19         '''
20         Subplot.
21 
22         Parameters
23         ----------
24         ax : matplotlib axis
25             The axis.
26         energies : 1D array-like
27             Energies (scaled between 0 and 1)
28         number_of_interactions : integer
29             Number of interactions in diagram.
30         title : string (optional)
31             Title of subplot. Default is empty string.
32         state_names: list of str (optional)
33             list of the names of the states
34         virtual: list of ints (optional)
35             list of indexes of any vitual energy states
36         state_font_size: numtype (optional)
37             font size for the state lables
38         state_text_buffer: numtype (optional)
39             space between the energy level bars and the state labels
40         '''
41         self.ax = ax
42         self.energies = energies
43         self.interactions = number_of_interactions
44         self.state_names = state_names
45 
46         # Plot Energy Levels
47         for i in range(len(self.energies)):
48             if i in virtual:
49                 linestyle = '--'
50             else:
51                 linestyle = '-'
52             self.ax.axhline(self.energies[i], color='k', linewidth=2,
53                             ls=linestyle, zorder=5)
54 
55         # add state names
56         if isinstance(state_names, list):
57             for i in range(len(self.energies)):
58                 if label_side == 'left':
59                     ax.text(-state_text_buffer, energies[i], state_names[i],
60                             fontsize=state_font_size, verticalalignment='center',
61                             horizontalalignment='center')
62                 elif label_side == 'right':
63                     ax.text(1 + state_text_buffer, energies[i], state_names[i],
64                             fontsize=state_font_size, verticalalignment='center',
65                             horizontalalignment='center')
66         # calculate interaction_positons
67         self.x_pos = np.linspace(0, 1, number_of_interactions)
68         # set limits
69         self.ax.set_xlim(-0.1, 1.1)
70         self.ax.set_ylim(-0.01, 1.01)
71         # remove guff
72         self.ax.axis('off')
73         # title
74         self.ax.set_title(title, fontsize=title_font_size)
75 
76     def add_arrow(self, index, between, kind, label='', head_length=0.075,
77                   head_aspect=2, font_size=7, color='k'):
78         '''
79         Add an arrow to the WMEL diagram.
80 
81         Parameters
82         ----------
83         index : integer
84             The interaction, or start and stop interaction for the arrow.
85         between : 2-element iterable of integers
86             The inital and final state of the arrow
87         kind : {'ket', 'bra'}
88             The kind of interaction.
89         label : string (optional)
90             Interaction label. Default is empty string.
91         head_length: number (optional)
92             size of arrow head
93         font_size : number (optional)
94             Label font size. Default is 14.
95         color : matplotlib color (optional)
96             Arrow color. Default is black.
97 
98         Returns
99         -------
100         [line,arrow_head,text]
101         '''
102         if hasattr(index, 'index'):
103             x_pos = list(index)
104         else:
105             x_pos = [index] * 2
106         x_pos = [np.linspace(0, 1, self.interactions)[i] for i in x_pos]
107 
108         # calculate arrow length
109         arrow_length = self.energies[between[1]] - self.energies[between[0]]
110         arrow_end = self.energies[between[1]]
111         if arrow_length > 0:
112             direction = 1
113             y_poss = [self.energies[between[0]], self.energies[between[1]] - head_length]
114         elif arrow_length < 0:
115             direction = -1
116             y_poss = [self.energies[between[0]], self.energies[between[1]] + head_length]
117         else:
118             raise ValueError('between invalid!')
119 
120         length = abs(y_poss[0] - y_poss[1])
121         if kind == 'ket':
122             line = self.ax.plot(x_pos, y_poss, linestyle='-', color=color,
123                                 linewidth=2, zorder=9)
124         elif kind == 'bra':
125             line = self.ax.plot(x_pos, y_poss, linestyle='--', color=color,
126                                 linewidth=2, zorder=9)
127         elif kind == 'out':
128             yi = np.linspace(y_poss[0], y_poss[1], 100)
129             xi = np.sin((yi - y_poss[0]) * int((1 / length) * 20)
130                         * 2 * np.pi * length) / 40 + x_pos[0]
131             line = self.ax.plot(xi, yi, linestyle='-', color=color,
132                                 linewidth=2, solid_capstyle='butt', zorder=9)
133         else:
134             raise ValueError('kind is not \'ket\', \'out\', or \'bra\'.')
135         # add arrow head
136         arrow_head = self.ax.arrow(x_pos[1], arrow_end - head_length * direction,
137                                    0, 0.0001 * direction,
138                                    head_width=head_length * head_aspect,
139                                    head_length=head_length,
140                                    fc=color, ec=color, linestyle='solid',
141                                    linewidth=0, zorder=10)
142         # add text
143         text = self.ax.text(np.mean(x_pos), -0.1, label, fontsize=font_size,
144                             horizontalalignment='center')
145         return line, arrow_head, text
146 
147 
148 ### artist ####################################################################
149 
150 
151 class Artist:
152 
153     def __init__(self, size, energies, state_names=None,
154                  number_of_interactions=4, virtual=[None],
155                  state_font_size=8, state_text_buffer=0.5):
156         '''
157         virtual a list of indicies
158         '''
159         # create figure
160         figsize = [int(size[0] * ((number_of_interactions + 1.) / 6.)), size[1] * 2.5]
161         fig, (subplots) = plt.subplots(size[1], size[0], figsize=figsize)
162         self.fig = fig
163         # wrap subplots if need be
164         if size == [1, 1]:
165             self.subplots = np.array([[subplots]])
166             plt.subplots_adjust(left=0.3)
167         elif size[1] == 1:
168             self.subplots = np.array([subplots])
169         else:
170             self.subplots = subplots
171         # add energy levels
172         self.energies = energies
173         for plot in self.subplots.flatten():
174             for i in range(len(self.energies)):
175                 if i in virtual:
176                     linestyle = '--'
177                 else:
178                     linestyle = '-'
179                 plot.axhline(energies[i], color='k', linewidth=2, linestyle=linestyle)
180         # add state names to leftmost plots
181         if state_names:
182             for i in range(size[1]):
183                 plot = self.subplots[i][0]
184                 for i in range(len(self.energies)):
185                     plot.text(-state_text_buffer, energies[i], state_names[i], fontsize=state_font_size,
186                               verticalalignment='center', horizontalalignment='center')
187         # calculate interaction_positons
188         self.x_pos = np.linspace(0, 1, number_of_interactions)
189         # plot cleans up a bunch - call it now as well as later
190         self.plot()
191 
192     def label_rows(self, labels, font_size=15, text_buffer=1.5):
193         for i in range(len(self.subplots)):
194             plot = self.subplots[i][-1]
195             plot.text(text_buffer, 0.5, labels[i], fontsize=font_size,
196                       verticalalignment='center', horizontalalignment='center')
197 
198     def label_columns(self, labels, font_size=15, text_buffer=1.15):
199         for i in range(len(labels)):
200             plot = self.subplots[0][i]
201             plot.text(0.5, text_buffer, labels[i], fontsize=font_size,
202                       verticalalignment='center', horizontalalignment='center')
203 
204     def clear_diagram(self, diagram):
205         plot = self.subplots[diagram[1]][diagram[0]]
206         plot.cla()
207 
208     def add_arrow(self, diagram, number, between, kind, label='',
209                   head_length=0.075, font_size=7, color='k'):
210         '''
211         kind one in [ket, bra, out] \n
212         returns [line, arrow_head, text]
213         '''
214         column, row = diagram
215         x_pos = self.x_pos[number]
216         # calculate arrow length
217         arrow_length = self.energies[between[1]] - self.energies[between[0]]
218         arrow_end = self.energies[between[1]]
219         if arrow_length > 0:
220             direction = 1
221             y_poss = [self.energies[between[0]], self.energies[between[1]] - head_length]
222         elif arrow_length < 0:
223             direction = -1
224             y_poss = [self.energies[between[0]], self.energies[between[1]] + head_length]
225         else:
226             print('between invalid!')
227             return
228         subplot = self.subplots[row][column]
229         # add line
230         length = abs(y_poss[0] - y_poss[1])
231         if kind == 'ket':
232             line = subplot.plot([x_pos, x_pos], y_poss, linestyle='-', color=color, linewidth=2)
233         elif kind == 'bra':
234             line = subplot.plot([x_pos, x_pos], y_poss, linestyle='--', color=color, linewidth=2)
235         elif kind == 'out':
236             yi = np.linspace(y_poss[0], y_poss[1], 100)
237             xi = np.sin((yi - y_poss[0]) * int((1 / length) * 20)
238                         * 2 * np.pi * length) / 40 + x_pos
239             line = subplot.plot(xi, yi, linestyle='-', color=color,
240                                 linewidth=2, solid_capstyle='butt')
241         # add arrow head
242         arrow_head = subplot.arrow(self.x_pos[number], arrow_end - head_length * direction,
243                                    0, 0.0001 * direction,
244                                    head_width=head_length * 2,
245                                    head_length=head_length,
246                                    fc=color, ec=color, linestyle='solid', linewidth=0)
247         # add text
248         text = subplot.text(self.x_pos[number], -0.1, label,
249                             fontsize=font_size, horizontalalignment='center')
250         return line, arrow_head, text
251 
252     def plot(self, save_path=None, close=False, bbox_inches='tight', pad_inches=1):
253         # final manipulations
254         for plot in self.subplots.flatten():
255             # set limits
256             plot.set_xlim(-0.1, 1.1)
257             plot.set_ylim(-0.1, 1.1)
258             # remove guff
259             plot.axis('off')
260         # save
261         if save_path:
262             plt.savefig(save_path, transparent=True, dpi=300,
263                         bbox_inches=bbox_inches, pad_inches=pad_inches)
264         # close
265         if close:
266             plt.close()
267 
268 
269 ### testing ###################################################################
270 
271 
272 if __name__ == '__main__':
273     # testing code
274 
275     plt.close('all')
276 
277     diagram = Artist(size=[6, 3],
278                      energies=[0., 0.4, 0.6, 1.],
279                      state_names=['g', 'a', 'b', 'a+b'])
280 
281     diagram.label_rows([r'$\mathrm{\alpha}$', r'$\mathrm{\beta}$', r'$\mathrm{\gamma}$'])
282     diagram.label_columns(['I', 'II', 'III', 'IV', 'V', 'VI'])
283 
284     # pw1 alpha
285     diagram.add_arrow([0, 0], 0, [0, 2], 'ket', '1')
286     diagram.add_arrow([0, 0], 1, [0, 1], 'bra', '-2')
287     diagram.add_arrow([0, 0], 2, [1, 0], 'bra', '2\'')
288     diagram.add_arrow([0, 0], 3, [3, 0], 'out')
289 
290     diagram.add_arrow([1, 0], 3, [2, 0], 'out')
291 
292     diagram.clear_diagram([2, 1])
293 
294     diagram.plot('WMEL_out.png')
295
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/WrightTools/diagrams/WMEL.py b/WrightTools/diagrams/WMEL.py
--- a/WrightTools/diagrams/WMEL.py
+++ b/WrightTools/diagrams/WMEL.py
@@ -15,9 +15,8 @@
 
     def __init__(self, ax, energies, number_of_interactions=4, title='',
                  title_font_size=16, state_names=None, virtual=[None],
-                 state_font_size=8, state_text_buffer=0.5, label_side='left'):
-        '''
-        Subplot.
+                 state_font_size=14, state_text_buffer=0.5, label_side='left'):
+        """ Subplot.
 
         Parameters
         ----------
@@ -37,7 +36,7 @@
             font size for the state lables
         state_text_buffer: numtype (optional)
             space between the energy level bars and the state labels
-        '''
+        """
         self.ax = ax
         self.energies = energies
         self.interactions = number_of_interactions
@@ -73,10 +72,9 @@
         # title
         self.ax.set_title(title, fontsize=title_font_size)
 
-    def add_arrow(self, index, between, kind, label='', head_length=0.075,
-                  head_aspect=2, font_size=7, color='k'):
-        '''
-        Add an arrow to the WMEL diagram.
+    def add_arrow(self, index, between, kind, label='', head_length=0.1,
+                  head_aspect=2, font_size=14, color='k'):
+        """ Add an arrow to the WMEL diagram.
 
         Parameters
         ----------
@@ -98,7 +96,7 @@
         Returns
         -------
         [line,arrow_head,text]
-        '''
+        """
         if hasattr(index, 'index'):
             x_pos = list(index)
         else:
@@ -140,7 +138,7 @@
                                    fc=color, ec=color, linestyle='solid',
                                    linewidth=0, zorder=10)
         # add text
-        text = self.ax.text(np.mean(x_pos), -0.1, label, fontsize=font_size,
+        text = self.ax.text(np.mean(x_pos), -0.15, label, fontsize=font_size,
                             horizontalalignment='center')
         return line, arrow_head, text
 
@@ -153,9 +151,7 @@
     def __init__(self, size, energies, state_names=None,
                  number_of_interactions=4, virtual=[None],
                  state_font_size=8, state_text_buffer=0.5):
-        '''
-        virtual a list of indicies
-        '''
+        """ virtual a list of indicies """
         # create figure
         figsize = [int(size[0] * ((number_of_interactions + 1.) / 6.)), size[1] * 2.5]
         fig, (subplots) = plt.subplots(size[1], size[0], figsize=figsize)
@@ -207,10 +203,10 @@
 
     def add_arrow(self, diagram, number, between, kind, label='',
                   head_length=0.075, font_size=7, color='k'):
-        '''
+        """
         kind one in [ket, bra, out] \n
         returns [line, arrow_head, text]
-        '''
+        """
         column, row = diagram
         x_pos = self.x_pos[number]
         # calculate arrow length

{"golden_diff": "diff --git a/WrightTools/diagrams/WMEL.py b/WrightTools/diagrams/WMEL.py\n--- a/WrightTools/diagrams/WMEL.py\n+++ b/WrightTools/diagrams/WMEL.py\n@@ -15,9 +15,8 @@\n \n     def __init__(self, ax, energies, number_of_interactions=4, title='',\n                  title_font_size=16, state_names=None, virtual=[None],\n-                 state_font_size=8, state_text_buffer=0.5, label_side='left'):\n-        '''\n-        Subplot.\n+                 state_font_size=14, state_text_buffer=0.5, label_side='left'):\n+        \"\"\" Subplot.\n \n         Parameters\n         ----------\n@@ -37,7 +36,7 @@\n             font size for the state lables\n         state_text_buffer: numtype (optional)\n             space between the energy level bars and the state labels\n-        '''\n+        \"\"\"\n         self.ax = ax\n         self.energies = energies\n         self.interactions = number_of_interactions\n@@ -73,10 +72,9 @@\n         # title\n         self.ax.set_title(title, fontsize=title_font_size)\n \n-    def add_arrow(self, index, between, kind, label='', head_length=0.075,\n-                  head_aspect=2, font_size=7, color='k'):\n-        '''\n-        Add an arrow to the WMEL diagram.\n+    def add_arrow(self, index, between, kind, label='', head_length=0.1,\n+                  head_aspect=2, font_size=14, color='k'):\n+        \"\"\" Add an arrow to the WMEL diagram.\n \n         Parameters\n         ----------\n@@ -98,7 +96,7 @@\n         Returns\n         -------\n         [line,arrow_head,text]\n-        '''\n+        \"\"\"\n         if hasattr(index, 'index'):\n             x_pos = list(index)\n         else:\n@@ -140,7 +138,7 @@\n                                    fc=color, ec=color, linestyle='solid',\n                                    linewidth=0, zorder=10)\n         # add text\n-        text = self.ax.text(np.mean(x_pos), -0.1, label, fontsize=font_size,\n+        text = self.ax.text(np.mean(x_pos), -0.15, label, fontsize=font_size,\n                             horizontalalignment='center')\n         return line, arrow_head, text\n \n@@ -153,9 +151,7 @@\n     def __init__(self, size, energies, state_names=None,\n                  number_of_interactions=4, virtual=[None],\n                  state_font_size=8, state_text_buffer=0.5):\n-        '''\n-        virtual a list of indicies\n-        '''\n+        \"\"\" virtual a list of indicies \"\"\"\n         # create figure\n         figsize = [int(size[0] * ((number_of_interactions + 1.) / 6.)), size[1] * 2.5]\n         fig, (subplots) = plt.subplots(size[1], size[0], figsize=figsize)\n@@ -207,10 +203,10 @@\n \n     def add_arrow(self, diagram, number, between, kind, label='',\n                   head_length=0.075, font_size=7, color='k'):\n-        '''\n+        \"\"\"\n         kind one in [ket, bra, out] \\n\n         returns [line, arrow_head, text]\n-        '''\n+        \"\"\"\n         column, row = diagram\n         x_pos = self.x_pos[number]\n         # calculate arrow length\n", "issue": "better WMEL defaults\n\n", "before_files": [{"content": "### import ####################################################################\n\n\nimport numpy as np\nimport matplotlib.pyplot as plt\n\n\n### define ####################################################################\n\n\n### subplot ###################################################################\n\n\nclass Subplot:\n\n    def __init__(self, ax, energies, number_of_interactions=4, title='',\n                 title_font_size=16, state_names=None, virtual=[None],\n                 state_font_size=8, state_text_buffer=0.5, label_side='left'):\n        '''\n        Subplot.\n\n        Parameters\n        ----------\n        ax : matplotlib axis\n            The axis.\n        energies : 1D array-like\n            Energies (scaled between 0 and 1)\n        number_of_interactions : integer\n            Number of interactions in diagram.\n        title : string (optional)\n            Title of subplot. Default is empty string.\n        state_names: list of str (optional)\n            list of the names of the states\n        virtual: list of ints (optional)\n            list of indexes of any vitual energy states\n        state_font_size: numtype (optional)\n            font size for the state lables\n        state_text_buffer: numtype (optional)\n            space between the energy level bars and the state labels\n        '''\n        self.ax = ax\n        self.energies = energies\n        self.interactions = number_of_interactions\n        self.state_names = state_names\n\n        # Plot Energy Levels\n        for i in range(len(self.energies)):\n            if i in virtual:\n                linestyle = '--'\n            else:\n                linestyle = '-'\n            self.ax.axhline(self.energies[i], color='k', linewidth=2,\n                            ls=linestyle, zorder=5)\n\n        # add state names\n        if isinstance(state_names, list):\n            for i in range(len(self.energies)):\n                if label_side == 'left':\n                    ax.text(-state_text_buffer, energies[i], state_names[i],\n                            fontsize=state_font_size, verticalalignment='center',\n                            horizontalalignment='center')\n                elif label_side == 'right':\n                    ax.text(1 + state_text_buffer, energies[i], state_names[i],\n                            fontsize=state_font_size, verticalalignment='center',\n                            horizontalalignment='center')\n        # calculate interaction_positons\n        self.x_pos = np.linspace(0, 1, number_of_interactions)\n        # set limits\n        self.ax.set_xlim(-0.1, 1.1)\n        self.ax.set_ylim(-0.01, 1.01)\n        # remove guff\n        self.ax.axis('off')\n        # title\n        self.ax.set_title(title, fontsize=title_font_size)\n\n    def add_arrow(self, index, between, kind, label='', head_length=0.075,\n                  head_aspect=2, font_size=7, color='k'):\n        '''\n        Add an arrow to the WMEL diagram.\n\n        Parameters\n        ----------\n        index : integer\n            The interaction, or start and stop interaction for the arrow.\n        between : 2-element iterable of integers\n            The inital and final state of the arrow\n        kind : {'ket', 'bra'}\n            The kind of interaction.\n        label : string (optional)\n            Interaction label. Default is empty string.\n        head_length: number (optional)\n            size of arrow head\n        font_size : number (optional)\n            Label font size. Default is 14.\n        color : matplotlib color (optional)\n            Arrow color. Default is black.\n\n        Returns\n        -------\n        [line,arrow_head,text]\n        '''\n        if hasattr(index, 'index'):\n            x_pos = list(index)\n        else:\n            x_pos = [index] * 2\n        x_pos = [np.linspace(0, 1, self.interactions)[i] for i in x_pos]\n\n        # calculate arrow length\n        arrow_length = self.energies[between[1]] - self.energies[between[0]]\n        arrow_end = self.energies[between[1]]\n        if arrow_length > 0:\n            direction = 1\n            y_poss = [self.energies[between[0]], self.energies[between[1]] - head_length]\n        elif arrow_length < 0:\n            direction = -1\n            y_poss = [self.energies[between[0]], self.energies[between[1]] + head_length]\n        else:\n            raise ValueError('between invalid!')\n\n        length = abs(y_poss[0] - y_poss[1])\n        if kind == 'ket':\n            line = self.ax.plot(x_pos, y_poss, linestyle='-', color=color,\n                                linewidth=2, zorder=9)\n        elif kind == 'bra':\n            line = self.ax.plot(x_pos, y_poss, linestyle='--', color=color,\n                                linewidth=2, zorder=9)\n        elif kind == 'out':\n            yi = np.linspace(y_poss[0], y_poss[1], 100)\n            xi = np.sin((yi - y_poss[0]) * int((1 / length) * 20)\n                        * 2 * np.pi * length) / 40 + x_pos[0]\n            line = self.ax.plot(xi, yi, linestyle='-', color=color,\n                                linewidth=2, solid_capstyle='butt', zorder=9)\n        else:\n            raise ValueError('kind is not \\'ket\\', \\'out\\', or \\'bra\\'.')\n        # add arrow head\n        arrow_head = self.ax.arrow(x_pos[1], arrow_end - head_length * direction,\n                                   0, 0.0001 * direction,\n                                   head_width=head_length * head_aspect,\n                                   head_length=head_length,\n                                   fc=color, ec=color, linestyle='solid',\n                                   linewidth=0, zorder=10)\n        # add text\n        text = self.ax.text(np.mean(x_pos), -0.1, label, fontsize=font_size,\n                            horizontalalignment='center')\n        return line, arrow_head, text\n\n\n### artist ####################################################################\n\n\nclass Artist:\n\n    def __init__(self, size, energies, state_names=None,\n                 number_of_interactions=4, virtual=[None],\n                 state_font_size=8, state_text_buffer=0.5):\n        '''\n        virtual a list of indicies\n        '''\n        # create figure\n        figsize = [int(size[0] * ((number_of_interactions + 1.) / 6.)), size[1] * 2.5]\n        fig, (subplots) = plt.subplots(size[1], size[0], figsize=figsize)\n        self.fig = fig\n        # wrap subplots if need be\n        if size == [1, 1]:\n            self.subplots = np.array([[subplots]])\n            plt.subplots_adjust(left=0.3)\n        elif size[1] == 1:\n            self.subplots = np.array([subplots])\n        else:\n            self.subplots = subplots\n        # add energy levels\n        self.energies = energies\n        for plot in self.subplots.flatten():\n            for i in range(len(self.energies)):\n                if i in virtual:\n                    linestyle = '--'\n                else:\n                    linestyle = '-'\n                plot.axhline(energies[i], color='k', linewidth=2, linestyle=linestyle)\n        # add state names to leftmost plots\n        if state_names:\n            for i in range(size[1]):\n                plot = self.subplots[i][0]\n                for i in range(len(self.energies)):\n                    plot.text(-state_text_buffer, energies[i], state_names[i], fontsize=state_font_size,\n                              verticalalignment='center', horizontalalignment='center')\n        # calculate interaction_positons\n        self.x_pos = np.linspace(0, 1, number_of_interactions)\n        # plot cleans up a bunch - call it now as well as later\n        self.plot()\n\n    def label_rows(self, labels, font_size=15, text_buffer=1.5):\n        for i in range(len(self.subplots)):\n            plot = self.subplots[i][-1]\n            plot.text(text_buffer, 0.5, labels[i], fontsize=font_size,\n                      verticalalignment='center', horizontalalignment='center')\n\n    def label_columns(self, labels, font_size=15, text_buffer=1.15):\n        for i in range(len(labels)):\n            plot = self.subplots[0][i]\n            plot.text(0.5, text_buffer, labels[i], fontsize=font_size,\n                      verticalalignment='center', horizontalalignment='center')\n\n    def clear_diagram(self, diagram):\n        plot = self.subplots[diagram[1]][diagram[0]]\n        plot.cla()\n\n    def add_arrow(self, diagram, number, between, kind, label='',\n                  head_length=0.075, font_size=7, color='k'):\n        '''\n        kind one in [ket, bra, out] \\n\n        returns [line, arrow_head, text]\n        '''\n        column, row = diagram\n        x_pos = self.x_pos[number]\n        # calculate arrow length\n        arrow_length = self.energies[between[1]] - self.energies[between[0]]\n        arrow_end = self.energies[between[1]]\n        if arrow_length > 0:\n            direction = 1\n            y_poss = [self.energies[between[0]], self.energies[between[1]] - head_length]\n        elif arrow_length < 0:\n            direction = -1\n            y_poss = [self.energies[between[0]], self.energies[between[1]] + head_length]\n        else:\n            print('between invalid!')\n            return\n        subplot = self.subplots[row][column]\n        # add line\n        length = abs(y_poss[0] - y_poss[1])\n        if kind == 'ket':\n            line = subplot.plot([x_pos, x_pos], y_poss, linestyle='-', color=color, linewidth=2)\n        elif kind == 'bra':\n            line = subplot.plot([x_pos, x_pos], y_poss, linestyle='--', color=color, linewidth=2)\n        elif kind == 'out':\n            yi = np.linspace(y_poss[0], y_poss[1], 100)\n            xi = np.sin((yi - y_poss[0]) * int((1 / length) * 20)\n                        * 2 * np.pi * length) / 40 + x_pos\n            line = subplot.plot(xi, yi, linestyle='-', color=color,\n                                linewidth=2, solid_capstyle='butt')\n        # add arrow head\n        arrow_head = subplot.arrow(self.x_pos[number], arrow_end - head_length * direction,\n                                   0, 0.0001 * direction,\n                                   head_width=head_length * 2,\n                                   head_length=head_length,\n                                   fc=color, ec=color, linestyle='solid', linewidth=0)\n        # add text\n        text = subplot.text(self.x_pos[number], -0.1, label,\n                            fontsize=font_size, horizontalalignment='center')\n        return line, arrow_head, text\n\n    def plot(self, save_path=None, close=False, bbox_inches='tight', pad_inches=1):\n        # final manipulations\n        for plot in self.subplots.flatten():\n            # set limits\n            plot.set_xlim(-0.1, 1.1)\n            plot.set_ylim(-0.1, 1.1)\n            # remove guff\n            plot.axis('off')\n        # save\n        if save_path:\n            plt.savefig(save_path, transparent=True, dpi=300,\n                        bbox_inches=bbox_inches, pad_inches=pad_inches)\n        # close\n        if close:\n            plt.close()\n\n\n### testing ###################################################################\n\n\nif __name__ == '__main__':\n    # testing code\n\n    plt.close('all')\n\n    diagram = Artist(size=[6, 3],\n                     energies=[0., 0.4, 0.6, 1.],\n                     state_names=['g', 'a', 'b', 'a+b'])\n\n    diagram.label_rows([r'$\\mathrm{\\alpha}$', r'$\\mathrm{\\beta}$', r'$\\mathrm{\\gamma}$'])\n    diagram.label_columns(['I', 'II', 'III', 'IV', 'V', 'VI'])\n\n    # pw1 alpha\n    diagram.add_arrow([0, 0], 0, [0, 2], 'ket', '1')\n    diagram.add_arrow([0, 0], 1, [0, 1], 'bra', '-2')\n    diagram.add_arrow([0, 0], 2, [1, 0], 'bra', '2\\'')\n    diagram.add_arrow([0, 0], 3, [3, 0], 'out')\n\n    diagram.add_arrow([1, 0], 3, [2, 0], 'out')\n\n    diagram.clear_diagram([2, 1])\n\n    diagram.plot('WMEL_out.png')\n", "path": "WrightTools/diagrams/WMEL.py"}], "after_files": [{"content": "### import ####################################################################\n\n\nimport numpy as np\nimport matplotlib.pyplot as plt\n\n\n### define ####################################################################\n\n\n### subplot ###################################################################\n\n\nclass Subplot:\n\n    def __init__(self, ax, energies, number_of_interactions=4, title='',\n                 title_font_size=16, state_names=None, virtual=[None],\n                 state_font_size=14, state_text_buffer=0.5, label_side='left'):\n        \"\"\" Subplot.\n\n        Parameters\n        ----------\n        ax : matplotlib axis\n            The axis.\n        energies : 1D array-like\n            Energies (scaled between 0 and 1)\n        number_of_interactions : integer\n            Number of interactions in diagram.\n        title : string (optional)\n            Title of subplot. Default is empty string.\n        state_names: list of str (optional)\n            list of the names of the states\n        virtual: list of ints (optional)\n            list of indexes of any vitual energy states\n        state_font_size: numtype (optional)\n            font size for the state lables\n        state_text_buffer: numtype (optional)\n            space between the energy level bars and the state labels\n        \"\"\"\n        self.ax = ax\n        self.energies = energies\n        self.interactions = number_of_interactions\n        self.state_names = state_names\n\n        # Plot Energy Levels\n        for i in range(len(self.energies)):\n            if i in virtual:\n                linestyle = '--'\n            else:\n                linestyle = '-'\n            self.ax.axhline(self.energies[i], color='k', linewidth=2,\n                            ls=linestyle, zorder=5)\n\n        # add state names\n        if isinstance(state_names, list):\n            for i in range(len(self.energies)):\n                if label_side == 'left':\n                    ax.text(-state_text_buffer, energies[i], state_names[i],\n                            fontsize=state_font_size, verticalalignment='center',\n                            horizontalalignment='center')\n                elif label_side == 'right':\n                    ax.text(1 + state_text_buffer, energies[i], state_names[i],\n                            fontsize=state_font_size, verticalalignment='center',\n                            horizontalalignment='center')\n        # calculate interaction_positons\n        self.x_pos = np.linspace(0, 1, number_of_interactions)\n        # set limits\n        self.ax.set_xlim(-0.1, 1.1)\n        self.ax.set_ylim(-0.01, 1.01)\n        # remove guff\n        self.ax.axis('off')\n        # title\n        self.ax.set_title(title, fontsize=title_font_size)\n\n    def add_arrow(self, index, between, kind, label='', head_length=0.1,\n                  head_aspect=2, font_size=14, color='k'):\n        \"\"\" Add an arrow to the WMEL diagram.\n\n        Parameters\n        ----------\n        index : integer\n            The interaction, or start and stop interaction for the arrow.\n        between : 2-element iterable of integers\n            The inital and final state of the arrow\n        kind : {'ket', 'bra'}\n            The kind of interaction.\n        label : string (optional)\n            Interaction label. Default is empty string.\n        head_length: number (optional)\n            size of arrow head\n        font_size : number (optional)\n            Label font size. Default is 14.\n        color : matplotlib color (optional)\n            Arrow color. Default is black.\n\n        Returns\n        -------\n        [line,arrow_head,text]\n        \"\"\"\n        if hasattr(index, 'index'):\n            x_pos = list(index)\n        else:\n            x_pos = [index] * 2\n        x_pos = [np.linspace(0, 1, self.interactions)[i] for i in x_pos]\n\n        # calculate arrow length\n        arrow_length = self.energies[between[1]] - self.energies[between[0]]\n        arrow_end = self.energies[between[1]]\n        if arrow_length > 0:\n            direction = 1\n            y_poss = [self.energies[between[0]], self.energies[between[1]] - head_length]\n        elif arrow_length < 0:\n            direction = -1\n            y_poss = [self.energies[between[0]], self.energies[between[1]] + head_length]\n        else:\n            raise ValueError('between invalid!')\n\n        length = abs(y_poss[0] - y_poss[1])\n        if kind == 'ket':\n            line = self.ax.plot(x_pos, y_poss, linestyle='-', color=color,\n                                linewidth=2, zorder=9)\n        elif kind == 'bra':\n            line = self.ax.plot(x_pos, y_poss, linestyle='--', color=color,\n                                linewidth=2, zorder=9)\n        elif kind == 'out':\n            yi = np.linspace(y_poss[0], y_poss[1], 100)\n            xi = np.sin((yi - y_poss[0]) * int((1 / length) * 20)\n                        * 2 * np.pi * length) / 40 + x_pos[0]\n            line = self.ax.plot(xi, yi, linestyle='-', color=color,\n                                linewidth=2, solid_capstyle='butt', zorder=9)\n        else:\n            raise ValueError('kind is not \\'ket\\', \\'out\\', or \\'bra\\'.')\n        # add arrow head\n        arrow_head = self.ax.arrow(x_pos[1], arrow_end - head_length * direction,\n                                   0, 0.0001 * direction,\n                                   head_width=head_length * head_aspect,\n                                   head_length=head_length,\n                                   fc=color, ec=color, linestyle='solid',\n                                   linewidth=0, zorder=10)\n        # add text\n        text = self.ax.text(np.mean(x_pos), -0.15, label, fontsize=font_size,\n                            horizontalalignment='center')\n        return line, arrow_head, text\n\n\n### artist ####################################################################\n\n\nclass Artist:\n\n    def __init__(self, size, energies, state_names=None,\n                 number_of_interactions=4, virtual=[None],\n                 state_font_size=8, state_text_buffer=0.5):\n        \"\"\" virtual a list of indicies \"\"\"\n        # create figure\n        figsize = [int(size[0] * ((number_of_interactions + 1.) / 6.)), size[1] * 2.5]\n        fig, (subplots) = plt.subplots(size[1], size[0], figsize=figsize)\n        self.fig = fig\n        # wrap subplots if need be\n        if size == [1, 1]:\n            self.subplots = np.array([[subplots]])\n            plt.subplots_adjust(left=0.3)\n        elif size[1] == 1:\n            self.subplots = np.array([subplots])\n        else:\n            self.subplots = subplots\n        # add energy levels\n        self.energies = energies\n        for plot in self.subplots.flatten():\n            for i in range(len(self.energies)):\n                if i in virtual:\n                    linestyle = '--'\n                else:\n                    linestyle = '-'\n                plot.axhline(energies[i], color='k', linewidth=2, linestyle=linestyle)\n        # add state names to leftmost plots\n        if state_names:\n            for i in range(size[1]):\n                plot = self.subplots[i][0]\n                for i in range(len(self.energies)):\n                    plot.text(-state_text_buffer, energies[i], state_names[i], fontsize=state_font_size,\n                              verticalalignment='center', horizontalalignment='center')\n        # calculate interaction_positons\n        self.x_pos = np.linspace(0, 1, number_of_interactions)\n        # plot cleans up a bunch - call it now as well as later\n        self.plot()\n\n    def label_rows(self, labels, font_size=15, text_buffer=1.5):\n        for i in range(len(self.subplots)):\n            plot = self.subplots[i][-1]\n            plot.text(text_buffer, 0.5, labels[i], fontsize=font_size,\n                      verticalalignment='center', horizontalalignment='center')\n\n    def label_columns(self, labels, font_size=15, text_buffer=1.15):\n        for i in range(len(labels)):\n            plot = self.subplots[0][i]\n            plot.text(0.5, text_buffer, labels[i], fontsize=font_size,\n                      verticalalignment='center', horizontalalignment='center')\n\n    def clear_diagram(self, diagram):\n        plot = self.subplots[diagram[1]][diagram[0]]\n        plot.cla()\n\n    def add_arrow(self, diagram, number, between, kind, label='',\n                  head_length=0.075, font_size=7, color='k'):\n        \"\"\"\n        kind one in [ket, bra, out] \\n\n        returns [line, arrow_head, text]\n        \"\"\"\n        column, row = diagram\n        x_pos = self.x_pos[number]\n        # calculate arrow length\n        arrow_length = self.energies[between[1]] - self.energies[between[0]]\n        arrow_end = self.energies[between[1]]\n        if arrow_length > 0:\n            direction = 1\n            y_poss = [self.energies[between[0]], self.energies[between[1]] - head_length]\n        elif arrow_length < 0:\n            direction = -1\n            y_poss = [self.energies[between[0]], self.energies[between[1]] + head_length]\n        else:\n            print('between invalid!')\n            return\n        subplot = self.subplots[row][column]\n        # add line\n        length = abs(y_poss[0] - y_poss[1])\n        if kind == 'ket':\n            line = subplot.plot([x_pos, x_pos], y_poss, linestyle='-', color=color, linewidth=2)\n        elif kind == 'bra':\n            line = subplot.plot([x_pos, x_pos], y_poss, linestyle='--', color=color, linewidth=2)\n        elif kind == 'out':\n            yi = np.linspace(y_poss[0], y_poss[1], 100)\n            xi = np.sin((yi - y_poss[0]) * int((1 / length) * 20)\n                        * 2 * np.pi * length) / 40 + x_pos\n            line = subplot.plot(xi, yi, linestyle='-', color=color,\n                                linewidth=2, solid_capstyle='butt')\n        # add arrow head\n        arrow_head = subplot.arrow(self.x_pos[number], arrow_end - head_length * direction,\n                                   0, 0.0001 * direction,\n                                   head_width=head_length * 2,\n                                   head_length=head_length,\n                                   fc=color, ec=color, linestyle='solid', linewidth=0)\n        # add text\n        text = subplot.text(self.x_pos[number], -0.1, label,\n                            fontsize=font_size, horizontalalignment='center')\n        return line, arrow_head, text\n\n    def plot(self, save_path=None, close=False, bbox_inches='tight', pad_inches=1):\n        # final manipulations\n        for plot in self.subplots.flatten():\n            # set limits\n            plot.set_xlim(-0.1, 1.1)\n            plot.set_ylim(-0.1, 1.1)\n            # remove guff\n            plot.axis('off')\n        # save\n        if save_path:\n            plt.savefig(save_path, transparent=True, dpi=300,\n                        bbox_inches=bbox_inches, pad_inches=pad_inches)\n        # close\n        if close:\n            plt.close()\n\n\n### testing ###################################################################\n\n\nif __name__ == '__main__':\n    # testing code\n\n    plt.close('all')\n\n    diagram = Artist(size=[6, 3],\n                     energies=[0., 0.4, 0.6, 1.],\n                     state_names=['g', 'a', 'b', 'a+b'])\n\n    diagram.label_rows([r'$\\mathrm{\\alpha}$', r'$\\mathrm{\\beta}$', r'$\\mathrm{\\gamma}$'])\n    diagram.label_columns(['I', 'II', 'III', 'IV', 'V', 'VI'])\n\n    # pw1 alpha\n    diagram.add_arrow([0, 0], 0, [0, 2], 'ket', '1')\n    diagram.add_arrow([0, 0], 1, [0, 1], 'bra', '-2')\n    diagram.add_arrow([0, 0], 2, [1, 0], 'bra', '2\\'')\n    diagram.add_arrow([0, 0], 3, [3, 0], 'out')\n\n    diagram.add_arrow([1, 0], 3, [2, 0], 'out')\n\n    diagram.clear_diagram([2, 1])\n\n    diagram.plot('WMEL_out.png')\n", "path": "WrightTools/diagrams/WMEL.py"}]}