shell:
    """
    python3 scripts/partition_strains_by_timepoint.py \
        {input.metadata} \
        {wildcards.timepoint} \
        {output} \
        --years-back {params.years_back} \
        {params.reference_strains}
    """

shell:
    """
    python3 scripts/extract_sequences.py \
        --sequences {input.sequences} \
        --samples {input.strains} \
        --output {output}
    """

shell:
    """
    augur refine \
        --tree {input.tree} \
        --alignment {input.alignment} \
        --metadata {input.metadata} \
        --output-tree {output.tree} \
        --output-node-data {output.node_data} \
        --timetree \
        --no-covariance \
        {params.clock_rate} \
        {params.clock_std_dev} \
        --coalescent {params.coalescent} \
        --date-confidence \
        --date-inference {params.date_inference} &> {log}
    """

shell:
    """
    augur frequencies \
        --method kde \
        --tree {input.tree} \
        --metadata {input.metadata} \
        --narrow-bandwidth {params.narrow_bandwidth} \
        --wide-bandwidth {params.wide_bandwidth} \
        --proportion-wide {params.proportion_wide} \
        --min-date {params.min_date} \
        --max-date {params.max_date} \
        --pivot-interval {params.pivot_frequency} \
        --output {output}
    """

    shell: """python3 scripts/frequencies.py {input.tree} {input.metadata} {output} \
--narrow-bandwidth {params.narrow_bandwidth} \
--wide-bandwidth {params.wide_bandwidth} \
--proportion-wide {params.proportion_wide} \
--pivot-frequency {params.pivot_frequency} \
--start-date {params.start_date} \
--end-date {wildcards.timepoint} \
--include-internal-nodes &> {log}"""

shell: """augur frequencies \
    --method diffusion \
    --tree {input.tree} \
    --metadata {input.metadata} \
    --output {output} \
    --include-internal-nodes \
    --stiffness {params.stiffness} \
    --inertia {params.inertia} \
    --pivot-interval {params.pivot_frequency} \
    --min-date {params.min_date} \
    --max-date {params.max_date} &> {log}"""

shell:
    """
    python3 scripts/frequencies_to_table.py \
        --tree {input.tree} \
        --frequencies {input.frequencies} \
        --method {params.method} \
        --output {output} \
        --annotations timepoint={wildcards.timepoint}
    """

shell:
    """
    python3 scripts/frequencies_to_table.py \
        --tree {input.tree} \
        --frequencies {input.frequencies} \
        --method {params.method} \
        --output {output} \
        --annotations timepoint={wildcards.timepoint}
    """

shell:
    """
    augur ancestral \
        --tree {input.tree} \
        --alignment {input.alignment} \
        --output {output.node_data} \
        --inference {params.inference} &> {log}
    """

shell:
    """
    augur translate \
        --tree {input.tree} \
        --ancestral-sequences {input.node_data} \
        --reference-sequence {input.reference} \
        --output {output.node_data} &> {log}
    """

shell:
    """
    augur reconstruct-sequences \
        --tree {input.tree} \
        --mutations {input.node_data} \
        --gene {wildcards.gene} \
        --output {output.aa_alignment} \
        --internal-nodes &> {log}
    """

shell:
    """
    python3 scripts/convert_translations_to_json.py \
        --tree {input.tree} \
        --alignment {input.translations} \
        --gene-names {params.gene_names} \
        --output {output.translations}
    """

shell:
    """
    python3 scripts/nonoverlapping_clades.py \
        --tree {input.tree} \
        --translations {input.translations} \
        --gene-names {params.gene_names} \
        --annotations timepoint={wildcards.timepoint} \
        --output {output.clades} \
        --output-tip-clade-table {output.tip_clade_table} &> {log}
    """

shell:
    """
    python3 scripts/calculate_delta_frequency.py \
        --tree {input.tree} \
        --frequencies {input.frequencies} \
        --frequency-method {params.method} \
        --delta-pivots {params.delta_pivots} \
        --output {output.delta_frequency} &> {log}
    """

shell:
    """
    augur traits \
        --tree {input.tree} \
        --metadata {input.metadata} \
        --output {output.node_data} \
        --columns {params.columns} \
        --confidence
    """

