feat(web): sort & run

packages/nextclade-cli/Cargo.toml

@@ -33,6 +33,7 @@ indexmap = { workspace = true }
 itertools = { workspace = true }
 lazy_static = { workspace = true }
 log = { workspace = true }
+maplit = { workspace = true }
 nextclade = { path = "../nextclade" }
 num_cpus = { workspace = true }
 ordered-float = { workspace = true }

packages/nextclade-cli/src/cli/nextclade_seq_sort.rs

@@ -4,14 +4,17 @@ use crate::io::http_client::HttpClient;
 use eyre::{Report, WrapErr};
 use itertools::Itertools;
 use log::{trace, LevelFilter};
+use maplit::btreemap;
 use nextclade::io::csv::CsvStructFileWriter;
 use nextclade::io::fasta::{FastaReader, FastaRecord, FastaWriter};
 use nextclade::io::fs::path_to_string;
-use nextclade::make_error;
 use nextclade::sort::minimizer_index::{MinimizerIndexJson, MINIMIZER_INDEX_ALGO_VERSION};
-use nextclade::sort::minimizer_search::{run_minimizer_search, MinimizerSearchDatasetResult, MinimizerSearchRecord};
+use nextclade::sort::minimizer_search::{
+  find_best_datasets, run_minimizer_search, MinimizerSearchDatasetResult, MinimizerSearchRecord,
+};
 use nextclade::utils::option::{OptionMapMutFallible, OptionMapRefFallible};
 use nextclade::utils::string::truncate;
+use nextclade::{make_error, make_internal_report};
 use ordered_float::OrderedFloat;
 use owo_colors::OwoColorize;
 use schemars::JsonSchema;
@@ -151,43 +154,122 @@ fn writer_thread(
   verbose: bool,
 ) -> Result<(), Report> {
   let NextcladeSortArgs {
+    input_fastas,
     output_dir,
     output_path,
     output_results_tsv,
     search_params,
     ..
   } = args;
 
-  let template = output_path.map_ref_fallible(move |output_path| -> Result<TinyTemplate, Report> {
-    let mut template = TinyTemplate::new();
-    template
-      .add_template("output", output_path)
-      .wrap_err_with(|| format!("When parsing template: '{output_path}'"))?;
-    Ok(template)
-  })?;
+  if search_params.global {
+    // NOTE(perf): this gathers all suggestions results and discards sequence data to make sure we don't store the
+    // whole thing in memory. We will have to read fasta again later to write some outputs.
+    let results: BTreeMap<usize, _> = result_receiver
+      .iter()
+      .map(|result| (result.fasta_record.index, result.result))
+      .collect();
+
+    // let seqs_with_no_hits = results
+    //   .iter()
+    //   .filter(|(_, r)| r.datasets.is_empty())
+    //   .map(|(fasta_index, _)| fasta_index)
+    //   .sorted_unstable()
+    //   .copied()
+    //   .collect_vec();
+
+    let best_datasets = find_best_datasets(&results, search_params)?;
+
+    let mut stats = StatsPrinter::new(
+      "Suggested datasets for each sequence (after global optimization)",
+      verbose,
+    );
+    let mut reader = FastaReader::from_paths(input_fastas)?;
+    let mut writer = DatasetSortWriter::new(output_path, output_dir, output_results_tsv)?;
+    loop {
+      let mut record = FastaRecord::default();
+      reader.read(&mut record)?;
+      if record.is_empty() {
+        break;
+      }
 
-  let mut writers = BTreeMap::new();
-  let mut stats = StatsPrinter::new(verbose);
+      let best_dataset = best_datasets
+        .iter()
+        .find(|best_dataset| best_dataset.qry_indices.contains(&record.index));
+
+      let datasets = if let Some(best_dataset) = best_dataset {
+        let dataset = results[&record.index]
+          .datasets
+          .iter()
+          .find(|d| d.name == best_dataset.name)
+          .ok_or_else(|| make_internal_report!("Unable to find dataset '{}'", best_dataset.name))?;
+        vec![dataset.clone()]
+      } else {
+        vec![]
+      };
+      stats.print_seq(&datasets, &record.seq_name);
+      writer.write_one(&record, &datasets)?;
+    }
+  } else {
+    let mut stats = StatsPrinter::new("Suggested datasets for each sequence", verbose);
+    let mut writer = DatasetSortWriter::new(output_path, output_dir, output_results_tsv)?;
+    for MinimizerSearchRecord { fasta_record, result } in result_receiver {
+      let datasets = {
+        if result.datasets.len() == 0 {
+          &[]
+        } else if search_params.all_matches {
+          &result.datasets[..]
+        } else {
+          &result.datasets[0..1]
+        }
+      };
+      stats.print_seq(datasets, &fasta_record.seq_name);
+      writer.write_one(&fasta_record, datasets)?;
+    }
+    stats.finish();
+  }
 
-  let mut results_csv =
-    output_results_tsv.map_ref_fallible(|output_results_tsv| CsvStructFileWriter::new(output_results_tsv, b'\t'))?;
+  Ok(())
+}
 
