add support for general 2D prior mask via mask_function argument

Author: cmbant

docs/plot_gallery.html

@@ -1,5 +1,5 @@
 __author__ = 'Antony Lewis'
-__version__ = "1.5.4"
+__version__ = "1.5.5"
 __url__ = "https://getdist.readthedocs.io"
 
 import os

docs/plot_gallery.ipynb

@@ -40,7 +40,7 @@ smooth_scale_1D =-1
 
 #0 is basic normalization correction
 #1 is linear boundary kernel (should get gradient correct)
-#2 is a higher order kernel, that also affects estimates way from the boundary (1D only)
+#2 is a higher order kernel, that also affects estimates away from the boundary (1D only)
 boundary_correction_order=1
 
 #Correct for (over-smoothing) biases using multiplicative bias correction

getdist/__init__.py

@@ -253,17 +253,19 @@ class Density2D(GridDensity):
     You can call it like a :class:`~scipy:scipy.interpolate.RectBivariateSpline` object to get interpolated values.
     """
 
-    def __init__(self, x, y, P=None, view_ranges=None):
+    def __init__(self, x, y, P=None, view_ranges=None, mask=None):
         """
         :param x: array of x values
         :param y: array of y values
         :param P: 2D array of density values at x, y
         :param view_ranges: optional ranges for viewing density
+        :param mask: optional 2D boolean array for non-trivial mask
         """
         self.x = x
         self.y = y
         self.axes = [y, x]
         self.view_ranges = view_ranges
+        self.mask = mask
         self.spacing = (self.x[1] - self.x[0]) * (self.y[1] - self.y[0])
         self.setP(P)
 

getdist/analysis_defaults.ini

@@ -1601,7 +1601,7 @@ def get1DDensityGridData(self, j, paramConfid=None, meanlikes=False, **kwargs):
 
         return density1D
 
-    def _setEdgeMask2D(self, parx, pary, prior_mask, winw, alledge=False):
+    def _setEdgeMask2D(self, parx, pary, prior_mask, winw):
         if parx.has_limits_bot:
             prior_mask[:, winw] /= 2
             prior_mask[:, :winw] = 0
@@ -1614,11 +1614,12 @@ def _setEdgeMask2D(self, parx, pary, prior_mask, winw, alledge=False):
         if pary.has_limits_top:
             prior_mask[-(winw + 1), :] /= 2
             prior_mask[-winw:, :] = 0
-        if alledge:
-            prior_mask[:, :winw] = 0
-            prior_mask[:, -winw:] = 0
-            prior_mask[:winw:] = 0
-            prior_mask[-winw:, :] = 0
+
+    def _setAllEdgeMask2D(self, prior_mask, winw):
+        prior_mask[:, :winw] = 0
+        prior_mask[:, -winw:] = 0
+        prior_mask[:winw:] = 0
+        prior_mask[-winw:, :] = 0
 
     def _getScaleForParam(self, par):
         # Also ensures that the 1D limits are initialized
@@ -1655,7 +1656,8 @@ def get2DDensity(self, x, y, normalized=False, **kwargs):
         return density
 
     # noinspection PyUnboundLocalVariable
-    def get2DDensityGridData(self, j, j2, num_plot_contours=None, get_density=False, meanlikes=False, **kwargs):
+    def get2DDensityGridData(self, j, j2, num_plot_contours=None, get_density=False, meanlikes=False,
+                             mask_function: callable = None, **kwargs):
         """
         Low-level function to get 2D plot marginalized density and optional additional plot data.
 