shell:
    """
    augur distance \
        --tree {input.tree} \
        --alignment {input.alignments} \
        --gene-names {params.genes} \
        --compare-to {params.comparisons} \
        --attribute-name {params.attribute_names} \
        --map {input.distance_maps} \
        --date-annotations {input.date_annotations} \
        --earliest-date {params.earliest_date} \
        --latest-date {params.latest_date} \
        --output {output}
    """

shell:
    """
    python3 scripts/pairwise_distances.py \
        --tree {input.tree} \
        --frequencies {input.frequencies} \
        --alignment {input.alignments} \
        --gene-names {params.genes} \
        --attribute-name {params.attribute_names} \
        --map {input.distance_maps} \
        --date-annotations {input.date_annotations} \
        --years-back-to-compare {params.years_back_to_compare} \
        --output {output} &> {log}
    """

shell:
    """
    python3 src/cross_immunity.py \
        --frequencies {input.frequencies} \
        --distances {input.distances} \
        --date-annotations {input.date_annotations} \
        --distance-attributes {params.distance_attributes} \
        --immunity-attributes {params.immunity_attributes} \
        --decay-factors {params.decay_factors} \
        --output {output}
    """

shell:
    """
    augur lbi \
        --tree {input.tree} \
        --branch-lengths {input.branch_lengths} \
        --output {output} \
        --attribute-names {params.names} \
        --tau {params.tau} \
        --window {params.window}
    """

shell:
    """
    augur lbi \
        --tree {input.tree} \
        --branch-lengths {input.branch_lengths} \
        --output {output} \
        --attribute-names {params.names} \
        --tau {params.tau} \
        --window {params.window} \
        --no-normalization
    """

shell:
    """
    augur titers sub \
        --titers {input.titers} \
        --alignment {input.alignments} \
        --tree {input.tree} \
        --gene-names {params.genes} \
        --allow-empty-model \
        --output {output.titers_model} &> {log}
    """

shell:
    """
    augur titers tree \
        --titers {input.titers} \
        --tree {input.tree} \
        --allow-empty-model \
        --output {output.titers_model} &> {log}
    """

shell:
    """
    augur titers tree \
        --titers {input.titers} \
        --tree {input.tree} \
        --allow-empty-model \
        --output {output.titers_model} &> {log}
    """

shell:
    """
    python3 scripts/rename_fields_in_fra_titer_models.py \
        --titers-model {input.titers_model} \
        --output {output.titers_model}
    """

shell:
    """
    python3 scripts/titer_model_to_distance_map.py \
        --model {input.model} \
        --output {output}
    """

shell:
    """
    python3 scripts/pairwise_distances.py \
        --tree {input.tree} \
        --frequencies {input.frequencies} \
        --alignment {input.alignments} \
        --gene-names {params.genes} \
        --attribute-name {params.attribute_names} \
        --map {input.distance_maps} \
        --date-annotations {input.date_annotations} \
        --years-back-to-compare {params.years_back_to_compare} \
        --output {output} &> {log}
    """

shell:
    """
    python3 scripts/pairwise_titer_tree_distances.py \
        --tree {input.tree} \
        --frequencies {input.frequencies} \
        --model {input.model} \
        --attribute-name {params.attribute_names} \
        --date-annotations {input.date_annotations} \
        --months-back-for-current-samples {params.months_back_for_current_samples} \
        --years-back-to-compare {params.years_back_to_compare} \
        --output {output} &> {log}
    """

shell:
    """
    python3 scripts/pairwise_titer_tree_distances.py \
        --tree {input.tree} \
        --frequencies {input.frequencies} \
        --model {input.model} \
        --model-attribute-name {params.model_attribute_name} \
        --attribute-name {params.attribute_names} \
        --date-annotations {input.date_annotations} \
        --months-back-for-current-samples {params.months_back_for_current_samples} \
        --years-back-to-compare {params.years_back_to_compare} \
        --output {output} &> {log}
    """

shell:
    """
    python3 src/cross_immunity.py \
        --frequencies {input.frequencies} \
        --distances {input.distances} \
        --date-annotations {input.date_annotations} \
        --distance-attributes {params.distance_attributes} \
        --immunity-attributes {params.immunity_attributes} \
        --decay-factors {params.decay_factors} \
        --output {output}
    """

