Python autograd.numpy.allclose() Examples

def check_random_tensor(demo, *args, **kwargs):
    leaves = demo.sampled_pops
    ranges = [list(range(n + 1)) for n in demo.sampled_n]

    config_list = momi.data.configurations.build_full_config_list(demo.sampled_pops, demo.sampled_n)

    esfs = expected_sfs(demo, config_list, *args, **kwargs)

    tensor_components = [np.random.normal(size=(1, n + 1)) for n in demo.sampled_n]
    #tensor_components_list = tuple(v[0,:] for _,v in sorted(tensor_components.iteritems()))

    #prod1 = sfs_tensor_prod(dict(list(zip(config_list,esfs))), tensor_components)
    # sfs = momi.site_freq_spectrum(demo.sampled_pops, [dict(zip((tuple(map(tuple,c)) for c in config_list),
    #                                                           esfs))])
    sfs = momi.site_freq_spectrum(demo.sampled_pops,
                                  [{tuple(map(tuple, c)): s for c, s in zip(config_list, esfs)}])
    #assert sfs.get_dict() == {tuple(map(tuple,c)): s for c,s in zip(config_list, esfs)}
    prod1 = sfs_tensor_prod(sfs, tensor_components)
    # prod1 = sfs_tensor_prod({tuple(map(tuple,c)): s for c,s in zip(config_list, esfs)},
    #                        tensor_components)
    prod2 = expected_sfs_tensor_prod(
        tensor_components, demo, sampled_pops=demo.sampled_pops)

    assert np.allclose(prod1, prod2) 

def test_getter():
    def fun(input_list):
        A = np.sum(input_list[0])
        B = np.sum(input_list[1])
        C = np.sum(input_list[1])
        return A + B + C

    d_fun = grad(fun)
    input_list = [npr.randn(5, 6),
                   npr.randn(4, 3),
                   npr.randn(2, 4)]

    result = d_fun(input_list)
    assert np.allclose(result[0], np.ones((5, 6)))
    assert np.allclose(result[1], 2 * np.ones((4, 3)))
    assert np.allclose(result[2], np.zeros((2, 4))) 

def EM(init_params, data, callback=None):
    def EM_update(params):
        natural_params = list(map(np.log, params))
        loglike, E_stats = vgrad(log_partition_function)(natural_params, data)  # E step
        if callback: callback(loglike, params)
        return list(map(normalize, E_stats))                                    # M step

    def fixed_point(f, x0):
        x1 = f(x0)
        while different(x0, x1):
            x0, x1 = x1, f(x1)
        return x1

    def different(params1, params2):
        allclose = partial(np.allclose, atol=1e-3, rtol=1e-3)
        return not all(map(allclose, params1, params2))

    return fixed_point(EM_update, init_params) 

def _test_hvp(func, optimized):
  np.random.seed(0)
  a = np.random.normal(scale=1, size=(300,)).astype('float32')
  v = a.ravel()

  modes = ['forward', 'reverse']
  for mode1 in modes:
    for mode2 in modes:
      if mode1 == mode2 == 'forward':
        continue
      df = tangent.autodiff(
          func,
          mode=mode1,
          motion='joint',
          optimized=optimized,
          check_dims=False)
      ddf = tangent.autodiff(
          df, mode=mode2, motion='joint', optimized=optimized, check_dims=False)
      dx = ddf(a, 1, v)
      hvp_ag = hessian_vector_product(func)
      dx_ag = hvp_ag(a, v)
      assert np.allclose(dx, dx_ag) 

