grapp 0.4__tar.gz → 0.5__tar.gz

grapp-0.5/MANIFEST.in

@@ -0,0 +1,4 @@
+include LICENSE
+include README.md
+include requirements.txt
+include popgen_requirements.txt

{grapp-0.4 → grapp-0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: grapp
-Version: 0.4
+Version: 0.5
 Summary: Statistical and population genetics methods on GRG
 Home-page: https://aprilweilab.github.io/
 Author: Drew DeHaas
@@ -13,19 +13,24 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: grg_pheno_sim>=1.4
+Requires-Dist: grg_pheno_sim<2.0,>=1.4
 Requires-Dist: numpy
 Requires-Dist: pandas
-Requires-Dist: pygrgl>=2.7
-Requires-Dist: pyigd>=1.3
+Requires-Dist: pyfaidx
+Requires-Dist: pygrgl<3.0,>=2.10
+Requires-Dist: pyigd<2.0,>=1.3
 Requires-Dist: scikit-learn
 Requires-Dist: scipy
+Provides-Extra: popgen
+Requires-Dist: pyfaidx; extra == "popgen"
+Requires-Dist: scikit-allel; extra == "popgen"
 Dynamic: author
 Dynamic: classifier
 Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: license-file
+Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: summary
 
@@ -36,7 +41,7 @@ on top of the Genotype Representation Graph (GRG) format. [GRG](https://github.c
 a file format and data structure that losslessly represents a genetic dataset. It has the advantage of
 compressing large datasets significantly, while also making calculations over that dataset extremely fast
 (see [the paper](https://www.nature.com/articles/s43588-024-00739-9) and the
-[core library](https://github.com/aprilweilab/grgl)).
+[main GRG documentation](https://grgl.readthedocs.io/en/latest/index.html)).
 
 Some notes on usage are below, and you can also browse [the Python documentation for grapp](https://grapp.readthedocs.io/en/latest/).
 
@@ -177,7 +182,7 @@ mutation filters:
 
 #### Library
 
-Library API reference document is [here](https://grapp.readthedocs.io/en/latest/). There are also some tutorials in the [GRG documentation](https://grgl.readthedocs.io/en/stable/tutorials/GWAS.html) (and correspond [Jupyter notebooks](https://github.com/aprilweilab/grgl/tree/main/doc/tutorials/notebooks)).
+Library API reference document is [here](https://grapp.readthedocs.io/en/latest/). There are also some tutorials in the [GRG documentation](https://grgl.readthedocs.io/en/latest/tutorials/GWAS.html) (and correspond [Jupyter notebooks](https://github.com/aprilweilab/grgl/tree/main/doc/tutorials/notebooks)).
 
 ### nn
 

{grapp-0.4 → grapp-0.5}/README.md

@@ -5,7 +5,7 @@ on top of the Genotype Representation Graph (GRG) format. [GRG](https://github.c
 a file format and data structure that losslessly represents a genetic dataset. It has the advantage of
 compressing large datasets significantly, while also making calculations over that dataset extremely fast
 (see [the paper](https://www.nature.com/articles/s43588-024-00739-9) and the
-[core library](https://github.com/aprilweilab/grgl)).
+[main GRG documentation](https://grgl.readthedocs.io/en/latest/index.html)).
 
 Some notes on usage are below, and you can also browse [the Python documentation for grapp](https://grapp.readthedocs.io/en/latest/).
 
@@ -146,7 +146,7 @@ mutation filters:
 
 #### Library
 
-Library API reference document is [here](https://grapp.readthedocs.io/en/latest/). There are also some tutorials in the [GRG documentation](https://grgl.readthedocs.io/en/stable/tutorials/GWAS.html) (and correspond [Jupyter notebooks](https://github.com/aprilweilab/grgl/tree/main/doc/tutorials/notebooks)).
+Library API reference document is [here](https://grapp.readthedocs.io/en/latest/). There are also some tutorials in the [GRG documentation](https://grgl.readthedocs.io/en/latest/tutorials/GWAS.html) (and correspond [Jupyter notebooks](https://github.com/aprilweilab/grgl/tree/main/doc/tutorials/notebooks)).
 
 ### nn
 

{grapp-0.4 → grapp-0.5}/grapp/assoc/__init__.py

@@ -1,11 +1,17 @@
 from grapp.linalg.ops_scipy import SciPyStdXOperator, SciPyXOperator
-from grapp.util.simple import allele_counts, _GenotypeDist, common_mut_dataframe
+from grapp.util.simple import (
+    allele_counts,
+    _GenotypeDist,
+    common_mut_dataframe,
+    _div_or_default,
+)
+from grapp.util.exceptions import UserInputError
 from grapp.grg_calculator import (
     GRGCalcInterface as _GRGCalcInterface,
     _wrap_grg,
 )
 from scipy.stats import t as t_distribution
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Tuple
 import itertools
 import math
 import numpy
@@ -16,54 +22,89 @@ import sklearn.linear_model
 import sys
 
 
-def _div_or_default(a, b, d):
-    """
-    y = a / b, unless b_i is 0, then y_i will be set to 0.
-
-    :param a: Numerator
-    :param b: Denominator
-    :param d: Default value for when denominator is 0.
-    """
-    result = numpy.full(a.shape, d)
-    return numpy.divide(a, b, out=result, where=(b != 0))
-
-
-def read_plink_covariates(covar_path: str) -> numpy.typing.NDArray:
+def read_plink_covariates(
+    covar_path: str, return_indivs: bool = False, verbose: bool = False
+) -> Union[numpy.typing.NDArray, Tuple[numpy.typing.NDArray, List[str]]]:
     """
-    Reads a PLINK-style covariate file (no headers) and returns a NumPy matrix of covariate values.
-    The first two columns (FID/IID) are ignored. Optionally adds an intercept column of 1s.
+    Reads a PLINK-style covariate file: Optional header line, FID and IID in first two columns,
+    covariates in remaining columns. The first two columns (FID/IID) are ignored. Does not allow
 
     :param path: Path to the covariate file.
     :type path: str
-    :return: A NumPy array of shape (n_samples, n_covariates [+1 if intercept]).
-    :rtype: numpy.ndarray
+    :param return_indivs: If True, return a list of individual identifiers (IIDs from the plink file).
+    :type return_indivs: bool
+    :param verbose: Emit warnings/information about the file if True. Default: True.
+    :type verbose: bool
+    :return: The covariates as a numpy array of shape (n_samples, n_covariates). If return_indivs is True, then also
+        return a list of individual IDs (strings).
+    :rtype: Union[numpy.typing.NDArray, Tuple[numpy.typing.NDArray, List[str]]]
     """
     rows = []
+    num_covar = 0
+    covar_names = []
+    iids = []
     with open(covar_path, "r") as f:
-        for line in f:
+        for i, line in enumerate(f):
             parts = line.strip().split()
             if not parts:
                 continue
-            # skip FID/IID, keep covariates
-            cov_vals = [float(v) for v in parts[2:]]
-            rows.append(cov_vals)
+            if i == 0 and parts[0] == "FID":
+                assert (
+                    len(parts) > 1 and parts[1] == "IID"
+                ), "Header line must include FID followed by IID."
+                covar_names = parts[2:]
+                num_covar = len(covar_names)
+            else:
+                cov_vals = [float(v) for v in parts[2:]]
+                if num_covar == 0:
+                    num_covar = len(cov_vals)
+                assert num_covar == len(
+                    cov_vals
+                ), f"Inconsistent number of columns in covariates file (line: {i}, cols: {len(cov_vals)})"
+                rows.append(cov_vals)
+                iids.append(parts[1])
+    assert num_covar > 0, "No covariates are present in the file (only two columns)."
 
     # stack into (n_samples × K)
     X = numpy.vstack(rows) if rows else numpy.empty((0, 0))
 
+    if verbose:
+        print(f"Loaded {num_covar} covariates for {len(rows)} individuals")
+        if covar_names:
+            print(f"Covariate names: {covar_names}")
+    if return_indivs:
+        return (X, iids)
     return X
 
 
-def read_pheno(filename: str, verbose: bool = True) -> numpy.typing.NDArray:
+def read_pheno(
+    filename: str, return_indivs: bool = False, verbose: bool = True
+) -> Union[numpy.typing.NDArray, Tuple[numpy.typing.NDArray, List[str]]]:
     """
-    Reads a PLINK/GCTA/GRG-style phenotype file and returns the phenotype vector.
+    Reads a PLINK/GCTA/GRG-style phenotype file and returns the phenotype vector. In all cases, the
+    row order of the file must match the individual indexing order of the GRG.
+
+    PLINK-style: Optional header "FID IID PHEN". The first column is the family ID (ignored by grapp),
+    the second column is the IID (used for validation by grapp), and the third column in the numerical
+    phenotype value. The IID order is checked against the GRG individual ID order, when IDs are present
+    in the GRG (and also not "NA" in the plink input file).
+
+    GRG-style: Optional header "person_id phenotypes". The first column is checked against the GRG individual
+    ID order, when IDs are present in GRG (and also not "NA" in the phenotype file). The second column is
+    the numerical phenotype value.
+
+    If there is no header, the number of columns determines the type (3 = plink, 2 = GRG).
+
+    In all cases, columns can be tab-separated or space separated.
 
     :param path: Path to the phenotype file.
     :type path: str
+    :param return_indivs: If True, return a list of individual identifiers (IIDs from the plink file).
+    :type return_indivs: bool
     :param verbose: Emit warnings/information about the file if True. Default: True.
     :type verbose: bool
     :return: A one-dimensional NumPy array of phenotype values.
-    :rtype: numpy.ndarray
+    :rtype: Union[numpy.typing.NDArray, Tuple[numpy.typing.NDArray, List[str]]]
     """
     header_line = None
     with open(filename, "r") as f:
@@ -76,33 +117,56 @@ def read_pheno(filename: str, verbose: bool = True) -> numpy.typing.NDArray:
             header_line = i
             break
         # grg_pheno_sim header
-        if re.match(r"^(person_id\tphenotypes)", line.strip(), re.IGNORECASE):
+        if re.match(
+            r"^((person_id|individual_id)\s+(phenotypes|phenotype))",
+            line.strip(),
+            re.IGNORECASE,
+        ):
             header_line = i
             break
 
     # Read data starting from the header (if present)
     if header_line is not None:
         df = pandas.read_csv(
-            filename, sep=r"\s+", skiprows=header_line, engine="python"
+            filename,
+            sep=r"\s+",
+            skiprows=header_line,
+            engine="python",
+            skipinitialspace=True,
         )
     else:
-        df = pandas.read_csv(filename, sep=r"\s+", header=None, engine="python")
+        df = pandas.read_csv(
+            filename,
+            sep=r"\s+",
+            header=None,
+            engine="python",
+            skipinitialspace=True,
+        )
 
     # Check column count
-    if df.shape[1] not in (2, 3):
-        raise ValueError(f"Expected 2 or 3 columns, but found {df.shape[1]}.")
+    if df.shape[1] > 3:
+        raise ValueError(f"Expected 1, 2, or 3 columns, but found {df.shape[1]}.")
 
     # Extract last column and make sure it's a number
     try:
         last_col = df.iloc[:, -1].astype(float).to_numpy()
     except ValueError:
-        raise ValueError("Last column contains non-numeric values.")
+        raise ValueError(
+            "Last column should be phenotype, but it contains non-numeric values."
+        )
     if verbose:
         print(
-            f"Using the column {df.shape[1]} (the last one) for phenotype.",
+            f"Using the {df.shape[1]}th column (the last one) for phenotype ({df.columns[-1]}).",
             file=sys.stderr,
         )
 
+    if return_indivs:
+        if len(df.columns) >= 2:
+            iids = df.iloc[:, -2].to_list()
+        else:
+            iids = ["NA"] * len(df)
+        return (last_col, iids)
+
     return last_col
 
 
@@ -122,6 +186,11 @@ def _computeDiagXTX(
         # diag(X^T @ X): for the non-standardized genotype matrix, GRG has a special initialization
         # method "xtx" which uses coalescence information to compute the diagonal in a single pass
         if dist == _GenotypeDist.SAMPLE.value:
+            if not grg.has_individual_coals:
+                raise UserInputError(
+                    "GRG does not have individual coalescences; try using the binomial approximation for variance"
+                )
+
             XX = grg.matmul(
                 numpy.ones((1, grg.num_samples), dtype=numpy.int32),
                 pygrgl.TraversalDirection.UP,
@@ -249,8 +318,9 @@ def linear_assoc_no_covar(
     s_tot = yy - (total_pheno**2) / n
     r2 = 1 - sse / s_tot
 
-    cdf_vals = t_distribution.cdf(t_stat, df=n - 2)
-    p_val = 2 * numpy.where(t_stat > 0, 1 - cdf_vals, cdf_vals)
+    # The CDF is symmetric around 0, so for numerical reasons we always choose the smaller
+    # probability (tail) and do cdf(-|t|) instead of (1 - cdf(|t|)).
+    p_val = 2 * t_distribution.cdf(-numpy.abs(t_stat), df=n - 2)
 
     return common_mut_dataframe(
         grg, COUNT=acount, BETA=beta, B0=b0, SE=se, R2=r2, T=t_stat, P=p_val
@@ -296,6 +366,7 @@ def linear_assoc_covar(
         GAMMA columns.
     :rtype: pandas.DataFrame
     """
+    # TODO: generalize this. We just require the use of binomial variance when doing non-diploid
     PLOIDY = 2
     assert grg.ploidy == PLOIDY, "GWAS is only supported on diploid individuals"
     assert method in ("QR", "regress"), 'Invalid "method" parameter'
@@ -383,10 +454,11 @@ def linear_assoc_covar(
     xadjTxadj = XX - diagonal
 
     # Compute (Xadj^TYadj)
-    xadjTyadj = Yadj @ X_op
+    xadjTyadj = numpy.where(acount < (PLOIDY * n_j), Yadj @ X_op, math.nan)
 
     beta = numpy.zeros(XX.size)
     beta = _div_or_default(xadjTyadj, xadjTxadj, math.nan)
+
     if only_beta:
         return pandas.DataFrame({"BETA": beta})
 
@@ -397,8 +469,9 @@ def linear_assoc_covar(
     se = numpy.sqrt(numpy.abs(_div_or_default(SSE, (df * xadjTxadj), math.nan)))
     t_vals = beta / se
 
-    cdf_vals = t_distribution.cdf(t_vals, df)
-    p = 2 * numpy.where(t_vals > 0, 1 - cdf_vals, cdf_vals)
+    # The CDF is symmetric around 0, so for numerical reasons we always choose the smaller
+    # probability (tail) and do cdf(-|t|) instead of (1 - cdf(|t|)).
+    p = 2 * t_distribution.cdf(-numpy.abs(t_vals), df=df)
 
     # Optional GAMMA calculation
     gamma_cols = {}