-  for record in result_receiver {
-    let datasets = &{
-      if search_params.all_matches {
-        record.result.datasets
-      } else {
-        record.result.datasets.into_iter().take(1).collect_vec()
-      }
-    };
+pub struct DatasetSortWriter<'t> {
+  writers: BTreeMap<PathBuf, FastaWriter>,
+  results_csv: Option<CsvStructFileWriter>,
+  output_dir: Option<PathBuf>,
+  template: Option<TinyTemplate<'t>>,
+}
 
-    stats.print_seq(datasets, &record.fasta_record.seq_name);
+impl<'t> DatasetSortWriter<'t> {
+  pub fn new(
+    output_path: &'t Option<String>,
+    output_dir: &Option<PathBuf>,
+    output_results_tsv: &Option<String>,
+  ) -> Result<Self, Report> {
+    let template = output_path.map_ref_fallible(move |output_path| -> Result<TinyTemplate<'t>, Report> {
+      let mut template = TinyTemplate::new();
+      template
+        .add_template("output", output_path)
+        .wrap_err_with(|| format!("When parsing template: '{output_path}'"))?;
+      Ok(template)
+    })?;
+
+    let results_csv =
+      output_results_tsv.map_ref_fallible(|output_results_tsv| CsvStructFileWriter::new(output_results_tsv, b'\t'))?;
+
+    Ok(Self {
+      writers: btreemap! {},
+      results_csv,
+      output_dir: output_dir.clone(),
+      template,
+    })
+  }
 
