omas_plot: New plot_equilibrium_summary_and_quality_and_core_profiles…

… in support of CAKE visualization

omas/machine_mappings/d3d.py

@@ -1505,19 +1505,22 @@ def core_profiles_profile_1d(ods, pulse, PROFILES_tree="OMFIT_PROFS"):
 @machine_mapping_function(__regression_arguments__, pulse=133221)
 def core_profiles_global_quantities_data(ods, pulse):
     from scipy.interpolate import interp1d
+    mpulse = pulse
+    if len(str(pulse))>8:
+        mpulse = int(str(pulse)[:6])
 
     ods1 = ODS()
-    unwrap(magnetics_hardware)(ods1, pulse)
+    unwrap(magnetics_hardware)(ods1, mpulse)
 
     with omas_environment(ods, cocosio=1):
         cp = ods['core_profiles']
         gq = ods['core_profiles.global_quantities']
         if 'time' not in cp:
-            m = mdsvalue('d3d', pulse=pulse, TDI="\\TOP.PROFILES.EDENSFIT", treename="ZIPFIT01")
+            m = mdsvalue('d3d', pulse=mpulse, TDI="\\TOP.PROFILES.EDENSFIT", treename="ZIPFIT01")
             cp['time'] = m.dim_of(1) * 1e-3
         t = cp['time']
 
-        m = mdsvalue('d3d', pulse=pulse, TDI=f"ptdata2(\"VLOOP\",{pulse})", treename=None)
+        m = mdsvalue('d3d', pulse=mpulse, TDI=f"ptdata2(\"VLOOP\",{mpulse})", treename=None)
 
         gq['v_loop'] = interp1d(m.dim_of(0) * 1e-3, m.data(), bounds_error=False, fill_value=np.nan)(t)
 

omas/omas_plot.py

@@ -836,7 +836,8 @@ def get2d(contour_quantity):
             except ValueError:
                 pass
             else:
-                ax.plot(xr, xz, **xkw)
+                if xr > 0:
+                    ax.plot(xr, xz, **xkw)
 
     # Internal flux surfaces w/ or w/o masking
     if wall is not None:
@@ -1138,6 +1139,232 @@ def equilibrium_summary(ods, time_index=None, time=None, fig=None, ggd_points_tr
                                     visible_x=True, **kw)
     return {'ax': axs}
 
