Allow the mask argument in min/max to be a function to be parametrize…

…d by data form the PlotData instance

masci_tools/vis/data.py

@@ -20,6 +20,7 @@
 import numpy as np
 import pandas as pd
 import copy
+import warnings
 
 try:
  from bokeh.models import ColumnDataSource
@@ -245,30 +246,28 @@ def get_values(self, data_key):
 
  return values
 
- def min(self, data_key, separate=False, mask=None):
+ def min(self, data_key, separate=False, mask=None, mask_data_key=None):
  """
  Get the minimum value for a given data column for all entries
 
  :param data_key: name of the data key to determine the minimum
  :param separate: bool if True the minimum will be determined and returned
  for all entries separately
  :param mask: optional mask to select specifc rows from the data entries
+ :param mask_data_key: optional data key to be used when ``mask`` is a function
 
  :returns: minimum value for all entries either combined or as a list
  """
  if data_key not in self._column_spec._fields:
  raise ValueError(f'Field {data_key} does not exist')
 
  if mask is not None:
- if len(mask) == len(self):
- mask_gen = (mask_indx for mask_indx in mask)
- else:
- mask_gen = (mask for i in self)
+ mask = self.get_mask(mask, data_key=mask_data_key if mask_data_key is not None else data_key)
  else:
- mask_gen = (None for i in self)
+ mask = [None] * len(self)
 
  min_val = []
- for (entry, source), mask_entry in zip(self.items(mappable=True), mask_gen):
+ for (entry, source), mask_entry in zip(self.items(mappable=True), mask):
 
  key = entry._asdict()[data_key]
 
@@ -287,30 +286,28 @@ def min(self, data_key, separate=False, mask=None):
  else:
  return min(min_val)
 
- def max(self, data_key, separate=False, mask=None):
+ def max(self, data_key, separate=False, mask=None, mask_data_key=None):
  """
  Get the maximum value for a given data column for all entries
 
  :param data_key: name of the data key to determine the maximum
  :param separate: bool if True the maximum will be determined and returned
  for all entries separately
  :param mask: optional mask to select specifc rows from the data entries
+ :param mask_data_key: optional data key to be used when ``mask`` is a function
 
  :returns: maximum value for all entries either combined or as a list
  """
  if data_key not in self._column_spec._fields:
  raise ValueError(f'Field {data_key} does not exist')
 
  if mask is not None:
- if len(mask) == len(self):
- mask_gen = (mask_indx for mask_indx in mask)
- else:
- mask_gen = (mask for i in range(len(self)))
+ mask = self.get_mask(mask, data_key=mask_data_key if mask_data_key is not None else data_key)
  else:
- mask_gen = (None for i in range(len(self)))
+ mask = [None] * len(self)
 
  max_val = []
- for (entry, source), mask_entry in zip(self.items(mappable=True), mask_gen):
+ for (entry, source), mask_entry in zip(self.items(mappable=True), mask):
 
  key = entry._asdict()[data_key]
 
@@ -390,6 +387,33 @@ def get_function_result(self, data_key, func, list_return=False, **kwargs):
  else:
  return result
 
+ def get_mask(self, mask, data_key=None):
+ """
+ Get mask list for use with the Data in this instance
+
+ :param mask: either list of callable, if it is callable it is used in
+ :py:meth:`get_function_result()` together with the ``data_key``
+ argument
+ :param data_key: str to be used for the data key if mask is a callable
+
+ :param
+ """
+ from collections.abc import Callable
+
+ if isinstance(mask, Callable):
+ if data_key is None:
+ raise ValueError('If mask is a function the data_key argument has to be given')
+ mask = self.get_function_result(data_key, mask, list_return=True)
+ else:
+ if len(mask) != len(self):
+ mask = [mask for i in range(len(self))]
+
+ for mask_entry in mask:
+ if not all(isinstance(val, bool) for val in mask_entry):
+ warnings.warn('Not all entries in the mask are booleans')
+
+ return mask
+
  def sort_data(self, by_data_keys, **kwargs):
  """
  Sort the data by the given data_key(s)

masci_tools/vis/plot_methods.py

@@ -1926,10 +1926,8 @@ def plot_bands(kpath,
  if 'y' in kwargs['limits']:
  ylimits = kwargs['limits']['y']
 
- data = plot_data.values(first=True)
- mask = np.logical_and(data.bands > ylimits[0], data.bands < ylimits[1])
-
- weight_max = plot_data.max('size', mask=mask)
+ mask = lambda bands, ylimits=tuple(ylimits): np.logical_and(bands > ylimits[0], bands < ylimits[1])
+ weight_max = plot_data.max('size', mask=mask, mask_data_key='bands')
  if 'vmax' not in kwargs:
  kwargs['vmax'] = weight_max
 
@@ -2081,9 +2079,9 @@ def plot_spinpol_bands(kpath,
  if 'y' in kwargs['limits']:
  ylimits = kwargs['limits']['y']
 
- data = plot_data.values()
- mask = [np.logical_and(col.bands > ylimits[0], col.bands < ylimits[1]) for col in data]
- weight_max = plot_data.max('size', mask=mask)
+
+ mask = lambda bands, ylimits=tuple(ylimits): np.logical_and(bands > ylimits[0], bands < ylimits[1])
+ weight_max = plot_data.max('size', mask=mask, mask_data_key='bands')
 
  if 'vmax' not in kwargs:
  kwargs['vmax'] = weight_max