From bec30452fa2dc1d163557b0b617d7e36bfa6d31d Mon Sep 17 00:00:00 2001 From: Thomas Gastine Date: Fri, 24 Nov 2023 11:53:33 +0100 Subject: [PATCH] add options to tune the grid in radialContour --- python/magic/plotlib.py | 66 +++++++++++++++++++++------------- python/magic/surf.py | 80 +++++++++++++++++++++++++++-------------- 2 files changed, 94 insertions(+), 52 deletions(-) diff --git a/python/magic/plotlib.py b/python/magic/plotlib.py index 7db58caa..9ce0d1ec 100644 --- a/python/magic/plotlib.py +++ b/python/magic/plotlib.py @@ -172,7 +172,7 @@ def cut(dat, vmax=None, vmin=None): def equatContour(data, radius, minc=1, label=None, levels=defaultLevels, cm=defaultCm, normed=None, vmax=None, vmin=None, cbar=True, - tit=True, normRad=False, deminc=True, bounds=True, + title=True, normRad=False, deminc=True, bounds=True, lines=False, linewidths=0.5, pcolor=False, rasterized=False): """ Plot the equatorial cut of a given field @@ -193,8 +193,8 @@ def equatContour(data, radius, minc=1, label=None, levels=defaultLevels, :type levels: int :param cm: name of the colormap ('jet', 'seismic', 'RdYlBu_r', etc.) :type cm: str - :param tit: display the title of the figure when set to True - :type tit: bool + :param title: display the title of the figure when set to True + :type title: bool :param cbar: display the colorbar when set to True :type cbar: bool :param vmax: maximum value of the contour levels @@ -236,7 +236,7 @@ def equatContour(data, radius, minc=1, label=None, levels=defaultLevels, maxS = np.sqrt(np.mean(data**2, axis=0)) data[:, maxS!=0.] /= maxS[maxS!=0.] - if tit and label is not None: + if title and label is not None: if cbar: fig = plt.figure(figsize=(6.5,5.5)) ax = fig.add_axes([0.01, 0.01, 0.76, 0.9]) @@ -331,7 +331,7 @@ def equatContour(data, radius, minc=1, label=None, levels=defaultLevels, pos = ax.get_position() l, b, w, h = pos.bounds if cbar: - if tit and label is not None: + if title and label is not None: cax = fig.add_axes([0.85, 0.46-0.7*h/2., 0.03, 0.7*h]) else: cax = fig.add_axes([0.85, 0.5-0.7*h/2., 0.03, 0.7*h]) @@ -348,7 +348,7 @@ def equatContour(data, radius, minc=1, label=None, levels=defaultLevels, def merContour(data, radius, label=None, levels=defaultLevels, cm=defaultCm, - normed=None, vmax=None, vmin=None, cbar=True, tit=True, + normed=None, vmax=None, vmin=None, cbar=True, title=True, fig=None, ax=None, bounds=True, lines=False, pcolor=False, linewidths=0.5, rasterized=False): """ @@ -365,8 +365,8 @@ def merContour(data, radius, label=None, levels=defaultLevels, cm=defaultCm, :type levels: int :param cm: name of the colormap ('jet', 'seismic', 'RdYlBu_r', etc.) :type cm: str - :param tit: display the title of the figure when set to True - :type tit: bool + :param title: display the title of the figure when set to True + :type title: bool :param cbar: display the colorbar when set to True :type cbar: bool :param vmax: maximum value of the contour levels @@ -401,7 +401,7 @@ def merContour(data, radius, label=None, levels=defaultLevels, cm=defaultCm, xx = rr * np.sin(tth) yy = rr * np.cos(tth) - if tit and label is not None: + if title and label is not None: if cbar: fsize = (5, 7.5) bb = [0.01, 0.01, 0.69, 0.91] @@ -421,7 +421,7 @@ def merContour(data, radius, label=None, levels=defaultLevels, cm=defaultCm, if ax is None: ax = fig.add_axes(bb) - if tit and label is not None: + if title and label is not None: ax.set_title(label, fontsize=24) if normed is None: @@ -502,7 +502,7 @@ def merContour(data, radius, label=None, levels=defaultLevels, cm=defaultCm, pos = ax.get_position() l, b, w, h = pos.bounds if cbar: - if tit and label is not None: + if title and label is not None: cax = fig.add_axes([0.82, 0.46-0.7*h/2., 0.04, 0.7*h]) else: cax = fig.add_axes([0.82, 0.5-0.7*h/2., 0.04, 0.7*h]) @@ -513,8 +513,9 @@ def merContour(data, radius, label=None, levels=defaultLevels, cm=defaultCm, def radialContour(data, rad=0.85, label=None, proj='hammer', lon_0=0., vmax=None, vmin=None, lat_0=30., levels=defaultLevels, cm=defaultCm, - normed=None, cbar=True, tit=True, lines=False, fig=None, - ax=None, linewidths=0.5, pcolor=False, rasterized=False): + normed=None, cbar=True, title=True, lines=False, fig=None, + ax=None, linewidths=0.5, pcolor=False, rasterized=False, + gridLineStyle=':', gridColor='k', gridLineWidth=0.7): """ Plot the radial cut of a given field @@ -533,8 +534,8 @@ def radialContour(data, rad=0.85, label=None, proj='hammer', lon_0=0., vmax=None :type levels: int :param cm: name of the colormap ('jet', 'seismic', 'RdYlBu_r', etc.) :type cm: str - :param tit: display the title of the figure when set to True - :type tit: bool + :param title: display the title of the figure when set to True + :type title: bool :param cbar: display the colorbar when set to True :type cbar: bool :param lines: when set to True, over-plot solid lines to highlight @@ -558,6 +559,14 @@ def radialContour(data, rad=0.85, label=None, proj='hammer', lon_0=0., vmax=None :param rasterized: when set to True, the rasterization for vector graphics is turned on :type rasterized: bool + :param gridColor: this is used to set the color of the grid + :type gridColor: str + :param gridLineStyle: this allows to set the line style of the grid + (':', '-', '--') + :type gridLineStyle: str + :param gridLineWidth: this is used to tune the thickness of the lines used + in the grid + :type gridLineWidth: float """ nphi, ntheta = data.shape @@ -574,7 +583,7 @@ def radialContour(data, rad=0.85, label=None, proj='hammer', lon_0=0., vmax=None if fig is None and ax is None: if proj == 'moll' or proj == 'hammer': - if tit and label is not None: + if title and label is not None: if cbar: fig = plt.figure(figsize=(9.1, 4.5)) ax = fig.add_axes([0.01, 0.01, 0.87, 0.87]) @@ -596,7 +605,7 @@ def radialContour(data, rad=0.85, label=None, proj='hammer', lon_0=0., vmax=None #verticalalignment='center', #transform = ax.transAxes) else: - if tit and label is not None: + if title and label is not None: if cbar: fig = plt.figure(figsize=(6,5.5)) ax = fig.add_axes([0.01, 0.01, 0.82, 0.9]) @@ -622,20 +631,27 @@ def radialContour(data, rad=0.85, label=None, proj='hammer', lon_0=0., vmax=None from mpl_toolkits.basemap import Basemap map = Basemap(projection=proj, lon_0=lon_0, lat_0=lat_0, resolution='c') - map.drawparallels([0.], dashes=[2, 3], linewidth=0.5) - map.drawparallels(circles, dashes=[2,3], linewidth=0.5) - map.drawmeridians(meridians, dashes=[2,3], linewidth=0.5) - map.drawmeridians([-180], dashes=[20,0], linewidth=0.5) - map.drawmeridians([180], dashes=[20,0], linewidth=0.5) + map.drawparallels([0.], dashes=[2, 3], linewidth=gridLineWidth, + color=gridColor) + map.drawparallels(circles, dashes=[2,3], linewidth=gridLineWidth, + color=gridColor) + map.drawmeridians(meridians, dashes=[2,3], linewidth=gridLineWidth, + color=gridColor) + map.drawmeridians([-180], dashes=[20,0], linewidth=gridLineWidth, + color=gridColor) + map.drawmeridians([180], dashes=[20,0], linewidth=gridLineWidth, + color=gridColor) x, y = list(map(lon2, lat2)) else: x, y = hammer2cart(ttheta, pphi) for lat0 in circles: x0, y0 = hammer2cart(lat0*np.pi/180., phi) - ax.plot(x0, y0, 'k:', linewidth=0.7) + ax.plot(x0, y0, ls=gridLineStyle, color=gridColor, + linewidth=gridLineWidth) for lon0 in meridians: x0, y0 = hammer2cart(theta, lon0*np.pi/180.) - ax.plot(x0, y0, 'k:', linewidth=0.7) + ax.plot(x0, y0, ls=gridLineStyle, color=gridColor, + linewidth=gridLineWidth) xxout, yyout = hammer2cart(theta, -np.pi)#-1e-3) xxin, yyin = hammer2cart(theta, np.pi)#+1e-3) ax.plot(xxin, yyin, 'k-') @@ -704,7 +720,7 @@ def radialContour(data, rad=0.85, label=None, proj='hammer', lon_0=0., vmax=None pos = ax.get_position() l, b, w, h = pos.bounds if cbar: - if tit and label is not None: + if title and label is not None: cax = fig.add_axes([0.9, 0.46-0.7*h/2., 0.03, 0.7*h]) else: cax = fig.add_axes([0.9, 0.51-0.7*h/2., 0.03, 0.7*h]) diff --git a/python/magic/surf.py b/python/magic/surf.py index 1499b99b..6a8d1d2f 100644 --- a/python/magic/surf.py +++ b/python/magic/surf.py @@ -70,7 +70,8 @@ def __init__(self, ivar=None, datadir='.', vort=False, ave=False, tag=None, def surf(self, field='Bphi', proj='hammer', lon_0=0., r=0.85, vmax=None, vmin=None, lat_0=30., levels=defaultLevels, cm=None, ic=False, - lon_shift=0, normed=None, cbar=True, tit=True, lines=False): + lon_shift=0, normed=None, cbar=True, title=True, lines=False, + pcolor=False): """ Plot the surface distribution of an input field at a given input radius (normalised by the outer boundary radius). @@ -80,7 +81,7 @@ def surf(self, field='Bphi', proj='hammer', lon_0=0., r=0.85, vmax=None, >>> s.surf(field='vr', r=0.95, levels=65, cm='seismic') >>> # Minimal plot (no cbar, not title) - >>> s.surf(field='entropyfluct', r=0.6, tit=False, cbar=False) + >>> s.surf(field='entropyfluct', r=0.6, title=False, cbar=False) >>> # Control the limit of the colormap from -1e3 to 1e3 >>> s.surf(field='vp', r=1., vmin=-1e3, vmax=1e3, levels=33) @@ -107,9 +108,9 @@ def surf(self, field='Bphi', proj='hammer', lon_0=0., r=0.85, vmax=None, :type lon_0: float :param lat_0: central latitude (only used with Basemap) :type lat_0: float - :param tit: display the title of the figure when set to True - :param tit: display the title of the figure when set to True - :type tit: bool + :param title: display the title of the figure when set to True + :param title: display the title of the figure when set to True + :type title: bool :param cbar: display the colorbar when set to True :type cbar: bool :param lines: when set to True, over-plot solid lines to highlight @@ -125,6 +126,8 @@ def surf(self, field='Bphi', proj='hammer', lon_0=0., r=0.85, vmax=None, :param lines: when set to True, over-plot solid lines to highlight the limits between two adjacent contour levels :type lines: bool + :param pcolor: when set to True, use pcolormesh instead of contourf + :type pcolor: bool """ if proj != 'ortho': @@ -374,11 +377,12 @@ def surf(self, field='Bphi', proj='hammer', lon_0=0., r=0.85, vmax=None, rprof = symmetrize(rprof, self.gr.minc) radialContour(rprof, rad, label, proj, lon_0, vmax, vmin, - lat_0, levels, cm, normed, cbar, tit, lines) + lat_0, levels, cm, normed, cbar, title, lines, + pcolor=pcolor) def equat(self, field='vr', levels=defaultLevels, cm=None, - normed=None, vmax=None, vmin=None, cbar=True, tit=True, - avg=False, normRad=False, ic=False): + normed=None, vmax=None, vmin=None, cbar=True, title=True, + avg=False, normRad=False, ic=False, pcolor=False): """ Plot the equatorial cut of a given field @@ -387,7 +391,7 @@ def equat(self, field='vr', levels=defaultLevels, cm=None, >>> s.equat(field='vortz', levels=65, cm='seismic') >>> # Minimal plot (no cbar, not title) - >>> s.equat(field='bphi', tit=False, cbar=False) + >>> s.equat(field='bphi', title=False, cbar=False) >>> # Control the limit of the colormap from -1e3 to 1e3 >>> s.equat(field='vr', vmin=-1e3, vmax=1e3, levels=33) @@ -409,8 +413,8 @@ def equat(self, field='vr', levels=defaultLevels, cm=None, :type levels: int :param cm: name of the colormap ('jet', 'seismic', 'RdYlBu_r', etc.) :type cm: str - :param tit: display the title of the figure when set to True - :type tit: bool + :param title: display the title of the figure when set to True + :type title: bool :param cbar: display the colorbar when set to True :type cbar: bool :param vmax: maximum value of the contour levels @@ -423,6 +427,8 @@ def equat(self, field='vr', levels=defaultLevels, cm=None, :param ic: when set to True, also display the contour levels in the inner core :type ic: bool + :param pcolor: when set to True, use pcolormesh instead of contourf + :type pcolor: bool """ phi = np.linspace(0., 2.*np.pi, self.gr.nphi) rr, pphi = np.meshgrid(self.gr.radius, phi) @@ -522,7 +528,7 @@ def equat(self, field='vr', levels=defaultLevels, cm=None, fig, xx, yy = equatContour(equator, self.gr.radius, self.gr.minc, label, levels, cm, normed, vmax, vmin, - cbar, tit, normRad) + cbar, title, normRad, pcolor=pcolor) ax = fig.get_axes()[0] if ic and data_ic is not None: @@ -539,7 +545,15 @@ def equat(self, field='vr', levels=defaultLevels, cm=None, vmax = max(abs(equator.max()), abs(equator.min())) vmin = -vmax cs = np.linspace(vmin, vmax, levels) - ax.contourf(xx_ic, yy_ic, equator_ic, cs, cmap=cm, extend='both') + if pcolor: + if normed: + ax.pcolormesh(xx_ic, yy_ic, equator_ic, cmap=cm, antialiased=True, + shading='gouraud', vmax=vmax, vmin=vmin) + else: + ax.pcolormesh(xx_ic, yy_ic, equator_ic, cmap=cm, antialiased=True, + shading='gouraud') + else: + ax.contourf(xx_ic, yy_ic, equator_ic, cs, cmap=cm, extend='both') # Variable conductivity: add a dashed line if hasattr(self.gr, 'nVarCond'): @@ -558,9 +572,9 @@ def equat(self, field='vr', levels=defaultLevels, cm=None, ax1.set_xlim(self.gr.radius.min(), self.gr.radius.max()) def avg(self, field='vphi', levels=defaultLevels, cm=None, - normed=None, vmax=None, vmin=None, cbar=True, tit=True, + normed=None, vmax=None, vmin=None, cbar=True, title=True, pol=False, tor=False, mer=False, merLevels=16, polLevels=16, - ic=False, lines=False): + ic=False, lines=False, pcolor=False): """ Plot the azimutal average of a given field. @@ -569,7 +583,7 @@ def avg(self, field='vphi', levels=defaultLevels, cm=None, >>> s.avg(field='vp', levels=65, cm='seismic') >>> # Minimal plot (no cbar, not title) - >>> s.avg(field='Br', tit=False, cbar=False) + >>> s.avg(field='Br', title=False, cbar=False) >>> # Axisymmetric Bphi + poloidal field lines >>> s.avg(field='Bp', pol=True, polLevels=8) @@ -583,8 +597,8 @@ def avg(self, field='vphi', levels=defaultLevels, cm=None, :type levels: int :param cm: name of the colormap ('jet', 'seismic', 'RdYlBu_r', etc.) :type cm: str - :param tit: display the title of the figure when set to True - :type tit: bool + :param title: display the title of the figure when set to True + :type title: bool :param cbar: display the colorbar when set to True :type cbar: bool :param vmax: maximum value of the contour levels @@ -614,6 +628,8 @@ def avg(self, field='vphi', levels=defaultLevels, cm=None, :param lines: when set to True, over-plot solid lines to highlight the limits between two adjacent contour levels :type lines: bool + :param pcolor: when set to True, use pcolormesh instead of contourf + :type pcolor: bool """ if pol: if ic: @@ -1136,7 +1152,8 @@ def avg(self, field='vphi', levels=defaultLevels, cm=None, cm = default_cmap(field) fig, xx, yy, im = merContour(phiavg, self.gr.radius, label, levels, cm, - normed, vmax, vmin, cbar, tit, lines=lines) + normed, vmax, vmin, cbar, title, lines=lines, + pcolor=pcolor) ax = fig.get_axes()[0] if ic: @@ -1156,8 +1173,17 @@ def avg(self, field='vphi', levels=defaultLevels, cm=None, vmax = max(abs(phiavg.max()), abs(phiavg.min())) vmin = -vmax cs = np.linspace(vmin, vmax, levels) - ax.contourf(xx_ic, yy_ic, phiavg_ic, cs, cmap=cm, - extend='both') + if pcolor: + if normed: + ax.pcolormesh(xx_ic, yy_ic, phiavg_ic, cmap=cm, + antialiased=True, shading='gouraud', + vmax=vmax, vmin=vmin) + else: + ax.pcolormesh(xx_ic, yy_ic, phiavg_ic, cmap=cm, + antialiased=True, shading='gouraud') + else: + ax.contourf(xx_ic, yy_ic, phiavg_ic, cs, cmap=cm, + extend='both') ax.plot([0, 0], [-ri, ri], 'k-') if pol: @@ -1187,7 +1213,7 @@ def avg(self, field='vphi', levels=defaultLevels, cm=None, ax.plot(radi*np.sin(th), radi*np.cos(th), 'k--') def slice(self, field='Bphi', lon_0=0., levels=defaultLevels, cm=None, - normed=None, vmin=None, vmax=None, cbar=True, tit=True, + normed=None, vmin=None, vmax=None, cbar=True, title=True, grid=False, nGridLevs=16, normRad=False, ic=False): """ Plot an azimuthal slice of a given field. @@ -1197,7 +1223,7 @@ def slice(self, field='Bphi', lon_0=0., levels=defaultLevels, cm=None, >>> s.slice(field='vp', lon_0=[0, 30, 60], levels=65, cm='seismic') >>> # Minimal plot (no cbar, not title) - >>> s.avg(field='vp', lon_0=32, tit=False, cbar=False) + >>> s.avg(field='vp', lon_0=32, title=False, cbar=False) >>> # Axisymmetric Bphi + poloidal field lines >>> s.avg(field='Bp', pol=True, polLevels=8) @@ -1214,8 +1240,8 @@ def slice(self, field='Bphi', lon_0=0., levels=defaultLevels, cm=None, :type levels: int :param cm: name of the colormap ('jet', 'seismic', 'RdYlBu_r', etc.) :type cm: str - :param tit: display the title of the figure when set to True - :type tit: bool + :param title: display the title of the figure when set to True + :type title: bool :param cbar: display the colorbar when set to True :type cbar: bool :param vmax: maximum value of the contour levels @@ -1504,7 +1530,7 @@ def slice(self, field='Bphi', lon_0=0., levels=defaultLevels, cm=None, phislice1 = data1[indPlot, ...] phislice = phislice + phislice1 - if tit: + if title: if cbar: fig = plt.figure(figsize=(5, 7.5)) ax = fig.add_axes([0.01, 0.01, 0.69, 0.91]) @@ -1564,7 +1590,7 @@ def slice(self, field='Bphi', lon_0=0., levels=defaultLevels, cm=None, pos = ax.get_position() l, b, w, h = pos.bounds if cbar: - if tit: + if title: cax = fig.add_axes([0.82, 0.46-0.7*h/2., 0.04, 0.7*h]) else: cax = fig.add_axes([0.82, 0.5-0.7*h/2., 0.04, 0.7*h])