Source code for batman.visualization.tree

"""
Tree for 2D
-----------
"""
import copy
from sklearn.preprocessing import MinMaxScaler
import numpy as np
from scipy.spatial.distance import cdist
from matplotlib import cm
import matplotlib.tri as tri
from .hdr import HdrBoxplot
from .kiviat import Arrow3D, Kiviat3D

np.set_printoptions(precision=3)


[docs]class Tree(Kiviat3D): """Tree. Extend principle of :class:`batman.visualization.Kiviat3D` but for 2D parameter space. Sample are represented by segments and an azimutal component encode the value from :class:`batman.visualization.HdrBoxplot`. Subclass :class:`batman.visualization.Kiviat3D` by overwriting :func:`batman.visualization.Kiviat3D._axis` and :func:`batman.visualization.Kiviat3D.plane`. """
[docs] def __init__(self, sample, data, bounds=None, plabels=None, range_cbar=None): """Prepare params for Tree plot. :param array_like sample: Sample of parameters of Shape (n_samples, n_params). :param array_like data: Sample of realization which corresponds to the sample of parameters :attr:`sample` (n_samples, n_features). :param array_like bounds: Boundaries to scale the colors shape ([min, n_features], [max, n_features]). :param list(str) plabels: Names of each parameters (n_features). :param array_like range_cbar: Minimum and maximum values for output function (2 values). """ super(Tree, self).__init__(sample, data, bounds, plabels, range_cbar) scaler = MinMaxScaler() self.sample = self.sample[:, :-1] if bounds is None: bounds = copy.deepcopy(self.sample) else: bounds = np.asarray(bounds) self.scale = scaler.fit(bounds) self.n_params = 2 self.alphas = [0, np.pi] self.plabels = ['x' + str(i) for i in range(self.n_params)]\ if plabels is None else plabels self.ticks = [0.2, 0.5, 0.8] self.ticks = np.tile(self.ticks, self.n_params).reshape(-1, len(self.ticks)).T self.ticks_values = self.scale.inverse_transform(self.ticks).T self.x_ticks = [-1, 1] self.y_ticks = [0] * 2 # Distance based on HDR hdr = HdrBoxplot(data) centroide = hdr.pca.transform(hdr.median.reshape(1, -1)) dists = cdist(centroide, hdr.pca.transform(data))[0] scaler = MinMaxScaler(feature_range=[0, np.pi / 4]) self.dists = scaler.fit_transform(dists.reshape(-1, 1)) hdr_90 = hdr.pca.transform(np.array(hdr.hdr_90)) dists = cdist(centroide, hdr_90)[0]
self.dist_max = max(scaler.transform(dists.reshape(-1, 1))) def _plane(self, ax, params, feval, idx, *args): """Create a Leaf in 2D. From a set of parameters and the corresponding function evaluation, a 2D segment is created: a leaf. This leaf has an azimutal component conresponding to the distance of the given sample from the mediane given by :class:`batman.visualization.HdrBoxplot`. :param ax: Matplotlib AxesSubplot instances to draw to. :param array_like params: Parameters of the plane (n_params). :param feval: Function evaluation corresponding to :attr:`params`. :param idx: *Z* coordinate of the plane. :return: List of artists added. :rtype: list. """ params = self.scale.transform(np.asarray(params).reshape(1, -1))[0] X = params * np.cos(self.dists[idx]) * [-1, 1] Y = params * np.sin(self.dists[idx]) * [-1, 1] Z = [idx] * (self.n_params) # Plot the surface color = self.cmap(self.scale_f(feval)) out = ax.plot(X, Y, Z, alpha=0.5, lw=3, c=color) return out def _axis(self, ax): """n-dimentions axis definition. Create axis arrows along with annotations with parameters name and ticks. :param ax: Matplotlib AxesSubplot instances to draw to. :return: List of artists added starting with [[axis, plabel, [tick, tick_label] * n_ticks] * n_features] and complemented with [vert_axis, [lim_HDR, label] * 2, [wedge] * 2]. :rtype: list. """ out = super(Tree, self)._axis(ax) # Vertical line on z to separate coordinates out.append(ax.add_artist( Arrow3D([0, 0], [0, 0], [self.z_offset, len(self.data)], mutation_scale=20, lw=2, arrowstyle="-", color="k"))) # Wedges for HDR x_ = np.cos(self.dist_max)[0] y_ = np.sin(self.dist_max)[0] # Separating plane # xx = np.array([-x_, -x_, x_, x_]) # yy = np.array([-y_, -y_, y_, y_]) # z = np.array([self.z_offset, len(self.data), # self.z_offset, len(self.data)]).reshape(-1, 1) # ax.plot_surface(xx, yy, z, alpha=0.2) # Limit HDR out.append(ax.add_artist( Arrow3D([x_ * 0.85, x_], [y_ * 0.85, y_], [self.z_offset, self.z_offset], mutation_scale=20, lw=2, arrowstyle="-", color="r"))) out.append(ax.text(1.2 * x_, 1.2 * y_, self.z_offset, '90% HDR', fontsize=10, ha='center', va='center', color='k')) out.append(ax.add_artist( Arrow3D([-x_ * 0.85, -x_], [-y_ * 0.85, -y_], [self.z_offset, self.z_offset], mutation_scale=20, lw=2, arrowstyle="-", color="r"))) out.append(ax.text(-1.2 * x_, -1.2 * y_, self.z_offset, '90% HDR', fontsize=10, ha='center', va='center', color='k')) def wedge(alpha_i, alpha_e): """Wedge in 3-dimensions.""" # First create the x and y coordinates of the points. n_angles = 50 n_radii = 2 min_radius = 0.9 radii = np.linspace(min_radius, 0.95, n_radii) angles = np.linspace(alpha_i, alpha_e, n_angles, endpoint=False) angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1) angles[:, 1::2] += np.pi/n_angles x = (radii * np.cos(angles)).flatten() y = (radii * np.sin(angles)).flatten() # z set to small scale to appear as flat scaler = MinMaxScaler(feature_range=[self.z_offset, self.z_offset + 1e-5]) z = scaler.fit_transform(angles.reshape(-1, 1)).flatten() # Create the Triangulation; no triangles so Delaunay triangulation triang = tri.Triangulation(x, y) # Mask off unwanted triangles. xmid = x[triang.triangles].mean(axis=1) ymid = y[triang.triangles].mean(axis=1) mask = np.where(xmid**2 + ymid**2 < min_radius**2, 1, 0) triang.set_mask(mask) return triang, z triang, z = wedge(0, np.pi / 4) out.append(ax.plot_trisurf(triang, z, cmap=cm.get_cmap('inferno'), edgecolor='none')) triang, z = wedge(np.pi, np.pi + np.pi / 4) out.append(ax.plot_trisurf(triang, z, cmap=cm.get_cmap('inferno'), edgecolor='none'))
return out