o Rŀgn@s.dZddlmZddlmZGdddZdS)zsControl GenePop through an easier interface. This interface is less efficient than the standard GenePopControler )GenePop)GenePopControllerc@seZdZdZd0ddZddZd1dd Z d2ddZ d3ddZddZ ddZ ddZ ddZ ddZ ddZdd Zd!d"Zd#d$Zd%d&Zd'd(Zd)d*Zd4d.d/ZdS)5EasyControllerz@Define a class for an easier interface with the GenePop program.NcCs ||_t||_i|_i|_dS)zInitialize the controller. genepop_dir is the directory where GenePop is. The binary should be called Genepop (capital G) N)_fnamer _controller_EasyController__fst_pair_locus!_EasyController__allele_frequency)selffname genepop_dirr U/var/www/html/myenv/lib/python3.10/site-packages/Bio/PopGen/GenePop/EasyController.py__init__s  zEasyController.__init__cCs@t|j }t|}Wdn1swY|j|jfS)z7Obtain the population list and loci list from the file.N)openrrreadpop_list loci_list)r frecr r rget_basic_info"s   zEasyController.get_basic_info probabilitycCsh|dkr |j|j}n|dkr|j|j}n |j|jd\}}}t|dD]}t|q)t|S)z2Perform Hardy-Weinberg test on the given position. deficiencyexcessz.Pr)rtest_pop_hz_deficiencyrtest_pop_hz_excesstest_pop_hz_probrangenext)r pop_pos test_typehw_resloci_res fisher_fullir r r test_hw_pop)s  zEasyController.test_hw_poprT'c CsT|dkr|j|j||||\}}}n|j|j||||\}}}t|t||fS)z0Perform Hardy-Weinberg global Heterozygote test.r)rtest_global_hz_deficiencyrtest_global_hz_excesslist) r r enum_testdememorizationbatches iterationspop_resloc_resallr r rtest_hw_global8s     zEasyController.test_hw_globalc CsZ|j|j|||d}|D]}|\\}} } ||kr| |ks&||kr*| |kr*| SqdS)zITest for linkage disequilibrium for each pair of loci in each population.rN)rtest_ldr) r locus1locus2r-r.r/all_ldld_casel1l2resultr r rtest_ld_all_pairLs   zEasyController.test_ld_all_paircCs|j|jS)z"Estimate Nm. Just a simple bridge.)r estimate_nmrr r r rr=XszEasyController.estimate_nmcCs2|j|j}|\}}t|}||d|dS)zReturn the heterozygosity info for a certain locus on a population. Returns (Expected homozygotes, observed homozygotes, Expected heterozygotes, observed heterozygotes) rrcalc_allele_genotype_freqsrr+r r locus_name geno_freqspop_iterloc_iterpopsr r rget_heterozygosity_info\sz&EasyController.get_heterozygosity_infocCs2|j|j}|\}}t|}||d|dS)z>Return the genotype counts for a certain population and locus.rrr?r rrBrCrDrEr r rget_genotype_countgsz!EasyController.get_genotype_countcCs6|j|j}|\}}t|}||d|ddS)aTReturn the Fis for a certain population and locus. Below CW means Cockerham and Weir and RH means Robertson and Hill. Returns a pair: - dictionary [allele] = (repetition count, frequency, Fis CW ) with information for each allele - a triple with total number of alleles, Fis CW, Fis RH rNr?rAr r rget_fisns zEasyController.get_fiscCs:|j|j}|\}}t|}t||d|dS)6Return the alleles for a certain population and locus.rrJ)rr@rr+keysrHr r r get_allelesszEasyController.get_allelescCs<|j|j}|\}}|D]}|d|kr|dSq dS)rLrrN)rr@r)r rBrCrDrE locus_infor r rget_alleles_all_popss  z#EasyController.get_alleles_all_popscCst|jdkr0|j|j}|\}}|D]}|ddur$d|j|d<q|dd|j|d<q|j|}|d|\}} } i} |d} t| D] \} }| | | |<qH| | fS)zCCalculate the allele frequency for a certain locus on a population.rN)NNr)lenr rr@r enumerate)r rrBrCrDrErOinfopop_namefreqstotal allele_freqallelesr$alleler r rget_allele_frequencys  z#EasyController.get_allele_frequencycC|j|jdS)zeReturn the multilocus F stats. Explain averaging. Returns Fis(CW), Fst, Fit rr calc_fst_allrr>r r rget_multilocus_f_statssz%EasyController.get_multilocus_f_statsc CsF|j|jd}|D]\}}}}}}||kr |||||fSq dS)zWReturn F stats for a locus. Returns Fis(CW), Fst, Fit, Qintra, Qinter rNr\) r rB loci_iternamefisfstfitqintraqinterr r r get_f_statss zEasyController.get_f_statscCr[)z$Calculate identity-base average Fis.r)r"calc_diversities_fis_with_identityrr>r r r get_avg_fiszEasyController.get_avg_fiscCr[)z@Calculate Allele size-base average Fis for all population pairs.r)r calc_fst_pairrr>r r rget_avg_fst_pairrizEasyController.get_avg_fst_paircCsFt|jdkr|j|jd}|D] }|d|j|d<q|j|S)zSCalculate Allele size-base average Fis for all population pairs of the given locus.rr)rQrrrjr)r locusiterrOr r rget_avg_fst_pair_locuss  z%EasyController.get_avg_fst_pair_locusaLogh㈵>cCs,|r |j|j|||S|j|j|||S)zBCalculate isolation by distance statistics for Diploid or Haploid.)rcalc_ibd_diplorcalc_ibd_haplo)r is_diplostatscalemin_distr r rcalc_ibdszEasyController.calc_ibd)N)r)rTr&r'r()r&r'r()Trorprq)__name__ __module__ __qualname____doc__rrr%r3r<r=rGrIrKrNrPrZr^rfrhrkrnrxr r r rrs4       rN)r| Bio.PopGenr Controllerrrr r r rs