import sys

sys.path

sys.path.insert(0, '/home/claudius/Downloads/dadi')

sys.path

import dadi

import pylab

pylab.rcParams['figure.figsize'] = [12.0, 10.0]

%matplotlib inline

"""
Explanation: Load $\delta$a$\delta$i
I have not installed dadi globally on huluvu. Instead, I left it in my Downloads directory '/home/claudius/Downloads/dadi'. In order for Python to find that module, I need to add that directory to the PYTHONPATH variable.
End of explanation
"""


% ll dadiExercises/

% cat dadiExercises/ERY.FOLDED.sfs.dadi_format

"""
Explanation: Load data
End of explanation
"""


fs_ery = dadi.Spectrum.from_file('dadiExercises/ERY.FOLDED.sfs.dadi_format')

%pdoc dadi.Spectrum.from_file

fs_ery

ns = fs_ery.sample_sizes
ns

fs_ery.pop_ids = ['ery'] # must be an array, otherwise leads to error later on

# the number of segregating sites in the spectrum

fs_ery.sum()

"""
Explanation: I have turned the 1D folded SFS's from realSFS into $\delta$d$\delta$i format by hand according to the description in section 3.1 of the manual.
Note, that the last line, indicating the mask, has length 37, but the folded spectrum has length 19. Dadi wants to mask counts from invariable sites. For an unfolded spectrum, i. e. polarised with respect to an inferred ancestral allele at each site, the first and the last count classes would correspond to invariable sites. In a folded spectrum, i. e. with counts of the minor allele at each site, the last count class corresponds to SNP's with minor sample allele frequency of $n/2$ (with even sample size).
End of explanation
"""


%pdoc dadi.Plotting.plot_1d_fs

pylab.rcParams['figure.figsize'] = [12.0, 10.0]

dadi.Plotting.plot_1d_fs(fs_ery, show=False)

"""
Explanation: According to the number of segregating sites, this spectrum should have good power to distinguish between alternative demographic models (see Adams2004). However, the noise in the data is extreme, as can be seen below, which might compromise this power and maybe even lead to false inferences.
Plot the data
End of explanation
"""


# show modules within dadi

dir(dadi)

dir(dadi.Demographics1D)

# show the source of the 'Demographics1D' method

%psource dadi.Demographics1D

"""
Explanation: Built-in 1D models
End of explanation
"""


# create link to method

func = dadi.Demographics1D.snm

# make the extrapolating version of the demographic model function

func_ex = dadi.Numerics.make_extrap_log_func(func)

# setting the smallest grid size slightly larger than the largest population sample size

pts_l = [40, 50, 60]

"""
Explanation: standard neutral model
End of explanation
"""


# calculate unfolded AFS under standard neutral model (up to a scaling factor theta)

model = func_ex(0, ns, pts_l)
model

dadi.Plotting.plot_1d_fs(model.fold()[:19], show=False)

"""
Explanation: The snm function does not take parameters to optimize. I can therefore get directly the expected model. The snm function does not take a fold argument. I am therefore going to calculated an unfolded expected spectrum and then fold.
End of explanation
"""


# get the source of the fold method, which is part of the Spectrum object

%psource dadi.Spectrum.fold

# get the docstring of the Spectrum object

%pdoc dadi.Spectrum

# retrieve the spectrum array from the Spectrum object

model.data

"""
Explanation: What's happening in the 18th count class?
End of explanation
"""


# reverse spectrum and add to itself

model_fold = model.data + model.data[::-1]
model_fold

# discard all count classes >n/2

model_fold = model_fold[:19]
model_fold

"""
Explanation: I am going to fold manually now.
End of explanation
"""


# divide highest sample frequency class by 2

model_fold[18] = model_fold[18]/2.0

model_fold

# create dadi Spectrum object from array, need to specify custom mask

model_folded = dadi.Spectrum(data=model_fold, mask_corners=False, mask= [1] + [0]*18)
model_folded

dadi.Plotting.plot_1d_fs(model_folded)