+@add_to__ODS__
+def equilibrium_summary_and_quality_and_core_profiles(ods, time_index=None, time=None,  fig=None, ggd_points_triangles=None, 
+                          ods_species=[-1,1], **kw):
+    """
+    Plot equilibrium cross section, pressure profile and constraints, and convergence error
+
+    :param ods: input ods
+
+    :param time_index: int, list of ints, or None
+        time slice to plot.  If None all timeslices are plotted.
+
+    :param time: float, list of floats, or None
+        time to plot. If None all timeslicess are plotted.
+        if not None, it takes precedence over time_index
+    """
+    from matplotlib import pyplot
+    from omas.omas_physics import get_plot_scale_and_unit
+
+    # caching of ggd data
+    if ggd_points_triangles is None and 'equilibrium.grids_ggd' in ods:
+        from .omas_physics import grids_ggd_points_triangles
+
+        ggd_points_triangles = grids_ggd_points_triangles(ods['equilibrium.grids_ggd[0].grid[0]'])
+
+
+    axs = kw.pop('ax', {})
+    if axs is None:
+        axs = {}
+    if not len(axs) and fig is None:
+        fig = pyplot.figure()
+
+    # time animation
+    time_index, time = handle_time(ods, 'equilibrium', time_index, time)
+    if isinstance(time_index, (list, numpy.ndarray)):
+        if len(time) == 1:
+            time_index = time_index[0]
+        else:
+            return ods_time_plot(
+                equilibrium_summary_and_quality, ods, time_index, time, fig=fig, ggd_points_triangles=ggd_points_triangles, ax=axs, ods_species=ods_species, **kw
+            )
+    ax = cached_add_subplot(fig, axs, 1, 4, 1)
+    tmp = equilibrium_CX(
+        ods, time_index=time_index, ax=ax, contour_quantity='psi', ggd_points_triangles=ggd_points_triangles, **kw
+    )
+    overlay(ods, ax=ax, thomson_scattering=True)
+
+    # Work with equilibrium - inspired by equilibrium_summary
+    eq = ods['equilibrium']['time_slice'][time_index]
+    raw_xName = 'psi'
+    def psi2psin(psix):
+        return (psix-eq['global_quantities']['psi_axis'])/(eq['global_quantities']['psi_boundary']-eq['global_quantities']['psi_axis'])
+    x = psi2psin(eq['profiles_1d']['psi'])
+    xName = r"$\Psi_\mathrm{n}$"
+
+    # pressure
+    ax = cached_add_subplot(fig, axs, 2, 4, 8)
+    pfac = 1.e-3
+    ax.plot(x, eq['profiles_1d']['pressure'] * pfac)
+    ytitle = 0.9
+    ax.set_title('Pressure [kPa]', y=ytitle, va='top')
+    ax.set_xlabel(xName)
+    kw.setdefault('color', ax.lines[-1].get_color())
+    p_x = psi2psin(eq['constraints.pressure.:.position.psi'])
+    p = eq['constraints.pressure.:.measured'] * pfac
+    p_e = eq['constraints.pressure.:.measured_error_upper'] * pfac
+    ax.errorbar(p_x, p, yerr=p_e, color='red', marker='', alpha=0.25)
+
+    # Convergence error - inspired by equilibrium_quality
+    ax = cached_add_subplot(fig, axs, 2, 4, 4)
+    ax.plot(ods['equilibrium.time'], ods['equilibrium.time_slice[:].convergence.grad_shafranov_deviation_value'])
+    ax.axvline(eq['time'],color='black')
+    ax.set_yscale('log')
+    ax.set_title('Convergence error vs time', y=ytitle, va='top')
+
+    # Core profiles (Mostly copy paste from core_profiles_summary)
+    kw.pop('color',None)
+    ptime_index, ptime = handle_time(ods, 'core_profiles', time_index, eq['time'])
+    prof1d = ods['core_profiles']['profiles_1d'][ptime_index]
+    x_axis = "psi_norm"
+    if x_axis == "psi_norm":
+        x = prof1d['grid.rho_pol_norm']**2
+        x_label = r"$\Psi_\mathrm{n}$"
+    else:
+        x = prof1d[f'grid.{x_axis}']
+        if "tor" in x_axis:
+            x_label = r'$\rho$'
+        elif "pol" in x_axis:
+            x_label = r'$\rho_\mathrm{pol}$'
+    # Determine subplot rows x colls
+    if ods_species is None:
+        ncols = len(prof1d['ion']) + 1
+        ods_species = [-1] + list(prof1d['ion'])
+    else:
+        ncols = len(ods_species)
+    nplots = 0
+    quantities = ['density_thermal', 'temperature']
+    for quant in quantities:
+        if 'density' in quant or 'temperature' in quant:
+            nplots += ncols
+        elif 'velocity' in quant:
+            nplots += ncols - 1
+        else:
+            nplots += 1
+    nrows = 2# int(numpy.ceil(nplots / ncols))
+    ncols = ncols+2 # Make room for equilibrium plots
+    nplots = nrows * ncols
+
+    # Generate species with corresponding name
+    species_in_tree = [f"ion.{i}" if i >= 0 else 'electrons' for i in ods_species]
+    names = [f"{prof1d[i]['label']} ion" if i != 'electrons' else "electron" for i in species_in_tree]
+
+    plotting_list = []
+    label_name = []
+    label_name_z = []
+    unit_list = []
+    data_list = []
+    for q in quantities:
+        if 'density' in q or 'temperature' in q or "velocity.toroidal" in q :
+            for index, specie in enumerate(species_in_tree):
+                #unit_list.append(omas_info_node(o2u(f"core_profiles.profiles_1d.0.{specie}.{q}"))['units'])
+                if q in prof1d[specie]:
+                    if "label" in prof1d[specie]:
+                        scale, unit = get_plot_scale_and_unit(q, prof1d[specie]["label"])
+                    else:
+                        scale, unit = get_plot_scale_and_unit(q)
+                    unit_list.append(unit)
+                    try:
+                        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:
+                            plotting_list.append(prof1d[specie][q]*scale)
+                    except Exception as _excp:
+                        print(f'Error in fetching {q}_error_upper')
+                        plotting_list.append(prof1d[specie][q]*scale)
+
+                    if x_axis == "psi_norm":
+                        try:
+                            qfit = q.replace('_thermal','')
+                            data_list.append([prof1d[specie][f"{qfit}_fit.psi_norm"],
+                                              prof1d[specie][f"{qfit}_fit.measured"]*scale,
+                                              prof1d[specie][f"{qfit}_fit.measured_error_upper"]*scale])
+                        except Exception as e:
+                            if 'has no data' not in str(e):
+                                print(f'Exception getting raw data: {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:
+            unit_list.append(omas_info_node(o2u(f"core_profiles.profiles_1d.0.{q}"))['units'])
+            plotting_list.append(prof1d[q])
+            label_name.append(q.capitalize())
+            data_list.append(None)
+    if "e_field.radial" in quantities:
+        try:
+            scale, unit = get_plot_scale_and_unit("e_field.radial")
+            unit_list.append(unit)
+            plotting_list.append(prof1d["e_field.radial"]*scale)
+            label_name_z.append("")
+            label_name.append('e_field.radial')
+            data_list.append(None)
+        except:
+            pass
+    last_quant = None
+    for index, y in enumerate(plotting_list):
+        if index >= len(label_name):
+            break
+        plot = index + 2
+        if plot > ncols-1:
+            plot = plot + 2
+
+        if index == 0:
+            sharey = None
+            sharex = None
+        try:
+            if last_quant.split(" ")[-1] == label_name[index].split(" ")[-1]:
+                sharex = ax
+                sharey = ax
+            else:
+                sharex = ax
+                sharey = None
+        except:
+            sharex = None
+            sharey = None
+        last_quant = label_name[index]
+        ax = cached_add_subplot(fig, axs, nrows, ncols, plot, sharex=sharex, sharey=sharey)
+        species_label = label_name[index].split()[0]
+        if data_list[index] is not None:
+            mask = numpy.ones(data_list[index][0].shape, dtype=bool)
+            # Remove NaNs
+            for j in range(3):
+                mask[numpy.isnan(data_list[index][j])] = False
+            # Remove measuremetns with 100% or more uncertainty
+            x_data = data_list[index][0][mask]
+            y_data = data_list[index][1][mask]
+            y_data_err = data_list[index][2][mask]
+            mask = mask[mask]
+            mask[numpy.abs(y_data_err[mask]) > numpy.abs(y_data[mask])] = False
+            if numpy.any(mask):
+                color = 'orange' if species_label=='electron' else 'green'
+                ax.errorbar(x_data[mask], y_data[mask], y_data_err[mask],
+                            linestyle='', marker=".", color=color, zorder=-10, **kw)
+        uband(x, y, ax=ax, **kw)
+
+        species_label = species_label.replace("electron", "e")
+        if "Temp" in label_name[index]:
+            title_label = r'$T_{{{}}}$'.format(species_label) + imas_units_to_latex(unit_list[index])
+        elif "Density" in label_name[index]:
+            title_label = r'$n_{{{}}}$'.format(species_label) + imas_units_to_latex(unit_list[index]) + label_name_z[index]
+        elif "e_field" in label_name[index].lower():
+            title_label = r'$E_\mathrm{r}$' + imas_units_to_latex(unit_list[index])
+        elif "Velocity" in label_name[index]:
+            title_label = r"$v_\mathrm{" + species_label[0] + r"}$" + imas_units_to_latex(unit_list[index])
+        else:
+            title_label = label_name[index][:10] + imas_units_to_latex(unit_list[index])
+        ax.set_title(title_label, y=ytitle, va='top')
+        if (nplots - plot) < ncols:
+            ax.set_xlabel(x_label)
+
+    if 'label' in kw:
+        ax.legend(loc='lower center')
+    #ax.set_xlim(0, 1)
+    return {'ax': axs}
+
 @add_to__ODS__
 def core_profiles_currents_summary(ods, time_index=None, time=None, ax=None, **kw):
     """