<commit_before>from __future__ import division, print_function
import sys, os, time, gzip, glob
from base.io_util import make_dir, remove_dir, tree_to_json, write_json, myopen
import json
from pdb import set_trace
from collections import defaultdict

def process_freqs(frequencies, num_dp):
    """ round frequency estimates to useful precision (reduces file size) """
    return [round(x, num_dp) for x in frequencies]

def export_frequency_json(process, prefix, indent):
    num_dp = 4;
    # construct a json file containing all frequency estimate
    # the format is region_protein:159F for mutations and region_clade:123 for clades
    if hasattr(process, 'pivots'):
        freq_json = {'pivots':process_freqs(process.pivots, num_dp)}
        if hasattr(process, 'mutation_frequencies'):
            freq_json['counts'] = {x:list(counts) for x, counts in process.mutation_frequency_counts.iteritems()}
            for (region, gene), tmp_freqs in process.mutation_frequencies.iteritems():
                for mut, freq in tmp_freqs.iteritems():
                    label_str =  region+"_"+ gene + ':' + str(mut[0]+1)+mut[1]
                    freq_json[label_str] = process_freqs(freq, num_dp)
        # repeat for clade frequencies in trees
        if hasattr(process, 'tree_frequencies'):
            for region in process.tree_frequencies:
                for clade, freq in process.tree_frequencies[region].iteritems():
                    label_str = region+'_clade:'+str(clade)
                    freq_json[label_str] = process_freqs(freq, num_dp)
        # repeat for named clades
        if hasattr(process, 'clades_to_nodes') and hasattr(process, 'tree_frequencies'):
            for region in process.tree_frequencies:
                for clade, node in process.clades_to_nodes.iteritems():
                    label_str = region+'_'+str(clade)
                    freq_json[label_str] = process_freqs(process.tree_frequencies[region][node.clade], num_dp)
        # write to one frequency json
        if hasattr(process, 'tree_frequencies') or hasattr(process, 'mutation_frequencies'):
            write_json(freq_json, prefix+'_frequencies.json', indent=indent)
    else:
        process.log.notify("Cannot export frequencies - pivots do not exist")

def export_tip_frequency_json(process, prefix, indent):
    if not (hasattr(process, 'pivots') and hasattr(process, 'tree_frequencies')):
        process.log.notify("Cannot export tip frequencies - pivots and/or tree_frequencies do not exist")
        return;

    num_dp = 6;
    freq_json = {'pivots':process_freqs(process.pivots, num_dp)}

    for n in process.tree.tree.get_terminals():
        freq_json["global_clade:"+str(n.clade)] = process_freqs(process.tree_frequencies["global"][n.clade], num_dp)
<commit_msg>final tip-frequencies spec change<commit_after>from __future__ import division, print_function
import sys, os, time, gzip, glob
from base.io_util import make_dir, remove_dir, tree_to_json, write_json, myopen
import json
from pdb import set_trace
from collections import defaultdict

def process_freqs(frequencies, num_dp):
    """ round frequency estimates to useful precision (reduces file size) """
    return [round(x, num_dp) for x in frequencies]

def export_frequency_json(process, prefix, indent):
    num_dp = 4;
    # construct a json file containing all frequency estimate
    # the format is region_protein:159F for mutations and region_clade:123 for clades
    if hasattr(process, 'pivots'):
        freq_json = {'pivots':process_freqs(process.pivots, num_dp)}
        if hasattr(process, 'mutation_frequencies'):
            freq_json['counts'] = {x:list(counts) for x, counts in process.mutation_frequency_counts.iteritems()}
            for (region, gene), tmp_freqs in process.mutation_frequencies.iteritems():
                for mut, freq in tmp_freqs.iteritems():
                    label_str =  region+"_"+ gene + ':' + str(mut[0]+1)+mut[1]
                    freq_json[label_str] = process_freqs(freq, num_dp)
        # repeat for clade frequencies in trees
        if hasattr(process, 'tree_frequencies'):
            for region in process.tree_frequencies:
                for clade, freq in process.tree_frequencies[region].iteritems():
                    label_str = region+'_clade:'+str(clade)
                    freq_json[label_str] = process_freqs(freq, num_dp)
        # repeat for named clades
        if hasattr(process, 'clades_to_nodes') and hasattr(process, 'tree_frequencies'):
            for region in process.tree_frequencies:
                for clade, node in process.clades_to_nodes.iteritems():
                    label_str = region+'_'+str(clade)
                    freq_json[label_str] = process_freqs(process.tree_frequencies[region][node.clade], num_dp)
        # write to one frequency json
        if hasattr(process, 'tree_frequencies') or hasattr(process, 'mutation_frequencies'):
            write_json(freq_json, prefix+'_frequencies.json', indent=indent)
    else:
        process.log.notify("Cannot export frequencies - pivots do not exist")

def export_tip_frequency_json(process, prefix, indent):
    if not (hasattr(process, 'pivots') and hasattr(process, 'tree_frequencies')):
        process.log.notify("Cannot export tip frequencies - pivots and/or tree_frequencies do not exist")
        return;

    num_dp = 6;
    freq_json = {'pivots':process_freqs(process.pivots, num_dp)}

    for n in process.tree.tree.get_terminals():
        freq_json[n.name] = {
            "frequencies" : process_freqs(process.tree_frequencies["global"][n.clade], num_dp),
            "weight": 1.0
        }

{'fraction_non_alphanumeric': 0.1004072959093324, 'fraction_numerical': 0.004250044271294493, 'mean_word_length': 3.285280728376328, 'pattern_counts': {'":': 1, '<': 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': '4402343', 'n_tokens_mistral': 1754, 'n_tokens_neox': 1694, 'n_words': 424}