def compute_stats(demo, sampled_sfs, true_sfs=None, true_branch_len=None):
    sampled_sfs = momi.site_freq_spectrum(demo.leafs, to_dict(sampled_sfs))
    demo.set_data(sampled_sfs, length=1)
    demo.set_mut_rate(1)

    exp_branch_len = demo.expected_branchlen()
    exp_sfs = demo.expected_sfs()

    configs = sorted([tuple(map(tuple, c)) for c in sampled_sfs.configs])
    exp_sfs = np.array([exp_sfs[c] for c in configs])

    # use ms units
    exp_branch_len = exp_branch_len / 4.0 / demo.N_e
    exp_sfs = exp_sfs / 4.0 / demo.N_e

    if true_sfs is not None:
        assert np.allclose(true_sfs, exp_sfs, rtol=1e-4)
    if true_branch_len is not None:
        assert np.allclose(true_branch_len, exp_branch_len, rtol=1e-4)

    return exp_sfs, exp_branch_len 

def test_getter():
    def fun(input_tuple):
        A = np.sum(input_tuple[0])
        B = np.sum(input_tuple[1])
        C = np.sum(input_tuple[1])
        return A + B + C

    d_fun = grad(fun)
    input_tuple = (npr.randn(5, 6),
                   npr.randn(4, 3),
                   npr.randn(2, 4))

    result = d_fun(input_tuple)
    assert np.allclose(result[0], np.ones((5, 6)))
    assert np.allclose(result[1], 2 * np.ones((4, 3)))
    assert np.allclose(result[2], np.zeros((2, 4))) 

def check_gradient(f, x):
    print(x, "\n", f(x))

    print("# grad2")
    grad2 = Gradient(f)(x)
    print("# building grad1")
    g = grad(f)
    print("# computing grad1")
    grad1 = g(x)

    print("gradient1\n", grad1, "\ngradient2\n", grad2)
    np.allclose(grad1, grad2)

    # check Hessian vector product
    y = np.random.normal(size=x.shape)
    gdot = lambda u: np.dot(g(u), y)
    hess1, hess2 = grad(gdot)(x), Gradient(gdot)(x)
    print("hess1\n", hess1, "\nhess2\n", hess2)
    np.allclose(hess1, hess2) 

def test_getter():
    def fun(input_dict):
        A = np.sum(input_dict['item_1'])
        B = np.sum(input_dict['item_2'])
        C = np.sum(input_dict['item_2'])
        return A + B + C

    d_fun = grad(fun)
    input_dict = {'item_1' : npr.randn(5, 6),
                  'item_2' : npr.randn(4, 3),
                  'item_X' : npr.randn(2, 4)}

    result = d_fun(input_dict)
    assert np.allclose(result['item_1'], np.ones((5, 6)))
    assert np.allclose(result['item_2'], 2 * np.ones((4, 3)))
    assert np.allclose(result['item_X'], np.zeros((2, 4))) 

def __init__(self, degree=3, gamma=None, coef0=1):
        """
        Polynomial kernel
          (gamma X^T Y + coef0)^degree

        degree: default 3
        gamma: default 1/dim
        coef0: float, default 1
        """
        self.degree = degree
        self.gamma = gamma
        self.coef0 = coef0
        if degree <= 0:
            raise ValueError("KPoly needs positive degree")
        if not np.allclose(degree, int(degree)):
            raise ValueError("KPoly needs integral degree") 

def _set_startprob(self, startprob):
        if startprob is None:
            startprob = np.tile(1.0 / self.n_components, self.n_components)
        else:
            startprob = np.asarray(startprob, dtype=np.float)

            normalize(startprob)

            if len(startprob) != self.n_components:
                if len(startprob) == self.n_unique:
                    startprob_split = np.copy(startprob) / (1.0+self.n_tied)
                    startprob = np.zeros(self.n_components)
                    for u in range(self.n_unique):
                        for t in range(self.n_chain):
                            startprob[u*(self.n_chain)+t] = \
                                startprob_split[u].copy()
                else:
                    raise ValueError("cannot match shape of startprob")

        if not np.allclose(np.sum(startprob), 1.0):
            raise ValueError('startprob must sum to 1.0')

        self._log_startprob = np.log(np.asarray(startprob).copy()) 

