<commit_before># vim: fdm=marker
'''
author:     <NAME>/<NAME>
date:       25/04/2015
content:    Data access module HIV patients.
'''
# Modules
import numpy as np
import pandas as pd
from Bio import SeqIO,Seq
from .samples import *
from .af_tools import *
from .sequence import alpha, alphaa



# Classes
class Patient(pd.Series):
    '''
    Class providing access to longitudinal sequencing data of HIV-1 populations
    in participants of the HIVEVO study. The class contains time-ordered samples
    and access methods to single nucleotide variants, pair frequencies, and genomic
    features of the the HIV poputions
    '''

    def __init__(self, *args, **kwargs):
        '''Initialize a patient with all his samples'''
        include_cell = kwargs.pop('include_cell', False)
        super(Patient, self).__init__(*args, **kwargs)
        self.samples = sorted(load_samples_sequenced(patients=[self.name]), 
                              key = lambda x:x.date)
        self._dates = [x.date for x in self.samples]
        self._cd4 = [x['CD4+ count'] for x in self.samples]
        self._viral_load = [x['viral load'] for x in self.samples]

        # We take 400 ul of serum
        # We typically have 6 reactions with that total volume (plus the F4 dilution
        # series, but each of those uses only 0.1x template which is very little)
        self._n_templates_viral_load = np.array([x*0.4/6.1 for x in self._viral_load], dtype = float)
        self._n_templates_dilutions = np.ma.masked_invalid([x.get_n_templates_dilutions() for x in self.samples])
        self._times = []
        self.reference = self.load_reference()
        # translate genbank encoded sequence features into a dictionary
        self.annotation = {x.qualifiers['note'][-1]:x for x in self.reference.features}
        self._initial_consensus_noinsertions()
        self.positions_to_features()


    @classmethod
    def load(cls, pname):
        from .filenames import get_table_filename
        patients = pd.read_excel(get_table_filename('patients'),
                                 'Patients',
                                 index_col=0)
        patients.index = pd.Index(map(str, patients.index))
        if pname in patients.index:
            return cls(patients.loc[pname])
        else:
            return cls(patients.loc[patients.code == pname].iloc[0])


    @property
    def _constructor(self):
        return Patient


    @property
    def transmission_date(self):
        '''The most likely time of transmission'''
        return self['infect date best']


    @property
    def dates(self):
        return self._dates


    @property
    def viral_load(self):
        return self._viral_load


    @property
    def cd4(self):
        return self._cd4


    @property
    def dsi(self):
        return self.times(unit='days')


    @property
    def msi(self):
        return self.times(unit='month')


    @property
    def ysi(self):
        return self.times(unit='year')
    
    
    def times(self, unit='days'):
        '''Get the times from transmission'''
        delta_days = [float(s['days since infection']) for s in self.samples]
        if unit.startswith('day'):
            return np.array(delta_days, dtype=float)
        elif unit.startswith('year'):
            return np.array(delta_days, dtype=float)/365.24
        elif unit.startswith('month'):
            return np.array(delta_days, dtype=float)/365.24*12
        else:
            raise ValueError("bad time unit")


    @property
    def n_templates_dilutions(self):
        '''Get the time course of the number of templates to PCR, limiting depth'''
        return self._n_templates_dilutions


    @property
    def n_templates_viral_load(self):
        '''Get the number of templates, estimated from the viral load'''
        return self._n_templates_viral_load


    @property
    def initial_sample(self):
        '''The initial sample used as a mapping reference'''
        return self.samples[0]


    def load_reference(self):
        from .filenames import get_initial_reference_filename
        return SeqIO.read(get_initial_reference_filename(self.name, "genomewide", format='gb'), 'gb')


    def _region_to_indices(self,region):
        '''returns a list of positions corresponding to a genomic region'''
        if region=='genomewide':
            return np.arange(len(self.reference))
        elif region in self.annotation:
            return np.array([int(x) for x in self.annotation[region]], dtype=int)
        else:
            raise ValueError('no annotation with name '+region)


    def _annotation_to_fragment_indices(self, anno):
        '''
        returns coordinates of a region specified in the annotation
        in terms of the fragments F1 to F5. This is needed to extract 
        region specific allele counts, frequencies etc.
        returns a dict containing 'length', 'start' (of the region in the genome)
        and for each fragment involved 'F1': (indices in the region of interest, indices on the fragment)
        '''
        coordinates = {}
        region_indices = self._region_to_indices(anno)
        coordinates['start'] = min(region_indices)
        coordinates['length'] = len(region_indices)
        fragments = ['F'+str(i) for i in xrange(1,7)]
        if anno not in fragments:
            for frag in fragments: # loop over fragments and extract the indices of the region on this fragment
                frag_ind = set(self._region_to_indices(frag))       # indices of the fragment
                region_indices_on_fragment = sorted(frag_ind.intersection(region_indices))  # intersection of region and fragment positions
                if len(region_indices_on_fragment): # attach indices in region and on fragment
                    anno_indices_self = np.arange(coordinates['length'])[np.in1d(region_indices, region_indices_on_fragment)]
                    coordinates[frag] = (anno_indices_self, 
                                       np.array(region_indices_on_fragment)- int(self.annotation[frag].location.start))
        else: # if requested region is a fragment, return only this fragment
            coordinates[anno] = (np.arange(coordinates['length']), np.arange(coordinates['length']))
        return coordinates


    def get_coverage_trajectories(self, region, **kwargs):
        '''Get coverage as a function of time'''
        coordinates = self._annotation_to_fragment_indices(region)
        cov = np.ma.array([tmp_sample.get_coverage(coordinates, **kwargs) for tmp_sample in self.samples])
        return cov


    def get_allele_count_trajectories(self, region, safe=False, type='nuc', **kwargs):
        '''Get the allele count trajectories from files
        
        Args:
          region (str): region to study, a fragment or a genomic feature (e.g. V3)
          type (str): 'nuc' for nucleotides, 'aa' for amino acids
          **kwargs: passed down to the function (VERBOSE, etc.).

        Note: the genomewide counts are currently saved to file.
        '''
        if type == 'nuc':
            coordinates = self._annotation_to_fragment_indices(region)
        elif type == 'aa':
            coordinates = {region: None}
        else:
            raise ValueError('Data type not understood')

        act = np.ma.array([tmp_sample.get_allele_counts(coordinates, type=type, **kwargs)
                           for tmp_sample in self.samples],
                          hard_mask=True,
                          shrink=False)
        if len(act.mask.shape) < 1:
            act.mask = np.zeros_like(act, dtype=bool)
        return act
