o Uݢgh@sddlZddlmZddlmZmZmZddlm Z ej j ddZ ddZdd Zd d Zd$d d ZddZd$ddZddZddZddZddZddZddZddZGd d!d!ZGd"d#d#ZdS)%N)optimize)betalngammalnloggamma*)seedcCs6t|d}tdd|D}tt|t|S)z=Inverse Simpson Index, or the effective diversity of power 2.css|]}|dVqdS)rN).0vr r _/oak/stanford/groups/akundaje/marinovg/programs/cellranger-9.0.1/lib/python/cellranger/stats.py sz&effective_diversity..)npsumtk_stats robust_dividefloat)counts numerator denominatorr r r effective_diversitysrcCstj|dd\}}t|||dk}t|}i}ttt||||D]\}}||vr5d||<q(||||<||d7<q(|S)a_Produces an incremental count vector from a categorial sample vector. Takes a set of drawn samples from a multinomial or similar distribution distribution and produces a vector of the same size, with each element the the number of times the sample at that position has been seen indices. I.e., it would turn the following vector of sample draws: [1, 1, 2, 1, 0, 2, 1, 3] into [1, 2, 1, 3, 1, 2, 4, 1] Args: sample_draws (np.ndarray(int)): A vector of sample indices. Returns: inc_counts (np.ndarray(int)): The incremental counts of those sample indices. T) return_countsr)runiqueisin ones_likeziparangelen) sample_drawsidxscmaskZ inc_countsrijr r r incremental_counts_from_samples $  r%cCstj||dS)anProduces a count of items in the sample draws. Takes an array of samples drawn from a distribution and returns a new array of counts of how many times each feature was seen in the draws. Args: sample_draws (np.ndarray(int)): A vector of sample indices. Returns: counts (np.ndarray(int): A vector of counts for each feature index. ) minlength)rbincount)r num_featuresr r r collapse_draws_to_counts5s r)c Cs|jd}tj|td}|durt|jddd}t|d}t|D]9}|j||j|d}}|j ||} |j ||} |j | ddd} ||t| d| | ||<q%|S)a(Computes the multinomial log-likelihood for many barcodes. Multinomial log-likelihood for a single barcode where the count of UMIs for feature i is x_i is: l = log(gamma(sum_i(x_i) + 1)) - sum_i(log(gamma(x_i + 1)) + sum_i(log(p_i) * x_i) Because the input matrix is sparse and zero counts do not contribute to the likelihood, we can rapidly remove zero elements from the computation. Args: matrix (scipy.sparse.csc_matrix): Matrix of UMI counts (feature x barcode) logp (np.ndarray(float)): The natural log of the multinomial probability vector across features n (int, optional): The total number of UMIs per barcode. Saves computation if it can be precomputed. Returns: loglk (np.ndarray(float)): Log-likelihood for each barcode rdtypeNraxisclipr-mode) shaperzerosrasarrayrrrangeindptrindicesdatatake) matrixlogpnnum_bcsloglkconstsr# idx_startidx_endr rowZ short_logpr r r eval_multinomial_loglikelihoodsEs   *rBcCsBt|}tdt|d}t|t|||}t|S)aComputes the cumulative multinomial log-likelihood for a vector. Given a vector of sample draws from a multinomial distribution, computes the log-likelihood of all of the draws up to each draw. This is done incrementally, by first pre-computing the rank R of each draw as the number of times that feature was seen when it was drawn. Then the incremental log-likelihood from that draw is dl = log(i) - log(R_i) + log(p_i) Args: sample_draws (np.ndarray(int)): A vector of sample draws logp (np.ndarray(float)): The log-probability of sampling each feature Returns: loglk (np.ndarray(float): The cumulative log-likelihood up to each draw r)r%rrrlogcumsum)rr:marginal_countsnvalsr=r r r )eval_multinomial_loglikelihood_cumulativehs rGc Cs|jd}t|}|durt|jddd}t|tt||}t|jdD]9}|j||j|d}}|j ||} |j ||} |j | ddd} ||t| t| | ||<q-|S)a?Computes the Dirichlet-multinomial log-likelihood for many barcodes. Dirichlet-Multinomial log-likelihood for a single barcode where the count of UMIs for feature i is x_i is: l = log(sum_i(x_i)) + log(beta(sum_i(a_i), sum_i(x_i))) - sum_i(log(x_i)) - sum_i(log(beta(a_i, x_i))) Because the input matrix is sparse and zero counts do not contribute to the likelihood, we can rapidly remove zero elements from the computation. Args: matrix (scipy.sparse.csc_matrix): Matrix of UMI counts (feature x barcode) alpha (np.ndarray(float)): The vector of Dirichlet parameters for each feature n (int, optional): The total number of UMIs per barcode. Saves computation if it can be precomputed. Returns: loglk (np.ndarray(float)): Log-likelihood for each barcode rNrr,r.r/) r1rr2r3rrCrr4r5r6r7r8) r9alphar;r<r=r>Zbc_indexr?r@r rAZ short_alphar r r )eval_dirichlet_multinomial_loglikelihoodss  *rIcCslt|}tdt|d}t|}t|t|t||dt|d||}t|S)aComputes the cumulative Dirichlet-multinomial log-likelihood for a vector. Given a vector of sample draws from a Dirichlet-multinomial distribution, computes the log-likelihood of all of the draws up to each draw. This is done incrementally, by first pre-computing the rank R of each draw as the number of times that feature was seen when it was drawn. Then the incremental log-likelihood from that draw is dl = log(i) - log(R_i) - log(i + sum_i(a_i) - 1) + log(R_i + a_i - 1) Args: sample_draws (np.ndarray(int)): A vector of sample draws alpha (np.ndarray(float)): The Dirichlet parameter for each feature Returns: loglk (np.ndarray(float): The cumulative log-likelihood up to each draw r)r%rrrrrCrD)rrHrErFZalpha_0r=r r r 3eval_dirichlet_multinomial_loglikelihood_cumulatives  rJcCs~|dd|f}t|jdd}dd}ddg}tj|||||fd}|js1td |jt d |j d ||j S) aXEstimates the best-fit overdispersion parameter for data. Uses a Dirichlet-multinomial to maximize the log-likelihood of the ambient barcode signal, given an input probability per feature that is scaled by a fixed overdispersion parameter. Args: matrix (scipy.sparse.csc_matrix): Matrix of UMI counts (feature x barcode) ambient_bcs (np.array): Array of barcode indexes to use for the ambient background p (np.ndarray(float)): The estimated multinomial probability vector across features Returns: alpha (float): Best-fit overdispersion for the data Nrr,cSstt||||d S)N)r;)rrrI)rHr9p umis_per_bcr r r ambient_loglksz8estimate_dirichlet_overdispersion..ambient_loglkgMbP?i')boundsargszCould not find valid alpha: zAlpha = zC: maximizes the likelihood of the ambient barcodes. Search bounds: ) rr3rflattenrminimize_scalarsuccess ValueErrormessageprintx)r9Z ambient_bcsrKrLrMrNresultr r r !estimate_dirichlet_overdispersionsrXcCstj|d}t||S)aReturns a fixed size array of multinomial draws from the input probabilities. Given a number of samples to draw and a probability vector for all features to sample from, produces an array of feature indices of the desired size, drawn according to a multinomial distribution. Because it is significantly faster to generate random numbers on the unit interval, we use that and then find the feature index using the cumulative probabilities of the probability vector. Args: num_draws (int): The number of samples to draw from the distribution p_cumulative (np.ndarray(float)): The cumulative probability of all possible features. The length of this array should be the number of features. It should be sorted from smallest to largest, and the last entry should be 1. Returns: sample_draws (np.ndarray(float)): A random set of draws from a multinomial distribution. )size)RNGrandomr searchsorted) num_draws p_cumulativeZrng_numsr r r draw_multinomial_samples  r_cCst|}t|t|S)aReturns an array of Dirichlet-multinomial draws from the input parameters. Given a number of samples to draw and a vector of alpha parameters for all features to sample from, produces an array of feature indices of the desired size, drawn according to a Dirichlet-multinomial distribution. Note that Dirichlet-multinomial draws are *not* independent and identically-distributed, so that all draws for a sample must be done with a single call or the sample variance will be under-stated. Args: num_draws (int): The number of samples to draw from the distribution alpha (np.ndarray(float)): The Dirichlet parameters for each feature. Normalized to a sum of 1, these are the mean probabilities of each feature over many draws of many samples. The length of this array should be the number of features. Returns: sample_draws (np.ndarray(float)): A random set of draws from a Dirichlet- multinomial distribution. )rZ dirichletr_rrD)r]rHprobsr r r !draw_dirichlet_multinomial_samples rbc Cstt|}t|}tjt||ftd}t|}t|}t |D]}t ||} t | ||d|dd|f<q%||fS)aSimulate draws from a multinomial distribution for many values of N. Note that the samples within each simulation are not independent; we generate N draws for the largest value of N and use subsamples of the largest simulation for smaller values of N. Args: p (np.ndarray(float)): Probability of observing each feature. umis_per_bc (np.ndarray(int)): UMI counts per barcode. num_sims (int): Number of simulations to perform. Returns: distinct_ns (np.ndarray(int)): an array containing the distinct N values that were simulated. log_likelihoods (np.ndarray(float)): a len(distinct_ns) x num_sims matrix containing the simulated log likelihoods. r*rN) r flatnonzeror'maxr2rrrDrCr4r_rG) rKrLnum_sims distinct_nmax_nr=r^r:r#drawr r r #simulate_multinomial_loglikelihoodss     ricCsltt|}t|}tjt||ftd}t|D]}t||}t |||d|dd|f<q||fS)aSimulate draws from a Dirichlet-multinomial distribution for many values of N. Note that the samples within each simulation are not independent; we generate N draws for the largest value of N and use subsamples of the largest simulation for smaller values of N. Args: alpha (np.ndarray(float)): Dirichlet parameter for each feature. umis_per_bc (np.ndarray(int)): UMI counts per barcode. num_sims (int): Number of simulations to perform. Returns: distinct_ns (np.ndarray(int)): an array containing the distinct N values that were simulated. log_likelihoods (np.ndarray(float)): a len(distinct_ns) x num_sims matrix containing the simulated log likelihoods. r*rN) rrcr'rdr2rrr4rbrJ)rHrLrerfrgr=r#rhr r r -simulate_dirichlet_multinomial_loglikelihoods8s  rjc Cst|t|ks J|jdt|ksJt||}|jd}t|}t|}t|D]}t|||ddf||k} td| d|||<q-|S)aCompute p-values for observed multinomial log-likelihoods. Args: umis_per_bc (nd.array(int)): UMI counts per barcode obs_loglk (nd.array(float)): Observed log-likelihoods of each barcode deriving from an ambient profile sim_n (nd.array(int)): Multinomial N for simulated log-likelihoods sim_loglk (nd.array(float)): Simulated log-likelihoods of shape (len(sim_n), num_simulations) Returns: pvalues (nd.array(float)): p-values rrN)rr1rr\r2r4rr) rLZ obs_loglkZsim_nZ sim_loglkZ sim_n_idxreZ num_barcodespvaluesr#Znum_lower_loglkr r r compute_ambient_pvaluesXs     "rlc@s(eZdZdZddZddZddZdS) CurveaCCurve information for plotting. Attributes: x (list of int): Curve x-axis (number of cells) y (list of float): Curve y-axis (expected value of number of unique clonotypes) y_std (list of float): Standard deviation of number of unique clonotypes in a curve y_ciu (list of float): Upper bound of 95% confidence interval for number of unique clonotypes in a curve y_cil (list of float): Lower bound of 95% confidence interval for number of unique clonotypes in a curve cCs"g|_g|_g|_g|_g|_dS)zInitialize an empty curve.N)rVyy_stdy_ciuy_cilselfr r r __init__s  zCurve.__init__cCs$t|jdkst|jdkrdSdS)zRCheck if curve is empty or calculated. Returns: bool rTF)rrnrVrrr r r is_emptyszCurve.is_emptycCs<dd|j|j|j|jfD}||dt|krdSdS)zcCheck if calculated curve is consistent (length match). Returns: bool cSg|]}t|qSr )r)r r#r r r z'Curve.is_consistent..rFT)rVrnrqrpcountr)rsZlen_listr r r is_consistentszCurve.is_consistentN)__name__ __module__ __qualname____doc__rtrurzr r r r rmus  rmc@seZdZdZddZddZdeeeffddZdefd d Z d d Z d dZ d'defddZ d(de dedefddZddZd)dedefddZddZd d!Z d*d"ed#edefd$d%Zd&S)+ DiversityaRepresents diversity with rarefaction and extrapolation curves. Attributes: sorted_hist (list of tuple): A histogram sorted by abundance (most abundant has lowest index) freq_counts (dict of int: int): Frequency counts table N (int): total number of counts in histogram rarefaction_curve (Curve): Rarefaction curve information extrapolation_curve (Curve): Extrapolation curve information cCs6t|dd|_||_||_t|_t|_dS)NT)reverse) sorted sorted_hist_calc_nN_get_freq_counts freq_countsrmrarefaction_curveextrapolation_curve)rshistr r r rts    zDiversity.__init__cCs|jdks |dvr dSdS)Nr)rFTr f_0_chao1rrr r r is_diversity_curve_possiblesz%Diversity.is_diversity_curve_possiblereturncCs6i}|jD]}||vrd||<q||d7<q|S)z]Get frequency counts (f_k in Colwell et al 2012). Returns: dict r)r)rsZret_dictZabundr r r rs   zDiversity._get_freq_countscCs t|jS)zmCalculates number of samples. (length of input histogram) Returns: int )rrrrr r r rs zDiversity._calc_ncCsj||j|kr dSt|j|dt|j|dt|jdt|j||d}tt|S)zCalculates alpha parameter in Colwell et al 2012. Args: n (int): rarefaction point n k (int): Frequency Returns: int rr)rrrexpreal)rsr;kZ log_alphar r r _alphas  zDiversity._alphacCstt|j}d}|jD]\}}|||||7}q||}d}|jD]\}}|d|||d|t|d|7}q*|t|fS)zCalculates rarefaction and standard deviation of rarefaction for a single point n. Args: n (int): rarefaction point n Returns: float, float rrr) rrrvaluesitemsrassemblage_size_estimatersqrt)rsr;Znum_clonZsum_helper_exp_valrryexp_valZsum_helper_std_devr r r _rarefactions  zDiversity._rarefaction( num_stepsc Cst|j|d}|dkrd}td|j|}t||kr%td|j||}d}|gt|}|gt|}|gt|}|gt|}t|D]'\} } || \|| <|| <|| d|| || <|| d|| || <qGt||j_||j_ ||j_ ||j_ ||j_ dS)zCalculates rarefaction curve for num_steps between 1 and N. Args: num_steps (int): Number of steps between 1 and N rr\(\?N) intrr4r enumeraterlistrrVrnrorqrp) rsr step_sizeZrc_x placeholderZrc_y_expZrc_y_stdZrc_y_ciuZrc_y_cilidxr;r r r calc_rarefaction_curves(   z Diversity.calc_rarefaction_curveToriginstdevc Cs||g}||jj|jjd|d|d|r?||jj|jjddd|jj|jjdddd|ddd|d d |S) a]Create a plotly curve for rarefaction. Args: origin (str): Origin string (see vdj inputs) color: color of the curve num_steps (int): number of extrapolation steps stdev (bool): flag for calculation of standard deviation (and confidence intarvals) Returns: [dict] Zscatterlines)rVrntypenamer0 line_colorNrz (Stdev)Ztoselfzrgba(255,255,255,0)g?)rVrnrrfillrZ fillcolorZopacity)rappendrrVrnrprq)rsrcolorrrrr r r calc_rarefaction_curve_plotlys0   z'Diversity.calc_rarefaction_curve_plotlyc Csttt|j}t|}t|jd}t|j}t||j}dd||||}|||}d} || fS)aZCalculates extrapolation for a single point (N + m). standard deviation is not implemented. Args: n_plus_m (int): Distance of extrapolation point to N (how many more samples collected) Returns: float: Expected value None: Standard deviation (not implemented) rg?N)rrrrrr) rsn_plus_mZs_obsf_0Zf_1rmZbracketsrZstd_devr r r _extrapolationIs   zDiversity._extrapolationrgc Cst||j|d}t|j||}t||kr!t|j|||}dgt|}dgt|}dgt|}dgt|}t|D]-\} } || \|| <|| <|| durn|| d|| || <|| d|| || <qAt||j_||j_ ||j_ ||j_ ||j_ dS)aCalculates extrapolation curve for num_steps between N and max_n. Args: num_steps (int): Number of steps between N and max_n max_n (int): The limit to which the function extrapolates Returns: None rrNr) rrr4rrrrrrVrnrorqrp) rsrrgrZec_xZec_y_expZec_y_stdZec_y_ciuZec_y_cilrrr r r calc_extrapolation_curvefs&     z"Diversity.calc_extrapolation_curvecCsfd|jvrd|jvrt|jddtd|jdSd|jvr1t|jd|jdddSdS)zEstimate f0 (number of clonotypes that we haven't seen yet) using chao1 estimator. Based on: Colwell et al 2012 Returns: float rrg@g)rrrrr r r rs $  zDiversity.f_0_chao1cCs|j|S)zEstimate assemblage size (number of total unique clonotypes) using chao1 estimator. Based on: Colwell et al 2012 Returns: float rrrr r r rsz"Diversity.assemblage_size_estimaterc_stepsec_stepsc CsL|jr ||||jr|||jrdS|jr%dSd}t|dp}||t|jj D]*\}}||jj ||jj ||jj |g} d dd| Ddgd} || q8t|jj D]*\}}||jj ||jj ||jj |g} d d d| Dd gd} || qiWdd S1swYd S) awCalculate rarefaction and extrapolation diversity curves. Save them in csv format in outfile. Args: outfile (str): Path to output file rc_steps (int): Rarefaction curve steps between 1 and N ec_steps (int): Extrapolation steps between N and max_n max_n (int): Maximum number of cells extrapolating to Nzx,y,y_cil,y_ciu,type w,cSrvr strr r;r r r rwrxz0Diversity.create_both_curves..Z rarefaction cSrvr rrr r r rwrxZ extrapolationr)rrurrrrzopenwriterrVrnrqrpjoin) rsoutfilerrrgheaderZ outhandlerrVnumbersliner r r create_both_curvess@                 zDiversity.create_both_curvesN)r)rT)rr)rrr)r{r|r}r~rtrdictrrrrrrrboolrrrrrrr r r r rs:    + r)N)numpyrscipyr scipy.specialrrr tenkit.statsstatsrr[ default_rngrZrr%r)rBrGrIrJrXr_rbrirjrlrmrr r r r s(   # "&  -