def _set_transmat(self, transmat_val):
        if transmat_val is None:
            transmat = np.tile(1.0 / self.n_components,
                               (self.n_components, self.n_components))
        else:
            transmat_val[np.isnan(transmat_val)] = 0.0
            normalize(transmat_val, axis=1)

            if (np.asarray(transmat_val).shape == (self.n_components,
                                                   self.n_components)):
                transmat = np.copy(transmat_val)
            elif transmat_val.shape[0] == self.n_unique:
                transmat = self._ntied_transmat(transmat_val)
            else:
                raise ValueError("cannot match shape of transmat")

        if not np.all(np.allclose(np.sum(transmat, axis=1), 1.0)):
            raise ValueError('Rows of transmat must sum to 1.0')
        self._log_transmat = np.log(np.asarray(transmat).copy())
        underflow_idx = np.isnan(self._log_transmat)
        self._log_transmat[underflow_idx] = NEGINF 

def test_make_diagonal():
    def fun(D):
        return np.make_diagonal(D, axis1=-1, axis2=-2)

    D = np.random.randn(4)
    A = np.make_diagonal(D, axis1=-1, axis2=-2)
    assert np.allclose(np.diag(A), D)
    check_grads(fun)(D)

    D = np.random.randn(3, 4)
    A = np.make_diagonal(D, axis1=-1, axis2=-2)
    assert all([np.allclose(np.diag(A[i]), D[i]) for i in range(3)])
    check_grads(fun)(D) 

def test_grad():
    p = .05
    def fun0(B, Bdims):
        return einsum2.einsum2(np.exp(B**2), Bdims, np.transpose(B), Bdims[::-1], [])
    def fun1(B, Bdims):
        if Bdims: Bdims = list(range(len(Bdims)))
        return np.einsum(np.exp(B**2), Bdims,
                         np.transpose(B), Bdims[::-1], [])
    grad0 = autograd.grad(fun0)
    grad1 = autograd.grad(fun1)
    B, Bdims = random_tensor(p)
    assert np.allclose(grad0(B, Bdims), grad1(B, Bdims)) 

def test_nograd():
    # we want this to raise non-differentiability error
    fun = lambda x: np.allclose(x, (x*3.0)/3.0)
    with pytest.raises(TypeError):
        with warnings.catch_warnings(record=True) as w:
            grad(fun)(np.array([1., 2., 3.])) 

def test_grad_identity():
    fun = lambda x : x
    df = grad(fun)
    ddf = grad(df)
    assert np.allclose(df(2.0), 1.0)
    assert np.allclose(ddf(2.0), 0.0) 

def test_grad_const():
    fun = lambda x : 1.0
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("ignore")
        df = grad(fun)
        assert np.allclose(df(2.0), 0.0) 

def test_jacobian_against_stacked_grads():
    scalar_funs = [
        lambda x: np.sum(x**3),
        lambda x: np.prod(np.sin(x) + np.sin(x)),
        lambda x: grad(lambda y: np.exp(y) * np.tanh(x[0]))(x[1])
    ]

    vector_fun = lambda x: np.array([f(x) for f in scalar_funs])

    x = npr.randn(5)
    jac = jacobian(vector_fun)(x)
    grads = [grad(f)(x) for f in scalar_funs]

    assert np.allclose(jac, np.vstack(grads)) 

def test_jacobian_scalar_to_vector():
    fun = lambda x: np.array([x, x**2, x**3])
    val = npr.randn()
    assert np.allclose(jacobian(fun)(val), np.array([1., 2*val, 3*val**2])) 

def test_jacobian_against_grad():
    fun = lambda x: np.sum(np.sin(x), axis=1, keepdims=True)
    A = npr.randn(1,3)
    assert np.allclose(grad(fun)(A), jacobian(fun)(A)) 