shell:
    """
    python3 src/cross_immunity.py \
        --frequencies {input.frequencies} \
        --distances {input.distances} \
        --date-annotations {input.date_annotations} \
        --distance-attributes {params.distance_attributes} \
        --immunity-attributes {params.immunity_attributes} \
        --decay-factors {params.decay_factors} \
        --output {output}
    """

shell:
    """
    python3 src/cross_immunity.py \
        --frequencies {input.frequencies} \
        --distances {input.distances} \
        --date-annotations {input.date_annotations} \
        --distance-attributes {params.distance_attributes} \
        --immunity-attributes {params.immunity_attributes} \
        --decay-factors {params.decay_factors} \
        --output {output}
    """

shell:
    """
    python3 scripts/normalize_fitness.py \
        --metadata {input.metadata} \
        --frequencies-table {input.frequencies} \
        --frequency-method {params.preferred_frequency_method} \
        --output {output.fitness}
    """

shell:
    """
    python3 scripts/distance_from_consensus.py \
        --sequences {input.sequences} \
        --frequencies {input.frequencies} \
        --output {output.distances}
    """

shell:
    """
    python3 scripts/node_data_to_table.py \
        --tree {input.tree} \
        --metadata {input.metadata} \
        --jsons {input.node_data} \
        --output {output} \
        {params.excluded_fields_arg} \
        --annotations timepoint={wildcards.timepoint} \
                      lineage={params.lineage} \
                      segment={params.segment}
    """

shell:
    """
    python3 scripts/merge_node_data_and_frequencies.py \
        --node-data {input.node_data} \
        --kde-frequencies {input.kde_frequencies} \
        --diffusion-frequencies {input.diffusion_frequencies} \
        --preferred-frequency-method {params.preferred_frequency_method} \
        --output {output.table}
    """

shell:
    """
    python3 scripts/collect_tables.py \
        --tables {input} \
        --output {output.attributes}
    """

shell:
    """
    python3 scripts/annotate_naive_tip_attribute.py \
        --tip-attributes {input.attributes} \
        --output {output.attributes}
    """

shell:
    """
    python3 scripts/calculate_target_distances.py \
        --tip-attributes {input.attributes} \
        --delta-months {params.delta_months} \
        --sequence-attribute-name {params.sequence_attribute_name} \
        --output {output}
    """

shell:
    """
    python3 src/weighted_distances.py \
        --tip-attributes {input.attributes} \
        --distances {input.distances} \
        --delta-months {params.delta_months} \
        --output {output}
    """

shell:
    """
    python3 src/fit_model.py \
        --tip-attributes {input.attributes} \
        --training-window {params.training_window} \
        --delta-months {params.delta_months} \
        --predictors {params.predictors} \
        --cost-function {params.cost_function} \
        --l1-lambda {params.l1_lambda} \
        --target distances \
        --distances {input.distances} \
        --errors-by-timepoint {output.errors} \
        --coefficients-by-timepoint {output.coefficients} \
        --include-scores \
        --output {output.model} &> {log}
    """

shell:
    """
    python3 scripts/extract_minimal_models_by_distances.py \
        --model {input.model} \
        --output {output.model}
    """

shell:
    """
    python3 scripts/annotate_model_tables.py \
        --tip-attributes {input.attributes} \
        --model {input.model} \
        --errors-by-timepoint {input.errors} \
        --coefficients-by-timepoint {input.coefficients} \
        --annotated-errors-by-timepoint {output.errors} \
        --annotated-coefficients-by-timepoint {output.coefficients} \
        --delta-months {params.delta_months} \
        --annotations type="{wildcards.type}" sample="{wildcards.sample}" error_type="{params.error_type}"
    """

shell:
    """
    python3 scripts/collect_tables.py \
        --tables {input} \
        --output {output.tip_clade_table}
    """

shell:
    """
    python3 scripts/select_clades.py \
        --tip-attributes {input.attributes} \
        --tips-to-clades {input.tips_to_clades} \
        --delta-months {params.delta_months} \
        --output {output} &> {log}
    """

