Source code for glow.gwas.lin_reg

import pandas as pd
import numpy as np
from nptyping import Float, NDArray
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union
from pyspark.sql import Column, DataFrame
from pyspark.sql.types import StringType, StructField, IntegerType
from scipy import stats
from typeguard import typechecked
from . import functions as gwas_fx
from .functions import _VALUES_COLUMN_NAME
from ..wgr.wgr_functions import _get_contigs_from_loco_df

__all__ = ['linear_regression']


[docs]@typechecked def linear_regression(genotype_df: DataFrame, phenotype_df: pd.DataFrame, covariate_df: pd.DataFrame = pd.DataFrame({}), offset_df: pd.DataFrame = pd.DataFrame({}), contigs: Optional[List[str]] = None, add_intercept: bool = True, values_column: Union[str, Column] = 'values', dt: type = np.float64, verbose_output: bool = False) -> DataFrame: ''' Uses linear regression to test for association between genotypes and one or more phenotypes. The implementation is a distributed version of the method used in regenie: https://www.nature.com/articles/s41588-021-00870-7 Implementation details: On the driver node, we decompose the covariate matrix into an orthonormal basis and use it to project the covariates out of the phenotype matrix. The orthonormal basis and the phenotype residuals are broadcast as part of a Pandas UDF. In each Spark task, we project the covariates out of a block of genotypes and then compute the regression statistics for each phenotype, taking into account the distinct missingness patterns of each phenotype. Examples: >>> np.random.seed(42) >>> n_samples, n_phenotypes, n_covariates = (710, 3, 3) >>> phenotype_df = pd.DataFrame(np.random.random((n_samples, n_phenotypes)), columns=['p1', 'p2', 'p3']) >>> covariate_df = pd.DataFrame(np.random.random((n_samples, n_phenotypes))) >>> genotype_df = (spark.read.format('vcf').load('test-data/1kg_sample.vcf') ... .select('contigName', 'start', 'genotypes')) >>> results = glow.gwas.linear_regression(genotype_df, phenotype_df, covariate_df, ... values_column=glow.genotype_states('genotypes')) >>> results.head() # doctest: +ELLIPSIS Row(contigName='1', start=904164, effect=0.0453..., stderror=0.0214..., tvalue=2.114..., pvalue=0.0348..., phenotype='p1') >>> phenotype_df = pd.DataFrame(np.random.random((n_samples, n_phenotypes)), columns=['p1', 'p2', 'p3']) >>> covariate_df = pd.DataFrame(np.random.random((n_samples, n_phenotypes))) >>> genotype_df = (spark.read.format('vcf').load('test-data/1kg_sample.vcf') ... .select('contigName', 'start', 'genotypes')) >>> contigs = ['1', '2', '3'] >>> offset_index = pd.MultiIndex.from_product([phenotype_df.index, contigs]) >>> offset_df = pd.DataFrame(np.random.random((n_samples * len(contigs), n_phenotypes)), ... index=offset_index, columns=phenotype_df.columns) >>> results = glow.gwas.linear_regression(genotype_df, phenotype_df, covariate_df, ... offset_df=offset_df, values_column=glow.genotype_states('genotypes')) Args: genotype_df : Spark DataFrame containing genomic data phenotype_df : Pandas DataFrame containing phenotypic data covariate_df : An optional Pandas DataFrame containing covariates offset_df : An optional Pandas DataFrame containing the phenotype offset, as output by GloWGR's RidgeRegression or Regenie step 1. The actual phenotype used for linear regression is the mean-centered, residualized and scaled ``phenotype_df`` minus the appropriate offset. The ``offset_df`` may have one or two levels of indexing. If one level, the index should be the same as the ``phenotype_df``. If two levels, the level 0 index should be the same as the ``phenotype_df``, and the level 1 index should be the contig name. The two level index scheme allows for per-contig offsets like LOCO predictions from GloWGR. contigs : When using LOCO offsets, this parameter indicates the contigs to analyze. You can use this parameter to limit the size of the broadcasted data, which may be necessary with large sample sizes. If this parameter is omitted, the contigs are inferred from the ``offset_df``. add_intercept : Whether or not to add an intercept column to the covariate DataFrame values_column : A column name or column expression to test with linear regression. If a column name is provided, ``genotype_df`` should have a column with this name and a numeric array type. If a column expression is provided, the expression should return a numeric array type. dt : The numpy datatype to use in the linear regression test. Must be ``np.float32`` or ``np.float64``. verbose_output: Whether or not to generate additional test statistics (n, sum_x, y_transpose_x) to the output DataFrame. These values are derived directly from phenotype_df and genotype_df, and does not reflect any standardization performed as part of the implementation of linear_regression. Returns: A Spark DataFrame that contains - All columns from ``genotype_df`` except the ``values_column`` and the ``genotypes`` column if one exists - ``effect``: The effect size estimate for the genotype - ``stderror``: The estimated standard error of the effect - ``tvalue``: The T statistic - ``pvalue``: P value estimated from a two sided T-test - ``phenotype``: The phenotype name as determined by the column names of ``phenotype_df`` - ``n``(int): (verbose_output only) number of samples with non-null phenotype - ``sum_x``(float): (verbose_output only) sum of genotype inputs - ``y_transpose_x``(float): (verbose_output only) dot product of phenotype response (missing values encoded as zeros) and genotype input, i.e. phenotype value * number of alternate alleles ''' gwas_fx._check_spark_version(genotype_df.sql_ctx.sparkSession) gwas_fx._validate_covariates_and_phenotypes(covariate_df, phenotype_df, is_binary=False) sql_type = gwas_fx._regression_sql_type(dt) genotype_df = gwas_fx._prepare_genotype_df(genotype_df, values_column, sql_type) C = covariate_df.to_numpy(dt, copy=True) if add_intercept: C = gwas_fx._add_intercept(C, phenotype_df.shape[0]) # Prepare covariate basis and phenotype residuals Q = np.linalg.qr(C)[0] Y = phenotype_df.to_numpy(dt, copy=True) Y_mask = (~np.isnan(Y)).astype(dt) Y = np.nan_to_num(Y, copy=False) Y_for_verbose_output = np.copy(Y) if verbose_output else None Y -= Y.mean(axis=0) # Mean-center Y = gwas_fx._residualize_in_place(Y, Q) * Y_mask # Residualize Y_scale = np.sqrt(np.sum(Y**2, axis=0) / (Y_mask.sum(axis=0) - Q.shape[1])) Y /= Y_scale[None, :] # Scale Y_state = _create_YState(Y, phenotype_df, offset_df, Y_mask, dt, contigs) dof = C.shape[0] - C.shape[1] - 1 return _generate_linreg_output(genotype_df, sql_type, Y_state, Y_mask, Y_scale, Q, dof, phenotype_df, Y_for_verbose_output, verbose_output)
@dataclass class YState: ''' Keeps track of state that varies per contig ''' Y: NDArray[(Any, Any), Float] YdotY: NDArray[(Any), Float] def _create_YState(Y: NDArray[(Any, Any), Float], phenotype_df: pd.DataFrame, offset_df: pd.DataFrame, Y_mask: NDArray[(Any, Any), Float], dt, contigs: Optional[List[str]]) -> Union[YState, Dict[str, YState]]: offset_type = gwas_fx._validate_offset(phenotype_df, offset_df) if offset_type != gwas_fx._OffsetType.LOCO_OFFSET: return _create_one_YState(Y, phenotype_df, offset_df, Y_mask, dt) if contigs is None: contigs = _get_contigs_from_loco_df(offset_df) return { contig: _create_one_YState(Y, phenotype_df, offset_df.xs(contig, level=1), Y_mask, dt) for contig in contigs } def _create_one_YState(Y: NDArray[(Any, Any), Float], phenotype_df: pd.DataFrame, offset_df: pd.DataFrame, Y_mask: NDArray[(Any, Any), Float], dt) -> YState: if not offset_df.empty: base_Y = pd.DataFrame(Y, phenotype_df.index, phenotype_df.columns) # Reindex so the numpy array maintains ordering after subtracting offset Y = (base_Y - offset_df).reindex(phenotype_df.index).to_numpy(dt) Y *= Y_mask return YState(Y, np.sum(Y * Y, axis=0)) # @typechecked -- typeguard does not support numpy array def _linear_regression_inner(genotype_pdf: pd.DataFrame, Y_state: YState, Y_mask: NDArray[(Any, Any), Float], Y_scale: NDArray[(Any, ), Float], Q: NDArray[(Any, Any), Float], dof: int, phenotype_names: pd.Series, Y_raw: Optional[NDArray[(Any, Any), Float]], verbose_output: Optional[bool]) -> pd.DataFrame: ''' Applies a linear regression model to a block of genotypes. We first project the covariates out of the genotype block and then perform single variate linear regression for each site. To account for samples with missing traits, we additionally accept a mask indicating which samples are missing for each phenotype. This mask is used to ensure that missing samples are not included when summing across individuals. Rather than use traditional matrix indices in the einsum expressions, we use semantic indices. s: sample g: genotype p: phenotype So, if a matrix's indices are `sg` (like the X matrix), it has one row per sample and one column per genotype. ''' X = np.column_stack(genotype_pdf[_VALUES_COLUMN_NAME].array) del genotype_pdf[_VALUES_COLUMN_NAME] num_genotypes = genotype_pdf.shape[0] out_df = pd.concat([genotype_pdf] * Y_state.Y.shape[1]) if verbose_output: out_df["n"] = list(Y_mask.sum(axis=0).repeat(num_genotypes)) out_df["sum_x"] = list(np.ravel(Y_mask.T @ X)) out_df["y_transpose_x"] = list(np.ravel(Y_raw.T @ X)) X = gwas_fx._residualize_in_place(X, Q) XdotY = Y_state.Y.T @ X XdotX_reciprocal = 1 / gwas_fx._einsum('sp,sg,sg->pg', Y_mask, X, X) betas = XdotY * XdotX_reciprocal standard_error = np.sqrt((Y_state.YdotY[:, None] * XdotX_reciprocal - betas * betas) / dof) T = betas / standard_error pvalues = 2 * stats.distributions.t.sf(np.abs(T), dof) Y_scale_mat = Y_scale[:, None] out_df['effect'] = list(np.ravel(betas * Y_scale_mat)) out_df['stderror'] = list(np.ravel(standard_error * Y_scale_mat)) out_df['tvalue'] = list(np.ravel(T)) out_df['pvalue'] = list(np.ravel(pvalues)) out_df['phenotype'] = phenotype_names.repeat(num_genotypes).tolist() return out_df def _generate_linreg_output(genotype_df, sql_type, Y_state, Y_mask, Y_scale, Q, dof, phenotype_df, Y_for_verbose_output, verbose_output) -> DataFrame: # Construct output schema result_fields = [ StructField('effect', sql_type), StructField('stderror', sql_type), StructField('tvalue', sql_type), StructField('pvalue', sql_type), StructField('phenotype', StringType()) ] if verbose_output: result_fields += ([ StructField('n', IntegerType()), StructField('sum_x', sql_type), StructField('y_transpose_x', sql_type) ]) result_struct = gwas_fx._output_schema(genotype_df.schema.fields, result_fields) def map_func(pdf_iterator): for pdf in pdf_iterator: yield gwas_fx._loco_dispatch(pdf, Y_state, _linear_regression_inner, Y_mask, Y_scale, Q, dof, phenotype_df.columns.to_series().astype('str'), Y_for_verbose_output, verbose_output) return genotype_df.mapInPandas(map_func, result_struct)