# -*- coding: utf-8 -*-
"""This module defines functions for supporting VMD plugin `Heatmapper`_
format files.
.. _Heatmapper: http://www.ks.uiuc.edu/Research/vmd/plugins/heatmapper/"""
__author__ = 'Ahmet Bakan, Anindita Dutta'
from numpy import arange, fromstring, array
from prody.utilities import openFile, intorfloat, startswith, addext
from prody import LOGGER
__all__ = ['parseHeatmap', 'writeHeatmap', 'showHeatmap']
HMTYPES = {
'max': float,
'min': float,
'numbering': lambda line: line.split(':'),
'title': str,
'xlabel': str,
'ylabel': str,
'xorigin': intorfloat,
'xstep': intorfloat,
}
[docs]def showHeatmap(heatmap, *args, **kwargs):
"""Show *heatmap*, which can be an two dimensional array or a Heat Mapper
:file:`.hm` file.
Heatmap is plotted using :func:`~matplotlib.pyplot.imshow` function.
Default values passed to this function are ``interpolation='nearest'``,
``aspect='auto'``, and ``origin='lower'``."""
try:
ndim, shape = heatmap.ndim, heatmap.shape
except AttributeError:
heatmap, headers = parseHeatmap(heatmap)
ndim, shape = heatmap.ndim, heatmap.shape
xorigin = headers.pop('xorigin', 0)
xextent = headers.pop('xstep', 1) * shape[0]
ylabel = kwargs.get('ylabel', '').lower()
indices = None
if ylabel:
for key in headers.get('numbering', []):
if startswith(ylabel, key.lower()):
indices = headers.get(key)
if indices is not None:
extent = [indices[0] - .5, indices[0] + len(indices) - .5,
xorigin - .5, xextent - .5]
else:
extent = [-.5, shape[1] * 2 - .5, xorigin - .5, xextent - .5]
kwargs.setdefault('extent', extent)
for key in ['numbering', 'min', 'max'] + headers.get('numbering', []):
headers.pop(key, None)
headers.update(kwargs)
kwargs = headers
if ndim != 2:
raise ValueError('mutinfo must be a 2D matrix')
kwargs.setdefault('interpolation', 'nearest')
kwargs.setdefault('origin', 'lower')
kwargs.setdefault('aspect', 'auto')
xlabel = kwargs.pop('xlabel', None)
ylabel = kwargs.pop('ylabel', None)
title = kwargs.pop('title', None)
import matplotlib.pyplot as plt
show = plt.imshow(heatmap, *args, **kwargs), plt.colorbar()
if xlabel is not None:
plt.xlabel(xlabel)
if ylabel is not None:
plt.ylabel(ylabel)
if title is not None:
plt.title(title)
return show
[docs]def parseHeatmap(heatmap, **kwargs):
"""Returns a two dimensional array and a dictionary with information parsed
from *heatmap*, which may be an input stream or an :file:`.hm` file in VMD
plugin Heat Mapper format."""
try:
readline, close = heatmap.readline, lambda: None
except AttributeError:
heatmap = openFile(heatmap)
readline, close = heatmap.readline, heatmap.close
meta = {}
arrs = []
line = readline()
while line:
if line.startswith('-'):
label, data = line[1:].split(None, 1)
data = data.strip()
if data[0] == data[-1] == '"':
data = data[1:-1]
label = label.strip()
try:
meta[label] = HMTYPES[label](data)
except KeyError:
LOGGER.warn('Unrecognized label encountered: {0}'
.format(repr(label)))
meta[label] = HMTYPES[label](data)
except TypeError:
LOGGER.warn('Could not parse data with label {0}.'
.format(repr(label)))
else:
arrs.append(line.rstrip())
line = readline()
close()
nnums = len(meta.get('numbering', ''))
heatmap = []
numbers = []
for arr in arrs:
if nnums:
items = arr.split(':', nnums + 1)
numbers.append(items[:nnums])
else:
items = [arr]
heatmap.append(fromstring(items[-1], float, sep=';'))
heatmap = array(heatmap)
if nnums:
numbering = meta['numbering']
try:
numbers = array(numbers, int)
except ValueError:
try:
numbers = array(numbers, float)
except ValueError:
LOGGER.warn('Numbering for y-axis could not be parsed.')
numbering = []
for i, label in enumerate(numbering):
meta[label] = numbers[:, i].copy()
return heatmap, meta
[docs]def writeHeatmap(filename, heatmap, **kwargs):
"""Returns *filename* that contains *heatmap* in Heat Mapper :file:`.hm`
file (extension is automatically added when not found). *filename* may
also be an output stream.
:arg title: title of the heatmap
:type title: str
:arg xlabel: x-axis lab, default is ``'unknown'``
:type xlabel: str
:arg ylabel: y-axis lab, default is ``'unknown'``
:type ylabel: str
:arg xorigin: x-axis origin, default is 0
:type xorigin: float
:arg xstep: x-axis step, default is 1
:type xstep: float
:arg min: minimum value, default is minimum in *heatmap*
:type min: float
:arg max: maximum value, default is maximum in *heatmap*
:type max: float
:arg format: number format, default is ``'%f'``
:type format: str
Other keyword arguments that are arrays with length equal to the y-axis
(second dimension of heatmap) will be considered as *numbering*."""
try:
ndim, shape = heatmap.ndim, heatmap.shape
except:
raise TypeError('heatmap must be an array object')
if ndim!= 2:
raise TypeError('heatmap must be a 2D array')
try:
write, close, stream = filename.write, lambda: None, filename
except AttributeError:
out = openFile(addext(filename, '.hm'), 'w')
write, close, stream = out.write, out.close, out
format = kwargs.pop('format', '%f')
write('-min "{0}"\n'.format(kwargs.pop('min', heatmap.min())))
write('-max "{0}"\n'.format(kwargs.pop('max', heatmap.max())))
for label, default in [
('title', 'unknown'),
('xlabel', 'unknown'),
('xorigin', 0),
('xstep', 1),
('ylabel', 'unknown'),
]:
write('-{0} "{1}"\n'.format(label, kwargs.pop(label, default)))
numbering = []
numlabels = []
for key, val in kwargs.items():
try:
length = len(val)
except TypeError:
LOGGER.warn('Keyword argument {0} is not used.'.format(key))
else:
if length == shape[0]:
numlabels.append(key)
numbering.append(val)
if not numbering:
numlabels.append('unknown')
numbering.append(arange(1, shape[0] + 1))
write('-numbering "{0}"\n'.format(':'.join(numlabels)))
for i, row in enumerate(heatmap):
write(':'.join(str(nums[i]) for nums in numbering) + ':')
row.tofile(stream, sep=';', format=format)
write(';\n')
close()
return filename
