Merge pull request #465 from VariantEffect/jstone-dev/clinical-heatmap

API endpoint for score set variant scores with mapped coordinates

Author: bencap

src/mavedb/lib/score_sets.py

@@ -5,10 +5,11 @@
 from operator import attrgetter
 from typing import Any, BinaryIO, Iterable, Optional, TYPE_CHECKING, Sequence, Literal
 
+from mavedb.models.mapped_variant import MappedVariant
 import numpy as np
 import pandas as pd
 from pandas.testing import assert_index_equal
-from sqlalchemy import Integer, cast, func, or_, select
+from sqlalchemy import Integer, and_, cast, func, or_, select
 from sqlalchemy.orm import Session, aliased, contains_eager, joinedload, selectinload
 
 from mavedb.lib.exceptions import ValidationError
@@ -17,6 +18,7 @@
     HGVS_NT_COLUMN,
     HGVS_PRO_COLUMN,
     HGVS_SPLICE_COLUMN,
+    REQUIRED_SCORE_COLUMN,
     VARIANT_COUNT_DATA,
     VARIANT_SCORE_DATA,
 )
@@ -57,6 +59,9 @@
 
 logger = logging.getLogger(__name__)
 
+HGVS_G_REGEX = re.compile(r"(^|:)g\.")
+HGVS_P_REGEX = re.compile(r"(^|:)p\.")
+
 
 class HGVSColumns:
     NUCLEOTIDE: str = "hgvs_nt"  # dataset.constants.hgvs_nt_column
@@ -402,26 +407,90 @@ def get_score_set_variants_as_csv(
     start: Optional[int] = None,
     limit: Optional[int] = None,
     drop_na_columns: Optional[bool] = None,
+    include_custom_columns: bool = True,
+    include_post_mapped_hgvs: bool = False,
 ) -> str:
+    """
+    Get the variant data from a score set as a CSV string.
+
+    Parameters
+    __________
+    db : Session
+        The database session to use.
+    score_set : ScoreSet
+        The score set to get the variants from.
+    data_type : {'scores', 'counts'}
+        The type of data to get. Either 'scores' or 'counts'.
+    start : int, optional
+        The index to start from. If None, starts from the beginning.
+    limit : int, optional
+        The maximum number of variants to return. If None, returns all variants.
+    drop_na_columns : bool, optional
+        Whether to drop columns that contain only NA values. Defaults to False.
+    include_custom_columns : bool, optional
+        Whether to include custom columns defined in the score set. Defaults to True.
+    include_post_mapped_hgvs : bool, optional
+        Whether to include post-mapped HGVS notations in the output. Defaults to False. If True, the output will include
+        columns for both post-mapped HGVS genomic (g.) and protein (p.) notations.
+
+    Returns
+    _______
+    str
+        The CSV string containing the variant data.
+    """
     assert type(score_set.dataset_columns) is dict
-    dataset_cols = "score_columns" if data_type == "scores" else "count_columns"
+    custom_columns_set = "score_columns" if data_type == "scores" else "count_columns"
     type_column = "score_data" if data_type == "scores" else "count_data"
 
-    count_columns = [str(x) for x in list(score_set.dataset_columns.get(dataset_cols, []))]
-    columns = ["accession", "hgvs_nt", "hgvs_splice", "hgvs_pro"] + count_columns
-
-    variants_query = (
-        select(Variant)
-        .where(Variant.score_set_id == score_set.id)
-        .order_by(cast(func.split_part(Variant.urn, "#", 2), Integer))
-    )
-    if start:
-        variants_query = variants_query.offset(start)
-    if limit:
-        variants_query = variants_query.limit(limit)
-    variants = db.scalars(variants_query).all()
+    columns = ["accession", "hgvs_nt", "hgvs_splice", "hgvs_pro"]
+    if include_post_mapped_hgvs:
+        columns.append("post_mapped_hgvs_g")
+        columns.append("post_mapped_hgvs_p")
+
+    if include_custom_columns:
+        custom_columns = [str(x) for x in list(score_set.dataset_columns.get(custom_columns_set, []))]
+        columns += custom_columns
+    elif data_type == "scores":
+        columns.append(REQUIRED_SCORE_COLUMN)
+
+    variants: Sequence[Variant] = []
+    mappings: Optional[list[Optional[MappedVariant]]] = None
+
+    if include_post_mapped_hgvs:
+        variants_and_mappings_query = (
+            select(Variant, MappedVariant)
+            .join(
+                MappedVariant,
+                and_(Variant.id == MappedVariant.variant_id, MappedVariant.current.is_(True)),
+                isouter=True,
+            )
+            .where(Variant.score_set_id == score_set.id)
+            .order_by(cast(func.split_part(Variant.urn, "#", 2), Integer))
+        )
+        if start:
+            variants_and_mappings_query = variants_and_mappings_query.offset(start)
+        if limit:
+            variants_and_mappings_query = variants_and_mappings_query.limit(limit)
+        variants_and_mappings = db.execute(variants_and_mappings_query).all()
+
+        variants = []
+        mappings = []
+        for variant, mapping in variants_and_mappings:
+            variants.append(variant)
+            mappings.append(mapping)
+    else:
+        variants_query = (
+            select(Variant)
+            .where(Variant.score_set_id == score_set.id)
+            .order_by(cast(func.split_part(Variant.urn, "#", 2), Integer))
+        )
+        if start:
+            variants_query = variants_query.offset(start)
+        if limit:
+            variants_query = variants_query.limit(limit)
+        variants = db.scalars(variants_query).all()
 
