protmap.py

# coding=utf-8
"""
                    ARIA Evolutionary Constraints Tools
"""
from __future__ import absolute_import, division, print_function

from ..core import ConversionTable as ConversionTable
import collections
import csv
import datetime
import itertools
import logging
import matplotlib; matplotlib.use("Agg")
import numpy as np
import operator
import os
import pandas as pd
import re
import seaborn as sns
import six
import sklearn.metrics as skm
import string
import textwrap
from collections import defaultdict
from copy import deepcopy
from matplotlib import pyplot as plt
from matplotlib.lines import Line2D

from ..core.legacy import AminoAcid as AminoAcid
from .common import (tickmin, tickrot, titleprint, addtup)
from .ndconv import net_deconv


LOG = logging.getLogger(__name__)
# TODO: check dataframe symmetry or always use unstack
# TODO: objet MapContainer contenant les differentes maps en attributs ! (et
# non en clef de dict)


class Map(pd.DataFrame):
    """
    Distance/contact matrix abstract class accepting only one type of value according to mtype arg

    Examples
    --------
    Score/distance matrix

    >>> d = Map(index=[0,1,2], columns=[0,1,2])
    >>> d
        0   1   2
    0 NaN NaN NaN
    1 NaN NaN NaN
    2 NaN NaN NaN
    >>> d = Map(index=[0,1,2], columns=[0,1,2], mtype='contact')
    >>> d
           0      1      2
    0  False  False  False
    1  False  False  False
    2  False  False  False

    """

    mtype_choices = {'contact': bool, 'distance': float, "score": float}
    _metadata = pd.DataFrame._metadata + ['_sort_list']

    def update(self, *args, **kwargs):
        """
        

        Parameters
        ----------
        args :
            param kwargs:
        *args :
            
        **kwargs :
            

        Returns
        -------

        
        """
        super(Map, self).update(*args, **kwargs)

    # def _constructor_expanddim(self):
    #     super(Map, self)._constructor_expanddim()

    def __init__(self, index=None, columns=None, mtype='distance',
                 duplicate_levels=False, data=None, dtype=None, sym=True,
                 desc="", path="", **kwargs):
        """
        
        Parameters
        ----------
        index
        columns
        mtype
        duplicate_levels
            Allow duplicate levels in dataframe
        data
        dtype
        sym
        desc
        path
        """
        if not dtype:
            dtype = self.check_type(mtype)
        # TODO: should probably init to np.NaN
        if data is None:
            data = False if dtype == bool else np.NaN
        super(Map, self).__init__(data=data, dtype=dtype, index=index,
                                  columns=columns)
        self.duplicate_levels = duplicate_levels
        self.mtype = mtype
        self.dtype = dtype
        if mtype == "score":
            self._sort_list = []
        self.sym = sym
        self.desc = desc
        self.path = path

    def __str__(self):
        return super(Map, self).__str__()

    @property
    def sort_list(self):
        return self._sort_list

    def sortedset(self, human_idx=False):
        """
        

        Parameters
        ----------
        human_idx :
            (Default value = False)

        Returns
        -------

        
        """
        # Remove duplicate in sort_list
        n = 1 if human_idx else 0
        if hasattr(self, "sort_list"):
            if self.sym:
                # Use OrderedDict to keep the order
                return [(x + n, y + n)
                        for x, y in
                        collections.OrderedDict.fromkeys(frozenset(x)
                                                         for x in
                                                         self.sort_list)]
            else:
                # Asym matrix, no need to remove duplicate
                return [(x + n, y + n) for x, y in self.sort_list]
        else:
            return None

    def check_type(self, mtype):
        """
        

        Parameters
        ----------
        mtype :
            return:

        Returns
        -------

        
        """
        value = self.mtype_choices.get(mtype)
        if value:
            return value
        else:
            LOG.error("Map type should be in list %s",
                      self.mtype_choices.keys())
            return None

    def reduce(self):
        """Low complexcity dataframe"""
        raise NotImplementedError

    def copy(self, **kwargs):
        """
        Copy the current map

        Parameters
        ----------
        **kwargs :
            

        Returns
        -------

        
        """
        raise NotImplementedError

    def remove(self, rm_list):
        """
        

        Parameters
        ----------
        rm_list :
            

        Returns
        -------

        
        """
        # Reset values at positions in rm_list
        value = False if self.dtype == bool else np.NaN
        for contact in rm_list:
            idx1, idx2 = self.index[contact[0]], self.index[contact[1]]
            self.set_value(idx1, idx2, value)
            # self.iat[(contact[0], contact[1])] = value
            # self.iat[(contact[1], contact[0])] = value
            if hasattr(self, 'sort_list'):
                # ! sort_list start at 1
                if (contact[0], contact[1]) in self.sort_list:
                    self.sort_list.remove((contact[0], contact[1]))
                if (contact[1], contact[0]) in self.sort_list and self.sym:
                    self.sort_list.remove((contact[1], contact[0]))

    def to_series(self):
        """Return panda series related to lower triangle values"""
        df = self.copy()
        df = df.astype(float)
        # remove values from upper triangle
        df.values[np.triu_indices_from(df, k=1)] = np.nan
        # pd.series with only lower triangle values
        return df.unstack().dropna()

    def topmap(self, scoremap, nb_topcontact):
        """
        

        Parameters
        ----------
        scoremap :
            param nb_topcontact:
        nb_topcontact :
            

        Returns
        -------

        
        """
        if self.dtype != bool:
            LOG.info("Error when retrieving top contact map. The type of "
                     "the given map is not a contact type!")
            return self
        self[:] = False
        if self.shape == scoremap.shape:
            pair_list = scoremap.sortedset()[:nb_topcontact]
            for contact in pair_list:
                self.iat[(contact[1], contact[0])] = True
                self.iat[(contact[0], contact[1])] = True
            return self
        else:
            LOG.error("Given scoremap has not the same dimension !")
            return None

    def subfill(self, pairdict, level=0):
        """
        Fill map with dict giving

        Parameters
        ----------
        pairdict :
            param level:
        level :
            (Default value = 0)

        Returns
        -------

        
        """
        pairdict = {k.upper(): v for k, v in pairdict.items()}
        if "def" in pairdict:
            self[:] = pairdict["def"]
        for idxval in pairdict:
            idx = self.index.get_level_values(level) == idxval
            self.loc[idx, idx] = pairdict[idxval]

    # TODO: remove sym attribute --> information is duplicated !!!!
    def set_value(self, index, col, value, **kwargs):
        """
        Assign value at [index, col] in the related dataframe

        >>> d = Map(index=[0,1,2], columns=[0,1,2])
        >>> d.set_value(1, 2, 2)
        >>> d.at[1, 2]
        2.0

        Parameters
        ----------
        index :
            param col:
        value :
            param kwargs:
        col :
            
        **kwargs :
            

        Returns
        -------

        
        """
        LOG.debug("Update {index}, {col} with {value} value {default}".format(index=index, col=col, value=value,
                                                                              default=self.at[index, col]))
        # super(Map, self).set_value(index, col, value, **kwargs)
        super(Map, self).at[index, col] = value
        if self.sym:
            super(Map, self).at[col, index] = value
            # super(Map, self).set_value(col, index, value, **kwargs)