+  pub fn write_one(&mut self, record: &FastaRecord, datasets: &[MinimizerSearchDatasetResult]) -> Result<(), Report> {
     if datasets.is_empty() {
-      results_csv.map_mut_fallible(|results_csv| {
+      self.results_csv.map_mut_fallible(|results_csv| {
         results_csv.write(&SeqSortCsvEntry {
-          index: record.fasta_record.index,
-          seq_name: &record.fasta_record.seq_name,
+          index: record.index,
+          seq_name: &record.seq_name,
           dataset: None,
           score: None,
           num_hits: None,
@@ -196,38 +278,34 @@ fn writer_thread(
     }
 
     for dataset in datasets {
-      results_csv.map_mut_fallible(|results_csv| {
+      self.results_csv.map_mut_fallible(|results_csv| {
         results_csv.write(&SeqSortCsvEntry {
-          index: record.fasta_record.index,
-          seq_name: &record.fasta_record.seq_name,
+          index: record.index,
+          seq_name: &record.seq_name,
           dataset: Some(&dataset.name),
           score: Some(dataset.score),
           num_hits: Some(dataset.n_hits),
         })
       })?;
     }
 
-    let names = datasets
+    let dataset_names = datasets
       .iter()
       .map(|dataset| get_all_prefix_names(&dataset.name))
       .collect::<Result<Vec<Vec<String>>, Report>>()?
       .into_iter()
       .flatten()
       .unique();
 
-    for name in names {
-      let filepath = get_filepath(&name, &template, output_dir)?;
-
+    for name in dataset_names {
+      let filepath = get_filepath(&name, &self.template, &self.output_dir)?;
       if let Some(filepath) = filepath {
-        let writer = get_or_insert_writer(&mut writers, filepath)?;
-        writer.write(&record.fasta_record.seq_name, &record.fasta_record.seq, false)?;
+        let writer = get_or_insert_writer(&mut self.writers, filepath)?;
+        writer.write(&record.seq_name, &record.seq, false)?;
       }
     }
+    Ok(())
   }
-
-  stats.finish();
-
-  Ok(())
 }
 
 pub fn get_all_prefix_names(name: impl AsRef<str>) -> Result<Vec<String>, Report> {
@@ -250,9 +328,9 @@ struct StatsPrinter {
 }
 
 impl StatsPrinter {
-  fn new(enabled: bool) -> Self {
+  fn new(title: impl AsRef<str>, enabled: bool) -> Self {
     if enabled {
-      println!("Suggested datasets for each sequence");
+      println!("{}", title.as_ref());
       println!("{}┐", "─".repeat(110));
       println!(
         "{:^40} │ {:^40} │ {:^10} │ {:^10} │",

packages/nextclade-cli/src/dataset/dataset_download.rs

@@ -150,6 +150,7 @@ pub fn dataset_zip_load(
   }
 
   Ok(NextcladeParams {
+    dataset_name: dataset_zip.to_str().unwrap().to_owned(),
     ref_record,
     gene_map,
     tree,
@@ -296,6 +297,7 @@ pub fn dataset_dir_load(
   }
 
   Ok(NextcladeParams {
+    dataset_name: dataset_dir.to_str().unwrap().to_owned(),
     ref_record,
     gene_map,
     tree,
@@ -344,14 +346,17 @@ pub fn dataset_json_load(
     }
 
     NextcladeParamsOptional {
+      dataset_name: dataset_json.to_str().map(ToOwned::to_owned),
       ref_record,
       gene_map,
       tree,
       virus_properties,
     }
   };
 
-  NextcladeParams::from_auspice(&auspice_json, &overrides, cdses)
+  // TODO: should we support multiple datasets here?
+  let mut datasets = NextcladeParams::from_auspice(&auspice_json, &overrides, cdses)?;
+  Ok(datasets.remove(0))
 }
 
 pub fn dataset_individual_files_load(
@@ -395,6 +400,12 @@ pub fn dataset_individual_files_load(
       }
 
       Ok(NextcladeParams {
+        dataset_name: run_args
+          .inputs
+          .input_pathogen_json
+          .as_ref()
+          .map(|s| s.to_str().unwrap().to_owned())
+          .unwrap_or_default(),
         ref_record,
         gene_map,
         tree,
@@ -470,6 +481,7 @@ pub fn dataset_str_download_and_load(
   }
 
   Ok(NextcladeParams {
+    dataset_name: name.to_owned(),
     ref_record,
     gene_map,
     tree,

packages/nextclade-web/.eslintrc.js

@@ -115,6 +115,7 @@ module.exports = {
     'jsx-a11y/label-has-associated-control': ['warn', { assert: 'either' }],
     'lodash/chaining': 'off',
     'lodash/import-scope': 'off',
+    'lodash/matches-prop-shorthand': 'off',
     'lodash/prefer-constant': 'off',
     'lodash/prefer-lodash-chain': 'off',
     'lodash/prefer-lodash-method': 'off',

packages/nextclade-web/src/components/Common/MutationBadge.tsx

@@ -1,6 +1,7 @@
 import React from 'react'
 import { useRecoilValue } from 'recoil'
 import { isNil } from 'lodash'
+import { viewedDatasetNameAtom } from 'src/state/dataset.state'
 import styled, { useTheme } from 'styled-components'
 import { shade } from 'polished'
 import type { AaSub, NucSub } from 'src/types'
@@ -52,12 +53,6 @@ export const PositionText = styled.span`
   color: ${(props) => props.theme.gray800};
 `
 
-export const VersionText = styled.span`
-  padding: 1px 2px;
-  background-color: ${(props) => props.theme.gray400};
-  color: ${(props) => props.theme.gray800};
-`
-
 export interface NucleotideMutationBadgeProps {
   mutation: NucSub
 }
@@ -99,7 +94,8 @@ export function AminoacidMutationBadge({ mutation }: AminoacidMutationBadgeProps
   const theme = useTheme()
 
   const { cdsName, refAa, qryAa, pos } = mutation
-  const cds = useRecoilValue(cdsAtom(cdsName))
+  const datasetName = useRecoilValue(viewedDatasetNameAtom)
+  const cds = useRecoilValue(cdsAtom({ datasetName, cdsName }))
 
   const geneBg = cds?.color ?? '#999'
   const refBg = getAminoacidColor(refAa)

packages/nextclade-web/src/components/Export/ExportPage.tsx

@@ -1,5 +1,6 @@
 import { useRouter } from 'next/router'
 import React, { useState } from 'react'
+import { ViewedDatasetSelector } from 'src/components/Main/ViewedDatasetSelector'
 import styled from 'styled-components'
 import { TabContent, TabLabel, TabNav, TabPane } from 'src/components/Common/TabsFull'
 import { ExportTabColumnConfig } from 'src/components/Export/ExportTabColumnConfig'
@@ -15,6 +16,11 @@ export function ExportPage() {
   return (
     <Layout>
       <Container>
+        <div>
+          <label> {t('Select dataset')}</label>
+          <ViewedDatasetSelector />
+        </div>
+
         <Header>
           <h4 className="mx-auto">{t('Download output files')}</h4>
         </Header>