<commit_msg>Amino acid count trajectories docs<commit_after># vim: fdm=marker
'''
author:     <NAME>/<NAME>
date:       25/04/2015
content:    Data access module HIV patients.
'''
# Modules
import numpy as np
import pandas as pd
from Bio import SeqIO,Seq
from .samples import *
from .af_tools import *
from .sequence import alpha, alphaa



# Classes
class Patient(pd.Series):
    '''
    Class providing access to longitudinal sequencing data of HIV-1 populations
    in participants of the HIVEVO study. The class contains time-ordered samples
    and access methods to single nucleotide variants, pair frequencies, and genomic
    features of the the HIV poputions
    '''

    def __init__(self, *args, **kwargs):
        '''Initialize a patient with all his samples'''
        include_cell = kwargs.pop('include_cell', False)
        super(Patient, self).__init__(*args, **kwargs)
        self.samples = sorted(load_samples_sequenced(patients=[self.name]), 
                              key = lambda x:x.date)
        self._dates = [x.date for x in self.samples]
        self._cd4 = [x['CD4+ count'] for x in self.samples]
        self._viral_load = [x['viral load'] for x in self.samples]

        # We take 400 ul of serum
        # We typically have 6 reactions with that total volume (plus the F4 dilution
        # series, but each of those uses only 0.1x template which is very little)
        self._n_templates_viral_load = np.array([x*0.4/6.1 for x in self._viral_load], dtype = float)
        self._n_templates_dilutions = np.ma.masked_invalid([x.get_n_templates_dilutions() for x in self.samples])
        self._times = []
        self.reference = self.load_reference()
        # translate genbank encoded sequence features into a dictionary
        self.annotation = {x.qualifiers['note'][-1]:x for x in self.reference.features}
        self._initial_consensus_noinsertions()
        self.positions_to_features()


    @classmethod
    def load(cls, pname):
        from .filenames import get_table_filename
        patients = pd.read_excel(get_table_filename('patients'),
                                 'Patients',
                                 index_col=0)
        patients.index = pd.Index(map(str, patients.index))
        if pname in patients.index:
            return cls(patients.loc[pname])
        else:
            return cls(patients.loc[patients.code == pname].iloc[0])


    @property
    def _constructor(self):
        return Patient


    @property
    def transmission_date(self):
        '''The most likely time of transmission'''
        return self['infect date best']


    @property
    def dates(self):
        return self._dates


    @property
    def viral_load(self):
        return self._viral_load


    @property
    def cd4(self):
        return self._cd4


    @property
    def dsi(self):
        return self.times(unit='days')


    @property
    def msi(self):
        return self.times(unit='month')


    @property
    def ysi(self):
        return self.times(unit='year')
    
    
    def times(self, unit='days'):
        '''Get the times from transmission'''
        delta_days = [float(s['days since infection']) for s in self.samples]
        if unit.startswith('day'):
            return np.array(delta_days, dtype=float)
        elif unit.startswith('year'):
            return np.array(delta_days, dtype=float)/365.24
        elif unit.startswith('month'):
            return np.array(delta_days, dtype=float)/365.24*12
        else:
            raise ValueError("bad time unit")


    @property
    def n_templates_dilutions(self):
        '''Get the time course of the number of templates to PCR, limiting depth'''
        return self._n_templates_dilutions


    @property
    def n_templates_viral_load(self):
        '''Get the number of templates, estimated from the viral load'''
        return self._n_templates_viral_load


    @property
    def initial_sample(self):
        '''The initial sample used as a mapping reference'''
        return self.samples[0]


    def load_reference(self):
        from .filenames import get_initial_reference_filename
        return SeqIO.read(get_initial_reference_filename(self.name, "genomewide", format='gb'), 'gb')


    def _region_to_indices(self,region):
        '''returns a list of positions corresponding to a genomic region'''
        if region=='genomewide':
            return np.arange(len(self.reference))
        elif region in self.annotation:
            return np.array([int(x) for x in self.annotation[region]], dtype=int)
        else:
            raise ValueError('no annotation with name '+region)


    def _annotation_to_fragment_indices(self, anno):
        '''
        returns coordinates of a region specified in the annotation
        in terms of the fragments F1 to F5. This is needed to extract 
        region specific allele counts, frequencies etc.
        returns a dict containing 'length', 'start' (of the region in the genome)
        and for each fragment involved 'F1': (indices in the region of interest, indices on the fragment)
        '''
        coordinates = {}
        region_indices = self._region_to_indices(anno)
        coordinates['start'] = min(region_indices)
        coordinates['length'] = len(region_indices)
        fragments = ['F'+str(i) for i in xrange(1,7)]
        if anno not in fragments:
            for frag in fragments: # loop over fragments and extract the indices of the region on this fragment
                frag_ind = set(self._region_to_indices(frag))       # indices of the fragment
                region_indices_on_fragment = sorted(frag_ind.intersection(region_indices))  # intersection of region and fragment positions
                if len(region_indices_on_fragment): # attach indices in region and on fragment
                    anno_indices_self = np.arange(coordinates['length'])[np.in1d(region_indices, region_indices_on_fragment)]
                    coordinates[frag] = (anno_indices_self, 
                                       np.array(region_indices_on_fragment)- int(self.annotation[frag].location.start))
        else: # if requested region is a fragment, return only this fragment
            coordinates[anno] = (np.arange(coordinates['length']), np.arange(coordinates['length']))
        return coordinates


    def get_coverage_trajectories(self, region, **kwargs):
        '''Get coverage as a function of time'''
        coordinates = self._annotation_to_fragment_indices(region)
        cov = np.ma.array([tmp_sample.get_coverage(coordinates, **kwargs) for tmp_sample in self.samples])
        return cov


    def get_allele_count_trajectories(self, region, safe=False, type='nuc', **kwargs):
        '''Get the allele count trajectories from files
        
        Args:
          region (str): region to study, a fragment or a genomic feature (e.g. V3)
          type (str): 'nuc' for nucleotides, 'aa' for amino acids
          **kwargs: passed down to the function (VERBOSE, etc.).

        Note: the genomewide counts are currently saved to file.

        Examples:

        1. Nucleotides from a certain region:

        patient.get_allele_count_trajectories('PR', type='nuc')

        2. Amino acids from a certain region:

        patient.get_allele_count_trajectories('PR', type='aa')
        '''
        if type == 'nuc':
            coordinates = self._annotation_to_fragment_indices(region)
        elif type == 'aa':
            coordinates = {region: None}
        else:
            raise ValueError('Data type not understood')

        act = np.ma.array([tmp_sample.get_allele_counts(coordinates, type=type, **kwargs)
                           for tmp_sample in self.samples],
                          hard_mask=True,
                          shrink=False)
        if len(act.mask.shape) < 1:
            act.mask = np.zeros_like(act, dtype=bool)
        return act

{'fraction_non_alphanumeric': 0.09438824817988532, 'fraction_numerical': 0.00631402615810837, 'mean_word_length': 2.4975214060387563, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2312155', 'n_tokens_mistral': 4375, 'n_tokens_neox': 4144, 'n_words': 1292}