split models into CatalogScoreVariant and ScoreVariant

Author: nebfield

pgscatalog.core/src/pgscatalog/core/lib/_normalise.py

@@ -9,7 +9,7 @@
 import pyliftover
 
 from .genomebuild import GenomeBuild
-from .scorevariant import EffectType, ScoreVariant, EffectAllele
+from .scorevariant import ScoreVariant, Allele
 from .pgsexceptions import LiftoverError
 
 logger = logging.getLogger(__name__)
@@ -45,8 +45,6 @@ def normalise(
     if drop_missing:
         variants = drop_hla(variants)
 
-    variants = assign_effect_type(variants)
-    variants = check_effect_weight(variants)
     variants = assign_other_allele(variants)
     variants = check_effect_allele(variants, drop_missing)
     variants = detect_complex(variants)
@@ -75,19 +73,11 @@ def check_duplicates(variants):
             seen_ids = {}
             current_accession = accession
 
-        # None other allele -> empty string
-        variant_id: str = ":".join(
-            [
-                str(getattr(variant, k) or "")
-                for k in ["chr_name", "chr_position", "effect_allele", "other_allele"]
-            ]
-        )
-
-        if variant_id in seen_ids:
+        if variant.variant_id in seen_ids:
             variant.is_duplicated = True
             n_duplicates += 1
 
-        seen_ids[variant_id] = True
+        seen_ids[variant.variant_id] = True
 
         yield variant
         n_variants += 1
@@ -110,9 +100,12 @@ def drop_hla(variants):
     []
     """
     n_dropped = 0
+    p = Allele(allele="P")
+    n = Allele(allele="N")
+
     for variant in variants:
         match variant:
-            case _ if variant.effect_allele in (EffectAllele("P"), EffectAllele("N")):
+            case _ if variant.effect_allele in (p, n):
                 n_dropped += 1
                 continue
             case _:
@@ -121,39 +114,15 @@ def drop_hla(variants):
     logger.warning(f"{n_dropped} HLA alleles detected and dropped")
 
 
-def check_effect_weight(variants):
-    """Check that effect weights are valid floats. Effect weights are intentionally
-    left as strings during processing.
-
-    >>> variant = ScoreVariant(**{"effect_allele": "A", "effect_weight": 5, "accession": "test", "row_nr": 0})
-    >>> list(check_effect_weight([variant])) # doctest: +ELLIPSIS
-    [ScoreVariant(effect_allele='A',effect_weight=5,...)]
-
-    >>> variant = ScoreVariant(**{"effect_allele": "A", "effect_weight": "potato", "accession": "test", "row_nr": 0})
-    >>> list(check_effect_weight([variant])) # doctest: +ELLIPSIS
-    Traceback (most recent call last):
-    ...
-    ValueError
-    """
-    for variant in variants:
-        try:
-            float(variant.effect_weight)
-        except ValueError as e:
-            logger.critical(f"{variant} has bad effect weight")
-            raise ValueError from e
-        else:
-            yield variant
-
-
 def assign_other_allele(variants):
     """Check if there's more than one possible other allele, remove if true
 
-    >>> variant = ScoreVariant(**{"effect_allele": "A", "effect_weight": 5, "accession": "test", "row_nr": 0, "other_allele": "A"})
+    >>> variant = ScoreVariant(**{"chr_position": 1, "rsID": None, "chr_name": "1", "effect_allele": "A", "effect_weight": 5, "other_allele": "A"})
     >>> list(assign_other_allele([variant])) # doctest: +ELLIPSIS
-    [ScoreVariant(effect_allele='A',...,other_allele='A',...)]
-    >>> variant = ScoreVariant(**{"effect_allele": "A", "effect_weight": 5, "accession": "test", "row_nr": 0, "other_allele": "A/C"})
+    [ScoreVariant(..., effect_allele='A', other_allele='A', ...)]
+    >>> variant = ScoreVariant(**{"chr_position": 1, "rsID": None, "chr_name": "1", "effect_allele": "A", "effect_weight": 5, "other_allele": "A/C"})
     >>> list(assign_other_allele([variant])) # doctest: +ELLIPSIS
-    [ScoreVariant(effect_allele='A',...,other_allele=None,...)]
+    [ScoreVariant(..., effect_allele='A', other_allele=None,...)]
     """
     n_dropped = 0
     for variant in variants:
@@ -171,51 +140,16 @@ def assign_other_allele(variants):
         logger.warning("Other allele for these variants is set to missing")
 
 
-def assign_effect_type(variants):
-    """Convert PGS Catalog effect type columns to EffectType enums
-
-    The most common type of effect type is additive:
-
-    >>> variant = ScoreVariant(**{"effect_allele": "A", "effect_weight": 5, "accession": "test", "row_nr": 0, "is_recessive": "False", "is_dominant": "False"})
-    >>> list(assign_effect_type([variant])) # doctest: +ELLIPSIS
-    [ScoreVariant(...,effect_type=EffectType.ADDITIVE,...)]
-    >>> variant = ScoreVariant(**{"effect_allele": "A", "effect_weight": 5, "accession": "test", "row_nr": 0, "is_recessive": "True", "is_dominant": "False"})
-    >>> list(assign_effect_type([variant])) # doctest: +ELLIPSIS
-    [ScoreVariant(...,effect_type=EffectType.RECESSIVE,...)]
-    >>> variant = ScoreVariant(**{"effect_allele": "A", "effect_weight": 5, "accession": "test", "row_nr": 0, "is_recessive": "False", "is_dominant": "True"})
-    >>> list(assign_effect_type([variant])) # doctest: +ELLIPSIS
-    [ScoreVariant(...,effect_type=EffectType.DOMINANT,...)]
-
-    is_recessive and is_dominant fields are parsed from strings to bools during __init__.
-    """
-    for variant in variants:
-        match (variant.is_recessive, variant.is_dominant):
-            case (None, None) | (False, False) | (None, False) | (False, None):
-                # none is OK because is_recessive or is_dominant column may be missing
-                # default value is already set to additive, so just yield the variant
-                pass
-            case (False, True) | (None, True):
-                # none is OK because is_recessive column may be missing
-                variant.effect_type = EffectType.DOMINANT
-            case (True, False) | (True, None):
-                # none is OK because is_dominant column may be missing
-                variant.effect_type = EffectType.RECESSIVE
-            case _:
-                logger.critical(f"Bad effect type setting: {variant}")
-                raise Exception
-        yield variant
-
-
 def remap_harmonised(variants, harmonised, target_build):
     """
     Overwrite key attributes with harmonised data, if available.
 
     In this case chr_name, chr_position, and other allele are missing.
     Perhaps authors submitted rsID and effect allele originally:
 
-    >>> variant = ScoreVariant(**{"effect_allele": "A", "effect_weight": 5, "accession": "test", "row_nr": 0, "hm_chr": 1, "hm_pos": 100, "hm_inferOtherAllele": "A"})
+    >>> variant = ScoreVariant(**{"chr_position": 1, "rsID": None, "chr_name": "2", "effect_allele": "A", "effect_weight": 5, "accession": "test", "hm_chr": 1, "hm_pos": 100, "hm_inferOtherAllele": "A"})
     >>> list(remap_harmonised([variant], harmonised=True, target_build=GenomeBuild.GRCh38)) # doctest: +ELLIPSIS
-    [ScoreVariant(...,chr_name=1,chr_position=100,...other_allele='A'...)]
+    [ScoreVariant(chr_name=1,chr_position=100,...other_allele='A'...)]
     """
     if harmonised:
         for variant in variants: