maplot.py

# coding=utf-8
"""
                            Input/Output aria_ec scripts
"""
from __future__ import absolute_import, division, print_function

import sys
import json
import logging
from .base import get_filename
from .reader import MapFileListReader
from .protmap import MapFilter
from .protein import Protein

LOG = logging.getLogger(__name__)


class AriaEcContactMap(object):
    """
    Contact maplot class
    """

    def __init__(self, settings):
        # TODO: check_type settings (AriaEcSettings)
        self.settings = settings
        self.protein = Protein(settings)
        self.file_reader = MapFileListReader(
            cont_def=settings.contactdef.config)
        self.filter = MapFilter(settings.setup.config)
        self.outprefix = ''
        self.allresmap = {}
        self.refmap = None
        self.refname = None
        self.reftype = ''

    def run(self):
        # Check input
        """
        Main method
        """
        LOG.debug("Settings:\n" + json.dumps(self.settings.contactmap.config,
                                             indent=4))
        LOG.debug("Args:\n" + json.dumps(self.settings.contactmap.args,
                                         indent=4))
        if not self.settings.contactmap.args.get("onlyreport", False):
            self.settings.make_infra()
        # ----------------------------- Input -------------------------------- #
        if self.settings.contactmap.args.get("prefix") and len(
                self.settings.contactmap.args.get("infiles")) == len(
                    self.settings.contactmap.args.get("prefix")):
                self.outprefix = self.settings.contactmap.args.get("prefix", "")
        else:
            self.outprefix = get_filename(self.settings.contactmap.args.get("seq",
                                                                            None))
        # Load Sequence file
        self.protein.set_aa_sequence(self.settings.contactmap.args.get("seq", None))
        # Load secondary structure prediction file
        self.protein.set_sec_struct(self.settings.contactmap.args.get("sspred",
                                                                      None),
                                    ssdist_filename=self.settings.ssdist,
                                    ssidx=self.settings.contactmap.args.get(
                                        "ssidx", False))
        # ---------------------------- Processing ---------------------------- #
        self.file_reader.read(self.settings.contactmap.args.get("infiles"),
                              filetypelist=self.settings.contactmap.args.get("contact_types"),
                              protein=self.protein,
                              groupby_method=self.settings.setup.config[
                                  'groupby_method'],
                              scsc=self.settings.scsc_min)

        for idx, fo in enumerate(self.file_reader.filelist):
            # fo need a contactmap in order to wite XML dist restraints
            # TODO: filter pour toutes les map de mapdict !! (fonction remove
            #  s'applique sur l'humanidx contenant les residus)
            if idx == 0:
                LOG.info("%s map set as reference", fo.filetype.capitalize())
                self.refmap = fo.mapdict
                self.reftype = fo.filetype
                self.refname = fo.filename if type(self.outprefix) != list \
                    else self.outprefix[idx]
            if (self.settings.contactmap.args.get("filter") and idx != 0) or \
                    (self.settings.contactmap.args.get("filter")
                     and len(self.file_reader.filelist) == 1):
                # Filtering all contactmap except the reference map
                LOG.info("Filtering %s map", fo.filetype)
                self.filter(fo.mapdict, fo.filetype, fo.contactlist,
                            self.protein, clashlist=fo.clashlist,
                            outprefix=self.outprefix[idx] if type(
                                self.outprefix) == list else self.outprefix,
                            outdir=self.settings.outdir)
            # else:
            #     Use only position filter
            #     self.filter(fo.mapdict, fo.filetype, fo.contactlist,
            #                 self.protein, clashlist=fo.clashlist,
            #                 outprefix=self.outprefix,
            #                 outdir=self.settings.outdir, mapfilters="pos")
            self.allresmap[(fo.filename if type(self.outprefix) != list else
                            self.outprefix[idx], fo.filetype,
                            fo.filepath)] = fo.mapdict

        try:
            refmap = self.refmap["contactmap"]
        except TypeError:
            LOG.error("First contact map should be a valid file")
            sys.exit(1)

        # nb_c = int(len(self.protein.aa_sequence.sequence) *
        #            self.settings.setup.config.get("n_factor"))
        plotparams = {k: self.settings.contactmap.config.get(k, None)
                      for k in ('size_fig', 'plot_ext', 'plot_dpi')}
        outdir = self.settings.outdir

        if self.settings.contactmap.config.get("save_fig") and not \
                self.settings.contactmap.args.get("onlyreport", False):
            refmap.saveplot(outdir=outdir,
                            outprefix=self.outprefix[0] if type(
                                self.outprefix) == list else self.outprefix,
                            **plotparams)

        if self.settings.contactmap.args.get("merge", None):
            # Combine merge maps with other maps
            mergelist = self.settings.contactmap.args.get("merge")
            for mergetype in mergelist:
                if mergetype in zip(self.allresmap.keys())[0]:
                    mergemaps = self.allresmap.pop(mergetype)
                    mergecontactmap = mergemaps.get("contactmap")
                    for mapname, mapt in self.allresmap.keys():
                        if mapt != self.reftype:
                            # TODO: DON'T WORK !!!!
                            LOG.info("Merging %s with %s map",
                                     mergetype, mapt)
                            up_map = self.allresmap[mapt]["contactmap"]
                            up_map[:] = up_map[:] + mergecontactmap[:]
                            mergekey = "%s_%s" % (mapt, mergetype)
                            self.allresmap[mergekey] = {}
                            self.allresmap[mergekey]["contactmap"] = up_map

        for mapname, mapt, mapath in self.allresmap.keys():

            prefix = "_".join((mapname, self.refname)).replace(".", "_")

            if mapname == self.refname:
                if self.settings.contactmap.args.get("onlyreport", False) is not False:
                    refmap.write_contacts(mapname,
                                          outdir=outdir,
                                          scoremap=self.refmap.get("scoremap",
                                                                   None))
                continue

            scoremap = self.allresmap[(mapname, mapt, mapath)].get(
                'scoremap', None)
            # if self.allresmap[mapt].get("contactmap") is not None and \
            #         self.allresmap[mapt].get("scoremap") is not None:
            #     Get top contact map/list
            #     cmpmap = self.allresmap[mapt]["contactmap"].topmap(
            #         self.allresmap[mapt]["scoremap"], nb_c)
            #     cmplist = self.allresmap[mapt]['scoremap'].sortedset(
            #         human_idx=True)[:nb_c]
            # elif self.allresmap[mapt].get("contactmap") is not None:
            #     If no score given, use all contact list
            cmpmap = self.allresmap[(mapname, mapt, mapath)]["contactmap"]
            cmplist = self.allresmap[(mapname, mapt, mapath)][
                'contactmap'].contact_list(human_idx=True)
            # else:
            #     LOG.warning("%s map doesn't have any score related. Can't "
            #                    "define top list related to this map", mapt)
            #     continue

            # TODO: only one function for output files
            # Write contact list in txt file
            refmap.report(cmpmap, scoremap=scoremap,
                          outdir=outdir,
                          plotag=not self.settings.contactmap.args.get(
                              "onlyreport"),
                          plotdir=self.settings.infra.get("graphics", outdir)
                          if not self.settings.contactmap.args.get(
                              "onlyreport") else self.settings.outdir)
            # Contact map comparison plot
            # TODO: elementwise error with compare method
            # Write cmp stats
            if not self.settings.contactmap.args.get("onlyreport", False):
                cmpmap.write_contacts(mapname,
                                      scoremap=scoremap,
                                      outdir=outdir)
                cmpmap.compare_contactmap(refmap, cmplist, prefix,
                                          distmap=self.refmap["distmap"],
                                          human_idx=True,
                                          outdir=outdir)
                refmap.compareplot(cmpmap, outprefix=prefix,
                                   outdir=outdir,
                                   save_fig=self.settings.contactmap.config.get(
                                       "save_fig"),
                                   alpha=self.settings.contactmap.config.get(
                                       "alpha"),
                                   **plotparams)
                # Contingency table
                # print(cmpmap.to_series())
                # LOG.info(pd.crosstab(cmpmap.values, refmap.values,
                #                         rownames=[mapt], colnames=[self.reftype]))