<commit_before>"""
infer frequencies of mutations or clades
"""
import json, os, sys
import numpy as np
from collections import defaultdict
from Bio import Phylo, AlignIO
from Bio.Align import MultipleSeqAlignment
from .utils import read_metadata, read_node_data, write_json, get_numerical_dates


def register_arguments(parser):
    parser.add_argument('--tree', '-t', type=str,
                        help="tree to estimate clade frequencies for")
    parser.add_argument('--alignments', type=str, nargs='+',
                        help="alignments to estimate mutations frequencies for")
    parser.add_argument('--metadata', type=str,
                        help="alignment to estimate mutations frequencies for")
    parser.add_argument('--gene-names', nargs='+', type=str,
                        help="names of the sequences in the alignment, same order assumed")
    parser.add_argument('--regions', type=str, nargs='+', default=['global'],
                        help="region to subsample to")
    parser.add_argument('--min-date', type=float,
                        help="minimal pivot value")
    parser.add_argument('--max-date', type=float,
                        help="maximal pivot value")
    parser.add_argument('--pivots-per-year', type=int, default=4,
                        help="number of pivots per year")
    parser.add_argument('--minimal-frequency', type=float, default=0.05,
                        help="minimal all-time frequencies for a trajectory to be estimates")
    parser.add_argument('--output', '-o', type=str,
                        help='JSON file to save estimated frequencies to')


def format_frequencies(freq):
    return list(freq)

def get_pivots(tps, dt,min_date, max_date):
    first_pivot = min_date if min_date else np.floor(np.min(tps)/dt)*dt
    last_pivot = max_date if max_date else np.ceil(np.max(tps)/dt)*dt
    return np.arange(first_pivot, last_pivot, dt)

def run(args):

    if args.metadata is None:
        print("ERROR: meta data with dates is required for frequency estimation")
        return 1

    metadata, columns = read_metadata(args.metadata)
    dates = get_numerical_dates(metadata, fmt='%Y-%m-%d')
    dt = 1.0/args.pivots_per_year
    stiffness = 5.0

    if args.tree:
        from .frequency_estimators import tree_frequencies
        tree = Phylo.read(args.tree, 'newick')
        tps = []
        for x in tree.get_terminals():
            x.num_date = np.mean(dates[x.name])
            tps.append(x.num_date)
            x.region = metadata[x.name]["region"]

        pivots = get_pivots(tps, dt, args.min_date, args.max_date)
        frequency_dict = {"pivots":format_frequencies(pivots)}

        # estimate tree frequencies
        # Omit strains sampled prior to the first pivot from frequency calculations.
        for region in args.regions:
            if region=='global':
                node_filter_func = lambda node: node.num_date >= first_pivot
            else:
                node_filter_func = lambda node: (node.region == region
                                                and node.num_date >= first_pivot)

            tree_freqs = tree_frequencies(tree, pivots, method='SLSQP',
                                          node_filter = node_filter_func,
                                          ws = max(2, tree.count_terminals()//10),
                                          stiffness = stiffness)

            tree_freqs.estimate_clade_frequencies()
            for x,val in tree_freqs.frequencies.items():
                frequency_dict["%s_clade:%d"%(region,x)] = format_frequencies(val)

            frequency_dict["%s:counts"%region] = [int(x) for x in tree_freqs.counts]

        json_success = write_json(frequency_dict, args.output)
        print("tree frequencies written to", args.output, file=sys.stdout)

    elif args.alignments:<commit_msg>augur/frequencies.py: fix problem in tree frequencies regarding undefined 'first_pivot'<commit_after>"""
infer frequencies of mutations or clades
"""
import json, os, sys
import numpy as np
from collections import defaultdict
from Bio import Phylo, AlignIO
from Bio.Align import MultipleSeqAlignment
from .utils import read_metadata, read_node_data, write_json, get_numerical_dates


def register_arguments(parser):
    parser.add_argument('--tree', '-t', type=str,
                        help="tree to estimate clade frequencies for")
    parser.add_argument('--alignments', type=str, nargs='+',
                        help="alignments to estimate mutations frequencies for")
    parser.add_argument('--metadata', type=str,
                        help="alignment to estimate mutations frequencies for")
    parser.add_argument('--gene-names', nargs='+', type=str,
                        help="names of the sequences in the alignment, same order assumed")
    parser.add_argument('--regions', type=str, nargs='+', default=['global'],
                        help="region to subsample to")
    parser.add_argument('--min-date', type=float,
                        help="minimal pivot value")
    parser.add_argument('--max-date', type=float,
                        help="maximal pivot value")
    parser.add_argument('--pivots-per-year', type=int, default=4,
                        help="number of pivots per year")
    parser.add_argument('--minimal-frequency', type=float, default=0.05,
                        help="minimal all-time frequencies for a trajectory to be estimates")
    parser.add_argument('--output', '-o', type=str,
                        help='JSON file to save estimated frequencies to')


def format_frequencies(freq):
    return list(freq)

def get_pivots(tps, dt,min_date, max_date):
    first_pivot = min_date if min_date else np.floor(np.min(tps)/dt)*dt
    last_pivot = max_date if max_date else np.ceil(np.max(tps)/dt)*dt
    return np.arange(first_pivot, last_pivot, dt)

def run(args):

    if args.metadata is None:
        print("ERROR: meta data with dates is required for frequency estimation")
        return 1

    metadata, columns = read_metadata(args.metadata)
    dates = get_numerical_dates(metadata, fmt='%Y-%m-%d')
    dt = 1.0/args.pivots_per_year
    stiffness = 5.0

    if args.tree:
        from .frequency_estimators import tree_frequencies
        tree = Phylo.read(args.tree, 'newick')
        tps = []
        for x in tree.get_terminals():
            x.num_date = np.mean(dates[x.name])
            tps.append(x.num_date)
            x.region = metadata[x.name]["region"]

        pivots = get_pivots(tps, dt, args.min_date, args.max_date)
        frequency_dict = {"pivots":format_frequencies(pivots)}

        # estimate tree frequencies
        # Omit strains sampled prior to the first pivot from frequency calculations.
        # (these tend to be reference strains included for phylogenetic context)
        for region in args.regions:
            if region=='global':
                node_filter_func = lambda node: node.num_date >= pivots[0]
            else:
                node_filter_func = lambda node: (node.region == region
                                                and node.num_date >= pivots[0])

            tree_freqs = tree_frequencies(tree, pivots, method='SLSQP',
                                          node_filter = node_filter_func,
                                          ws = max(2, tree.count_terminals()//10),
                                          stiffness = stiffness)

            tree_freqs.estimate_clade_frequencies()
            for x,val in tree_freqs.frequencies.items():
                frequency_dict["%s_clade:%d"%(region,x)] = format_frequencies(val)

            frequency_dict["%s:counts"%region] = [int(x) for x in tree_freqs.counts]

        json_success = write_json(frequency_dict, args.output)
        print("tree frequencies written to", args.output, file=sys.stdout)

    elif args.alignments:
{'fraction_non_alphanumeric': 0.10653138871274571, 'fraction_numerical': 0.003297400126823082, 'mean_word_length': 2.3586030664395228, 'pattern_counts': {'":': 2, '<': 3, '<?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': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '311022', 'n_tokens_mistral': 2290, 'n_tokens_neox': 2173, 'n_words': 559}