shell:
    """
    python3 src/fit_model.py \
        --tip-attributes {input.attributes} \
        --final-clade-frequencies {input.final_clade_frequencies} \
        --training-window {params.training_window} \
        --delta-months {params.delta_months} \
        --predictors {params.predictors} \
        --cost-function {params.cost_function} \
        --l1-lambda {params.l1_lambda} \
        --pseudocount {params.pseudocount} \
        --target clades \
        --output {output} &> {log}
    """

shell:
    """
    python3 scripts/plot_tree.py {input} {output} &> {log}
    """

shell: "gs -dBATCH -dNOPAUSE -q -sDEVICE=pdfwrite -sOutputFile={output} {input}"

shell:
    """
    python3 scripts/calculate_target_distances.py \
        --tip-attributes {input.attributes} \
        --delta-months {params.delta_months} \
        --output {output}
    """

shell:
    """
    python3 src/forecast_model.py \
        --tip-attributes {input.attributes} \
        --distances {input.distances} \
        --frequencies {input.frequencies} \
        --model {input.model} \
        --delta-months {params.delta_months} \
        --output-node-data {output.node_data} \
        --output-frequencies {output.frequencies}
    """

shell:
    """
    augur export \
        --tree {input.tree} \
        --metadata {input.metadata} \
        --node-data {input.node_data} {input.forecasts} \
        --colors {input.colors} \
        --auspice-config {input.auspice_config} \
        --output-tree {output.auspice_tree} \
        --output-meta {output.auspice_metadata} \
        --panels {params.panels} \
        --minify-json
    """

shell:
    """
    python3 src/forecast_model.py \
        --tip-attributes {input.attributes} \
        --distances {input.distances} \
        --model {input.model} \
        --delta-months {params.delta_months} \
        --output-table {output.table}
    """

shell:
    """
    python3 src/fit_model.py \
        --tip-attributes {input.attributes} \
        --target distances \
        --distances {input.distances} \
        --fixed-model {input.model} \
        --errors-by-timepoint {output.errors} \
        --coefficients-by-timepoint {output.coefficients} \
        --include-scores \
        --output {output.model} &> {log}
    """

shell:
    """
    python3 scripts/annotate_model_tables.py \
        --tip-attributes {input.attributes} \
        --model {input.model} \
        --errors-by-timepoint {input.errors} \
        --coefficients-by-timepoint {input.coefficients} \
        --annotated-errors-by-timepoint {output.errors} \
        --annotated-coefficients-by-timepoint {output.coefficients} \
        --delta-months {params.delta_months} \
        --annotations type="{wildcards.type}" sample="{params.sample}" error_type="test"
    """

shell:
    """
    python3 scripts/plot_validation_figure_by_population.py \
        --tip-attributes {input.attributes} \
        --tips-to-clades {input.tips_to_clades} \
        --forecasts {input.forecasts} \
        --model-errors {input.model_errors} \
        --population {wildcards.type} \
        --sample {wildcards.sample} \
        --predictors {wildcards.predictors} \
        --output {output.figure} \
        --output-clades-table {output.clades} \
        --output-ranks-table {output.ranks}
    """

shell:
    """
    cd data/simulated/{wildcards.sample} && java -jar {SNAKEMAKE_DIR}/dist/santa.jar -seed={params.seed} {SNAKEMAKE_DIR}/{input.simulation_config}
    """

shell:
    """
    augur parse \
        --sequences {input.sequences} \
        --output-sequences {output.sequences} \
        --output-metadata {output.metadata} \
        --fields {params.fasta_fields}
    """

shell:
    """
    python3 scripts/standardize_simulated_sequence_dates.py \
        --metadata {input.metadata} \
        --start-year {params.start_year} \
       --generations-per-year {params.generations_per_year} \
        --output {output.metadata}
    """

shell:
    """
    augur filter \
        --sequences {input.sequences} \
        --metadata {input.metadata} \
        --min-date {params.min_date} \
        --group-by {params.group_by} \
        --sequences-per-group {params.viruses_per_month} \
        --output {output}
    """

shell:
    """
    python3 scripts/filter_simulated_metadata.py \
        --sequences {input.sequences} \
        --metadata {input.metadata} \
        --output {output.metadata}
    """

