Include MLR lineage fitness

This include a script to grab live estimate of MLR lineage fitness and apply this as a tip-level coloring in the tree.

Author: trvrb

scripts/fetch_mlr_lineage_fitness.py

@@ -0,0 +1,69 @@
+fromaugur.ioimportread_metadata
+fromaugur.utilsimportwrite_json
+importrequests
+importjson
+importpandasaspd
+importargparse
+importmath
+
+# Set up argument parser
+parser=argparse.ArgumentParser(description="Process metadata and growth advantage data.")
+parser.add_argument("--metadata", required=True, help="Path to the metadata file (TSV or compressed .tsv.xz format).")
+parser.add_argument("--metadata-id-columns", default=["strain", "name", "Virus name"], nargs="+", help="List of columns to use as identifiers in the metadata file.")
+parser.add_argument("--metadata-clade-attribute", default="Nextclade_pango", help="Matched attribute to MLR variants.")
+parser.add_argument("--mlr-url", default="https://data.nextstrain.org/files/workflows/forecasts-ncov/gisaid/pango_lineages/global/mlr/latest_results.json", help="URL to fetch the forecasts JSON data.")
+parser.add_argument("--output-node-data", required=True, help="Path to save the output JSON node data.")
+
+args=parser.parse_args()
+
+deffetch_growth_advantages(mlr_url):
+try:
+response=requests.get(mlr_url)
+response.raise_for_status() # Raise an exception for HTTP errors
+json_data=response.json() # Parse the JSON content
+data=json_data["data"]
+
+growth_advantages= {}
+forentryindata:
+ifall(keyinentryforkeyin ["location", "site", "variant", "value", "ps"]):
+ifentry["location"] =="hierarchical"andentry["site"] =="ga"andentry["ps"] =="median":
+growth_advantages[entry["variant"]] =entry["value"]
+returngrowth_advantages
+exceptExceptionase:
+print(f"Error fetching the JSON file: {e}")
+returnNone
+
+try:
+# Fetch the growth advantages
+growth_advantages=fetch_growth_advantages(args.mlr_url)
+
+# Load the local metadata file
+metadata_file=args.metadata
+metadata=read_metadata(
+metadata_file,
+id_columns=args.metadata_id_columns
+ )
+
+# Match Nextclade_pango entries to the growth advantage
+ifgrowth_advantages:
+metadata[args.metadata_clade_attribute] =metadata[args.metadata_clade_attribute].map(growth_advantages)
+else:
+metadata[args.metadata_clade_attribute] =math.nan
+
+# Output rows with matched data
+print(metadata.head()) # Display the first few rows as an example
+
+# Create a node data object with growth advantages
+node_data= {}
+forindex, recordinmetadata.iterrows():
+node_data[index] = {
+"mlr_lineage_fitness": record[args.metadata_clade_attribute]
+ }
+
+# Save node data
+write_json({"nodes": node_data}, args.output_node_data)
+
+exceptFileNotFoundErrorase:
+print(f"Error reading metadata file: {e}")
+exceptExceptionase:
+print(f"An unexpected error occurred: {e}")

workflow/snakemake_rules/export_for_nextstrain.smk

@@ -192,6 +192,11 @@ rule auspice_config:
 "title": "Mutational Fitness",
 "type": "continuous"
  },
+ {
+"key": "mlr_lineage_fitness",
+"title": "MLR lineage fitness",
+"type": "continuous"
+ },
  {
 "key": "region",
 "title": "Region",

workflow/snakemake_rules/main_workflow.smk

@@ -1170,6 +1170,27 @@ rule mutational_fitness:
  --output {output} 2>&1 | tee {log}
  """
 
+rulemlr_lineage_fitness:
+input:
+metadata="results/{build_name}/metadata_adjusted.tsv.xz",
+output:
+node_data="results/{build_name}/mlr_lineage_fitness.json",
+benchmark:
+"benchmarks/mlr_lineage_fitness_{build_name}.txt",
+conda:
+config["conda_environment"],
+log:
+"logs/mlr_lineage_fitness_{build_name}.txt",
+params:
+metadata_id_columns=config["sanitize_metadata"]["metadata_id_columns"],
+shell:
+r"""
+ python3 scripts/fetch_mlr_lineage_fitness.py \
+ --metadata {input.metadata} \
+ --metadata-id-columns {params.metadata_id_columns:q} \
+ --output-node-data {output.node_data} 2>&1 | tee {log}
+ """
+
 rulecalculate_epiweeks:
 input:
 metadata="results/{build_name}/metadata_adjusted.tsv.xz",
@@ -1255,6 +1276,7 @@ def _get_node_data_by_wildcards(wildcards):
 rules.traits.output.node_data,
 rules.logistic_growth.output.node_data,
 rules.mutational_fitness.output.node_data,
+rules.mlr_lineage_fitness.output.node_data,
 rules.distances.output.node_data,
 rules.calculate_epiweeks.output.node_data,
  ]