add gallery example comparing ioneq

examples/idl_comparisons/ioneq_comparison.py

@@ -25,38 +25,39 @@ def plot_idl_comparison(x, y_idl, y_python, fig, n_rows, i_row, quantity_name, t
     idx = np.where(y_idl>=y_idl.max()*thresh)
     x_thresh_min, x_thresh_max = x[idx[0][[0,-1]]]
     # Direct comparison
-    ax = fig.add_subplot(n_rows, 3, i_row+1)
-    ax.plot(x, y_python, label='fiasco')
-    ax.plot(x, y_idl, color='k', marker='o', ls='--', markevery=2, label='IDL')
-    ax.axvspan(x[0], x_thresh_min, color='r', alpha=0.25)
-    ax.axvspan(x_thresh_max, x[-1], color='r', alpha=0.25)
-    ax.set_xscale('log')
-    ax.set_yscale('log')
+    ax1 = fig.add_subplot(n_rows, 3, i_row+1)
+    ax1.plot(x, y_python, label='fiasco')
+    ax1.plot(x, y_idl, color='k', marker='o', ls='--', markevery=2, label='IDL')
+    ax1.axvspan(x[0], x_thresh_min, color='r', alpha=0.25)
+    ax1.axvspan(x_thresh_max, x[-1], color='r', alpha=0.25)
+    ax1.set_xscale('log')
+    ax1.set_yscale('log')
     if i_row == 0:
-        ax.set_title('Ionization Fraction')
-    ax.set_xlim(x[[0,-1]])
-    ax.set_ylim(y_python.max()*np.array([thresh,2]))
-    ax.set_ylabel(quantity_name)
-    plt.legend()
+        ax1.set_title('Ionization Fraction')
+    ax1.set_xlim(x[[0,-1]])
+    ax1.set_ylim(y_python.max()*np.array([thresh,2]))
+    ax1.set_ylabel(quantity_name)
+    ax1.legend()
     # Ratio
-    ax = fig.add_subplot(n_rows, 3, i_row+2, sharex=ax)
-    plt.plot(x, (y_python/y_idl).decompose())
-    ax.axvspan(x[0], x_thresh_min, color='r', alpha=0.25)
-    ax.axvspan(x_thresh_max, x[-1], color='r', alpha=0.25)
-    ax.axhline(y=1, color='k', ls=':')
+    ax2 = fig.add_subplot(n_rows, 3, i_row+2, sharex=ax1)
+    ax2.plot(x, (y_python/y_idl).decompose())
+    ax2.axvspan(x[0], x_thresh_min, color='r', alpha=0.25)
+    ax2.axvspan(x_thresh_max, x[-1], color='r', alpha=0.25)
+    ax2.axhline(y=1, color='k', ls=':')
     if i_row == 0:
-        ax.set_title('fiasco / IDL')
+        ax2.set_title('fiasco / IDL')
     diff_limit = .001
-    ax.set_ylim(1-diff_limit, 1+diff_limit)
+    ax2.set_ylim(1-diff_limit, 1+diff_limit)
     # Normalized difference
-    ax = fig.add_subplot(n_rows, 3, i_row+3, sharex=ax)
-    ax.plot(x,  ((y_python - y_idl)/y_idl).decompose())
-    ax.axvspan(x[0], x_thresh_min, color='r', alpha=0.25)
-    ax.axvspan(x_thresh_max, x[-1], color='r', alpha=0.25)
-    ax.axhline(y=0, color='k', ls=':')
+    ax3 = fig.add_subplot(n_rows, 3, i_row+3, sharex=ax1)
+    ax3.plot(x,  ((y_python - y_idl)/y_idl).decompose())
+    ax3.axvspan(x[0], x_thresh_min, color='r', alpha=0.25)
+    ax3.axvspan(x_thresh_max, x[-1], color='r', alpha=0.25)
+    ax3.axhline(y=0, color='k', ls=':')
     if i_row == 0:
-        ax.set_title('(fiasco - IDL)/IDL')
-    ax.set_ylim(-diff_limit, diff_limit)
+        ax3.set_title('(fiasco - IDL)/IDL')
+    ax3.set_ylim(-diff_limit, diff_limit)
+    return ax1, ax2, ax3
 
 
 ################################################
@@ -82,8 +83,12 @@ def plot_idl_comparison(x, y_idl, y_python, fig, n_rows, i_row, quantity_name, t
     ion = fiasco.Ion((idl_result['Z'], idl_result['iz']),
                      idl_result['temperature'],
                      ioneq_filename=idl_result['ioneq_filename'])
+    element = fiasco.Element(ion.atomic_symbol, ion.temperature)
+    ioneq = element.equilibrium_ionization
     print(f'IDL code to produce ioneq result for {ion.ion_name_roman}:')
     print(Template(idl_result['idl_script']).render(**idl_result))
-    plot_idl_comparison(ion.temperature, idl_result['ioneq'], ion.ioneq,
-                        fig, len(ioneq_files), 3*i, f'{ion.ion_name_roman}')
+    axes = plot_idl_comparison(ion.temperature, idl_result['ioneq'], ion.ioneq,
+                               fig, len(ioneq_files), 3*i, f'{ion.ion_name_roman}')
+    axes[0].plot(element.temperature, ioneq[:, ion.charge_state],
+                 label='fiasco, from rates', color='C0', ls='--')
 plt.show()