<commit_before>from __future__ import division, print_function
from collections import defaultdict
import sys
sys.path.append('')  # need to import from base
from base.process import process
import numpy as np
from datetime import datetime, timedelta
from base.io_util import myopen
from Bio import SeqIO
from glob import glob
import os

regions = ['africa', 'south_asia', 'europe', 'china', 'north_america',
           'china', 'south_america', 'japan_korea', 'oceania', 'southeast_asia']

region_groups = {'NA':'north_america',
                 'AS':['china', 'japan_korea', 'south_asia', 'southeast_asia'],
                 'OC':'oceania', 'EU':'europe'}

attribute_nesting = {'geographic location':['region', 'country', 'city'],}

color_options = {
    "country":{"key":"country", "legendTitle":"Country", "menuItem":"country", "type":"discrete"},
    "region":{"key":"region", "legendTitle":"Region", "menuItem":"region", "type":"discrete"},
    "num_date":{"key":"num_date", "legendTitle":"Sampling date", "menuItem":"date", "type":"continuous"},
    "ep":{"key":"ep", "legendTitle":"Epitope Mutations", "menuItem":"epitope mutations", "type":"continuous"},
    "ne":{"key":"ne", "legendTitle":"Non-epitope Mutations", "menuItem":"nonepitope mutations", "type":"continuous"},
    "rb":{"key":"rb", "legendTitle":"Receptor Binding Mutations", "menuItem":"RBS mutations", "type":"continuous"},
    "gt":{"key":"genotype", "legendTitle":"Genotype", "menuItem":"genotype", "type":"discrete"}}
panels = ['tree', 'entropy', 'frequencies']


class dengue_process(process):

    def __init__(self, *args, **kwargs):
        super(dengue_process, self).__init__(*args, **kwargs)
        print("Initializing dengue_process")
        self.titer_fname = kwargs['titer_fname']
        self.strains_fname = kwargs['strain_fname']


    def subsample(self, sampling_threshold, **kwargs):

        def sampling_category(x):
            '''
            Subsample per region, per month.
            '''
            return (x.attributes['region'],
                    x.attributes['date'].year,
                    x.attributes['date'].month)

        # load titer count to prioritize sequences
        titer_count = {}
        forced_strains = []
        with myopen(self.strains_fname,'r') as ifile:
            for line in ifile:
                strain, count = line.strip().split()
                titer_count[strain]=int(count)
                if strain.startswith('DENV%s'%serotype):
                    forced_strains.append(strain)

        def sampling_priority(seq):
            '''
            Prefers more titer measurements, longer sequences. Penalizes ambiguity codes.
            '''
            sname = seq.attributes['strain']
            if sname in titer_count:
                pr = titer_count[sname]
            else:
                pr = 0

            if seq.id in forced_strains:
                pr += 10.0

            return pr + len(seq.seq)*0.00005 - 0.01*np.sum([seq.seq.count(nuc) for nuc in 'NRWYMKSHBVD'])

        self.seqs.subsample(category = sampling_category,
                            threshold=sampling_threshold,
                            priority=sampling_priority,
                            forced_strains=forced_strains,
                            **kwargs)<commit_msg>this was missing in prev commit<commit_after>from __future__ import division, print_function
from collections import defaultdict
import sys
sys.path.append('')  # need to import from base
from base.process import process
import numpy as np
from datetime import datetime, timedelta
from base.io_util import myopen
from Bio import SeqIO
from glob import glob
import os

regions = ['africa', 'south_asia', 'europe', 'china', 'north_america',
           'china', 'south_america', 'japan_korea', 'oceania', 'southeast_asia']

region_groups = {'NA':'north_america',
                 'AS':['china', 'japan_korea', 'south_asia', 'southeast_asia'],
                 'OC':'oceania', 'EU':'europe'}

attribute_nesting = {'geographic location':['region', 'country', 'city'],}

color_options = {
    "country":{"key":"country", "legendTitle":"Country", "menuItem":"country", "type":"discrete"},
    "region":{"key":"region", "legendTitle":"Region", "menuItem":"region", "type":"discrete"},
    "num_date":{"key":"num_date", "legendTitle":"Sampling date", "menuItem":"date", "type":"continuous"},
    "ep":{"key":"ep", "legendTitle":"Epitope Mutations", "menuItem":"epitope mutations", "type":"continuous"},
    "ne":{"key":"ne", "legendTitle":"Non-epitope Mutations", "menuItem":"nonepitope mutations", "type":"continuous"},
    "rb":{"key":"rb", "legendTitle":"Receptor Binding Mutations", "menuItem":"RBS mutations", "type":"continuous"},
    "gt":{"key":"genotype", "legendTitle":"Genotype", "menuItem":"genotype", "type":"discrete"}}
panels = ['tree', 'entropy', 'frequencies']


class dengue_process(process):

    def __init__(self, *args, **kwargs):
        super(dengue_process, self).__init__(*args, **kwargs)
        print("Initializing dengue_process")
        self.titer_fname = kwargs['titer_fname']
        self.strains_fname = kwargs['strain_fname']


    def subsample(self, sampling_threshold, **kwargs):

        def sampling_category(x):
            '''
            Subsample per region, per month.
            '''
            return (x.attributes['region'],
                    x.attributes['date'].year,
                    x.attributes['date'].month)

        # load titer count to prioritize sequences
        titer_count = {}
        forced_strains = []
        with myopen(self.strains_fname,'r') as ifile:
            for line in ifile:
                strain, count = line.strip().split()
                titer_count[strain]=int(count)
                if strain.startswith('DENV%s'%serotype):
                    forced_strains.append(strain)

        def sampling_priority(seq):
            '''
            Prefers more titer measurements, longer sequences. Penalizes ambiguity codes.
            '''
            sname = seq.attributes['strain']
            if sname in titer_count:
                pr = titer_count[sname]
            else:
                pr = 0

            if seq.id in forced_strains:
                pr += 10.0

            return pr + len(seq.seq)*0.00005 - 0.01*np.sum([seq.seq.count(nuc) for nuc in 'NRWYMKSHBVD'])

        self.seqs.subsample(category = sampling_category,
                            threshold=sampling_threshold,
                            priority=sampling_priority,
                            **kwargs)
{'fraction_non_alphanumeric': 0.15035763201947952, 'fraction_numerical': 0.0039567797899863035, 'mean_word_length': 2.823152995927865, 'pattern_counts': {'":': 70, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6475565', 'n_tokens_mistral': 1995, 'n_tokens_neox': 1837, 'n_words': 411}