"""
Explanation: When the sample size is even, then highest sample frequency class corresponds to just one unfolded class (18). This has been added to itself and those SNP's are counted twice at the moment. I need to divide this class by 2 to get the correct count for this folded class.
End of explanation
"""


# fold the unfolded model

model_folded = model.fold()
#model_folded = model_folded[:(ns[0]+1)]
model_folded.pop_ids = ['ery'] # be sure to give an array, not a scalar string
model_folded

ll_model_folded = dadi.Inference.ll_multinom(model_folded, fs_ery)

print 'The log composite likelihood of the observed ery spectrum given a standard neutral model is {0:.3f}.'.format(ll_model_folded)

"""
Explanation: The folded expected spectrum is correct. Also, see figure 4.5 in Wakeley2009.
How to fold an unfolded spectrum
End of explanation
"""


theta = dadi.Inference.optimal_sfs_scaling(model_folded, fs_ery)

print 'The optimal value of theta is {0:.3f}.'.format(theta)

"""
Explanation: $\theta$ and implied $N_{ref}$
End of explanation
"""


mu = 3e-9
L = fs_ery.data.sum() # this sums over all entries in the spectrum, including masked ones, i. e. also contains invariable sites
print "The total sequence length is " + str(L)
N_ref = theta/L/mu/4
print "The effective ancestral population size (in number of diploid individuals) implied by this theta is: {0}.".format(int(N_ref))

"""
Explanation: This theta estimate is a little bit higher than what I estimated with curve fitting in Fist_Steps_with_dadi.ipynb, which was 10198.849.
What effective ancestral population size would that imply?
According to section 4.4 in the dadi manual:
$$
\theta = 4 N_{ref} \mu_{L} \qquad \text{L: sequence length}
$$
Let's assume the mutation rate per nucleotide site per generation is $3\times 10^{-9}$ (see e. g. Liu2017). Then
$$
\mu_{L} = \mu_{site} \times L
$$
So
$$
\theta = 4 N_{ref} \mu_{site} \times L
$$
and
$$
N_{ref} = \frac{\theta}{4 \mu_{site} L}
$$
End of explanation
"""


x = pylab.arange(0, 100)
y = 0.5**(x)
pylab.plot(x, y)

x[:10] * y[:10]

sum(x * y)

"""
Explanation: This effective population size is consistent with those reported in Lynch2016 for other insect species.

Begin Digression:
End of explanation
"""


model_folded * theta

pylab.semilogy(model_folded * theta, "bo-", label='SNM')
pylab.plot(fs_ery, "ro-", label='ery')
pylab.legend()

%psource dadi.Plotting.plot_1d_comp_Poisson

# compare model prediction and data visually with dadi function

dadi.Plotting.plot_1d_comp_multinom(model_folded[:19], fs_ery[:19], residual='linear')

"""
Explanation: End Digression
End of explanation
"""


fs_ery

# make copy of spectrum array
data_abc = fs_ery.data.copy()

# resize the array to the unfolded length

data_abc.resize((37,))
data_abc

fs_ery_ext = dadi.Spectrum(data_abc)
fs_ery_ext

fs_ery_ext.fold()

fs_ery_ext = fs_ery_ext.fold()
fs_ery_ext.pop_ids = ['ery']
fs_ery_ext

fs_ery_ext.sample_sizes