shell:
    """
    python3 {path_to_fauna}/vdb/download.py \
        --database vdb \
        --virus flu \
        --fasta_fields {params.fasta_fields} \
        --resolve_method split_passage \
        --select locus:{params.segment} lineage:seasonal_{params.lineage} \
        --path data/natural/{wildcards.sample} \
        --fstem original_sequences
    """

shell:
    """
    python3 {path_to_fauna}/tdb/download.py \
        --database {params.databases} \
        --virus flu \
        --subtype {params.lineage} \
        --select assay_type:{params.assay} \
        --path data/natural/{wildcards.sample} \
        --fstem complete_fra \
        --ftype json
    """

shell:
    """
    python3 scripts/get_titers_by_passage.py \
        --titers {input.titers} \
        --passage-type {params.passage} \
        --output {output.titers}
    """

shell:
    """
    python3 scripts/get_titers_by_passage.py \
        --titers {input.titers} \
        --passage-type {params.passage} \
        --output {output.titers}
    """

shell:
    """
    augur parse \
        --sequences {input.sequences} \
        --output-sequences {output.sequences} \
        --output-metadata {output.metadata} \
        --fields {params.fasta_fields}
    """

shell:
    """
    augur filter \
        --sequences {input.sequences} \
        --metadata {input.metadata} \
        --min-length {params.min_length} \
        --exclude {params.exclude} \
        --exclude-where country=? region=? passage=egg \
        --output {output}
    """

shell:
    """
    python3 scripts/filter_strains_with_ambiguous_dates.py \
        --metadata {input.metadata} \
        --output {output.metadata}
    """

shell:
    """
    python3 scripts/select_strains.py \
        --sequences {input.sequences} \
        --metadata {input.metadata} \
        --segments {params.segment} \
        --include {input.include} \
        --lineage {params.lineage} \
        --time-interval {params.start_date} {params.end_date} \
        --viruses_per_month {params.viruses_per_month} \
        --titers {input.titers} \
        --output {output.strains}
    """

shell:
    """
    python3 scripts/filter_metadata_by_strains.py \
        --metadata {input.metadata} \
        --strains {input.strains} \
        --output {output.metadata}
    """

import argparse
import pandas as pd


if __name__ == "__main__":
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--tip-attributes", required=True, help="table of tip attributes from one or more timepoints")
    parser.add_argument("--output", required=True, help="table of tip attributes annotated with a 'naive' predictor")

    args = parser.parse_args()

    # Annotate a predictor for a naive model with no growth.
    df = pd.read_csv(args.tip_attributes, sep="\t")
    df["naive"] = 0.0
    df.to_csv(args.output, sep="\t", index=False)

import argparse
import numpy as np
import pandas as pd


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Calculate pairwise distances between samples at adjacent timepoints (t and t - delta months)",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument("--tip-attributes", required=True, help="a tab-delimited file describing tip attributes at one or more timepoints")
    parser.add_argument("--delta-months", required=True, nargs="+", type=int, help="number of months between timepoints to be compared")
    parser.add_argument("--output", help="tab-delimited file of pairwise distances between tips in timepoints separate by the given delta time", required=True)
    parser.add_argument("--sequence-attribute-name", default="aa_sequence", help="attribute name of sequences to compare")
    args = parser.parse_args()

    # Load tip attributes.
    tips = pd.read_csv(args.tip_attributes, sep="\t", parse_dates=["timepoint"])

    # Open output file handle to enable streaming distances to disk instead of
    # storing them in memory.
    output_handle = open(args.output, "w")
    output_handle.write("\t".join(["sample", "other_sample", "distance"]) + "\n")

    # Calculate pairwise distances between all tips within a timepoint and at
    # the next timepoint as defined by the given delta months.
    for timepoint, timepoint_df in tips.groupby("timepoint"):
        current_tips = [
            tuple(values)
            for values in timepoint_df.loc[:, ["strain", args.sequence_attribute_name]].values.tolist()
        ]
        comparison_tips = current_tips

        for delta_month in args.delta_months:
            future_timepoint_df = tips[tips["timepoint"] == (timepoint + pd.DateOffset(months=delta_month))]
            future_tips = [
                tuple(values)
                for values in future_timepoint_df.loc[:, ["strain", args.sequence_attribute_name]].values.tolist()
            ]
            comparison_tips = comparison_tips + future_tips

        comparison_tips = list(set(comparison_tips))
        for current_tip, current_tip_sequence in current_tips:
            current_tip_sequence_array = np.frombuffer(current_tip_sequence.encode(), dtype="S1")

            for future_tip, future_tip_sequence in comparison_tips:
                future_tip_sequence_array = np.frombuffer(future_tip_sequence.encode(), dtype="S1")
                distance = (current_tip_sequence_array != future_tip_sequence_array).sum()
                output_handle.write("\t".join([current_tip, future_tip, str(distance)]) + "\n")

    # Close output.
    output_handle.close()