def test_einsum2_str():
    p = .5
    A, Adims = random_tensor(p)
    B, Bdims = random_tensor(p)
    Cdims = np.random.permutation([k for k in set(Adims + Bdims) if np.random.uniform() <= p])

    dim_idxs = {k:i for i,k in enumerate(set(Adims + Bdims))}
    Adims, Bdims, Cdims = [''.join(dims)
                           for dims in (Adims,Bdims,Cdims)]
    subs = Adims + "," + Bdims + "->" + Cdims
    assert np.allclose(einsum2.einsum2(subs, A, B),
                       einsum2.einsum2(A, Adims,
                                       B, Bdims, Cdims)) 

def test_einsum2():
    p = .5
    A, Adims = random_tensor(p)
    B, Bdims = random_tensor(p)
    Cdims = np.random.permutation([k for k in set(Adims + Bdims) if np.random.uniform() <= p])

    dim_idxs = {k:i for i,k in enumerate(set(Adims + Bdims))}
    assert np.allclose(einsum2.einsum2(A, Adims, B, Bdims, Cdims),
                       np.einsum(A, [dim_idxs[s] for s in Adims],
                                 B, [dim_idxs[s] for s in Bdims],
                                 [dim_idxs[s] for s in Cdims])) 

def test_batched_dot():
    I,J,K,L = np.random.randint(1, 4, size=4)

    A = np.random.normal(size=(I,J,K))
    B = np.random.normal(size=(I,K,L))

    assert np.allclose(einsum2.batched_dot(A,B), A @ B) 

def test_count_subsets():
    demo = simple_admixture_demo()
    #data = momi.simulate_ms(ms_path, demo.demo_hist,
    #                        sampled_pops=demo.pops,
    #                        sampled_n=demo.n,
    #                        num_loci=1000, mut_rate=1.0)
    num_bases = 1000
    mu = 1.0
    num_replicates = 100
    data = demo.simulate_data(
        muts_per_gen=mu/num_bases,
        recoms_per_gen=0,
        length=num_bases,
        num_replicates=num_replicates,
        sampled_n_dict={"b":2,"a":3}).extract_sfs(None)

    subconfig = []
    for n in data.sampled_n:
        sub_n = random.randrange(n+1)
        d = random.randrange(sub_n+1)
        subconfig.append([sub_n-d, d])

    subconfig_probs = np.ones(len(data.configs))
    for i, (a, d) in enumerate(subconfig):
        #subconfig_probs *= scipy.special.comb(
        #    data.configs.value[:, i, :], [a, d]).prod(axis=1)
        #subconfig_probs /= scipy.special.comb(
        #    data.configs.value[:, i, :].sum(axis=1), a+d)
        subconfig_probs *= scipy.stats.hypergeom.pmf(
            d, data.configs.value[:, i, :].sum(axis=1),
            data.configs.value[:, i, 1], a+d)

    assert np.allclose(subconfig_probs,
                       data.configs.subsample_probs(subconfig)) 

def test_batches():
    demo = simple_five_pop_demo()

    sampled_n_dict = dict(zip(demo.leafs, [10]*5))
    num_bases=1000
    sfs = demo.simulate_data(
        length=num_bases,
        muts_per_gen=.1/num_bases,
        recoms_per_gen=0,
        num_replicates=1000,
        sampled_n_dict=sampled_n_dict)._sfs


    sfs_len = sfs.n_nonzero_entries

    print("total entries", sfs_len)
    print("total snps", sfs.n_snps())

    assert sfs_len > 30

    demo = demo._get_demo(sampled_n_dict)
    assert np.isclose(SfsLikelihoodSurface(sfs, batch_size=5).log_lik(demo),
                      momi.likelihood._composite_log_likelihood(sfs, demo))

    assert np.allclose(SfsLikelihoodSurface(sfs, batch_size=5).log_lik(demo, vector=True),
                      momi.likelihood._composite_log_likelihood(sfs, demo, vector=True))



# TODO does this test still make sense?
# uses obsolete style of demo functions which I think makes it slow 