-    rows_data = variants_to_csv_rows(variants, columns=columns, dtype=type_column)  # type: ignore
+    rows_data = variants_to_csv_rows(variants, columns=columns, dtype=type_column, mappings=mappings)  # type: ignore
     if drop_na_columns:
         rows_data, columns = drop_na_columns_from_csv_file_rows(rows_data, columns)
 
@@ -462,7 +531,63 @@ def is_null(value):
     return null_values_re.fullmatch(value) or not value
 
 
-def variant_to_csv_row(variant: Variant, columns: list[str], dtype: str, na_rep="NA") -> dict[str, Any]:
+def hgvs_from_vrs_allele(allele: dict) -> str:
+    """
+    Extract the HGVS notation from the VRS allele.
+    """
+    try:
+        # VRS 2.X
+        return allele["expressions"][0]["value"]
+    except KeyError:
+        raise ValueError("VRS 1.X format not supported.")
+        # VRS 1.X. We don't want to allow this.
+        # return allele["variation"]["expressions"][0]["value"]
+
+
+def get_hgvs_from_mapped_variant(post_mapped_vrs: Any) -> Optional[str]:
+    if post_mapped_vrs["type"] == "Haplotype":  # type: ignore
+        variations_hgvs = [hgvs_from_vrs_allele(allele) for allele in post_mapped_vrs["members"]]
+    elif post_mapped_vrs["type"] == "CisPhasedBlock":  # type: ignore
+        variations_hgvs = [hgvs_from_vrs_allele(allele) for allele in post_mapped_vrs["members"]]
+    elif post_mapped_vrs["type"] == "Allele":  # type: ignore
+        variations_hgvs = [hgvs_from_vrs_allele(post_mapped_vrs)]
+    else:
+        return None
+
+    if len(variations_hgvs) == 0:
+        return None
+        # raise ValueError(f"No variations found in variant {variant_urn}.")
+    if len(variations_hgvs) > 1:
+        return None
+        # raise ValueError(f"Multiple variations found in variant {variant_urn}.")
+
+    return variations_hgvs[0]
+
+
+# TODO (https://github.com/VariantEffect/mavedb-api/issues/440) Temporarily, we are using these functions to distinguish
+# genomic and protein HGVS strings produced by the mapper. Using hgvs.parser.Parser is too slow, and we won't need to do
+# this once the mapper extracts separate g., c., and p. post-mapped HGVS strings.
+def is_hgvs_g(hgvs: str) -> bool:
+    """
+    Check if the given HGVS string is a genomic HGVS (g.) string.
+    """
+    return bool(HGVS_G_REGEX.search(hgvs))
+
+
+def is_hgvs_p(hgvs: str) -> bool:
+    """
+    Check if the given HGVS string is a protein HGVS (p.) string.
+    """
+    return bool(HGVS_P_REGEX.search(hgvs))
+
+
+def variant_to_csv_row(
+    variant: Variant,
+    columns: list[str],
+    dtype: str,
+    mapping: Optional[MappedVariant] = None,
+    na_rep="NA",
+) -> dict[str, Any]:
     """
     Format a variant into a containing the keys specified in `columns`.
 
@@ -491,6 +616,18 @@ def variant_to_csv_row(variant: Variant, columns: list[str], dtype: str, na_rep=
             value = str(variant.hgvs_splice)
         elif column_key == "accession":
             value = str(variant.urn)
+        elif column_key == "post_mapped_hgvs_g":
+            hgvs_str = get_hgvs_from_mapped_variant(mapping.post_mapped) if mapping and mapping.post_mapped else None
+            if hgvs_str is not None and is_hgvs_g(hgvs_str):
+                value = hgvs_str
+            else:
+                value = ""
+        elif column_key == "post_mapped_hgvs_p":
+            hgvs_str = get_hgvs_from_mapped_variant(mapping.post_mapped) if mapping and mapping.post_mapped else None
+            if hgvs_str is not None and is_hgvs_p(hgvs_str):
+                value = hgvs_str
+            else:
+                value = ""
         else:
             parent = variant.data.get(dtype) if variant.data else None
             value = str(parent.get(column_key)) if parent else na_rep
@@ -502,7 +639,11 @@ def variant_to_csv_row(variant: Variant, columns: list[str], dtype: str, na_rep=
 
 
 def variants_to_csv_rows(
-    variants: Sequence[Variant], columns: list[str], dtype: str, na_rep="NA"
+    variants: Sequence[Variant],
+    columns: list[str],
+    dtype: str,
+    mappings: Optional[Sequence[Optional[MappedVariant]]] = None,
+    na_rep="NA",
 ) -> Iterable[dict[str, Any]]:
     """
     Format each variant into a dictionary row containing the keys specified in `columns`.
