Add data types for the control data.

To help with type-checking, parsing, and generating correct control
files.

1 files changed, 164 insertions, 0 deletions

diff --git a/r_qtl/r_qtl2/types.py b/r_qtl/r_qtl2/types.py
new file mode 100644
index 0000000..13b8db4
--- /dev/null
+++ b/r_qtl/r_qtl2/types.py
@@ -0,0 +1,164 @@
+"""The types used by our R/qtl2 system."""
+from enum import Enum
+from pathlib import Path
+from dataclasses import field, dataclass
+from typing import (Union,
+                    Literal,
+                    Optional,
+                    Sequence,
+                    TypeAlias,
+                    Annotated)
+
+
+_ALLELE_COUNTS_BY_CROSSTYPE = {
+    "bc": 2, "f2": 2, "riself": 2, "risib": 2, "dh": 2, "haploid": 2, "ail": 2,
+    "ail3": 3,
+    "riself4": 4, "risib4": 4,
+    "dh6": 6,
+    "hs": 8, "do": 8, "riself8": 8, "risib8": 8,
+    "dof1": 9, # 8 DO founders + 1 base inbred line
+    "riself16": 16,
+    "magic19": 19
+}
+
+
+
+PathLike = Union[Path, str]
+
+AlleleLabel: TypeAlias = str
+
+GenoCodeAlleleValue = Union[
+    Literal[-1, 0, 1],
+    Literal[1, 2, 3],# T0D0: reconcile this to genocodes of 0, 1, 2 instead
+    Annotated[float, "Value must be between 0.0 and 2.0"]
+]
+
+GenoCode = dict[AlleleLabel, GenoCodeAlleleValue]
+
+# Sex information
+SexKey = Union[str, int]# codes used in the covariate file
+@dataclass(frozen=True)
+class ControlDataSex:
+    """Sex information in the control file."""
+    covar: Optional[str] = None
+    file: Optional[PathLike] = None
+    female: Optional[SexKey] = None
+    male: Optional[SexKey] = None
+
+
+# Cross information
+CrossType = Literal[
+    # T0D0: Look at https://kbroman.org/qtl2/assets/vignettes/input_files.html#Backcross to implement appropriate QC dependent on the crosstype.
+    "bc",# Backcross
+    "f2",# F2 intercross
+    "riself",# RIL by self -- RIL: Recombinant Inbred Lines
+    "risib",# RIL by sib
+    "dh",# Double haploid
+    "haploid",# Haploid
+    "ail",# AIL -- Advanced Inbred Lines
+    "hs",# Heterogeneous stock
+    "do",# Diversity outbreds
+    "riself4",# Multi-parent recombinant inbred lines
+    "riself8",# Multi-parent recombinant inbred lines
+    "riself16",# Multi-parent recombinant inbred lines
+    "risib4",# Multi-parent recombinant inbred lines
+    "risib8",# Multi-parent recombinant inbred lines
+    "magic19",# 19-way MAGIC lines
+    "dh6",# 6-way doubled haploids
+    "dof1",# DOF1
+    "ail3",# 3-way advanced intercross lines
+    "genail",# General advanced intercross lines
+    "genril",# General recombinant inbred lines
+]
+
+class CrossDirection(Enum):
+    FORWARD = 0
+    REVERSE = 1
+
+
+@dataclass(frozen=True)
+class ShortCrossInfo:
+    """Cross information in the control file."""
+    # "covar": indicates the name of the column in the covariate data
+    covar: Optional[str] = None
+    file: Optional[PathLike] = None
+    forward_crosses: Sequence[str] = field(default_factory=tuple)
+    reverse_crosses: Sequence[str] = field(default_factory=tuple)
+
+    def __post_init__(self):
+        """Validate the cross info."""
+        # A valid short config MUST have either a `covar` or a `file` source
+        # but not both.
+        if not (self.covar and self.file):
+            raise ValueError("ShortCrossInfo requires either a 'covar' or a "
+                             "'file' parameter.")
+        if self.covar and self.file:
+            raise ValueError("ShortCrossInfo cannot define both 'covar' and "
+                             "'file' simultaneously.")
+
+
+LongCrossInfo: TypeAlias = PathLike
+CrossInfo = Union[LongCrossInfo, ShortCrossInfo]
+
+
+def expected_allele_count_for_crosstype(
+        alleles: Sequence[str],
+        crosstype: CrossType,
+        expected_counts: dict[CrossType, int]
+) -> bool:
+    """Check that the number of alleles matches what crosstype expects."""
+    def __validate__(expected, alleles):
+        if expected is not None and len(alleles) != expected:
+            raise ValueError(
+                f"Cross type '{crosstype}' expects exactly {expected} allele "
+                f"labels. Received {len(alleles)}: {alleles}")
+
+    if crosstype.startswith(("genail", "genril")):
+        # genail and genril are dynamic (e.g., genail8)
+        # Extract the trailing digits to determine expected founder count
+        num_part = "".join(filter(str.isdigit, crosstype))
+        if num_part:
+            return __validate__(int(num_part), alleles)
+
+    return __validate__(expected_counts.get(crosstype), alleles)
+            
+
+
+@dataclass(frozen=True)
+class ControlData:
+    """Class for the R/qtl2 control data."""
+    # File names: Force listings, rather than singular strings
+    geno: Sequence[PathLike] = field(default_factory=tuple)
+    founder_geno: Sequence[PathLike] = field(default_factory=tuple)
+    pheno: Sequence[PathLike] = field(default_factory=tuple)
+    covar: Sequence[PathLike] = field(default_factory=tuple)
+    phenocovar: Sequence[PathLike] = field(default_factory=tuple)
+    gmap: Sequence[PathLike] = field(default_factory=tuple)
+    pmap: Sequence[PathLike] = field(default_factory=tuple)
+
+    # X Chromosome
+    x_chr: Optional[str] = None
+
+    # Allele labels
+    alleles: Sequence[str] = field(default_factory=tuple)
+
+    # Genotype codes
+    genotypes: GenoCode = field(default_factory=dict)
+
+    # sex
+    sex: Optional[ControlDataSex] = None
+
+    # Cross info
+    crosstype: Optional[CrossType] = None
+    cross_info: Optional[CrossInfo] = None
+
+
+    # CSV fields
+    na_strings: Sequence[str] = ("-", "NA", "N/A")
+    sep: str = ","
+    comment_char: Optional[str] = "#"
+
+    def __post_init__(self):
+        if self.alleles and self.crosstype:
+            expected_allele_count_for_crosstype(
+                self.alleles, self.crosstype, _ALLELE_COUNTS_BY_CROSSTYPE)