source: trunk/mzcms/models.py

"""Models for mzcms.

"""
import time
from datetime import datetime
import textwrap
from itertools import combinations

from persistent import Persistent
from repoze.folder import Folder
from repoze.catalog.catalog import Catalog
from repoze.catalog.document import DocumentMap
from repoze.catalog.indexes.field import CatalogFieldIndex
from repoze.catalog.indexes.path2 import CatalogPathIndex2
from repoze.bfg.traversal import model_path

from parsers import parse_dats
from msquant import parse_msq_files


# XXX: move constant masses to another module
PROTON = 1.00727646677

class Experiment(Folder):
    """A proteomics experiment"""
    __parent__ = __name__ = None
    def __init__(self):
        self.catalog = Catalog()
        self.catalog['max_score'] = CatalogFieldIndex(get_max_score)
        self.catalog['path'] = CatalogPathIndex2(get_path)
        self.catalog.document_map = DocumentMap()
        super(Experiment, self).__init__()

    def update_indexes(self):
        catalog = self.catalog

        for data in self.values():
            docid = catalog.document_map.add(model_path(data))
            catalog.index_doc(docid, data)

        for prot in self['proteins'].values():
            docid = catalog.document_map.add(model_path(prot))
            catalog.index_doc(docid, prot)

        for pep in self['peptides'].values():
            docid = catalog.document_map.add(model_path(pep))
            catalog.index_doc(docid, pep)

# Test of catalog for peptides. This will be moved somewhere else.
def get_max_score(object, default):
    return getattr(object, 'max_score', float())

def get_path(object, default):
    return model_path(object)

class Proteins(Folder):
    """A protein container"""
    def update(self, protein_container):
        """Updates the protein container so that if the protein is
           already in the container only the peptide references are
           added.
        """
        for prot_id, prot in protein_container.items():
            if prot_id in self:
                # XXX: Test this is done correctly
                self[prot_id].pep_refs.update(prot.pep_refs)
            else:
                self[prot_id] = prot

class Protein(Persistent):
    """A Protein"""
    def __init__(self, native_id, sequence, pep_refs,
                 ext_ids=None, description=None, ratios=None,
                 max_score=None, coverage=None, covered_seq=None):
        self.native_id = str(native_id)
        self.sequence = str(sequence)
        self.ext_ids = ext_ids
        self.description = description
        self.pep_refs = pep_refs
        self.ratios = ratios
        self.coverage = coverage
        self.covered_seq = None

    def update_coverage(self):
        cov_seq = self.sequence[:]
        for pep in self.pep_refs:
            new_seq = list()
            seq = self.sequence.replace(pep.naked_seq,
                    'x' * len(pep.naked_seq))
            for s, c in zip(seq, cov_seq):
                if s == 'x':
                    new_seq.append('x')
                else:
                    new_seq.append(c)
            cov_seq = ''.join(new_seq)
        self.covered_seq = cov_seq

    @property
    def coverage_seq(self):
        coverage = float(self.covered_seq.count('x'))/len(self.sequence)
        return round(coverage * 100, 2)

    def update_pep_num(self):
        self.pep_num = len(self.pep_refs)

    def update_psm_num(self):
        psm_count = int()
        for pep in self.pep_refs:
            psm_count += pep.psm_num
        self.psm_num = psm_count

    def formatted_seq(self):
        # TODO: Make sure this doesn't screw up in windows
        # Maybe this is better in the view
        return '\n'.join(textwrap.wrap(self.sequence))

    def update_ratios(self):
        ratios = list()
        for pep in self.pep_refs:
            for psm in pep.psm_refs:
                if psm.ratios:
                    ratios.append(psm.ratios)
        if len(ratios) > 0:
            il_ratio = sum(r[0] for r in ratios)/len(ratios)
            lh_ratio = sum(r[1] for r in ratios)/len(ratios)
            hi_ratio = sum(r[2] for r in ratios)/len(ratios)
            self.ratios = il_ratio, lh_ratio, hi_ratio

    def update_max_score(self):
        """Checks what is the maximum score in all peptide references.
        """
        max_score = float()
        for pep in self.pep_refs:
            if pep.max_score > max_score:
                max_score = pep.max_score
        self.max_score = max_score

class Peptide(Persistent):
    """A peptide"""
    def __init__(self, sequence, mass, psm_refs=None, max_score=None):
        self.sequence = tuple(sequence)
        self.naked_seq = ''.join((x[0] for x in sequence[1:-1]))
        self.mass = float(mass)
        self.length = len(sequence)
        self.psm_refs = psm_refs
        self.max_score = max_score

    def update_psm_num(self):
        self.psm_num = len(self.psm_refs)

    def update_max_score(self):
        self.max_score = max((p.score for p in self.psm_refs))

    def formatted_seq(self):
        seq = str()
        for aa, m in self.sequence:
            if m:
                f_aa = ''.join((m[1], aa))
            else:
                f_aa = aa
            seq = ''.join((seq, f_aa))
        return seq

    def was_quanted(self):
        for psm in self.psm_refs:
            if len(psm.quant_feat) >= 1:
                return True
        return False

    def formatted_ptms(self):
        return ', '.join((x[1][1] for x in self.sequence
                                    if x[1] is not None))