@@ -522,7 +663,12 @@ def variants_to_csv_rows(
     -------
     list[dict[str, Any]]
     """
-    return map(lambda v: variant_to_csv_row(v, columns, dtype, na_rep), variants)
+    if mappings is not None:
+        return map(
+            lambda pair: variant_to_csv_row(pair[0], columns, dtype, mapping=pair[1], na_rep=na_rep),
+            zip(variants, mappings),
+        )
+    return map(lambda v: variant_to_csv_row(v, columns, dtype, na_rep=na_rep), variants)
 
 
 def find_meta_analyses_for_score_sets(db: Session, urns: list[str]) -> list[ScoreSet]:

src/mavedb/routers/score_sets.py

@@ -225,6 +225,95 @@ async def show_score_set(
     return score_set.ScoreSet.from_orm(item).copy(update={"experiment": enriched_experiment})
 
 
+@router.get(
+    "/score-sets/{urn}/variants/data",
+    status_code=200,
+    responses={
+        200: {
+            "content": {"text/csv": {}},
+            "description": """Variant data in CSV format, with four fixed columns (accession, hgvs_nt, hgvs_pro,"""
+            """ and hgvs_splice), plus score columns defined by the score set.""",
+        }
+    },
+)
+def get_score_set_variants_csv(
+    *,
+    urn: str,
+    start: int = Query(default=None, description="Start index for pagination"),
+    limit: int = Query(default=None, description="Maximum number of variants to return"),
+    drop_na_columns: Optional[bool] = None,
+    db: Session = Depends(deps.get_db),
+    user_data: Optional[UserData] = Depends(get_current_user),
+) -> Any:
+    """
+    Return tabular variant data from a score set, identified by URN, in CSV format.
+
+    This differs from get_score_set_scores_csv() in that it returns only the HGVS columns, score column, and mapped HGVS
+    string.
+
+    TODO (https://github.com/VariantEffect/mavedb-api/issues/446) We may want to turn this into a general-purpose CSV
+    export endpoint, with options governing which columns to include.
+
+    Parameters
+    __________
+
+    Parameters
+    __________
+    urn : str
+        The URN of the score set to fetch variants from.
+    start : Optional[int]
+        The index to start from. If None, starts from the beginning.
+    limit : Optional[int]
+        The maximum number of variants to return. If None, returns all variants.
+    drop_na_columns : bool, optional
+        Whether to drop columns that contain only NA values. Defaults to False.
+    db : Session
+        The database session to use.
+    user_data : Optional[UserData]
+        The user data of the current user. If None, no user-specific permissions are checked.
+
+    Returns
+    _______
+    str
+        The CSV string containing the variant data.
+    """
+    save_to_logging_context(
+        {
+            "requested_resource": urn,
+            "resource_property": "scores",
+            "start": start,
+            "limit": limit,
+            "drop_na_columns": drop_na_columns,
+        }
+    )
+
+    if start and start < 0:
+        logger.info(msg="Could not fetch scores with negative start index.", extra=logging_context())
+        raise HTTPException(status_code=400, detail="Start index must be non-negative")
+    if limit is not None and limit <= 0:
+        logger.info(msg="Could not fetch scores with non-positive limit.", extra=logging_context())
+        raise HTTPException(status_code=400, detail="Limit must be positive")
+
+    score_set = db.query(ScoreSet).filter(ScoreSet.urn == urn).first()
+    if not score_set:
+        logger.info(msg="Could not fetch the requested scores; No such score set exists.", extra=logging_context())
+        raise HTTPException(status_code=404, detail=f"score set with URN '{urn}' not found")
+
+    assert_permission(user_data, score_set, Action.READ)
+
+    csv_str = get_score_set_variants_as_csv(
+        db,
+        score_set,
+        "scores",
+        start,
+        limit,
+        drop_na_columns,
+        include_custom_columns=False,
+        include_post_mapped_hgvs=True,
+    )
+    return StreamingResponse(iter([csv_str]), media_type="text/csv")
+
+
 @router.get(
     "/score-sets/{urn}/scores",
     status_code=200,