import argparse
import numpy as np
import pandas as pd
import re
import sys


class BadTransformException(Exception):
    pass


def parse_transform(transform):
    """Parse a transform string into its corresponding new column name, Python function, and original column name.

    Return `None` for the Python function if the requested function is not valid.

    Parameters
    ----------
    transform : str
        transformation definition string (e.g., "log_lbi=log(lbi)")

    Returns
    -------
    str, callable, str
        new column name, transformation function, and original column name

    >>> parse_transform("log_lbi=log(lbi)")
    ('log_lbi', <ufunc 'log'>, 'lbi')
    >>> parse_transform("fake_col=fake(col)")
    Traceback (most recent call last):
        ...
    collect_tables.BadTransformException: the requested function was invalid

    >>> parse_transform("bad_transform")
    Traceback (most recent call last):
        ...
    collect_tables.BadTransformException: the requested transform was malformed

    """
    match = re.match(r"(?P<new_column>\w+)=(?P<function>\w+)\((?P<column>\w+)\)", transform)
    if match is None:
        raise BadTransformException("the requested transform was malformed")

    new_column, function_string, column = match.groups()
    function = getattr(np, function_string, None)
    if function is None:
        raise BadTransformException("the requested function was invalid")

    return new_column, function, column


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Collect two or more data frame tables",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--tables", nargs="+", required=True, help="tab-delimited files with the same columns to be collected into a single file")
    parser.add_argument("--transforms", nargs="+", help="a list of new columns to create by transformation of existing columns (e.g., 'log_lbi=log(lbi)')")
    parser.add_argument("--output", required=True, help="tab-delimited output file collecting the given input tables")
    args = parser.parse_args()

    # Concatenate tip attributes across all timepoints.
    df = pd.concat([pd.read_table(table) for table in args.tables], ignore_index=True)

    # Apply transformations.
    if args.transforms:
        for transform in args.transforms:
            try:
                new_column, transform_function, column = parse_transform(transform)
                df[new_column] = transform_function(df[column])
            except BadTransformException as e:
                print(f"Error: Could not apply transformation '{transform}' because {e}", file=sys.stderr)
            except Exception as e:
                print(f"Error: Failed to apply transformation '{transform}' ({e})", file=sys.stderr)

    df.to_csv(args.output, sep="\t", index=False)

import argparse
import pandas as pd


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--tables", nargs="+", help="tables to concatenate")
    parser.add_argument("--separator", default="\t", help="separator between columns in the given tables")
    parser.add_argument("--output", help="concatenated table")

    args = parser.parse_args()

    # Concatenate tables.
    df = pd.concat([
        pd.read_csv(table_file, sep=args.separator)
        for table_file in args.tables
    ], ignore_index=True, sort=True)

    df.to_csv(args.output, sep=args.separator, header=True, index=False)