"""
Explanation: The lower plot is for the scaled Poisson residuals. 
$$
residuals = (model - data)/\sqrt{model}
$$
The model is the expected counts in each frequency class. If these counts are Poisson distributed, then their variance is equal to their expectation. The differences between model and data are therefore scaled by the expected standard deviation of the model counts.
The observed counts deviate by up to 30 standard deviations from the model!
What could be done about this? 
The greatest deviations are seen for the first two frequency classes, the ones that should provide the greatest amount of information (Fu1994) for theta and therefore probably also other parameters. Toni has suggested that the doubleton class is inflated due to "miscalling" heterozygotes as homozygotes. When they contain a singleton they will be "called" as homozygote and therefore contribute to the doubleton count. This is aggravated by the fact that the sequenced individuals are all male which only possess one X chromosome. The X chromosome is the fourth largest of the 9 chromosomes of these grasshoppers (8 autosomes + X) (see Gosalvez1988, fig. 2). That is, about 1/9th of the sequenced RAD loci are haploid but ANGSD assumes all loci to be diploid. The genotype likelihoods it calculates are all referring to diploid genotypes.
I think one potential reason for the extreme deviations is that the genotype likelihoods are generally biased toward homozygote genotypes (i. e. also for autosomal loci) due to PCR duplicates (see eq. 1 in Nielsen2012). So, one potential improvement would be to remove PCR duplicates. 
Another potential improvement could be found by subsampling 8/9th to 8/10th of the contigs in the SAF files and estimating an SFS from these. Given enough subsamples, one should eventually be found that maximally excludes loci from the X chromosome. This subsample is expected to produce the least squared deviations from an expected SFS under the standard neutral model. However, one could argue that this attempt to exclude problematic loci could also inadvertently remove loci that strongly deviate from neutral expectations due to non-neutral evolution, again reducing power to detect deviations from the standard neutral model. I think one could also just apply the selection criterion of the second MAF class to be lower than the first and just save all contig subsamples and SFS's that fulfill that criterioin, since that should be true for all demographic scenarios.
Exponential growth
Creating a folded spectrum exactly how dadi expects it
As seen above in the folded model spectrum, dadi just masks out entries that are not sensical in a folded spectrum, but keeps the length of the spectrum the same as the unfolded. That way the sample size (i. e. number of chromosomes) is determined correctly. Let's create a correct folded spectrum object for ery.
End of explanation
"""


% cat dadiExercises/ERY.FOLDED.sfs.dadi_format

# load the spectrum that was created from folded SAF's

fs_ery_folded_by_Angsd = dadi.Spectrum.from_file('dadiExercises/ERY.FOLDED.sfs.dadi_format')
fs_ery_folded_by_Angsd

# extract unmasked entries of the SFS

m = fs_ery_folded_by_Angsd.mask
fs_ery_folded_by_Angsd[m == False]

"""
Explanation: Now, the reported sample size is correct and we have a Spectrum object that dadi can handle correctly.
To fold or not to fold by ANGSD
Does estimating an unfolded spectrum with ANGSD and then folding yield a sensible folded SFS when the sites are not polarised with respect to an ancestral allele but with respect to the reference allele? Matteo Fumagalli thinks that this is sensible.
Load SFS folded by ANGSD
End of explanation
"""


% ll ../ANGSD/SFS/ERY/

"""
Explanation: Load unfolded SFS
End of explanation
"""


% ll

% cat ERY.unfolded.sfs

# load unfolded spectrum

fs_ery_unfolded_by_ANGSD = dadi.Spectrum.from_file('ERY.unfolded.sfs')
fs_ery_unfolded_by_ANGSD

# fold unfolded spectrum

fs_ery_unfolded_by_Angsd_folded = fs_ery_unfolded_by_ANGSD.fold()
fs_ery_unfolded_by_Angsd_folded

# plot the two spectra

pylab.rcParams['figure.figsize'] = [12.0, 10.0]

pylab.plot(fs_ery_folded_by_Angsd, 'ro-', label='folded by ANGSD')
pylab.plot(fs_ery_unfolded_by_Angsd_folded, 'bo-', label='folded by DADI')
pylab.legend()
pylab.savefig('ery_fold_comp.png')

%psource dadi.Plotting.plot_1d_comp_Poisson

dadi.Plotting.plot_1d_comp_Poisson(fs_ery_folded_by_Angsd[:19], fs_ery_unfolded_by_Angsd_folded[:19], \
                                   residual='linear')

"""
Explanation: I have copied the unfolded SFS into the current directory.
End of explanation
"""


% ll dadiExercises/

% cat dadiExercises/PAR.FOLDED.sfs.dadi_format

# load the spectrum folded by ANGSD

fs_par_folded_by_Angsd = dadi.Spectrum.from_file('dadiExercises/PAR.FOLDED.sfs.dadi_format')
fs_par_folded_by_Angsd

% cat PAR.unfolded.sfs

# load spectrum that has been created from unfolded SAF's

fs_par_unfolded_by_Angsd = dadi.Spectrum.from_file('PAR.unfolded.sfs')
fs_par_unfolded_by_Angsd

fs_par_unfolded_by_Angsd_folded = fs_par_unfolded_by_Angsd.fold()
fs_par_unfolded_by_Angsd_folded

dadi.Plotting.plot_1d_comp_Poisson(fs_par_folded_by_Angsd[:19], fs_par_unfolded_by_Angsd_folded[:19], \
                                   residual='linear')

#pylab.subplot(2,1,1)
pylab.plot(fs_par_folded_by_Angsd[:19], 'ro-', label='folded by ANGSD')

#pylab.subplot(2,1,2)
pylab.plot(fs_par_unfolded_by_Angsd_folded, 'bo-', label='folded by DADI')
pylab.legend()
pylab.savefig('par_fold_comp.png')

"""
Explanation: The sizes of the residuals (scaled by the Poisson standard deviations) indicate that the two versions of the folded SFS of ery are significantly different.
Now, what does the parallelus data say?
End of explanation
"""


%psource dadi.Spectrum.from_data_dict

"""
Explanation: The unfolded spectrum folded by dadi seems to be a bit better behaved than the one folded by ANGSD. I really wonder whether folding in ANGSD is needed.
The folded 2D spectrum from ANGSD is a 19 x 19 matrix. This is not a format that dadi can understand.
End of explanation
"""


# show the source of the 'Demographics1D' method

%psource dadi.Demographics1D.growth

# create link to function that specifies a simple growth or decline model

func = dadi.Demographics1D.growth

# create extrapolating version of the function

func_ex = dadi.Numerics.make_extrap_log_func(func)

# set lower and upper bounds to nu and T

upper_bound = [100, 3]
lower_bound = [1e-2, 0]

# set starting value

p0 = [1, 1] # corresponds to constant population size

%pdoc dadi.Misc.perturb_params

# perturb starting values by up to a factor of 2

p0 = dadi.Misc.perturb_params(p0, fold=1, upper_bound=upper_bound, lower_bound=lower_bound)

p0

%psource dadi.Inference.optimize_log

# run optimisation of paramters

popt = dadi.Inference.optimize_log(p0=p0, data=fs_ery, model_func=func_ex, pts=pts_l, \
                                   lower_bound=lower_bound, upper_bound=upper_bound, \
                                   verbose=0, maxiter=100, full_output=False)

popt

"""
Explanation: See this thread on the dadi forum.

Exponential growth model
End of explanation
"""


from ipyparallel import Client

cl = Client()

cl.ids

"""
Explanation: Parallelised $\delta$a$\delta$i
I need to run the simulation with different starting values to check convergence.
I would like to do these runs in parallel. I have 12 cores available on huluvu.
End of explanation
"""


# create load balanced view of the engines

lbview = cl.load_balanced_view()

lbview.block

# create direct view of all engines

dview = cl[:]

"""
Explanation: I now have connections to 11 engines. I started the engines with ipcluster start -n 11 &amp; in the terminal.
End of explanation
"""


# set starting value for all engines

dview['p0'] = [1, 1]
dview['p0']

# set lower and upper bounds to nu and T for all engines

dview['upper_bound'] = [100, 3]
dview['lower_bound'] = [1e-2, 0]

dview['fs_ery'] = fs_ery
cl[0]['fs_ery']

dview['func_ex'] = func_ex
dview['pts_l'] = pts_l

"""
Explanation: import variables to namespace of engines
End of explanation
"""


with dview.sync_imports():
    import sys

dview.execute('sys.path.insert(0, \'/home/claudius/Downloads/dadi\')')

cl[0]['sys.path']

with dview.sync_imports():
    import dadi

"""
Explanation: import dadi on all engines
End of explanation
"""


@lbview.parallel(block=True)
def run_dadi(x): # for the function to be called with map, it needs to have one input variable
    # perturb starting values by up to a factor of 2
    p1 = dadi.Misc.perturb_params(p0, fold=1, upper_bound=upper_bound, lower_bound=lower_bound)
    # run optimisation of paramters
    popt = dadi.Inference.optimize_log(p0=p1, data=fs_ery, model_func=func_ex, pts=pts_l, \
                                   lower_bound=lower_bound, upper_bound=upper_bound, \
                                   verbose=0, maxiter=100, full_output=False)
    return popt

run_dadi.map(range(20))

popt

# set starting value
p0 = [1, 1]

# perturb starting values by up to a factor of 2
p0 = dadi.Misc.perturb_params(p0, fold=1, upper_bound=upper_bound, lower_bound=lower_bound)

# run optimisation of paramters
popt = dadi.Inference.optimize_log(p0=p0, data=fs_ery_ext, model_func=func_ex, pts=pts_l, \
                                   lower_bound=lower_bound, upper_bound=upper_bound, \
                                   verbose=0, maxiter=100, full_output=False)
popt

"""
Explanation: create parallel function to run dadi
End of explanation
"""


def exp_growth(x):
    p0 = [1, 1]

    # perturb starting values by up to a factor of 2
    p0 = dadi.Misc.perturb_params(p0, fold=1, upper_bound=upper_bound, lower_bound=lower_bound)

    # run optimisation of paramters
    popt = dadi.Inference.optimize_log(p0=p0, data=fs_ery_ext, model_func=func_ex, pts=pts_l, \
                                   lower_bound=lower_bound, upper_bound=upper_bound, \
                                   verbose=0, maxiter=100, full_output=False)
    return popt

popt = map(exp_growth, range(10))

# this will run a few minutes

# popt

import ipyparallel as ipp

c = ipp.Client()
c.ids

%%time

dview = c[:]

popt = dview.map_sync(exp_growth, range(10))

"""
Explanation: 
End of explanation
"""


popt

"""
Explanation: Unfortunately, parallelisation is not as straightforward as it should be.
End of explanation
"""


ns = fs_ery_ext.sample_sizes
ns

print popt[0]
print popt[9]

"""
Explanation: Except for the last iteration, the two parameter estimates seem to have converged.
End of explanation
"""


model_one = func_ex(popt[0], ns, pts_l)
ll_model_one = dadi.Inference.ll_multinom(model_one, fs_ery_ext)
ll_model_one

model_two = func_ex(popt[9], ns, pts_l)
ll_model_two = dadi.Inference.ll_multinom(model_two, fs_ery_ext)
ll_model_two

"""
Explanation: What is the log likelihood of the model given these two different parameter sets?
End of explanation
"""


print 'The model suggests that exponential decline in population size started {0:.0f} generations ago.'.format(popt[0][1] * 2 * N_ref)

"""
Explanation: The lower log-likelihood for the last set of parameters inferred indicates that the optimisation got trapped in a local minimum in the last run of the optimisation.
What the majority of the parameter sets seem to indicate is that at about time $0.007 \times 2 N_{ref}$ generations in the past the ancestral population started to shrink exponentially, reaching a population size of about $0.14 \times N_{ref}$ at present.
End of explanation
"""


dir(dadi.Demographics1D)

%psource dadi.Demographics1D.two_epoch

"""
Explanation: Two epoch model
End of explanation
"""


func = dadi.Demographics1D.two_epoch
func_ex = dadi.Numerics.make_extrap_log_func(func)

upper_bound = [10, 3]
lower_bound = [1e-3, 0]
pts_l = [40, 50, 60]

