Merge pull request #104 from SMTG-Bham/bug-fix

bug fixes and wideband update

tests/test_plot/test_heatmap.py

@@ -111,7 +111,10 @@ def test_max(self, mock_colourbar):
 
  @patch.object(plt, 'colorbar')
  def test_colourmap(self, mock_colourbar):
- cmap = mpl.cm.get_cmap('viridis')
+ try:
+ cmap = mpl.cm.get_cmap('viridis')
+ except AttributeError:
+ cmap = mpl.colormaps['viridis']
  cbar = heatmap.add_heatmap(self.ax, self.x, self.y, self.c,
  xinterp=None, yinterp=None, xscale='linear',
  yscale='linear', cscale='linear',

tp/cli/cli.py

@@ -879,7 +879,10 @@ def avg_rates(mesh_h5, mfp, total, x, crt, exclude, doping, direction,
  colour = colour[0]
 
  try:
- colours = mpl.cm.get_cmap(colour)(np.linspace(0, 1, nlines))
+ try:
+ colours = mpl.cm.get_cmap(colour)(np.linspace(0, 1, nlines))
+ except AttributeError:
+ colours = mpl.colormaps[colour](np.linspace(0, 1, nlines))
  colours = [c for c in colours]
  except ValueError:
  if isinstance(colour, str) and colour == 'skelton':
@@ -1369,7 +1372,10 @@ def kappa(kfile, efile, component, direction, tmin, tmax, dtype, doping,
  legend_title = defleg['title']
 
  try:
- colours = mpl.cm.get_cmap(colour[0])(np.linspace(0, 1, len(data)))
+ try:
+ colours = mpl.cm.get_cmap(colour[0])(np.linspace(0, 1, len(data)))
+ except AttributeError:
+ colours = mpl.colormaps[colour[0]](np.linspace(0, 1, len(data)))
  colours = [c for c in colours]
  except ValueError:
  if isinstance(colour[0], str) and colour[0] == 'skelton':
@@ -1513,12 +1519,7 @@ def kappa_target(transport_file, zt, direction, interpolate, kind, colour,
 @plot.command('phonons', no_args_is_help=True)
 @adminsitrative_options
 @inputs_function('band_yaml')
-@click.option('--bandmin',
- help='Minimum band index.',
- type=click.IntRange(0))
-@click.option('--bandmax',
- help='Maximum band index.',
- type=click.IntRange(0))
+@bandrange_options
 
 @click.option('-c', '--colour',
  help='Colourmap name or min and max colours or list of '
@@ -1959,7 +1960,11 @@ def transport(transport_file, kfile, quantity, direction, tmin, tmax, dtype,
  q: edata2[q]})
 
  try:
- colours = mpl.cm.get_cmap(colour[0])(np.linspace(0, 1, lendata))
+ try:
+ colours = mpl.cm.get_cmap(colour[0])(np.linspace(0, 1, lendata))
+ except AttributeError:
+ colours = mpl.colormaps[colour[0]](np.linspace(0, 1, lendata))
+
  colours = [c for c in colours]
  except ValueError:
  if isinstance(colour[0], str) and colour[0] == 'skelton':
@@ -2203,6 +2208,7 @@ def waterfall(kappa_hdf5, y, x, projected, direction, temperature, colour, alpha
 @adminsitrative_options
 @inputs_function('band_yaml', nargs=1)
 @inputs_function('kappa_hdf5', nargs=1)
+@bandrange_options
 
 @temperature_option
 @click.option('-p', '--poscar',
@@ -2228,9 +2234,9 @@ def waterfall(kappa_hdf5, y, x, projected, direction, temperature, colour, alpha
 @plot_io_function('tp-wideband')
 @verbose_option
 
-def wideband(band_yaml, kappa_hdf5, temperature, poscar, colour, smoothing,
- style, xmin, xmax, ymin, ymax, large, save, show, extension,
- output, verbose):
+def wideband(band_yaml, kappa_hdf5, bandmin, bandmax, temperature, poscar,
+ colour, smoothing, style, xmin, xmax, ymin, ymax, large, save,
+ show, extension, output, verbose):
  """Plots a broadened phonon dispersion."""
 
  axes = tp.axes.large if large else tp.axes.small
@@ -2244,9 +2250,11 @@ def wideband(band_yaml, kappa_hdf5, temperature, poscar, colour, smoothing,
 
  fig, ax, _ = axes.one(style)
 
- tp.plot.phonons.add_wideband(ax, kdata, pdata, temperature=temperature,
- poscar=poscar, smoothing=smoothing,
- colour=colour, verbose=verbose)
+ tp.plot.phonons.add_wideband(ax, kdata, pdata, bandmin=bandmin,
+ bandmax=bandmax, ymin=ymin, ymax=ymax,
+ temperature=temperature, poscar=poscar,
+ smoothing=smoothing, colour=colour,
+ verbose=verbose)
 
  if xmin is not None:
  if xmax is not None:

tp/cli/options.py

@@ -7,6 +7,8 @@
 # function for help and version display.
 # axes_limit_function
 # function for setting the axis limits.
+# bandrange_options
+# option for range of bands plotted.
 # direction_function:
 # function for picking the --direction (-d).
 # doping_type_option:
@@ -127,6 +129,33 @@ def axes_limit_options(f):
  return f
  return axes_limit_options
 
+def bandrange_options(f):
+ """Options for specifying bands plotted.
+ 
+ Options
+ -------
+
+ --bandmin : float, optional
+ minimum band index (0-indexed, inclusive).
+ --bandmax : float, optional
+ maximum band index (0-indexed, exclusive).
+
+ Returns
+ -------
+
+ decorator
+ band range options decorator.
+ """
+
+ f = click.option('--bandmin',
+ help="Minimum band index.",
+ type=click.IntRange(0))(f)
+ f = click.option('--bandmax',
+ help="Maximum band index.",
+ type=click.IntRange(0))(f)
+
+ return f
+
 def direction_function(multiple=False):
  """Function to create direction options.
  
@@ -334,7 +363,27 @@ def dos_options(f):
  return dos_options
 
 def heatmap_options(f):
- """Options for heatmaps."""
+ """Options for heatmaps.
+ 
+ Options
+ -------
+
+ --discrete/--continuous : bool, optional
+ discretise the colourmap. Default: --continuous
+ -l, --levels : int or array-like, optional
+ boundaries for discrete plots. Lists specify actual values
+ while integers specify maximum-1 number of boundaries.
+ --contours : float or array-like, optional
+ contour line values.
+ --contourcolours: str or array-like, optional
+ contour colours. Default: black.
+
+ Returns
+ -------
+
+ decorator
+ heatmap options decorator.
+ """
 
  f = click.option('--discrete/--continuous',
  help='Discretise colourmap. [default: continuous]',

tp/plot/frequency.py

@@ -702,7 +702,10 @@ def add_waterfall(ax, data, quantity, xquantity='frequency', temperature=300,
 
  s = np.shape(data[xquantity])
  try:
- colour = mpl.cm.get_cmap(colour)
+ try:
+ colour = mpl.cm.get_cmap(colour)
+ except AttributeError:
+ colour = mpl.colormaps[colour]
  colours = [colour(i) for i in np.linspace(0, 1, s[1])]
  colours = np.tile(colours, (s[0], 1))
  except ValueError:

tp/plot/heatmap.py

@@ -190,7 +190,10 @@ def add_heatmap(ax, x, y, c, xinterp=None, yinterp=None, kind='linear',
  # #rrggbb colour as the highlight colour for a tp.plot.colour.uniform.
 
  try:
- colours = copy(mpl.cm.get_cmap(colour))
+ try:
+ colours = copy(mpl.cm.get_cmap(colour))
+ except AttributeError:
+ colours = copy(mpl.colormaps[colour])
  except ValueError:
  if isinstance(colour, mpl.colors.ListedColormap):
  colours = colour
@@ -282,7 +285,10 @@ def add_heatmap(ax, x, y, c, xinterp=None, yinterp=None, kind='linear',
  cmap = None
  try:
  cmap = contourcolours
- contourcolours = mpl.cm.get_cmap(contourcolours)(ctnorm)
+ try:
+ contourcolours = mpl.cm.get_cmap(contourcolours)(ctnorm)
+ except AttributeError:
+ contourcolours = mpl.colormaps[contourcolours](ctnorm)
  plt.contour(x[:-1], y[:-1], c, contours, cmap=cmap,
  **contourkwargs)
  except ValueError:

tp/plot/phonons.py

@@ -53,8 +53,8 @@
 workers = tp.settings.get_workers()
 
 def add_dispersion(ax, data, sdata=None, bandmin=None, bandmax=None, main=True,
- label=None, colour='#800080', linestyle='solid',
- marker=None, xmarkkwargs={}, **kwargs):
+ label=None, colour='#800080', linestyle='solid', marker=None,
+ xmarkkwargs={}, **kwargs):
  """Adds a phonon band structure to a set of axes.
 
  Labels, colours and linestyles can be given one for the whole
@@ -195,6 +195,8 @@ def add_dispersion(ax, data, sdata=None, bandmin=None, bandmax=None, main=True,
  if main:
  if round(np.amin(f), 1) == 0:
  ax.set_ylim(bottom=0)
+ else:
+ ax.set_ylim(bottom=np.amin(f))
  if sdata is None: sdata = data
  formatting(ax, sdata, 'frequency', **xmarkkwargs)
 
@@ -291,7 +293,10 @@ def add_multi(ax, data, bandmin=None, bandmax=None, main=True, label=None,
  # line appearance
 
  try:
- colours = mpl.cm.get_cmap(colour)(np.linspace(0, 1, len(data)))
+ try:
+ colours = mpl.cm.get_cmap(colour)(np.linspace(0, 1, len(data)))
+ except AttributeError:
+ colours = mpl.colormaps[colour](np.linspace(0, 1, len(data)))
  colours = [[c] for c in colours]
  except ValueError:
  if isinstance(colour, mpl.colors.ListedColormap):
@@ -812,6 +817,8 @@ def add_projected_dispersion(ax, data, pdata, quantity, bandmin=None,
  if main:
  if round(np.amin(f), 1) == 0:
  ax.set_ylim(bottom=0)
+ else:
+ ax.set_ylim(bottom=np.amin(f))
  formatting(ax, pdata, 'frequency', **xmarkkwargs)
 
  return cbar
@@ -1058,9 +1065,10 @@ def add_alt_projected_dispersion(ax, data, pdata, quantity, projected,
  return cbar
 
 @tp.docstring_replace(workers=str(workers))
-def add_wideband(ax, kdata, pdata, temperature=300, poscar='POSCAR', main=True,
- smoothing=5, colour='viridis', workers=workers,
- xmarkkwargs={}, verbose=False, **kwargs):
+def add_wideband(ax, kdata, pdata, temperature=300, bandmin=None, bandmax=None,
+ ymin=None, ymax=None, poscar='POSCAR', main=True, smoothing=5,
+ colour='viridis', workers=workers, xmarkkwargs={},
+ verbose=False, **kwargs):
  """Plots a phonon dispersion with broadened bands.
 
  Requires a POSCAR.
@@ -1096,6 +1104,14 @@ def add_wideband(ax, kdata, pdata, temperature=300, poscar='POSCAR', main=True,
 
  temperature : float, optional
  approximate temperature in K (finds closest). Default: 300.
+ bandmin : int, optional
+ zero-indexed minimum band index to plot. Default: None.
+ bandmax : int, optional
+ zero-indexed maximum band index to plot. Default: None.
+ ymin : float, optional
+ minimum y-value plotted.
+ ymax : float, optional
+ maximum y-value plotted.
 
  poscar : str, optional
  VASP POSCAR filepath. Default: POSCAR.
@@ -1180,6 +1196,19 @@ def add_wideband(ax, kdata, pdata, temperature=300, poscar='POSCAR', main=True,
  c = np.array(kdata['gamma'])
  qk = kdata['qpoint']
 
+ # data selection
+
+ if bandmin is None:
+ bandmin = 0
+ else:
+ bandmin = np.amax([0, bandmin])
+ if bandmax is None:
+ bandmax = len(pdata['frequency'][0])
+ else:
+ bandmax = np.amin([len(pdata['frequency'][0]), bandmax])
+
+ c = c[:,bandmin:bandmax]
+
  # Phonopy data formatting
 
  qp = pdata['qpoint']
@@ -1199,20 +1228,18 @@ def add_wideband(ax, kdata, pdata, temperature=300, poscar='POSCAR', main=True,
  min_id = pool.map(geq, qpi)
  c2 = c[min_id, :]
 
- x, indices = np.unique(x, return_index=True)
- f = np.array(pdata['frequency'])[indices]
- where = np.where(c2 == np.amax(c2))
-
  # interpolate
 
+ x, indices = np.unique(x, return_index=True)
+ f = np.array(pdata['frequency'])[indices,bandmin:bandmax]
  x2 = np.linspace(min(x), max(x), 2500)
  finterp = interp1d(x, f, kind='cubic', axis=0)
  f = finterp(x2)
 
  cinterp = interp1d(xi, c2, kind='cubic', axis=0)
  c2 = np.abs(cinterp(x2))
- fmax = np.amax(np.add(f, c2))
- fmin = np.amin(np.subtract(f, c2))
+ fmax = np.amax(np.add(f, c2)) if ymax is None else ymax
+ fmin = np.amin(np.subtract(f, c2)) if ymin is None else ymin
  c2 = np.where(c2==0, np.nanmin(c2[np.nonzero(c2)]), c2)
  f2 = np.linspace(fmin, fmax, 2500)
 
@@ -1231,7 +1258,10 @@ def add_wideband(ax, kdata, pdata, temperature=300, poscar='POSCAR', main=True,
  # the min and max values.
 
  try:
- cmap = mpl.cm.get_cmap(colour)
+ try:
+ cmap = mpl.cm.get_cmap(colour)
+ except AttributeError:
+ cmap = mpl.colormaps[colour]
  except ValueError:
  if isinstance(colour, mpl.colors.ListedColormap):
  cmap = colour
@@ -1256,6 +1286,8 @@ def add_wideband(ax, kdata, pdata, temperature=300, poscar='POSCAR', main=True,
  if main:
  if round(np.amin(f), 1) == 0:
  ax.set_ylim(bottom=0)
+ else:
+ ax.set_ylim(bottom=fmin)
  formatting(ax, pdata, 'frequency', **xmarkkwargs)
 
  return

tp/plot/utilities.py

@@ -141,7 +141,10 @@ def parse_colours(colour):
  from copy import copy
 
  try:
- cmap = copy(mpl.cm.get_cmap(colour))
+ try:
+ cmap = copy(mpl.cm.get_cmap(colour))
+ except AttributeError:
+ cmap = copy(mpl.colormaps[colour])
  except ValueError:
  if isinstance(colour, mpl.colors.ListedColormap):
  cmap = copy(colour)