#!/usr/bin/env python

import pydigree
import sys
import time
import itertools

def sharepairs(ped, inds, loc):
    from pydigree import count,ibs
    npairs = float( len(inds) * (len(inds)-1)/2 ) 
    r = [ibs(j.get_genotype(loc), k.get_genotype(loc)) > 0 for j,k
         in itertools.combinations(inds,2)]
    return count(True,r) / npairs


pop = pydigree.Population(5000)
                
ped = pydigree.io.read_ped(sys.argv[1], pop)[sys.argv[2]]
# Clear the genotypes, if present
ped.clear_genotypes()

c = pydigree.Chromosome()
c.add_genotype(0.5, 0)
print c
ped.add_chromosome(c)


affs = [x for x in ped if x.phenotypes['affected']]


niter = 1000
print "%s simulations" % niter

sim_share = []

t = time.time()
for x in range(niter):
    if x % 250 == 0:
        print 'Simulation %s' % x
    ped.simulate_ibd_states()
    s = sharepairs(ped, affs, (0,0))
    ped.clear_genotypes()
    sim_share.append(s)
t2 = time.time()-t

print "Maximum simulated allele sharing: %s" % max(sim_share)
print "Empiric P: %s" % (len([x for x in sim_share if x >= 1.0])/float(niter))
print "Time: %s (time per pedigree: %s)" % (t2,t2/niter)


ofilename = 'null.dist'
print "Outputting distribution to %s" % ofilename
with open(ofilename,'w') as of:
    for s in sim_share:
        of.write(str(s))
        of.write('\n')

#!/usr/bin/env python

import pydigree
import sys
import time
import itertools
import argparse

from pydigree import ibs

parser = argparse.ArgumentParser()
parser.add_argument('-f', '--file', required=True,
                    help='Pedigree file for simulations')
parser.add_argument('-n', '--niter', type=int, default=10000,
                    help='Number of simulations per pedigree')
parser.add_argument('--only', metavar='PED',
                    help='Only simulate specified pedigrees',
                    nargs='*', dest='onlypeds')
parser.add_argument('--obs', type=float, default=1.0,
                    help='Maximum observed value of test statistic')
parser.add_argument('--writedist', action='store_true',
                    help='Write null distribution to file')
args = parser.parse_args()


def sbool(ped, inds, loc):
    npairs = float(len(inds) * (len(inds) - 1) / 2)
    r = [ibs(j.get_genotype(loc, checkhasgeno=False),
             k.get_genotype(loc, checkhasgeno=False),
             checkmissing=False) > 0 for j, k
         in itertools.combinations(inds, 2)]
    return sum(r) / npairs


def genedrop(ped, affs, scorer, iteration):
    if iteration % (args.niter / 10) == 0:
        print 'Simulation %s' % x
    ped.simulate_ibd_states(inds=affs)
    s = scorer(ped, affs, (0, 0))
    return s


scorefunction = sbool

pop = pydigree.Population(5000)
peds = pydigree.io.read_ped(args.file, pop)
for ped in peds:
    if args.onlypeds and ped.label not in args.onlypeds:
        continue

    # Clear the genotypes, if present
    ped.clear_genotypes()

    c = pydigree.Chromosome()
    c.add_genotype(0.5, 0)
    ped.add_chromosome(c)

    affs = [x for x in ped if x.phenotypes['affected']]

    print "Pedigree %s, %s simulations" % (ped.label, args.niter)

    sim_share = [genedrop(ped, affs, scorefunction, x)
                 for x in xrange(args.niter)]

    print "Maximum simulated allele sharing: %s" % max(sim_share)
    print "Empiric P: %s" % (len([x for x in sim_share if x >= args.obs]) / float(args.niter))

    if not args.writedist:
        continue
    ofilename = '%s.null.dist' % ped.label
    print "Outputting distribution to %s" % ofilename
    with open(ofilename, 'w') as of:
        for s in sim_share:
            of.write(str(s))
            of.write('\n')

optimize and make run like a real application and not a one-off script