def stepwise_pop_change(x):
    # set initial values
    p0 = [1, 1]

    # perturb initial parameter values randomly by up to 2 * fold
    p0 = dadi.Misc.perturb_params(p0, fold=1.5, \
                                  upper_bound=upper_bound, lower_bound=lower_bound)
    
    # run optimisation
    popt = dadi.Inference.optimize_log(p0, fs_ery_ext, func_ex, pts_l, \
                                      upper_bound=upper_bound, lower_bound=lower_bound,
                                      verbose=0, maxiter=10)
    
    return popt

stepwise_pop_change(1)

stepwise_pop_change(1)

popt = map(stepwise_pop_change, range(10))

popt

"""
Explanation: This model specifies a stepwise change in population size some time ago. It assumes that the population size has stayed constant since the change.
End of explanation
"""


nu = [i[0] for i in popt]
nu

T = [i[1] for i in popt]
T

pylab.rcParams['font.size'] = 14.0

pylab.loglog(nu, T, 'bo')
pylab.xlabel(r'$\nu$')
pylab.ylabel('T')

"""
Explanation: This model does not converge on a set of parameter values.
End of explanation
"""


%psource dadi.Demographics1D

"""
Explanation: Both parameters seem to be correlated. With the available data, it may not be possible to distinguish between a moderate reduction in population size a long time ago (topright in the above figure) and a drastic reduction in population size a short time ago (bottomleft in the above figure).
Bottleneck then exponential growth
End of explanation
"""


func = dadi.Demographics1D.bottlegrowth
func_ex = dadi.Numerics.make_extrap_log_func(func)

upper_bound = [100, 100, 3]
lower_bound = [1e-3, 1e-3, 0]
pts_l = [40, 50, 60]

def bottleneck_growth(x):
    p0 = [1, 1, 1] # corresponds to constant population size
    
        # perturb initial parameter values randomly by up to 2 * fold
    p0 = dadi.Misc.perturb_params(p0, fold=1.5, \
                                  upper_bound=upper_bound, lower_bound=lower_bound)
    
    # run optimisation
    popt = dadi.Inference.optimize_log(p0, fs_ery_ext, func_ex, pts_l, \
                                      upper_bound=upper_bound, lower_bound=lower_bound,
                                      verbose=0, maxiter=10)
    
    return popt

%%time

popt = map(bottleneck_growth, range(10))

popt

"""
Explanation: This model has three parameters. $\nu_B$ is the ratio of the population size (with respect to the ancestral population size $N_{ref}$) after the first stepwise change at time T in the past. The population is then asumed to undergo exponential growth/decline to a ratio of population size $\nu_F$ at present.
End of explanation
"""


func = dadi.Demographics1D.three_epoch
func_ex = dadi.Numerics.make_extrap_log_func(func)

%psource dadi.Demographics1D.three_epoch

"""
Explanation: There is no convergence of parameters estimates. The parameter combinations stand for vastly different demographic scenarios. Most seem to suggest a population increase (up to 100 times the ancestral population size), followed by exponential decrease to about the ancestral population size.
Three epochs
End of explanation
"""


upper_bound = [100, 100, 3, 3]
lower_bound = [1e-3, 1e-3, 0, 0]
pts_l = [40, 50, 60]

def opt_three_epochs(x):
    p0 = [1, 1, 1, 1] # corresponds to constant population size
    
        # perturb initial parameter values randomly by up to 2 * fold
    p0 = dadi.Misc.perturb_params(p0, fold=1.5, \
                                  upper_bound=upper_bound, lower_bound=lower_bound)
    
    # run optimisation
    popt = dadi.Inference.optimize_log(p0, fs_ery_ext, func_ex, pts_l, \
                                      upper_bound=upper_bound, lower_bound=lower_bound,
                                      verbose=0, maxiter=10)
    
    return popt

%%time

popt = map(opt_three_epochs, range(10))

popt

"""
Explanation: This model tries to estimate three parameters. The populations is expected to undergo a stepwise population size change (bottleneck) at time TF + TB. At time TF it is expected to recover immediately to the current population size.
End of explanation
"""