import argparse
from augur.reconstruct_sequences import load_alignments
from augur.utils import write_json
import Bio.Phylo


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Convert translation FASTA to a node data JSON",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument("--tree", required=True, help="Newick file for the tree used to construct the given node data JSONs")
    parser.add_argument("--alignment", nargs="+", help="sequence(s) to be used, supplied as FASTA files", required=True)
    parser.add_argument('--gene-names', nargs="+", type=str, help="names of the sequences in the alignment, same order assumed", required=True)
    parser.add_argument("--output", help="JSON file with translated sequences by node", required=True)
    parser.add_argument("--include-internal-nodes", action="store_true", help="include data associated with internal nodes in the output JSON")
    parser.add_argument("--attribute-name", default="aa_sequence", help="name of attribute to store the complete amino acid sequence of each node")
    args = parser.parse_args()

    # Load tree.
    tree = Bio.Phylo.read(args.tree, "newick")

    # Load sequences.
    alignments = load_alignments(args.alignment, args.gene_names)

    # Concatenate translated sequences into a single sequence indexed by sample name.
    is_node_terminal = {node.name: node.is_terminal() for node in tree.find_clades()}

    translations = {}
    for gene in args.gene_names:
        alignment = alignments[gene]

        for record in alignment:
            if is_node_terminal[record.name] or args.include_internal_nodes:
                # Initialize new samples by name with an empty string.
                if record.name not in translations:
                    translations[record.name] = {args.attribute_name:  ""}

                # Append the current gene's amino acid sequence to the current
                # string for this sample.
                translations[record.name][args.attribute_name] += str(record.seq)

    # Write out the node annotations.
    write_json({"nodes": translations}, args.output)

import argparse
import json


if __name__ == "__main__":
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--model", required=True, help="JSON for a complete fitness model output")
    parser.add_argument("--output", required=True, help="JSON for a minimal fitness model (coefficients only)")

    args = parser.parse_args()

    with open(args.model, "r") as fh:
        model = json.load(fh)

    if "scores" in model:
        del model["scores"]

    with open(args.output, "w") as oh:
        json.dump(model, oh, indent=1)

import argparse
import Bio
import Bio.SeqIO


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Extract sample sequences by name",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("--sequences", required=True, help="FASTA file of all sample sequences")
    parser.add_argument("--samples", required=True, help="text file of samples names with one name per line")
    parser.add_argument("--output", required=True, help="FASTA file of extracted sample sequences")
    args = parser.parse_args()

    with open(args.samples) as infile:
        samples = set([line.strip() for line in infile])

    with open(args.output, 'w') as outfile:
        for seq in Bio.SeqIO.parse(args.sequences, 'fasta'):
            if seq.name in samples:
                Bio.SeqIO.write(seq, outfile, 'fasta')

import argparse
import pandas as pd


if __name__ == "__main__":
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--metadata", required=True, help="table of metadata to be filtered based on a date column")
    parser.add_argument("--strains", required=True, help="text file with one strain per line that should be included in the output")
    parser.add_argument("--output", required=True, help="table of filtered metadata")

    args = parser.parse_args()

    metadata = pd.read_table(args.metadata)
    strains = pd.read_table(args.strains, header=None, names=["strain"])

    selected_metadata = strains.merge(metadata, how="left", on="strain")
    selected_metadata.to_csv(args.output, sep="\t", index=False)

import argparse
import Bio.SeqIO
import pandas as pd


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--sequences", help="simulated sequences that have already been filtered")
    parser.add_argument("--metadata", help="original metadata table for simulated sequences")
    parser.add_argument("--output", help="filtered metadata where only samples present in the given sequences are included")

    args = parser.parse_args()

    # Get a list of all samples that passed the sequence filtering step.
    sequences = Bio.SeqIO.parse(args.sequences, "fasta")
    sample_ids = [sequence.id for sequence in sequences]

    # Load all metadata.
    metadata = pd.read_csv(args.metadata, sep="\t")
    filtered_metadata = metadata[metadata["strain"].isin(sample_ids)].copy()

    # Save only the metadata records that have entries in the filtered sequences.
    filtered_metadata.to_csv(args.output, sep="\t", header=True, index=False)

import argparse
import pandas as pd


if __name__ == "__main__":
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--metadata", required=True, help="table of metadata to be filtered based on a date column")
    parser.add_argument("--date-field", default="date", help="name of date column in the metadata")
    parser.add_argument("--output", required=True, help="table of filtered metadata")

    args = parser.parse_args()

    df = pd.read_csv(args.metadata, sep="\t")

    # Exclude strains with ambiguous collection dates.
    df[~df[args.date_field].str.contains("XX")].to_csv(args.output, sep="\t", header=True, index=False)

