Robustness upgrades to pass omas_test_plot

gafusion · Dec 11, 2023 · 77bad43 · 77bad43
1 parent b627204
commit 77bad43
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/omas/omas_plot.py b/omas/omas_plot.py
@@ -762,7 +762,10 @@ def get2d(contour_quantity):
     if levels is None and value_1d is not None:
         if contour_quantity == 'q':
             max_q = int(numpy.round(omas_interp1d(0.95, x_value_1d, value_1d)))
-            levels = numpy.arange(max_q)
+            levels = numpy.arange(numpy.abs(max_q))
+            if max_q < 0:
+                levels *= -1
+                levels = levels[::-1]
         else:
             levels = numpy.linspace(numpy.min(value_1d), numpy.max(value_1d), 11)[1:-1]
             levels = numpy.hstack((levels, levels[-1] + (levels[1] - levels[0]) * numpy.arange(100)[1:]))
@@ -799,8 +802,10 @@ def get2d(contour_quantity):
             r = scipy.ndimage.zoom(r, sf)
             z = scipy.ndimage.zoom(z, sf)
             value_2d = scipy.ndimage.zoom(value_2d, sf)
-
-        cs = ax.contour(r, z, value_2d, levels, **kw)
+        if levels is not None:
+            cs = ax.contour(r, z, value_2d, levels, **kw)
+        else:
+            cs = ax.contour(r, z, value_2d, **kw)
 
         if label_contours or ((label_contours is None) and (contour_quantity == 'q')):
             ax.clabel(cs)
@@ -1280,7 +1285,8 @@ def core_profiles_summary(ods, time_index=None, time=None, fig=None,
                     #     plotting_list.append(prof1d[specie][q]*scale * prof1d[specie]['element[0].z_n'])
                     #     label_name_z.append(r'$\times$' + f" {int(prof1d[specie]['element[0].z_n'])}")
                     # else:
-                    if len(prof1d[specie][q]) == len(prof1d[specie][q + "_error_upper"]):
+
+                    if q + "_error_upper" in prof1d[specie] and len(prof1d[specie][q]) == len(prof1d[specie][q + "_error_upper"]):
                         plotting_list.append(unumpy.uarray(prof1d[specie][q]*scale,
                                              prof1d[specie][q + "_error_upper"]*scale))
                     else:
@@ -1292,6 +1298,8 @@ def core_profiles_summary(ods, time_index=None, time=None, fig=None,
                                               prof1d[specie][q + "_fit.measured_error_upper"]*scale])
                         except Exception as e:
                             data_list.append(None)
+                    else:
+                        data_list.append(None)
                     label_name_z.append("")
                     label_name.append(f'{names[index]} {q.capitalize()}')
         elif "e_field.radial" not in q: