Remove the remaining duplicate implementations: annotation, snpdb, cross-app #1801

analysis/fake_variant_tags.py

@@ -26,8 +26,7 @@
 
 from analysis.models import VariantTag
 from analysis.models.enums import TagLocation
-from annotation.models import VariantAnnotation, VariantAnnotationVersion
-from genes.models import TranscriptVersion
+from annotation.fake_data import get_variant_ids_by_gene, zipf_weight
 from library.guardian_utils import all_users_group
 from snpdb.models import Allele, AlleleConversionTool, AlleleOrigin, GenomeBuild, GlobalSettings, Lab, Organization, \
     Tag, TagColor, TagColorsCollection, VariantAllele
@@ -219,13 +218,13 @@ def _create_tags(self):
     def _pick_variants(self, genome_build: GenomeBuild, group: str, genes: list[str],
                        num_variants: int) -> list[FakeVariant]:
         """ Real variants, so the gene cards have real symbols to group on and the links all work """
-        variant_ids_by_gene = _variant_ids_by_gene(genome_build, genes)
+        variant_ids_by_gene = get_variant_ids_by_gene(genome_build, genes)
         tags = [t for t in FAKE_TAGS if t.group in (group, BOTH)]
 
         fake_variants = []
         for i, gene_symbol in enumerate(genes):
             available = variant_ids_by_gene.get(gene_symbol, [])
-            wanted = int(num_variants * _zipf_weight(i) / sum(_zipf_weight(j) for j in range(len(genes))))
+            wanted = int(num_variants * zipf_weight(i) / sum(zipf_weight(j) for j in range(len(genes))))
             if len(available) < wanted:
                 self.stdout.write(f"{gene_symbol}: only {len(available)} annotated variants, wanted {wanted}")
                 wanted = len(available)
@@ -333,22 +332,6 @@ def _tag_colors_collection() -> TagColorsCollection:
     return collection
 
 
-def _variant_ids_by_gene(genome_build: GenomeBuild, genes: list[str]) -> dict[str, list[int]]:
-    gene_symbol_field = "transcript_version__gene_version__gene_symbol_id"
-    transcript_versions_qs = TranscriptVersion.objects.filter(genome_build=genome_build,
-                                                              gene_version__gene_symbol__in=genes)
-    variant_annotation_qs = VariantAnnotation.objects.filter(version=VariantAnnotationVersion.latest(genome_build),
-                                                             transcript_version__in=transcript_versions_qs)
-    variant_ids_by_gene = {}
-    for gene_symbol, variant_id in variant_annotation_qs.values_list(gene_symbol_field, "variant_id"):
-        variant_ids_by_gene.setdefault(gene_symbol, []).append(variant_id)
-    return variant_ids_by_gene
-
-
-def _zipf_weight(rank: int) -> float:
-    return 1 / (rank + 1) ** 0.8
-
-
 def _allocate_events(fake_variants: list[FakeVariant], num_events: int):
     """ Everything is tagged once, then the rest are dealt out by weight - the tail is where re-tagging lives """
     extra = num_events - len(fake_variants)

annotation/external_annotation.py

@@ -152,6 +152,59 @@ def parse_dump_metadata(path) -> dict:
     return meta
 
 