import argparse
from augur.utils import read_metadata, get_numerical_dates
from augur.frequencies import TreeKdeFrequencies
import Bio
import Bio.Phylo
import datetime
import json
import numpy as np
import os
import sys


def get_time_interval_as_floats(time_interval):
    """
    Converts the given datetime interval to start and end floats.

    Returns:
        start_date (float): the start of the given time interval
        end_date (float): the end of the given time interval
    """
    start_date = time_interval[1].year + (time_interval[1].month - 1) / 12.0
    end_date = time_interval[0].year + (time_interval[0].month - 1) / 12.0
    return start_date, end_date


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("tree", help="Newick tree")
    parser.add_argument("metadata", help="tab-delimited metadata for tips in the given tree including a date field")
    parser.add_argument("frequencies", help="JSON with frequencies estimated from the given tree and used to estimate the given parameters")
    parser.add_argument("--narrow-bandwidth", type=float, default=1 / 12.0, help="the bandwidth for the narrow KDE")
    parser.add_argument("--wide-bandwidth", type=float, default=3 / 12.0, help="the bandwidth for the wide KDE")
    parser.add_argument("--proportion-wide", type=float, default=0.2, help="the proportion of the wide bandwidth to use in the KDE mixture model")
    parser.add_argument("--pivot-frequency", type=int, default=1, help="number of months between pivots")
    parser.add_argument("--start-date", help="the start of the interval to estimate frequencies across")
    parser.add_argument("--end-date", help="the end of the interval to estimate frequencies across")
    parser.add_argument("--include-internal-nodes", action="store_true", help="calculate frequencies for internal nodes as well as tips")
    parser.add_argument("--weights", help="a dictionary of key/value mappings in JSON format used to weight tip frequencies")
    parser.add_argument("--weights-attribute", help="name of the attribute on each tip whose values map to the given weights dictionary")

    parser.add_argument("--precision", type=int, default=6, help="number of decimal places to retain in frequency estimates")
    parser.add_argument("--censored", action="store_true", help="calculate censored frequencies at each pivot")

    args = parser.parse_args()

    # Load tree.
    tree = Bio.Phylo.read(args.tree, "newick")

    # Load metadata.
    metadata, columns = read_metadata(args.metadata)
    dates = get_numerical_dates(metadata, fmt='%Y-%m-%d')

    # Annotate tree with dates and other metadata.
    for tip in tree.find_clades(terminal=True):
        tip.attr = {"num_date": np.mean(dates[tip.name])}

        # Annotate tips with metadata to enable filtering and weighting of
        # frequencies by metadata attributes.
        for key, value in metadata[tip.name].items():
            tip.attr[key] = value

    # Convert start and end dates to floats from time interval format.
    if args.start_date is not None and args.end_date is not None:
        # Convert the string time interval to a datetime instance and then to floats.
        time_interval = [
            datetime.datetime.strptime(time, "%Y-%m-%d")
            for time in (args.end_date, args.start_date)
        ]

        start_date, end_date = get_time_interval_as_floats(time_interval)
    else:
        start_date = end_date = None

    # Load weights if they have been provided.
    if args.weights:
        with open(args.weights, "r") as fh:
            weights = json.load(fh)

        weights_attribute = args.weights_attribute
    else:
        weights = None
        weights_attribute = None

    # Estimate frequencies.
    frequencies = TreeKdeFrequencies(
        sigma_narrow=args.narrow_bandwidth,
        sigma_wide=args.wide_bandwidth,
        proportion_wide=args.proportion_wide,
        pivot_frequency=args.pivot_frequency,
        start_date=start_date,
        end_date=end_date,
        weights=weights,
        weights_attribute=weights_attribute,
        include_internal_nodes=args.include_internal_nodes,
        censored=args.censored
    )
    frequencies.estimate(tree)

    # Export frequencies to JSON.
    json_frequencies = frequencies.to_json()

    # Set precision of frequency estimates.
    for clade in json_frequencies["data"]["frequencies"]:
        json_frequencies["data"]["frequencies"][clade] = np.around(
            np.array(
                json_frequencies["data"]["frequencies"][clade]
            ),
            args.precision
        ).tolist()

    with open(args.frequencies, "w") as oh:
        json.dump(json_frequencies, oh, indent=1, sort_keys=True)