Re-implement multi-target mapping

Author: sallybg

src/api/routers/map.py

@@ -1,7 +1,6 @@
 """"Provide mapping router"""
 from pathlib import Path
 
-from cool_seq_tool.schemas import AnnotationLayer
 from fastapi import APIRouter, HTTPException
 from fastapi.responses import JSONResponse
 from requests import HTTPError
@@ -140,46 +139,60 @@ async def map_scoreset(urn: str, store_path: Path | None = None) -> ScoresetMapp
             error_message="No annotated variant mappings available for this score set",
         )
 
-    # TODO this will need to be changed to support multi-target score sets.
-    # This version works for accession based score sets.
-    # Not implementing multi-target changes because this will require corresponding changes on mavedb-api and we want to get this on staging quickly right now.
-    # For now, only accept single-target score sets so that we don't need to change structure of JSON output.
-    target_gene = list(metadata.target_genes.keys())[0]  # noqa: RUF015
     try:
         raw_metadata = get_raw_scoreset_metadata(urn, store_path)
-        preferred_layers = {
-            _set_scoreset_layer(urn, annotated_vrs_results[target_gene]),
-        }
-
-        reference_sequences = {
-            layer: {
-                "computed_reference_sequence": None,
-                "mapped_reference_sequence": None,
+        reference_sequences: dict[str, dict] = {}
+        mapped_scores: list[ScoreAnnotation] = []
+        for target_gene in annotated_vrs_results:
+            preferred_layers = {
+                _set_scoreset_layer(urn, annotated_vrs_results[target_gene]),
             }
-            for layer in AnnotationLayer
-        }
-        # sometimes Nonetype layers show up in preferred layers dict; remove these
-        preferred_layers.discard(None)
-        for layer in preferred_layers:
-            reference_sequences[layer][
-                "computed_reference_sequence"
-            ] = _get_computed_reference_sequence(
-                metadata.target_genes[target_gene], layer, transcripts[target_gene]
-            )
-            reference_sequences[layer][
-                "mapped_reference_sequence"
-            ] = _get_mapped_reference_sequence(
-                metadata.target_genes[target_gene],
-                layer,
-                transcripts[target_gene],
-                alignment_results[target_gene],
-            )
+            reference_sequences[target_gene] = {
+                layer: {
+                    "computed_reference_sequence": None,
+                    "mapped_reference_sequence": None,
+                }
+                for layer in preferred_layers
+            }
+            # sometimes Nonetype layers show up in preferred layers dict; remove these
+            preferred_layers.discard(None)
+            for layer in preferred_layers:
+                reference_sequences[target_gene][layer][
+                    "computed_reference_sequence"
+                ] = _get_computed_reference_sequence(
+                    metadata.target_genes[target_gene], layer, transcripts[target_gene]
+                )
+                reference_sequences[target_gene][layer][
+                    "mapped_reference_sequence"
+                ] = _get_mapped_reference_sequence(
+                    metadata.target_genes[target_gene],
+                    layer,
+                    transcripts[target_gene],
+                    alignment_results[target_gene],
+                )
+
+            for m in annotated_vrs_results[target_gene]:
+                if m.pre_mapped is None:
+                    mapped_scores.append(ScoreAnnotation(**m.model_dump()))
+                elif m.annotation_layer in preferred_layers:
+                    # drop annotation layer from mapping object
+                    mapped_scores.append(ScoreAnnotation(**m.model_dump()))
+
+            # drop Nonetype reference sequences
+            for target_gene in reference_sequences:
+                for layer in list(reference_sequences[target_gene].keys()):
+                    if (
+                        reference_sequences[target_gene][layer][
+                            "mapped_reference_sequence"
+                        ]
+                        is None
+                        and reference_sequences[target_gene][layer][
+                            "computed_reference_sequence"
+                        ]
+                        is None
+                    ) or layer is None:
+                        del reference_sequences[target_gene][layer]
 
-        mapped_scores: list[ScoreAnnotation] = []
-        for m in annotated_vrs_results[target_gene]:
-            if m.annotation_layer in preferred_layers:
-                # drop annotation layer from mapping object
-                mapped_scores.append(ScoreAnnotation(**m.model_dump()))
     except Exception as e:
         return JSONResponse(
             content=ScoresetMapping(
@@ -190,18 +203,7 @@ async def map_scoreset(urn: str, store_path: Path | None = None) -> ScoresetMapp
     return JSONResponse(
         content=ScoresetMapping(
             metadata=raw_metadata,
-            computed_protein_reference_sequence=reference_sequences[
-                AnnotationLayer.PROTEIN
-            ]["computed_reference_sequence"],
-            mapped_protein_reference_sequence=reference_sequences[
-                AnnotationLayer.PROTEIN
-            ]["mapped_reference_sequence"],
-            computed_genomic_reference_sequence=reference_sequences[
-                AnnotationLayer.GENOMIC
-            ]["computed_reference_sequence"],
-            mapped_genomic_reference_sequence=reference_sequences[
-                AnnotationLayer.GENOMIC
-            ]["mapped_reference_sequence"],
+            reference_sequences=reference_sequences,
             mapped_scores=mapped_scores,
         ).model_dump(exclude_none=True)
     )

src/dcd_mapping/annotate.py

@@ -603,8 +603,7 @@ def save_mapped_output_json(
                 "computed_reference_sequence": None,
                 "mapped_reference_sequence": None,
             }
-            # TODO change this back after reimplementing multi-target mapping
-            for layer in AnnotationLayer
+            for layer in preferred_layers
         }
         # sometimes Nonetype layers show up in preferred layers dict; remove these
         preferred_layers.discard(None)
@@ -630,43 +629,21 @@ def save_mapped_output_json(
                 # drop annotation layer from mapping object
                 mapped_scores.append(ScoreAnnotation(**m.model_dump()))
 
-        # TODO drop this "continue" after reimplementing multi-target mapping
-        continue
-
-        # TODO add this back after reimplementing multi-target mapping
         # drop Nonetype reference sequences
-        # for target_gene in reference_sequences:
-        #     for layer in list(reference_sequences[target_gene].keys()):
-        #         if (
-        #             reference_sequences[target_gene][layer]["mapped_reference_sequence"]
-        #             is None
-        #             and reference_sequences[target_gene][layer][
-        #                 "computed_reference_sequence"
-        #             ]
-        #             is None
-        #         ) or layer is None:
-        #             del reference_sequences[target_gene][layer]
-
-        # TODO drop this "continue" after reimplementing multi-target mapping
-        continue
-    # TODO drop this after reimplementing multi-target mapping
-    reference_sequences = reference_sequences.popitem()[1]  # get only value in dict
-    # TODO change this back after reimplementing multi-target mapping
-    # this only works for --prefer_genomic right now, which is fine because we're going to change it back after reimplementing multi-target mapping
+        for target_gene in reference_sequences:
+            for layer in list(reference_sequences[target_gene].keys()):
+                if (
+                    reference_sequences[target_gene][layer]["mapped_reference_sequence"]
+                    is None
+                    and reference_sequences[target_gene][layer][
+                        "computed_reference_sequence"
+                    ]
+                    is None
+                ) or layer is None:
+                    del reference_sequences[target_gene][layer]
+
     output = ScoresetMapping(
         metadata=metadata.model_dump(),
-        computed_protein_reference_sequence=reference_sequences[
-            AnnotationLayer.PROTEIN
-        ]["computed_reference_sequence"],
-        mapped_protein_reference_sequence=reference_sequences[AnnotationLayer.PROTEIN][
-            "mapped_reference_sequence"
-        ],
-        computed_genomic_reference_sequence=reference_sequences[
-            AnnotationLayer.GENOMIC
-        ]["computed_reference_sequence"],
-        mapped_genomic_reference_sequence=reference_sequences[AnnotationLayer.GENOMIC][
-            "mapped_reference_sequence"
-        ],
         reference_sequences=reference_sequences,
         mapped_scores=mapped_scores,
     )

src/dcd_mapping/mavedb_data.py

@@ -187,9 +187,6 @@ def get_scoreset_metadata(
     metadata = get_raw_scoreset_metadata(scoreset_urn, dcd_mapping_dir)
     target_genes = {}
     multi_target = len(metadata["targetGenes"]) > 1
-    if multi_target:
-        msg = f"Multiple target genes for {scoreset_urn}. Multi-target score sets are not currently supported."
-        raise ScoresetNotSupportedError(msg)
 
     for gene in metadata["targetGenes"]:
         if not _metadata_response_is_human(metadata):

src/dcd_mapping/schemas.py

@@ -205,17 +205,16 @@ class ScoresetMapping(BaseModel):
     mapped_date_utc: str = Field(
         default=datetime.datetime.now(tz=datetime.UTC).isoformat()
     )
-    # TODO re-implement metadata change later to support multi-target score sets. will require corresponding changes in mavedb-api
-    # reference_sequences: dict[
-    #     str,
-    #     dict[
-    #         AnnotationLayer,
-    #         dict[str, ComputedReferenceSequence | MappedReferenceSequence | None],
-    #     ],
-    # ] | None = None
-    computed_protein_reference_sequence: ComputedReferenceSequence | MappedReferenceSequence | None = None
-    mapped_protein_reference_sequence: MappedReferenceSequence | None = None
-    computed_genomic_reference_sequence: ComputedReferenceSequence | MappedReferenceSequence | None = None
-    mapped_genomic_reference_sequence: MappedReferenceSequence | None = None
+    reference_sequences: dict[
+        str,
+        dict[
+            AnnotationLayer,
+            dict[str, ComputedReferenceSequence | MappedReferenceSequence | None],
+        ],
+    ] | None = None
+    # computed_protein_reference_sequence: ComputedReferenceSequence | MappedReferenceSequence | None = None
+    # mapped_protein_reference_sequence: MappedReferenceSequence | None = None
+    # computed_genomic_reference_sequence: ComputedReferenceSequence | MappedReferenceSequence | None = None
+    # mapped_genomic_reference_sequence: MappedReferenceSequence | None = None
     mapped_scores: list[ScoreAnnotation] | None = None
     error_message: str | None = None