def plot_scatter_4D(data, x_key, y_key, size_key=None, color_key=None,
x_label=None, y_label=None, size_label=None, color_fn=None,
add_marker_text=False, ax=None):
"""y_key can be a list..."""
colors = np.asarray(['b','r','g','y'])
# Massage inputs
if isinstance(y_key, string_types):
y_key = [y_key]
if isinstance(size_key, string_types):
size_key = [size_key]
if isinstance(color_key, string_types):
color_key = [color_key]
if x_label is None:
x_label = compute_scatter_label(x_key)
if y_label is None:
y_label = compute_scatter_label(y_key)
if size_label is None:
size_label = compute_scatter_label(s_key)
if ax is None:
ax = plt.figure(figsize=(11, 10.5)).gca()
# Now get all the data, and manipulate as needed
kwargs = {
'x': compute_key_data(data.data_dict, x_key),
'y': [compute_key_data(data.data_dict, k) for k in y_key]}
if size_key is not None:
kwargs['s'] = np.asarray([compute_key_data(data.data_dict, k)
for k in size_key])
if color_key is not None:
kwargs['c'] = np.asarray([compute_key_data(data.data_dict, k)
for k in color_key])
elif color_fn is not None:
kwargs['c'] = np.asarray([{k: color_fn(k, v) for k, v in kwargs['x'].items()}])
# Make sure everybody has the same keys
common_keys = [get_nonhemi_key(k)
for k in kwargs['x'].keys()
if not is_bad_key(k) and ~np.all(np.isnan(kwargs['x'][k]))]
if len(common_keys) == 0:
raise ValueError('Your x key has an issue.')
for key in list(set(kwargs.keys()) - set(['x'])):
# Loop over
cur_keys = [get_nonhemi_key(k)
for ddata in kwargs[key]
for k in ddata.keys()
if ~np.all(np.isnan(ddata[k]))]
common_keys = [k for k in common_keys if k in cur_keys]
if len(common_keys) == 0:
raise ValueError('Your x and y keys have no overlap.')
# Finally, we're safe to convert all of the data to numpy arrays,
# then massage the data.
# NOTE: Loop over common keys, so all are ordered similarly
# BUT the actual keys in each dict is NOT the common_key,
# but some measure-specific version of it.
#
gmc = get_measure_key
kwargs['x'] = np.asarray([kwargs['x'][gmc(ck, kwargs['x'].keys())]
for ck in common_keys])
for key in list(set(kwargs.keys()) - set(['x'])):
kwargs[key] = np.asarray([sdata[gmc(ck, sdata.keys())]
for sdata in kwargs[key]
for ck in common_keys])
if 's' in kwargs:
kwargs['s'] = 1000 * kwargs['s'] / np.abs(kwargs['s']).mean()
if 'c' in kwargs:
kwargs['c'] = colors[kwargs['c']].ravel()
# Now plot it, and annotate it!
ax.scatter(**kwargs)
ax.tick_params(labelsize=16)
if x_label:
ax.set_xlabel(x_label, fontsize=18)
if y_label:
ax.set_ylabel(y_label, fontsize=18)
if size_label:
if 'thickness' in size_label: # hack
loc='upper left'
else:
loc='upper right'
ax.legend([size_label], loc=loc)
if add_marker_text:
# Interesting if it's outside of some range of values
is_interesting = lambda v, varr, dist: np.abs(varr.mean() - v) >= dist * varr.std()
for label, x, y, s, c in zip(common_keys, kwargs['x'], kwargs['y'], kwargs['s'], kwargs['c']):
annotations = [key for key, sval in zip(['x', 'y', 's'], [1.35, 1.5, 2])
if is_interesting(locals()[key], kwargs[key], sval)]
if len(annotations) > 0:
plt.annotate(
'%s (%s)' % (get_anatomical_name(get_nonhemi_key(label)), ', '.join(annotations)),
xy = (x, y), xytext = (25, 25),
textcoords = 'offset points', ha = 'right', va = 'bottom',
bbox = dict(boxstyle = 'round,pad=0.5', fc = 'yellow', alpha = 0.5),
arrowprops = dict(arrowstyle = '->', connectionstyle = 'arc3,rad=0'),
fontsize=16)
plt.axis('equal') #ax.set_aspect('equal')
if np.any(kwargs['x'] <= 0) and np.any(kwargs['x'] >= 0):
ax.plot([0, 0], ax.get_ylim(), 'k--') # y-axis
if np.any(kwargs['y'] <= 0) and np.any(kwargs['y'] >= 0):
ax.plot(ax.get_xlim(), [0, 0], 'k--') # x-axis
plt.axis('tight')
return ax
