N_CODON = len(CODON_TO_INDEX) CODON_NONSYN = numpy.full((N_CODON, N_CODON), False, dtype='bool') CODON_TRANSITION = numpy.full((N_CODON, N_CODON), False, dtype='bool') CODON_SINGLEMUT = numpy.full((N_CODON, N_CODON), False, dtype='bool') ALMOST_ZERO = 1.0e-6 CODON_NT_INDEX = numpy.full((3, N_CODON), -1, dtype='int') STOP_CODON_TO_NT_INDICES = [] N_STOP = 0 N_NT = len(INDEX_TO_NT) N_AA = len(INDEX_TO_AA) INDEX_TO_NT = dict(enumerate('ACGT')) CODON_NT_MUT = numpy.full((N_NT, N_CODON, N_CODON), False, dtype='bool') CODON_NT_COUNT = numpy.zeros((N_NT, N_CODON), dtype='int') CODON_TO_AA = [] INDEX_TO_AA = dict(enumerate('ACDEFGHIKLMNPQRSTVWY')) AA_TO_INDEX = {aa: i for (i, aa) in INDEX_TO_AA.items()} CODON_NT = numpy.full((3, N_NT, N_CODON), False, dtype='bool')
import copy
import six
import abc
import warnings
import numpy
import scipy.misc
import scipy.optimize
import scipy.linalg
import scipy.stats
import scipy.special
    import doctest
from phydmslib.numutils import (broadcastMatrixMultiply,
                                broadcastMatrixVectorMultiply,
                                broadcastGetCols)
from phydmslib.constants import (N_CODON, CODON_NONSYN, CODON_TRANSITION,
                                 CODON_SINGLEMUT, ALMOST_ZERO, CODON_NT_INDEX,
                                 STOP_CODON_TO_NT_INDICES, N_STOP, N_NT, N_AA,
                                 INDEX_TO_NT, CODON_NT_MUT, CODON_NT_COUNT,
                                 CODON_TO_AA, INDEX_TO_AA, AA_TO_INDEX,
                                 CODON_NT)
        _checkParam('e_pw', e_pw, self.PARAMLIMITS, self.PARAMTYPES)
        self.e_pw = e_pw.copy()
        self.phi = self._calculate_correctedF3X4()
        assert (numpy.allclose(self.phi.sum(axis=1),
                               numpy.ones(3, dtype='float'),
                               atol=1e-4, rtol=5e-3)), "The `phi` values do \
                               not sum to 1 for all `p`"

        self.Phi_x = numpy.ones(N_CODON, dtype='float')
        self._calculate_Phi_x()
        self._nsites = nsites
        assert self._nsites > 0, "There must be more than 1 site in the gene"

        # check allowed params
        assert all(map(lambda x: x in self.ALLOWEDPARAMS, freeparams)), (
              "Invalid entry in freeparams\nGot: {0}\nAllowed: {1}"
              .format(', '.join(freeparams), ', '.join(self.ALLOWEDPARAMS)))
        self._freeparams = list(freeparams)  # `freeparams` is a property

        # set attributes to calling params
        self._mu = mu  # underscore as `mu` is property
        self.kappa = kappa
        self.omega = omega
        for (name, value) in [('kappa', self.kappa), ('omega', self.omega),
                              ('mu', self.mu)]:
            _checkParam(name, value, self.PARAMLIMITS, self.PARAMTYPES)

        # define other params, initialized appropriately
        # single site dimension to be carried through the calcs added here
        self.Pxy = numpy.zeros((1, N_CODON, N_CODON), dtype='float')
        self.Pxy_no_omega = numpy.zeros((1, N_CODON, N_CODON), dtype='float')
        self.D = numpy.zeros((1, N_CODON), dtype='float')
        self.A = numpy.zeros((1, N_CODON, N_CODON), dtype='float')
        self.Ainv = numpy.zeros((1, N_CODON, N_CODON), dtype='float')
        self.dPxy = {}
        self.B = {}
        for param in self.freeparams:
            if param in self.ALLOWEDPARAMS:
                self.dPxy[param] = numpy.zeros((1, N_CODON, N_CODON),
                                               dtype='float')
                self.B[param] = numpy.zeros((1, N_CODON, N_CODON),
                                            dtype='float')
            else:
                raise ValueError("Unrecognized param {0}".format(param))

        # indexes diagonals in square matrices
        self._diag_indices = numpy.diag_indices(N_CODON)
        self.updateParams({}, update_all=True)

    @property
    def stationarystate(self):
        """See docs for `Model` abstract base class."""
        # repeat the single state by the number sites
        return numpy.tile(self.Phi_x, (self.nsites, 1))

    def dstationarystate(self, param):
        """See docs for `Model` abstract base class."""
        return numpy.zeros((self.nsites, N_CODON), dtype='float')

    @property
    def paramsReport(self):
        """See docs for `Model` abstract base class."""
        report = {}
        for param in self._REPORTPARAMS:
            pvalue = getattr(self, param)
            if isinstance(pvalue, float):
                report[param] = pvalue
            elif (isinstance(pvalue, numpy.ndarray) and
                  pvalue.shape == (3, N_NT)):
                for p in range(3):
                    for w in range(N_NT - 1):
                        report['{0}{1}{2}'.format(param, p, INDEX_TO_NT[w])] =\
                                pvalue[p][w]
            else:
                raise ValueError("Unexpected param: {0}".format(param))
        return report

    @property
    def branchScale(self):
        """See docs for `Model` abstract base class."""
        bs = -(self.Phi_x * numpy.diagonal(self.Pxy[0])).sum() * self.mu
        assert bs > 0
        return bs

    @property
    def nsites(self):
        """See docs for `Model` abstract base class."""
        return self._nsites

    @property
    def freeparams(self):
        """See docs for `Model` abstract base class."""
        return self._freeparams

    @property
    def mu(self):
        """See docs for `Model` abstract base class."""
        return self._mu

    @mu.setter
    def mu(self, value):
        """Set new `mu` value."""
        self._mu = value

    def _calculate_correctedF3X4(self):
        """Calculate `phi` based on the empirical `e_pw` values"""
        def F(phi):
            phi_reshape = phi.reshape((3, N_NT))
            functionList = []
            stop_frequency = []

            for x in range(N_STOP):
                codonFrequency = STOP_CODON_TO_NT_INDICES[x] * phi_reshape
                codonFrequency = numpy.prod(codonFrequency.sum(axis=1))
                stop_frequency.append(codonFrequency)
            C = numpy.sum(stop_frequency)

            for p in range(3):
                for w in range(N_NT):
                    s = 0

                    for x in range(N_STOP):