N_CODON = len(CODON_TO_INDEX) NT_TO_INDEX = {nt: i for (i, nt) in INDEX_TO_NT.items()} CODON_TO_AA = [] INDEX_TO_CODON = {} INDEX_TO_AA = dict(enumerate('ACDEFGHIKLMNPQRSTVWY')) CODON_SINGLEMUT = numpy.full((N_CODON, N_CODON), False, dtype='bool') CODON_TRANSITION = numpy.full((N_CODON, N_CODON), False, dtype='bool') CODON_NONSYN = numpy.full((N_CODON, N_CODON), False, dtype='bool') N_NT = len(INDEX_TO_NT) AA_TO_INDEX = {aa: i for (i, aa) in INDEX_TO_AA.items()} N_AA = len(INDEX_TO_AA)
import random
import unittest
import numpy
import scipy.linalg
import scipy.optimize
import sympy
import phydmslib.models
from phydmslib.constants import (N_CODON, NT_TO_INDEX, CODON_TO_AA,
                                 INDEX_TO_CODON, INDEX_TO_AA, CODON_SINGLEMUT,
                                 CODON_TRANSITION, CODON_NONSYN, N_NT,
                                 AA_TO_INDEX, N_AA)
"""Tests `phydmslib.models.ExpCM` class.

Written by Jesse Bloom.

Uses `sympy` to make sure attributes and derivatives of attributes
are correct for `ExpCM` implemented in `phydmslib.models`.
"""




class testExpCM(unittest.TestCase):
    """Tests ``ExpCM`` model."""

    def test_ExpCM(self):
        """Initialize `ExpCM`, test values, update, test again."""
        # create preferences
        random.seed(1)
        numpy.random.seed(1)
        self.nsites = 2
        self.prefs = []
        minpref = 0.01
        for _r in range(self.nsites):
            rprefs = numpy.random.dirichlet([0.5] * N_AA)
            rprefs[rprefs < minpref] = minpref
            rprefs /= rprefs.sum()
            self.prefs.append(dict(zip(sorted(AA_TO_INDEX.keys()), rprefs)))

        # create initial ExpCM
        phi = numpy.random.dirichlet([2] * N_NT)
        omega = 0.7
        kappa = 2.5
        beta = 1.9
        self.expcm = phydmslib.models.ExpCM(
            self.prefs, phi=phi, omega=omega, kappa=kappa, beta=beta)
        self.assertTrue(numpy.allclose(phi, self.expcm.phi))
        self.assertTrue(numpy.allclose(omega, self.expcm.omega))
        self.assertTrue(numpy.allclose(kappa, self.expcm.kappa))
        self.assertTrue(numpy.allclose(beta, self.expcm.beta))

        self.assertTrue(
            numpy.allclose(
                numpy.repeat(1.0, self.nsites),
                self.expcm.stationarystate.sum(axis=1)))

        # now check ExpCM attributes / derivates, updating several times
        for _update in range(2):
            self.params = {
                "omega": random.uniform(*self.expcm.PARAMLIMITS["omega"]),
                "kappa": random.uniform(*self.expcm.PARAMLIMITS["kappa"]),
                "beta": random.uniform(0.5, 2.5),
                "eta": numpy.array(
                    [random.uniform(*self.expcm.PARAMLIMITS["eta"])
                     for i in range(N_NT - 1)]),
                "mu": random.uniform(0.05, 3.0)}
            self.expcm.updateParams(self.params)
            self.check_ExpCM_attributes()
            self.check_ExpCM_derivatives()
            self.check_ExpCM_matrix_exponentials()

    def check_ExpCM_attributes(self):
        """Make sure `ExpCM` has the expected attribute values."""
        self.assertEqual(self.nsites, self.expcm.nsites)

        # make sure Prxy has rows summing to zero
        self.assertFalse(numpy.isnan(self.expcm.Prxy).any())
        self.assertFalse(numpy.isinf(self.expcm.Prxy).any())
        diag = numpy.eye(N_CODON, dtype="bool")
        for r in range(self.nsites):
            self.assertTrue(numpy.allclose(0, numpy.sum(self.expcm.Prxy[r],
                            axis=1)))
            self.assertTrue(numpy.allclose(0, self.expcm.Prxy[r].sum()))
            self.assertTrue((self.expcm.Prxy[r][diag] <= 0).all())
            self.assertTrue((self.expcm.Prxy[r][~diag] >= 0).all())

        # make sure prx sums to 1 for each r
        self.assertTrue((self.expcm.prx >= 0).all())
        for r in range(self.nsites):
            self.assertTrue(numpy.allclose(1, self.expcm.prx[r].sum()))

        # prx is eigenvector or Prxy for the same r, but not different r
        for r in range(self.nsites):
            self.assertTrue(
                numpy.allclose(0, numpy.dot(self.expcm.prx[r],
                               self.expcm.Prxy[r])))
            if r > 0:
                self.assertFalse(
                    numpy.allclose(
                        0,
                        numpy.dot(self.expcm.prx[r], self.expcm.Prxy[r - 1])))

    def check_ExpCM_derivatives(self):
        """Use `sympy` to check values & derivatives of `ExpCM` attributes."""
        (Prxy, Qxy, phiw, beta, omega,
         eta0, eta1, eta2, kappa) = sympy.symbols("Prxy, Qxy, phiw, beta, "
                                                  "omega, eta0, eta1, eta2, "
                                                  "kappa")

        values = {"beta": self.params["beta"],
                  "omega": self.params["omega"],
                  "kappa": self.params["kappa"],
                  "eta0": self.params["eta"][0],
                  "eta1": self.params["eta"][1],
                  "eta2": self.params["eta"][2]}

        # check Prxy
        for r in range(self.nsites):
            for x in range(N_CODON):

                pirAx = self.prefs[r][INDEX_TO_AA[CODON_TO_AA[x]]]