+class _ExternalDump:
+    """ One dump directory's worth of runs. The callers differ only in how they get hold of their
+        AnnotationRuns - creating them up front, or adopting existing ones - then each run goes through
+        add_run() and the dump is closed off with finish(). """
+
+    def __init__(self, variant_annotation_version: VariantAnnotationVersion, output_dir: str,
+                 pipeline_type, min_variants: int, described_as: str):
+        self.variant_annotation_version = variant_annotation_version
+        self.output_dir = output_dir
+        self.pipeline_type = pipeline_type
+        self.min_variants = min_variants
+        self.described_as = described_as  # "external" or "existing", for the log lines
+        # Every run in this dump shares the same VAV, so compute the version identity (which runs VEP) once,
+        # lazily on the first run so a no-op dump never invokes VEP.
+        self._identity = None
+        self.annotation_runs = []
+        self.reverted = []
+        os.makedirs(output_dir, exist_ok=True)
+
+    def add_run(self, annotation_run: AnnotationRun):
+        # If already counted off-thread (AnnotationRun.count, #1646), reject a too-small run before dumping.
+        # A counted run here is always non-empty (the count lane finishes count==0 runs itself).
+        if annotation_run.count is not None and annotation_run.count < self.min_variants:
+            _revert_too_small_run(annotation_run, annotation_run.count, self.min_variants, self.reverted)
+            return
+        dump_count = get_runner(self.pipeline_type).dump(annotation_run, dump_dir=self.output_dir)
+        # Range locks are sized without pipeline_type, so an SV dump leaves most locks with zero SVs; such a
+        # run is already FINISHED (get_status: dump_count == 0) with nothing to annotate, so skip its sidecar
+        # rather than emit a no-op .meta.json (thousands of them, for a handful of SVs).
+        if dump_count == 0:
+            return
+        if dump_count < self.min_variants:  # uncounted run whose dump turned out too small
+            _revert_too_small_run(annotation_run, dump_count, self.min_variants, self.reverted)
+            return
+        if self._identity is None:
+            self._identity = variant_annotation_version_identity(self.variant_annotation_version)
+        meta_filename = write_dump_metadata(annotation_run, dump_dir=self.output_dir, identity=self._identity)
+        logging.info("Dumped %s AnnotationRun %s: %d variants -> %s (meta %s)",
+                     self.described_as, annotation_run.pk, dump_count, annotation_run.vcf_dump_filename,
+                     meta_filename)
+        self.annotation_runs.append(annotation_run)
+
+    def finish(self) -> list[AnnotationRun]:
+        write_snakemake_bundle(self.output_dir, self.variant_annotation_version,
+                               pipeline_type=self.pipeline_type)
+        if self.reverted:
+            dispatch_annotation_runs.si(self.variant_annotation_version.pk).apply_async()
+        logging.info("Dumped %d %s annotation run(s) for %s into %s (reverted %d too-small run(s) to local)",
+                     len(self.annotation_runs), self.described_as, self.variant_annotation_version,
+                     self.output_dir, len(self.reverted))
+        return self.annotation_runs
+
+
 def dump_external_annotation_runs(variant_annotation_version: VariantAnnotationVersion,
                                   output_dir: str,
                                   pipeline_type=VariantAnnotationPipelineType.STANDARD,
@@ -167,47 +220,16 @@ def dump_external_annotation_runs(variant_annotation_version: VariantAnnotationV
         A run holding fewer than `min_variants` variants is reverted to the local pipeline instead of parked
         external (see DEFAULT_MIN_EXTERNAL_VARIANTS) - the external round-trip is not worth it for a tiny run. """
     _require_sv_offload_supported(pipeline_type)
-    os.makedirs(output_dir, exist_ok=True)
-    # Every run in this dump shares the same VAV, so compute the version identity (which runs VEP) once,
-    # lazily on the first run so a no-op dump never invokes VEP.
-    identity = None
-    annotation_runs = []
-    reverted = []
+    dump = _ExternalDump(variant_annotation_version, output_dir, pipeline_type, min_variants, "external")
     while True:
         range_lock, _unannotated_count = get_annotation_range_lock_and_unannotated_count(
             variant_annotation_version, settings.ANNOTATION_VEP_BATCH_MIN, settings.ANNOTATION_VEP_BATCH_MAX)
         if range_lock is None:
             break
         range_lock.save()
-        annotation_run = AnnotationRun.objects.create(annotation_range_lock=range_lock,
-                                                      pipeline_type=pipeline_type, external=True)
-        # If already counted off-thread (AnnotationRun.count, #1646), reject a too-small run before dumping.
-        # A counted run here is always non-empty (the count lane finishes count==0 runs itself).
-        if annotation_run.count is not None and annotation_run.count < min_variants:
-            _revert_too_small_run(annotation_run, annotation_run.count, min_variants, reverted)
-            continue
-        dump_count = get_runner(pipeline_type).dump(annotation_run, dump_dir=output_dir)
-        # Range locks are sized without pipeline_type, so an SV dump leaves most locks with zero SVs; such a
-        # run is already FINISHED (get_status: dump_count == 0) with nothing to annotate, so skip its sidecar
-        # rather than emit a no-op .meta.json (thousands of them, for a handful of SVs).
-        if dump_count == 0:
-            continue
-        if dump_count < min_variants:  # uncounted run whose dump turned out too small
-            _revert_too_small_run(annotation_run, dump_count, min_variants, reverted)
-            continue
-        if identity is None:
-            identity = variant_annotation_version_identity(variant_annotation_version)
-        meta_filename = write_dump_metadata(annotation_run, dump_dir=output_dir, identity=identity)
-        logging.info("Dumped external AnnotationRun %s: %d variants -> %s (meta %s)",
-                     annotation_run.pk, dump_count, annotation_run.vcf_dump_filename, meta_filename)
-        annotation_runs.append(annotation_run)
-
-    write_snakemake_bundle(output_dir, variant_annotation_version, pipeline_type=pipeline_type)
-    if reverted:
-        dispatch_annotation_runs.si(variant_annotation_version.pk).apply_async()
-    logging.info("Dumped %d external annotation run(s) for %s into %s (reverted %d too-small run(s) to local)",
-                 len(annotation_runs), variant_annotation_version, output_dir, len(reverted))
-    return annotation_runs
+        dump.add_run(AnnotationRun.objects.create(annotation_range_lock=range_lock,
+                                                  pipeline_type=pipeline_type, external=True))
+    return dump.finish()
 
 
 def dump_existing_annotation_runs(variant_annotation_version: VariantAnnotationVersion,
@@ -227,8 +249,8 @@ def dump_existing_annotation_runs(variant_annotation_version: VariantAnnotationV
     if leave < 0:
         raise ValueError(f"leave must be >= 0, got {leave}")
     _require_sv_offload_supported(pipeline_type)
+    dump = _ExternalDump(variant_annotation_version, output_dir, pipeline_type, min_variants, "existing")
 
-    os.makedirs(output_dir, exist_ok=True)
     now = timezone.now()
     # Mirror the dispatcher's dispatchable filter (annotation_scheduler_task._dispatchable_runs_qs) and its
     # lowest-min-variant-first order, so we adopt exactly the runs it would otherwise launch.
@@ -246,11 +268,6 @@ def dump_existing_annotation_runs(variant_annotation_version: VariantAnnotationV
     logging.info("dump_existing: %d dispatchable run(s); leaving %d on the local pipeline, dumping %d",
                  len(kept) + len(candidate_ids), len(kept), len(candidate_ids))
 
-    # Every run in this dump shares the same VAV, so compute the version identity (which runs VEP) once,
-    # lazily on the first claimed run so a no-op dump never invokes VEP.
-    identity = None
-    annotation_runs = []
-    reverted = []
     for pk in candidate_ids:
         # Atomically claim as external only while still dispatchable (same filter as the dispatcher) so we
         # never adopt a run it just leased. If we lose the race (0 rows updated) skip it; if we win, the run
@@ -266,33 +283,9 @@ def dump_existing_annotation_runs(variant_annotation_version: VariantAnnotationV
                             "scheduler?)", pk)
             continue
 
-        annotation_run = AnnotationRun.objects.get(pk=pk)
-        # If already counted off-thread (AnnotationRun.count, #1646), reject a too-small run before dumping.
-        # A counted run here is always non-empty (the count lane finishes count==0 runs itself).
-        if annotation_run.count is not None and annotation_run.count < min_variants:
-            _revert_too_small_run(annotation_run, annotation_run.count, min_variants, reverted)
-            continue
-        dump_count = get_runner(pipeline_type).dump(annotation_run, dump_dir=output_dir)
-        # A zero-count run is already FINISHED (get_status: dump_count == 0) with nothing to annotate, so
-        # skip its sidecar rather than emit a no-op .meta.json (see dump_external_annotation_runs).
-        if dump_count == 0:
-            continue
-        if dump_count < min_variants:  # uncounted run whose dump turned out too small
-            _revert_too_small_run(annotation_run, dump_count, min_variants, reverted)
-            continue
-        if identity is None:
-            identity = variant_annotation_version_identity(variant_annotation_version)
-        meta_filename = write_dump_metadata(annotation_run, dump_dir=output_dir, identity=identity)
-        logging.info("Dumped existing AnnotationRun %s: %d variants -> %s (meta %s)",
-                     annotation_run.pk, dump_count, annotation_run.vcf_dump_filename, meta_filename)
-        annotation_runs.append(annotation_run)
-
-    write_snakemake_bundle(output_dir, variant_annotation_version, pipeline_type=pipeline_type)
-    if reverted:
-        dispatch_annotation_runs.si(variant_annotation_version.pk).apply_async()
-    logging.info("Dumped %d existing annotation run(s) for %s into %s (reverted %d too-small run(s) to local)",
-                 len(annotation_runs), variant_annotation_version, output_dir, len(reverted))
-    return annotation_runs
+        dump.add_run(AnnotationRun.objects.get(pk=pk))
+
+    return dump.finish()
 
 
 def verify_annotated_vcf_variant_ids(annotation_run: AnnotationRun, meta: dict):

annotation/fake_data.py

@@ -0,0 +1,24 @@
+"""
+Helpers shared by the 'manage.py create_fake_data' subcommands, which live in the app whose data they
+create (@see snpdb.management.commands.create_fake_data).
+"""
+from annotation.models import VariantAnnotation, VariantAnnotationVersion
+from genes.models import TranscriptVersion
+from snpdb.models import GenomeBuild
+
+
+def get_variant_ids_by_gene(genome_build: GenomeBuild, genes: list[str]) -> dict[str, list[int]]:
+    gene_symbol_field = "transcript_version__gene_version__gene_symbol_id"
+    transcript_versions_qs = TranscriptVersion.objects.filter(genome_build=genome_build,
+                                                              gene_version__gene_symbol__in=genes)
+    variant_annotation_qs = VariantAnnotation.objects.filter(version=VariantAnnotationVersion.latest(genome_build),
+                                                             transcript_version__in=transcript_versions_qs)
+    variant_ids_by_gene = {}
+    for gene_symbol, variant_id in variant_annotation_qs.values_list(gene_symbol_field, "variant_id"):
+        variant_ids_by_gene.setdefault(gene_symbol, []).append(variant_id)
+    return variant_ids_by_gene
+
+
+def zipf_weight(rank: int) -> float:
+    """ Long tail - a few genes/records get most of the action, the rest get a little """
+    return 1 / (rank + 1) ** 0.8

classification/fake_reclassifications.py

@@ -29,13 +29,12 @@
 from django.utils.timezone import get_current_timezone, now
 from guardian.models import GroupObjectPermission
 
-from annotation.models import VariantAnnotation, VariantAnnotationVersion
+from annotation.fake_data import get_variant_ids_by_gene, zipf_weight
 from classification.enums import (ClinicalSignificance, CriteriaEvaluation, ShareLevel, SpecialEKeys,
                                   SubmissionSource)
 from classification.models import (Classification, ClassificationModification, ImportedAlleleInfo,
                                    ReclassificationEvent, ReclassificationEventBuilder, ResolvedVariantInfo)
 from classification.models.classification_variant_info_models import ImportedAlleleInfoStatus
-from genes.models import TranscriptVersion
 from library.guardian_utils import all_users_group
 from snpdb.models import GenomeBuild, GenomeBuildPatchVersion, Lab, Organization
 
@@ -241,7 +240,7 @@ def _create_labs_and_curators(self) -> tuple[dict[str, Lab], list[User]]:
     def _pick_records(self, genome_build: GenomeBuild, labs: dict[str, Lab],
                       num_classifications: int, years: int) -> list[FakeRecord]:
         """ Real variants, so the gene chart has real symbols and the records link somewhere sensible """
-        variant_ids_by_gene = _variant_ids_by_gene(genome_build, GENES)
+        variant_ids_by_gene = get_variant_ids_by_gene(genome_build, GENES)
         lab_weights = [fake_lab.weight for fake_lab in FAKE_LABS]
         significances = list(SIGNIFICANCE_BEHAVIOUR)
         significance_weights = [SIGNIFICANCE_BEHAVIOUR[s].weight for s in significances]
@@ -251,8 +250,8 @@ def _pick_records(self, genome_build: GenomeBuild, labs: dict[str, Lab],
         records = []
         for index, gene_symbol in enumerate(GENES):
             available = variant_ids_by_gene.get(gene_symbol, [])
-            wanted = round(num_classifications * _zipf_weight(index)
-                           / sum(_zipf_weight(i) for i in range(len(GENES))))
+            wanted = round(num_classifications * zipf_weight(index)
+                           / sum(zipf_weight(i) for i in range(len(GENES))))
             if len(available) < wanted:
                 self.stdout.write(f"{gene_symbol}: only {len(available)} annotated variants, wanted {wanted}")
                 wanted = len(available)
@@ -408,22 +407,6 @@ def _period_start(years: int) -> datetime:
     return now().astimezone(get_current_timezone()) - timedelta(days=years * 365)
 
 
-def _zipf_weight(rank: int) -> float:
-    return 1 / (rank + 1) ** 0.8
-
-
-def _variant_ids_by_gene(genome_build: GenomeBuild, genes: list[str]) -> dict[str, list[int]]:
-    gene_symbol_field = "transcript_version__gene_version__gene_symbol_id"
-    transcript_versions_qs = TranscriptVersion.objects.filter(genome_build=genome_build,
-                                                              gene_version__gene_symbol__in=genes)
-    variant_annotation_qs = VariantAnnotation.objects.filter(version=VariantAnnotationVersion.latest(genome_build),
-                                                             transcript_version__in=transcript_versions_qs)
-    variant_ids_by_gene = {}
-    for gene_symbol, variant_id in variant_annotation_qs.values_list(gene_symbol_field, "variant_id"):
-        variant_ids_by_gene.setdefault(gene_symbol, []).append(variant_id)
-    return variant_ids_by_gene
-
-
 def _summary(step: FakeStep) -> dict:
     return {"clinical_significance": step.significance, "date": step.curated_on.isoformat()}
 

library/graphs/chromosomes_graph.py

@@ -4,6 +4,7 @@
 
 from collections import defaultdict
 
+import numpy as np
 import pandas as pd
 
 from library.genomics import format_chrom
@@ -81,3 +82,30 @@ def plot_chromosomes(ax, cytoband_filename, has_chr=False, **kwargs):
     ax.set_yticklabels(yticklabels)
     ax.set_xticks([])
     return chrom_ranges
+
+
+def plot_chromosome_bin_values(ax, chrom_ranges, bin_values_by_chrom, bin_size, cmap, vmax,
+                               padding, alpha):
+    """ Draw a heatmap band down each chromosome from its per-bin values, all sharing one colour scale.
+
+        Returns one of the QuadMeshes, for the figure to hang a colorbar off. """
+    quadmesh = None
+    for chrom, (_xranges, yranges) in chrom_ranges.items():
+        bins = bin_values_by_chrom[chrom]
+        num_bins = len(bins) + 1
+        x_pos = np.arange(num_bins) * bin_size
+
+        # Get into right dimensions
+        bins = [bins]
+        # pcolor says x,y should have dimensions 1 greater than colors
+        x_pos = [x_pos, x_pos]
+        y_top = np.empty(num_bins)
+        y_top.fill(yranges[0])
+        y_bottom = np.empty(num_bins)
+        y_bottom.fill(yranges[0] + yranges[1])
+        y_pos = [y_top + padding, y_bottom - padding]
+
+        masked_c = np.ma.masked_invalid(np.array(bins))
+        quadmesh = ax.pcolormesh(np.array(x_pos), np.array(y_pos), masked_c, cmap=cmap, alpha=alpha)
+        quadmesh.set_clim(vmin=0, vmax=vmax)
+    return quadmesh