@@ -1665,6 +1667,9 @@ def get2DDensityGridData(self, j, j2, num_plot_contours=None, get_density=False,
         :param get_density: only get the 2D marginalized density, don't calculate confidence level members
         :param meanlikes: calculate mean likelihoods as well as marginalized density
                           (returned as array in density.likes)
+        :param mask_function: optional function, mask_function(minx, miny,  stepx, stepy, mask),
+                which which sets mask to zero for values of parameters that are excluded by prior. Note this is not
+                needed for standard min, max bounds aligned with axes, as they are handled by default.
         :param kwargs: optional settings to override instance settings of the same name (see `analysis_settings`):
 
             - **fine_bins_2D**
@@ -1689,7 +1694,7 @@ def get2DDensityGridData(self, j, j2, num_plot_contours=None, get_density=False,
         mult_bias_correction_order = kwargs.get('mult_bias_correction_order', self.mult_bias_correction_order)
         smooth_scale_2D = float(kwargs.get('smooth_scale_2D', self.smooth_scale_2D))
 
-        has_prior = parx.has_limits or pary.has_limits
+        has_prior = parx.has_limits or pary.has_limits or mask_function
 
         corr = self.getCorrelationMatrix()[j2][j]
         actual_corr = corr
@@ -1761,7 +1766,7 @@ def get2DDensityGridData(self, j, j2, num_plot_contours=None, get_density=False,
         logging.debug('time 2D binning and bandwidth: %s ; bins: %s', time.time() - start, fine_bins_2D)
         start = time.time()
         cache = {}
-        convolvesize = xsize + 2 * winw + Win.shape[0]
+        convolvesize = xsize + 2 * winw + Win.shape[0]  # larger than needed for selecting fft pixel count
         bins2D = convolve2D(histbins, Win, 'same', largest_size=convolvesize, cache=cache)
 
         if meanlikes:
@@ -1779,15 +1784,24 @@ def get2DDensityGridData(self, j, j2, num_plot_contours=None, get_density=False,
         else:
             bin2Dlikes = None
 
+        if has_prior and boundary_correction_order >= 0 or mult_bias_correction_order or mask_function:
+            prior_mask = np.ones((ysize + 2 * winw, xsize + 2 * winw))
+            if mask_function:
+                mask_function(xbinmin - winw * finewidthx, ybinmin - winw * finewidthy, finewidthx, finewidthy,
+                              prior_mask)
+                bool_mask = prior_mask[winw:-winw, winw:-winw] < 1e-8
+
         if has_prior and boundary_correction_order >= 0:
             # Correct for edge effects
-            prior_mask = np.ones((ysize + 2 * winw, xsize + 2 * winw))
             self._setEdgeMask2D(parx, pary, prior_mask, winw)
             a00 = convolve2D(prior_mask, Win, 'valid', largest_size=convolvesize, cache=cache)
             ix = a00 * bins2D > np.max(bins2D) * 1e-8
             a00 = a00[ix]
             normed = bins2D[ix] / a00
-            if boundary_correction_order == 1:
+            if boundary_correction_order == 0:
+                # simple boundary correction by normalization
+                bins2D[ix] = normed
+            elif boundary_correction_order == 1:
                 # linear boundary correction
                 indexes = np.arange(-winw, winw + 1)
                 y = np.empty(Win.shape)
@@ -1811,26 +1825,28 @@ def get2DDensityGridData(self, j, j2, num_plot_contours=None, get_density=False,
                 Ay = a01 * a20 - a10 * a11
                 corrected = (bins2D[ix] * A + xP * Ax + yP * Ay) / denom
                 bins2D[ix] = normed * np.exp(np.minimum(corrected / normed, 4) - 1)
-            elif boundary_correction_order == 0:
-                # simple boundary correction by normalization
-                bins2D[ix] = normed
             else:
                 raise SettingError('unknown boundary_correction_order (expected 0 or 1)')
 
         if mult_bias_correction_order:
-            prior_mask = np.ones((ysize + 2 * winw, xsize + 2 * winw))
-            self._setEdgeMask2D(parx, pary, prior_mask, winw, alledge=True)
+            self._setAllEdgeMask2D(prior_mask, winw)
             a00 = convolve2D(prior_mask, Win, 'valid', largest_size=convolvesize, cache=cache, cache_args=[2])
             for _ in range(mult_bias_correction_order):
                 box = histbins.copy()
                 ix2 = bins2D > np.max(bins2D) * 1e-8
                 box[ix2] /= bins2D[ix2]
                 bins2D *= convolve2D(box, Win, 'same', largest_size=convolvesize, cache=cache, cache_args=[2])
-                bins2D /= a00
+                if mask_function:
+                    bins2D[~bool_mask] /= a00[~bool_mask]
+                else:
+                    bins2D /= a00
+
+        if mask_function:
+            bins2D[bool_mask] = 0
 
         x = np.linspace(xbinmin, xbinmax, xsize)
         y = np.linspace(ybinmin, ybinmax, ysize)
-        density = Density2D(x, y, bins2D,
+        density = Density2D(x, y, bins2D, mask=None if not mask_function else np.asarray(bool_mask),
                             view_ranges=[(parx.range_min, parx.range_max), (pary.range_min, pary.range_max)])
         density.normalize('max', in_place=True)
         if get_density:

getdist/densities.py

@@ -1024,7 +1024,8 @@ def _is_color_like(self, color):
             return False
 
     def add_2d_contours(self, root, param1=None, param2=None, plotno=0, of=None, cols=None, contour_levels=None,
-                        add_legend_proxy=True, param_pair=None, density=None, alpha=None, ax=None, **kwargs):
+                        add_legend_proxy=True, param_pair=None, density=None, alpha=None, ax=None,
+                        mask_function: callable = None, **kwargs):
         """
         Low-level function to add 2D contours to plot for samples with given root name and parameters
 
@@ -1043,6 +1044,9 @@ def add_2d_contours(self, root, param1=None, param2=None, plotno=0, of=None, col
         :param alpha: alpha for the contours added
         :param ax: optional :class:`~matplotlib:matplotlib.axes.Axes` instance (or y,x subplot coordinate)
                    to add to (defaults to current plot or the first/main plot if none)
+        :param mask_function: optional function, mask_function(minx, miny,  stepx, stepy, mask),
+                which which sets mask to zero for values of parameter name that are excluded by prior.
+                See the example in the plot gallery.
         :param kwargs: optional keyword arguments:
 
                - **filled**: True to make filled contours
@@ -1055,7 +1059,13 @@ def add_2d_contours(self, root, param1=None, param2=None, plotno=0, of=None, col
         if density is None:
             param1, param2 = self.get_param_array(root, param_pair or [param1, param2])
             ax.getdist_params = (param1, param2)
-            if isinstance(root, MixtureND):
+            if mask_function is not None:
+                samples = self.samples_for_root(root)
+                density = samples.get2DDensityGridData(param1.name, param2.name,
+                                                       mask_function=mask_function,
+                                                       num_plot_contours=self.settings.num_plot_contours,
+                                                       meanlikes=self.settings.shade_meanlikes)
+            elif isinstance(root, MixtureND):
                 density = root.marginalizedMixture(params=[param1, param2]).density2D()
             else:
                 density = self.sample_analyser.get_density_grid(root, param1, param2,
@@ -1086,6 +1096,7 @@ def add_2d_contours(self, root, param1=None, param2=None, plotno=0, of=None, col
         def clean_args(_args):
             return {k: v for k, v in _args.items() if k not in ('color', 'ls', 'lw')}
 
+        z = density.P if density.mask is None else np.ma.masked_where(density.mask, density.P)
         if kwargs.get('filled'):
             if cols is None:
                 color = kwargs.get('color')
@@ -1098,13 +1109,13 @@ def clean_args(_args):
                 else:
                     cols = color
             levels = sorted(np.append([density.P.max() + 1], contour_levels))
-            cs = ax.contourf(density.x, density.y, density.P, levels, colors=cols, alpha=alpha, **clean_args(kwargs))
+            cs = ax.contourf(density.x, density.y, z, levels, colors=cols, alpha=alpha, **clean_args(kwargs))
 
             fc = tuple(cs.to_rgba(cs.cvalues[-1], cs.alpha))
             if proxy_ix >= 0:
                 self.contours_added[proxy_ix] = (
                     matplotlib.patches.Rectangle((0, 0), 1, 1, fc=fc))
-            ax.contour(density.x, density.y, density.P, levels[:1], colors=(fc,),
+            ax.contour(density.x, density.y, z, levels[:1], colors=(fc,),
                        linewidths=self._scaled_linewidth(self.settings.linewidth_contour
                                                          if kwargs.get('lw') is None else kwargs['lw']),
                        linestyles=kwargs.get('ls'),
@@ -1116,7 +1127,7 @@ def clean_args(_args):
             lws = args['lw']  # note linewidth_contour is only used for filled contours
             kwargs = self._get_plot_args(plotno, **kwargs)
             kwargs['alpha'] = alpha
-            cs = ax.contour(density.x, density.y, density.P, sorted(contour_levels), colors=cols, linestyles=linestyles,
+            cs = ax.contour(density.x, density.y, z, sorted(contour_levels), colors=cols, linestyles=linestyles,
                             linewidths=lws, **clean_args(kwargs))
             dashes = args.get('dashes')
             if dashes:
@@ -1658,14 +1669,15 @@ def plot_1d(self, roots, param, marker=None, marker_color=None, label_right=Fals
             self.finish_plot()
 
     def plot_2d(self, roots, param1=None, param2=None, param_pair=None, shaded=False,
-                add_legend_proxy=True, line_offset=0, proxy_root_exclude=(), ax=None, **kwargs):
+                add_legend_proxy=True, line_offset=0, proxy_root_exclude=(), ax=None,
+                mask_function: callable = None, **kwargs):
         """
         Create a single 2D line, contour or filled plot.
 
         :param roots: root name or :class:`~.mcsamples.MCSamples` instance (or list of any of either of these) for
                       the samples to plot
         :param param1: x parameter name
-        :param param2:  y parameter name
+        :param param2: y parameter name
         :param param_pair: An [x,y] pair of params; can be set instead of param1 and param2
         :param shaded: True or integer if plot should be a shaded density plot, where the integer specifies
                        the index of which contour is shaded (first samples shaded if True provided instead
@@ -1675,6 +1687,15 @@ def plot_2d(self, roots, param1=None, param2=None, param_pair=None, shaded=False
         :param proxy_root_exclude: any root names not to include when adding to the legend proxy
         :param ax: optional :class:`~matplotlib:matplotlib.axes.Axes` instance (or y,x subplot coordinate)
                    to add to (defaults to current plot or the first/main plot if none)
+        :param mask_function: Function that defines regions in the 2D parameter space to exclude from the plot.
+                Must have signature mask_function(minx, miny, stepx, stepy, mask), where:
+                - minx, miny: minimum values of x and y parameters
+                - stepx, stepy: step sizes in x and y directions
+                - mask: 2D boolean numpy array (modified in-place)
+                The function should set mask values to 0 where points should be excluded by the prior.
+                Useful for implementing non-rectangular prior boundaries not aligned with parameter axes,
+                - see the example in the plot gallery.
+                Note it should not include simple axis-aligned range priors that are accounted for automatically.
         :param kwargs: additional optional arguments:
 
                 * **filled**: True for filled contours
@@ -1711,6 +1732,7 @@ def plot_2d(self, roots, param1=None, param2=None, param_pair=None, shaded=False
         contour_args = self._make_contour_args(len(roots), **kwargs)
         for i, root in enumerate(roots):
             res = self.add_2d_contours(root, param_pair[0], param_pair[1], line_offset + i, of=len(roots), ax=ax,
+                                       mask_function=mask_function,
                                        add_legend_proxy=add_legend_proxy and root not in proxy_root_exclude,
                                        **contour_args[i])
             xbounds, ybounds = self._update_limits(res, xbounds, ybounds)