# TODO: Matrices PosAaAtmMap, AaAtmMap, AtmMap
class ProteinMap(Map):
    """
    Abstract class for protein contact map objects
    """
    # TODO: matrix should be not only for heavy atoms
    # Matrix only for heavy atoms.
    heavy_reg = re.compile(r"[CNOS][ABGDEZH][0-9]?")
    all_reg = re.compile(r"^((?!cns|dyana).*)$")
    # TODO: Autre methodes de dist
    distance_method = 'euclidean'
    _metadata = Map._metadata + ['_evflags', '_maplot', '_sequence',
                                 '_atom_groups']

    def __init__(self, sequence, flaglist=None, atom_groups="min", **kwargs):
        if kwargs.get("index") is None or kwargs.get("columns") is None:
            kwargs["index"], kwargs["columns"] = self.create_index(
                sequence, **kwargs)
        super(ProteinMap, self).__init__(**kwargs)
        self._evflags = flaglist
        self._maplot = None
        self._sequence = sequence
        self._atom_groups = atom_groups

    # def _constructor_expanddim(self):
    #    return self._constructor_expanddim()
    @property
    def sequence(self):
        """
        Specific sequence related to the map in a string
        Returns
        -------
        str

        """
        raise NotImplementedError

    @property
    def contact_flags(self):
        return self._evflags

    def create_heatmap(self):
        """
        Generate heatmap

        Parameters
        ----------

        Returns
        -------


        """
        raise NotImplementedError

    def contact_map(self, contactdef, scsc_min=None):
        """
        Generate contact map

        Parameters
        ----------
        contactdef :
            param scsc_min:
        scsc_min :
            (Default value = None)

        Returns
        -------


        """
        raise NotImplementedError

    def create_index(self, sequence, **kwargs):
        """

        Parameters
        ----------
        sequence
        kwargs

        Returns
        -------

        """
        # Indexation matrice (tous les atomes ou tous les residus)
        raise NotImplementedError

    def reduce(self):
        """Lower index level if multi index"""
        # TODO: check if multiindex
        columns = ['-'.join(idx) for idx in self.columns]
        index = ['-'.join(idx) for idx in self.index]
        return getattr(self, '__init__')(self.sequence,
                                         path=self.path,
                                         data=self.as_matrix(),
                                         index=index, desc=self.desc,
                                         sym=self.sym, mtype=self.mtype,
                                         columns=columns, dtype=self.dtype)

    def copy(self, **kwargs):
        """
        Copy dataframe and related attributes of the map

        Parameters
        ----------
        **kwargs :
            

        Returns
        -------

        
        """
        df = super(Map, self).copy()
        return ProteinMap(
            sequence=self.sequence, path=self.path, data=df, desc=self.desc,
            sym=self.sym, mtype=self.mtype, dtype=self.dtype)

    def plotflush(self):
        """Flush contact map plot"""
        plt.clf()
        plt.cla()
        plt.close()
        self._maplot = None

    @property
    def maplot(self, linewidths=0.0):
        """
        Contact map plot
        :return:

        Parameters
        ----------
        linewidths :
            (Default value = 0.0)

        Returns
        -------

        
        """
        # Contact map Plot
        if not self._maplot:
            # Flush matplot
            LOG.debug("Build maplot")
            minticks = tickmin(self, shift=1)  # Nb graduations

            self._maplot = sns.heatmap(self, square=True, cbar=False,
                                       linewidths=linewidths, vmax=1, vmin=-1,
                                       cmap=sns.diverging_palette(20, 220, n=7,
                                                                  as_cmap=True),
                                       xticklabels=minticks[0],
                                       yticklabels=minticks[1],
                                       linecolor="grey")
        return self._maplot

    def saveplot(self, outdir='', outprefix="protein", size_fig=10,
                 plot_ext="pdf", plot_dpi=200):
        """
        Save plot

        Parameters
        ----------
        outdir :
            param outprefix: (Default value = '')
        size_fig :
            param plot_ext: (Default value = 10)
        plot_dpi :
            (Default value = 200)
        outprefix :
            (Default value = "protein")
        plot_ext :
            (Default value = "pdf")

        Returns
        -------


        """
        plotpath = os.path.join(outdir, "%s.maplot.%s" % (
            outprefix, plot_ext))
        LOG.info("Generate contact map plot (%s)", plotpath)
        f, ax = plt.subplots()
        tickrot(self.maplot.axes, self.maplot.figure)
        self.maplot.figure.set_size_inches(size_fig, size_fig)
        map_title = "%s contacts map" % outprefix
        self.maplot.set_title(map_title)
        self.maplot.figure.tight_layout()

        f.tight_layout()
        self.maplot.figure.savefig(plotpath, dpi=plot_dpi)
        plt.close('all')

    def contactset(self, human_idx=False):
        """
        Remove duplicate in contact_list

        Parameters
        ----------
        human_idx :
            return: (Default value = False)

        Returns
        -------
        list
        
        """
        contacts = self.contact_list(human_idx)
        if contacts:
            contacts = sorted(set((tuple(sorted((x, y)))
                                   for x, y in contacts)))
        return contacts

    def contact_list(self, human_idx=False):
        """
        Return contact list

        Parameters
        ----------
        human_idx :
            return: (Default value = False)

        Returns
        -------
        list
        
        """
        contact_list = []
        n = 1 if human_idx else 0
        if self.dtype is bool:
            # for irow, row in enumerate(self):
            for irow, row in enumerate(self.index):
                for icol, value in enumerate(self[row]):
                    if value:
                        contact_list.append((irow + n, icol + n))
        return contact_list

    def compareplot(self, protmap, save_fig=True, alpha=None, **kwargs):
        """
        Compare 2 contact map and plot the differences

        Parameters
        ----------
        protmap :
        save_fig :
        alpha :
        kwargs :

        Returns
        -------

        """
        self.plotflush()
        # Contact map plot
        if getattr(protmap, "shape") and self.shape != protmap.shape:
            logging.error("Cant't compare %s map with %s" % (
                protmap.__class__.__name__, self.__class__.__name__))
            return None
        else:
            cmplist = protmap.contact_list(human_idx=True)

            ymax = len(self.sequence) - 1
            if protmap.contact_flags:
                flags = set(protmap.contact_flags.values())
                # Color palette
                pal = sns.color_palette("hls", len(flags))
                for i, flag in enumerate(flags):
                    conlist = [contact for contact in protmap.contact_flags if
                               protmap.contact_flags[contact] == flag]
                    xind = [x - .5 for x in
                            zip(*conlist)[0] + zip(*conlist)[1]]
                    yind = [ymax - y + 1.5 for y in
                            zip(*conlist)[1] + zip(*conlist)[0]]
                    color = pal[i]
                    mark = Line2D.filled_markers[i]
                    for x, y in zip(xind, yind):
                        self.maplot.axes.scatter(x, y, s=8, c=color,
                                                 linewidths=0.1, alpha=alpha,
                                                 marker=mark)
            else:
                LOG.info("Contact list: %s", cmplist)
                xind = [x - .5 for x in
                        zip(*cmplist)[0] + zip(*cmplist)[1]]
                yind = [ymax - y + 1.5 for y in
                        zip(*cmplist)[1] + zip(*cmplist)[0]]
                LOG.debug("Xind: %s", xind)
                LOG.debug("Yind: %s", yind)
                color = "red"
                # width = [0.3 for _ in xind]
                # for x, y, h in zip(xind, yind, width):
                for x, y in zip(xind, yind):
                    self.maplot.axes.scatter(x, y, s=30, c=color,
                                             linewidths=0,
                                             alpha=alpha)
            if save_fig:
                self.saveplot(**kwargs)

    @staticmethod
    def _write_report(reportpath, **scores):
        """
        Write report file from score dict

        Parameters
        ----------
        reportpath :
            
        scores :
            
        **scores :
            

        Returns
        -------

        
        """
        with open(reportpath, 'w') as reportf:
            msg = string.Formatter().vformat("""\
## Report {map1name} vs. {map2name}
##
## Date: {date}
##
## Plots: {outdir}
## Reference map: {map1path}
## Contact map: {map2path}
##
## -----------------------------------------------------------------------------
##
## Sequence:
## {seq}
## Protein length: {protlen}
##
## -----------------------------------------------------------------------------
##
## Matthews correlation coefficient (MCC): {mcc}
## F1 score: {f1s}
## F2 score: {f2s}
## F0.5 score: {f05s}
##
## Precision: {precision}
## Recall (Sensibility): {recall}
## Accuracy: {accuracy}
##
## Hamming loss: {hamm}
## Hinge loss: {hin}
##
## -----------------------------------------------------------------------------
##
##                                  Plot scores
##
## ROC Area Under Curve: {roc_auc}
## Average precision score: {aver_prec}
##
## -----------------------------------------------------------------------------
##
##                          Precision recall curve
##
## Precision values:
## {allprec}
## Recall values:
## {allrec}
## Score tresholds ({map2name}):
## {prthres}
##
## -----------------------------------------------------------------------------
##
##                               ROC curve
##
## True Positive Rate (Sensibility) values:
## {alltpr}
## False Positive Rate (1 - Specificity) values:
## {allfpr}
## Score tresholds ({map2name}):
## {rocthres}""", (), defaultdict(str, **scores))
            LOG.debug("\n" + msg)
            reportf.write(msg)

    @staticmethod
    def classification_metrics(y_true, y_pred, y_scores=None):
        """
        Compute classification metrics

        Parameters
        ----------
        y_true : numpy array of true values
            
        y_pred : numpy array of predicted values
            
        y_scores : numpy array of prediction scores
            (Default value = None)

        Returns
        -------

        
        """
        metrics = {}

        if 1 in y_pred and 1 in y_true:
            metrics.update({
                'accuracy': skm.accuracy_score(y_true, y_pred),
                'precision': skm.precision_score(y_true, y_pred),
                'recall': skm.recall_score(y_true, y_pred),
                'mcc': skm.matthews_corrcoef(y_true, y_pred),
                'f1s': skm.f1_score(y_true, y_pred),
                'f2s': skm.fbeta_score(y_true, y_pred, 2),
                'f05s': skm.fbeta_score(y_true, y_pred, 0.5),
                'hamm': skm.hamming_loss(y_true, y_pred),
                'hin': skm.hinge_loss(y_true, y_pred)
            })

        if y_scores:
            # ROC plot
            # Replace nan values in y_scores
            y_scores[np.isnan(y_scores)] = np.floor(np.nanmin(y_scores))
            metrics.update({
                'roc_auc': skm.roc_auc_score(y_true, y_scores),
                'aver_prec': skm.average_precision_score(y_true, y_scores)
            })
            metrics['allfpr'], metrics['alltpr'], metrics['rocthres'] = \
                skm.roc_curve(y_true, y_scores, pos_label=1)
            metrics['allprec'], metrics['allrec'], metrics['prthres'] = \
                skm.precision_recall_curve(y_true, y_scores)

        return metrics

    @staticmethod
    def _metrics_plot(metrics, mapnames=('', ''), plotdir="", outprefix="",
                      plot_ext="pdf"):
        """
        

        Parameters
        ----------
        metrics :
            
        mapnames :
            (Default value = ('')
        '' :
            

        Returns
        -------

        
        """
        outprefix = outprefix if outprefix else "maplot"
        csv_roc = os.path.join(plotdir, "%s.roc.csv" % outprefix)
        LOG.info("Generate roc file (%s)", csv_roc)
        with open(csv_roc, "w") as f:
            f.write("TPR,FPR,Treshold\n")
            writer = csv.writer(f)
            writer.writerows(zip(metrics['alltpr'], metrics['allfpr'],
                                 metrics['rocthres']))

        plotpath = os.path.join(plotdir, "%s.roc.%s" % (outprefix,
                                                        plot_ext))
        LOG.info("Generate roc plot (%s)", plotpath)
        plt.figure()
        plt.plot(metrics['allfpr'], metrics['alltpr'],
                 label='ROC curve (area = %0.2f)' % metrics['roc_auc'])
        plt.plot([0, 1], [0, 1], 'k--')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.0])
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title('Receiver operating characteristic %s vs. %s' % (
            mapnames[0], mapnames[1]))
        plt.legend(loc="lower right")
        plt.savefig(plotpath)

        csv_precall = os.path.join(plotdir, "%s.roc.csv" % outprefix)
        LOG.info("Generate precall file (%s)", csv_precall)
        with open(csv_precall, "w") as f:
            f.write("Precision,Recall,Treshold\n")
            writer = csv.writer(f)
            writer.writerows(zip(metrics['allprec'], metrics['allrec'],
                                 metrics['prthres']))

        plotpath = os.path.join(plotdir, "%s.precall.%s" % (outprefix,
                                                            plot_ext))
        LOG.info("Generate precall plot (%s)", plotpath)
        # Precision recall curve
        plt.clf()
        plt.plot(metrics['allrec'], metrics['allprec'],
                 label='Precision-Recall curve')
        plt.xlabel('Recall')
        plt.ylabel('Precision')
        plt.ylim([0.0, 1.0])
        plt.xlim([0.0, 1.0])
        plt.title('Precision-Recall {1} vs. {2}: AUC={0:0.2f}'.format(
            metrics['aver_prec'], mapnames[0], mapnames[1]))
        plt.legend(loc="lower left")
        plt.savefig(plotpath)

    def report(self, cmpmap, scoremap=None, outprefix="", outdir="", plotdir="",
               # plot_ext="pdf", plotag=True, n_factors=(0.5, 1.0, 1.5)):
               plot_ext="pdf", plotag=True):
        """
        Generate contact map report file

        Parameters
        ----------
        cmpmap :
            
        scoremap :
            (Default value = None)
        outprefix :
            (Default value = "")
        outdir :
            (Default value = "")
        plotdir :
            (Default value = "")
        plot_ext :
            (Default value = "pdf")
        plotag :
            (Default value = True)

        Returns
        -------

        """
        filename = ".".join((outprefix, "mapreport")) if outprefix else \
            "mapreport"
        reportpath = "%s/%s" % (outdir, filename)
        LOG.info("Generate map report file (%s)", reportpath)

        # TODO: evaluate map statistics for several treshold (use submap var)
        # nb_c = int(len(self.sequence) * float(n_factors[0]))
        # submap = cmpmap.copy()
        # print(np.array_equal(cmpmap.values.astype(int).flat,
        #                      submap.topmap(scoremap, nb_c).values.astype(int).flat))

        y_true = list(self.values.astype(int).flat)
        y_pred = list(cmpmap.values.astype(int).flat)
        y_scores = scoremap.values.astype(float).flat if scoremap is not None else None

        # Compute classification metrics for the entire map
        metrics = self.classification_metrics(y_true, y_pred, y_scores)

        self._write_report(
            reportpath, map1name=self.desc, map2name=cmpmap.desc,
            map1path=self.path, map2path=cmpmap.path,
            seq="\n## ".join(textwrap.wrap(self.sequence, width=77)),
            protlen=len(self.sequence),
            date=datetime.date.today().strftime("%A %d. %B %Y"),
            outdir=outdir, **metrics)

        if plotag and scoremap is not None:
            self._metrics_plot(metrics, mapnames=(self.desc, cmpmap.desc),
                               outprefix=outprefix, plot_ext=plot_ext,
                               plotdir=plotdir)

    def compare_contactmap(self, cmpmap, contactlist, outprefix,
                           outdir="", distmap=None, human_idx=True):
        """
        Compare 2 contact map and plot differences

        Parameters
        ----------
        cmpmap :
            param contactlist:
        outprefix :
            param outdir:
        distmap :
            param human_idx: (Default value = None)
        contactlist :
            
        outdir :
            (Default value = "")
        human_idx :
            (Default value = True)

        Returns
        -------

        
        """
        # CSV file giving TP/FP contacts
        outpath = "%s/%s.contactcmp.csv" % (outdir,
                                            outprefix if outprefix else "maplot")
        LOG.info("Generate stat file (%s)", outpath)
        with open(outpath, 'w') as outfile:
            offset = 1 if human_idx else 0
            extra_header = "" if distmap is None else ",dmin"
            outfile.write("#resid1,resid2,res1,res2,TP/FP%s\n" % extra_header)
            for x in contactlist:
                if extra_header:
                    dmin = "," + str(distmap.iat[(int(x[0]) - offset, int(x[1]) -
                                                  offset)])
                else:
                    dmin = ""
                contact = self.iat[(int(x[0]) - offset, int(x[1]) - offset)]
                cmpcontact = cmpmap.iat[(int(x[0]) - offset, int(x[1]) - offset)]
                if contact and cmpcontact:
                    eq = "TP"
                elif contact and not cmpcontact:
                    eq = "FP"
                elif not contact and cmpcontact:
                    eq = "FN"
                else:
                    eq = "TN"
                msg = "%s,%s,%s,%s,%s%s\n" % (x[0], x[1],
                                            self.sequence[int(x[0]) - offset],
                                            self.sequence[int(x[1]) - offset],
                                            eq,
                                            dmin)
                outfile.write(msg)

    def write_contacts(self, filename, outdir="", prefix="", human_idx=True,
                       scoremap=None):
        """
        

        Parameters
        ----------
        prefix
        filename :
            param outdir:
        human_idx :
            param scoremap: (Default value = True)
        outdir :
            (Default value = "")
        scoremap :
            (Default value = None)

        Returns
        -------

        
        """
        prefix = prefix + "_" if prefix else prefix
        filepath = "%s/%s%s.contact.txt" % (outdir, prefix, filename.replace(".", "_"))
        LOG.info("Generate contact file (%s)", filepath)
        with open(filepath, 'w') as outfile:
            offset = 1 if human_idx else 0
            # contacts = [sorted(contact) for contact in self.contactset()]
            # for contact in sorted(contacts):
            for contact in self.contactset():
                contact = sorted([contact[0] + offset, contact[1] + offset])
                if scoremap is not None:
                    score = scoremap.iat[(int(contact[0]) - offset,
                                          int(contact[1]) - offset)]
                    outfile.write("%d %d %.4f\n" %
                                  (contact[0], contact[1], score))
                else:
                    outfile.write("%d %d\n" % (contact[0], contact[1]))


class ResAtmMap(ProteinMap):
    """
    Protein distance/contact matrix for all atom pairs. If no sequence given,
    protein distance/contact matrix for all amino acids

    Attributes
    ----------

    Examples
    --------

    """

    # def _constructor_expanddim(self):
    #     super(ResAtmMap, self)._constructor_expanddim()

    def __init__(self, sequence, flaglist=None, atom_groups="min", **kwargs):
        # Sequence: 1L string or MultiIndex object
        # Dataframe is in 3L code
        if not sequence:
            sequence = ConversionTable.ConversionTable().table['AMINO_ACID'][
                'iupac'].keys()
        # Super call will initialize index and columns with self.create_index()