def _apply_error_matrices(vecs, error_matrices):
    if not all([np.allclose(np.sum(err, axis=0), 1.0) for err in error_matrices]):
        raise Exception("Columns of error matrix should sum to 1")

    return [np.dot(v, err) for v, err in zip(vecs, error_matrices)]

# inverse of a PSD matrix 

def check_probs_matrix(x):
    x = truncate0(x)
    rowsums = np.sum(x, axis=1)
    assert np.allclose(rowsums, 1.0)
    return np.einsum('ij,i->ij', x, 1.0 / rowsums) 

def check_psd(X, **tol_kwargs):
    X = check_symmetric(X)
    d, U = np.linalg.eigh(X)
    d = truncate0(d, **tol_kwargs)
    ret = np.dot(U, np.dot(np.diag(d), np.transpose(U)))
    #assert np.allclose(ret, X)
    return np.array(ret, ndmin=2) 

def check_symmetric(X):
    Xt = np.transpose(X)
    assert np.allclose(X, Xt)
    return 0.5 * (X + Xt) 

def add_data(self, S, F=None):
        """
        Add a data set to the list of observations.
        First, filter the data with the impulse response basis,
        then instantiate a set of parents for this data set.

        :param S: a TxK matrix of of event counts for each time bin
                  and each process.
        """
        assert isinstance(S, np.ndarray) and S.ndim == 2 and S.shape[1] == self.K \
               and np.amin(S) >= 0 and S.dtype == np.int, \
               "Data must be a TxK array of event counts"

        T = S.shape[0]

        if F is None:
            # Filter the data into a TxKxB array
            Ftens = self.basis.convolve_with_basis(S)

            # Flatten this into a T x (KxB) matrix
            # [F00, F01, F02, F10, F11, ... F(K-1)0, F(K-1)(B-1)]
            F = Ftens.reshape((T, self.K * self.B))
            assert np.allclose(F[:,0], Ftens[:,0,0])
            if self.B > 1:
                assert np.allclose(F[:,1], Ftens[:,0,1])
            if self.K > 1:
                assert np.allclose(F[:,self.B], Ftens[:,1,0])

            # Prepend a column of ones
            F = np.hstack((np.ones((T,1)), F))

        for k,node in enumerate(self.nodes):
            node.add_data(F, S[:,k]) 

def _build_event_tree(G):
    # def node_time(v):
    #     return G.node[v]['sizes'][0]['t']

    eventEdgeList = []
    currEvents = {k: (k,) for k, v in list(dict(G.out_degree()).items()) if v == 0}
    eventDict = {e: {'subpops': (v,), 'parent_pops': (
        v,), 'child_pops': {}} for v, e in list(currEvents.items())}
    for e in G.graph['events_as_edges']:
        # get the population edges forming the event
        parent_pops, child_pops = list(map(set, list(zip(*e))))
        child_events = set([currEvents[c] for c in child_pops])

        sub_pops = set(itertools.chain(
            *[eventDict[c]['subpops'] for c in child_events]))
        sub_pops.difference_update(child_pops)
        sub_pops.update(parent_pops)

        # try:
        #     times = [t for t in map(node_time, parent_pops)]
        #     assert np.allclose(times, times[0])
        # except TypeError:
        #     ## autograd sometimes raise TypeError for this assertion
        #     pass

        eventDict[e] = {'parent_pops': tuple(parent_pops), 'subpops': tuple(
            sub_pops), 'child_pops': {c: currEvents[c] for c in child_pops}}
        currEvents.update({p: e for p in sub_pops})
        for p in child_pops:
            del currEvents[p]
        eventEdgeList += [(e, c) for c in child_events]
    ret = nx.DiGraph(eventEdgeList)
    for e in eventDict:
        ret.add_node(e, **(eventDict[e]))

    assert len(currEvents) == 1
    root, = [v for k, v in list(currEvents.items())]
    ret.root = root

    return ret

# methods for constructing demography from string