class Spectrum(Persistent):
    """A Spectrum"""
    def __init__(self, scan, run, peaks, prec_mz, prec_charge,
                 psm_refs=None):
        self.scan = int(scan)
        self.run = str(run)
        self.peaks = tuple(peaks)
        self.prec_mz = float(prec_mz)
        self.prec_charge = int(prec_charge)
        self.psm_refs = psm_refs or set()

class Psm(Persistent):
    """A Peptide to Spectrum Match"""
    def __init__(self, score, rank, delta_mass, pep_ref, spec_ref,
                 params_ref=None, mme=None, quant_feat=None, xics=None,
                 ratios=None):
        self.score = float(score)
        self.rank = int(rank)
        self.delta_mass = float(delta_mass)
        # mass measurement error
        self.mme = mme
        #self.delta_score = float(delta_score)
        #self.sequence_rank2 = str(sequence_rank2)
        self.pep_ref = pep_ref
        self.spec_ref = spec_ref
        self.params_ref = params_ref
        self.quant_feat = quant_feat or list()
        self.xics = xics
        self.ratios = ratios

    def update_mme(self):
        charge= self.spec_ref.prec_charge 
        theo_mass = self.pep_ref.mass
        theo_mz = (theo_mass + charge * PROTON) / charge
        exp_mz = self.spec_ref.prec_mz
        ratio = exp_mz/theo_mz
        if ratio > 1.0:
            mme = 10**6 * (ratio - 1.0)
        elif ratio < 1.0:
            mme = -10**6 * (1.0 - ratio)
        elif ratio == 1.0:
            mme = 0
        self.mme = mme

    def was_quanted(self):
        if len(self.quant_feat) >= 1:
            return True
        return False

    def update_ratios(self):
        ratios = list()
        if self.xics:
            for comb in combinations(self.xics, 2):
                try:
                # IL, HL, HI
                    ratios.append(comb[1]/comb[0])
                except ZeroDivisionError:
                    ratios.append(float())
            self.ratios = tuple(ratios)

class Ptm(Persistent):
    """A Post-tranlational modification"""
    def __init__(self, mass, name):
        self.name = name
        self.mass = float(mass)

class SerParameters(Persistent):
    """A set of Search Engine Results parameters"""
    def __init__(self,
                 frag_mode='CID',
                 db_name=None,
                 pkl_fn=None,
                 enzyme='Trypsin',
                 prec_tol=(0.6, 'Da'),
                 frag_tol=(50.0, 'ppm'),
                 quantitation=None,
                 program='Mascot',
                 version='2.2.04',
                 date=time.time(),
                 ):
        self.frag_mode=frag_mode
        self.db_name = db_name
        self.pkl_fn = pkl_fn
        self.enzyme = enzyme
        self.prec_tol = prec_tol
        self.frag_tol = frag_tol
        self.quantitation = quantitation
        self.program = program
        self.version = version
        self.date = datetime.fromtimestamp(date)

def update_psms(psms):
    for psm in psms.values():
        psm.update_mme()
        psm.update_ratios()

def update_peptides(peptides):
    """Updates the number of psms in all peptides"""
    for peptide in peptides.values():
      peptide.update_psm_num()
      peptide.update_max_score()


#TODO: Move this to another module
def update_proteins(proteins):
    """Updates the number of psms and peptides in all proteins.
    """
    for protein in proteins.values():
      protein.update_pep_num()
      protein.update_psm_num()
      protein.update_coverage()
      protein.update_ratios()
      protein.update_max_score()

def appmaker(zodb_root):
    if not 'app_root' in zodb_root:
        p_t = time.time()
        app_root = Experiment()
        proteins = Proteins()
        peptides = Folder()
        spectra = Folder()
        psms = Folder()
        zodb_root['app_root'] = app_root
        app_root['proteins'] = proteins
        app_root['peptides'] = peptides
        app_root['spectra'] = spectra
        app_root['psms'] = psms

        parse_dats('./min-dats',
                   proteins_container=proteins,
                   peptides_container=peptides,
                   spectra_container=spectra,
                   psms_container=psms,
                   protein_factory=Protein,
                   peptide_factory=Peptide,
                   spectrum_factory=Spectrum,
                   psm_factory=Psm,
                   serparam_factory=SerParameters,
                   )
        import transaction
        transaction.commit()
        m_t = time.time()
        print(u"Time to parse dat files {0} secs.".format((m_t - p_t)))
        p_t = m_t

        parse_msq_files(u"./min-dats", psms_container=psms,
                        spectra_container=spectra)
        transaction.commit()
        m_t = time.time()
        print(u"Time to parse MsQuant output: {0} secs.".format((m_t - p_t)))
        p_t = m_t

        update_psms(psms)
        transaction.commit()
        m_t = time.time()
        print(u"Time to update PSMs: {0} secs.".format((m_t - p_t)))
        p_t = m_t

        update_peptides(peptides)
        transaction.commit()
        m_t = time.time()
        print(u"Time to update peptides: {0} secs.".format((m_t - p_t)))
        p_t = m_t

        update_proteins(proteins)
        transaction.commit()
        m_t = time.time()
        print(u"Time to update proteins: {0} secs.".format((m_t - p_t)))
        p_t = m_t

        app_root.update_indexes()
        transaction.commit()
        m_t = time.time()
        print(u"Time to update indexes: {0} secs.".format((m_t - p_t)))
        p_t = m_t

        print(u"Finished building")

    return zodb_root['app_root']