Python numpy.ix_() Examples

The following are 30 code examples for showing how to use numpy.ix_(). These examples are extracted from open source projects. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example.

You may check out the related API usage on the sidebar.

You may also want to check out all available functions/classes of the module numpy , or try the search function .

Example 1
Project: pyscf   Author: pyscf   File: kmp2.py    License: Apache License 2.0 6 votes vote down vote up
def padded_mo_energy(mp, mo_energy):
    """
    Pads energies of active MOs.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        mo_energy (ndarray): original non-padded molecular energies;

    Returns:
        Padded molecular energies.
    """
    frozen_mask = get_frozen_mask(mp)
    padding_convention = padding_k_idx(mp, kind="joint")
    nkpts = mp.nkpts

    result = np.zeros((nkpts, mp.nmo), dtype=mo_energy[0].dtype)
    for k in range(nkpts):
        result[np.ix_([k], padding_convention[k])] = mo_energy[k][frozen_mask[k]]

    return result 
Example 2
Project: pyscf   Author: pyscf   File: kmp2.py    License: Apache License 2.0 6 votes vote down vote up
def padded_mo_coeff(mp, mo_coeff):
    """
    Pads coefficients of active MOs.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        mo_coeff (ndarray): original non-padded molecular coefficients;

    Returns:
        Padded molecular coefficients.
    """
    frozen_mask = get_frozen_mask(mp)
    padding_convention = padding_k_idx(mp, kind="joint")
    nkpts = mp.nkpts

    result = np.zeros((nkpts, mo_coeff[0].shape[0], mp.nmo), dtype=mo_coeff[0].dtype)
    for k in range(nkpts):
        result[np.ix_([k], np.arange(result.shape[1]), padding_convention[k])] = mo_coeff[k][:, frozen_mask[k]]

    return result 
Example 3
Project: pyscf   Author: pyscf   File: kump2.py    License: Apache License 2.0 6 votes vote down vote up
def padded_mo_energy(mp, mo_energy):
    """
    Pads energies of active MOs.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        mo_energy (ndarray): original non-padded molecular energies;

    Returns:
        Padded molecular energies.
    """
    frozen_mask = get_frozen_mask(mp)
    padding_convention = padding_k_idx(mp, kind="joint")
    nkpts = mp.nkpts

    result = (np.zeros((nkpts, mp.nmo), dtype=mo_energy[0][0].dtype),
              np.zeros((nkpts, mp.nmo), dtype=mo_energy[0][0].dtype))
    for spin in [0, 1]:
        for k in range(nkpts):
            result[spin][np.ix_([k], padding_convention[k])] = mo_energy[spin][k][frozen_mask[k]]

    return result 
Example 4
Project: pyscf   Author: pyscf   File: kump2.py    License: Apache License 2.0 6 votes vote down vote up
def padded_mo_coeff(mp, mo_coeff):
    """
    Pads coefficients of active MOs.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        mo_coeff (ndarray): original non-padded molecular coefficients;

    Returns:
        Padded molecular coefficients.
    """
    frozen_mask = get_frozen_mask(mp)
    padding_convention = padding_k_idx(mp, kind="joint")
    nkpts = mp.nkpts

    result = (np.zeros((nkpts, mo_coeff[0][0].shape[0], mp.nmo[0]), dtype=mo_coeff[0][0].dtype),
              np.zeros((nkpts, mo_coeff[1][0].shape[0], mp.nmo[1]), dtype=mo_coeff[0][0].dtype))
    for spin in [0, 1]:
        for k in range(nkpts):
            result[spin][np.ix_([k], np.arange(result[spin].shape[1]), padding_convention[spin][k])] = \
                mo_coeff[spin][k][:, frozen_mask[spin][k]]

    return result 
Example 5
Project: pyscf   Author: pyscf   File: eom_kccsd_ghf.py    License: Apache License 2.0 6 votes vote down vote up
def mask_frozen_ip(eom, vector, kshift, const=LARGE_DENOM):
    '''Replaces all frozen orbital indices of `vector` with the value `const`.'''
    r1, r2 = eom.vector_to_amplitudes(vector, kshift=kshift)
    nkpts = eom.nkpts
    nocc, nmo = eom.nocc, eom.nmo
    kconserv = eom.kconserv

    # Get location of padded elements in occupied and virtual space
    nonzero_opadding, nonzero_vpadding = eom.nonzero_opadding, eom.nonzero_vpadding

    new_r1 = const * np.ones_like(r1)
    new_r2 = const * np.ones_like(r2)

    new_r1[nonzero_opadding[kshift]] = r1[nonzero_opadding[kshift]]
    for ki in range(nkpts):
        for kj in range(nkpts):
            kb = kconserv[ki, kshift, kj]
            idx = np.ix_([ki], [kj], nonzero_opadding[ki], nonzero_opadding[kj], nonzero_vpadding[kb])
            new_r2[idx] = r2[idx]

    return eom.amplitudes_to_vector(new_r1, new_r2, kshift, kconserv) 
Example 6
Project: pyscf   Author: pyscf   File: eom_kccsd_ghf.py    License: Apache License 2.0 6 votes vote down vote up
def mask_frozen_ea(eom, vector, kshift, const=LARGE_DENOM):
    '''Replaces all frozen orbital indices of `vector` with the value `const`.'''
    r1, r2 = eom.vector_to_amplitudes(vector, kshift=kshift)
    kconserv = eom.kconserv
    nkpts = eom.nkpts
    nocc, nmo = eom.nocc, eom.nmo

    # Get location of padded elements in occupied and virtual space
    nonzero_opadding, nonzero_vpadding = eom.nonzero_opadding, eom.nonzero_vpadding

    new_r1 = const * np.ones_like(r1)
    new_r2 = const * np.ones_like(r2)

    new_r1[nonzero_vpadding[kshift]] = r1[nonzero_vpadding[kshift]]
    for kj in range(nkpts):
        for ka in range(nkpts):
            kb = kconserv[kshift, ka, kj]
            idx = np.ix_([kj], [ka], nonzero_opadding[kj], nonzero_vpadding[ka], nonzero_vpadding[kb])
            new_r2[idx] = r2[idx]

    return eom.amplitudes_to_vector(new_r1, new_r2, kshift, kconserv) 
Example 7
Project: pyscf   Author: pyscf   File: test_krhf_slow_supercell.py    License: Apache License 2.0 6 votes vote down vote up
def test_eri(self):
        """Tests all ERI implementations: with and without symmetries."""
        for eri in (ktd.PhysERI, ktd.PhysERI4, ktd.PhysERI8):
            if not eri == ktd.PhysERI8 or self.test8:
                try:
                    e = eri(self.model_krhf)
                    m = e.tdhf_full_form()

                    # Test matrix vs ref
                    testing.assert_allclose(m, retrieve_m_hf(e), atol=1e-11)

                    # Test matrix vs pyscf
                    testing.assert_allclose(self.ref_m, m[numpy.ix_(self.ov_order, self.ov_order)], atol=1e-5)
                except Exception:
                    print("When testing {} the following exception occurred:".format(eri))
                    raise 
Example 8
Project: pyscf   Author: pyscf   File: test_proxy.py    License: Apache License 2.0 6 votes vote down vote up
def test_eri(self):
        """Tests all ERI implementations: with and without symmetries."""
        e = PhysERI(self.model_rks, "dft")
        mk, _ = e.tdhf_mk_form()
        testing.assert_allclose(self.ref_m, mk, atol=1e-13)

        # Test frozen
        for frozen in (1, [0, -1]):
            try:
                e = PhysERI(self.model_rks, "dft", frozen=frozen)
                mk, _ = e.tdhf_mk_form()
                ov_mask = tdhf_frozen_mask(e, kind="1ov")
                ref_m = self.ref_m[numpy.ix_(ov_mask, ov_mask)]
                testing.assert_allclose(ref_m, mk, atol=1e-13)

            except Exception:
                print("When testing with frozen={} the following exception occurred:".format(repr(frozen)))
                raise 
Example 9
Project: pyscf   Author: pyscf   File: test_proxy.py    License: Apache License 2.0 6 votes vote down vote up
def test_eri(self):
        """Tests all ERI implementations: with and without symmetries."""
        e = PhysERI(self.model_rhf, "hf")
        m = e.tdhf_full_form()
        testing.assert_allclose(self.ref_m, m, atol=1e-13)

        # Test frozen
        for proxy in (None, self.td_model_rhf):
            for frozen in (1, [0, -1]):
                try:
                    e = PhysERI(self.model_rhf, "dft", frozen=frozen)
                    m = e.tdhf_full_form()
                    ov_mask = tdhf_frozen_mask(e, kind="ov")
                    ref_m = self.ref_m[numpy.ix_(ov_mask, ov_mask)]
                    testing.assert_allclose(ref_m, m, atol=1e-13)

                except Exception:
                    print("When testing with proxy={} and frozen={} the following exception occurred:".format(
                        repr(proxy),
                        repr(frozen)
                    ))
                    raise 
Example 10
Project: lambda-packs   Author: ryfeus   File: factorization_ops_test_utils.py    License: MIT License 6 votes vote down vote up
def remove_empty_rows_columns(np_matrix):
  """Simple util to remove empty rows and columns of a matrix.

  Args:
    np_matrix: A numpy array.
  Returns:
    A tuple consisting of:
    mat: A numpy matrix obtained by removing empty rows and columns from
      np_matrix.
    nz_row_ids: A numpy array of the ids of non-empty rows, such that
      nz_row_ids[i] is the old row index corresponding to new index i.
    nz_col_ids: A numpy array of the ids of non-empty columns, such that
      nz_col_ids[j] is the old column index corresponding to new index j.
  """
  nz_row_ids = np.where(np.sum(np_matrix, axis=1) != 0)[0]
  nz_col_ids = np.where(np.sum(np_matrix, axis=0) != 0)[0]
  mat = np_matrix[np.ix_(nz_row_ids, nz_col_ids)]
  return mat, nz_row_ids, nz_col_ids 
Example 11
Project: lambda-packs   Author: ryfeus   File: test_regression.py    License: MIT License 6 votes vote down vote up
def test_large_fancy_indexing(self, level=rlevel):
        # Large enough to fail on 64-bit.
        nbits = np.dtype(np.intp).itemsize * 8
        thesize = int((2**nbits)**(1.0/5.0)+1)

        def dp():
            n = 3
            a = np.ones((n,)*5)
            i = np.random.randint(0, n, size=thesize)
            a[np.ix_(i, i, i, i, i)] = 0

        def dp2():
            n = 3
            a = np.ones((n,)*5)
            i = np.random.randint(0, n, size=thesize)
            a[np.ix_(i, i, i, i, i)]

        self.assertRaises(ValueError, dp)
        self.assertRaises(ValueError, dp2) 
Example 12
Project: SparseSC   Author: microsoft   File: sub_matrix_inverse.py    License: MIT License 6 votes vote down vote up
def subinv(x,eps=None):
    """ Given an matrix (x), calculate all the inverses of leave-one-out sub-matrices.
    
    :param x: a square matrix for which to find the inverses of all it's leave one out sub-matrices.
    :param eps: If not None, used to assert that the each calculated
           sub-matrix-inverse is within eps of the brute force calculation.
           Testing only, this slows the process way down since the inverse of
           each sub-matrix is calculated by the brute force method. Typically
           set to a multiple of `np.finfo(float).eps`
    """
    # handy constant for indexing
    xi = x.I
    N = x.shape[0]
    rng = np.arange(N)
    out = [None,] * N
    for k in range(N):
        k_rng = rng[rng != k]
        out[k] = xi[np.ix_(k_rng,k_rng)] - xi[k_rng,k].dot(xi[k,k_rng])/xi[k,k]
        if eps is not None:
            if not (abs(out[k] - x[np.ix_(k_rng,k_rng)].I) < eps).all():
                raise RuntimeError("Fast and brute force methods were not within epsilon (%s) for sub-matrix k = %s; max difference = %s" % 
                                   (eps, k,  abs(out[k] - x[np.ix_(k_rng,k_rng)].I).max(), ) )
    return out 
Example 13
Project: QCElemental   Author: MolSSI   File: align.py    License: BSD 3-Clause "New" or "Revised" License 5 votes vote down vote up
def align_hessian(self, hess) -> Array:
        blocked_hess = blockwise_expand(hess, (3, 3), False)
        alhess = np.zeros_like(blocked_hess)

        nat = blocked_hess.shape[0]
        for iat in range(nat):
            for jat in range(nat):
                alhess[iat, jat] = (self.rotation.T).dot(blocked_hess[iat, jat].dot(self.rotation))

        alhess = alhess[np.ix_(self.atommap, self.atommap)]

        alhess = blockwise_contract(alhess)
        return alhess 
Example 14
Project: pyscf   Author: pyscf   File: test_kgw_slow_supercell.py    License: Apache License 2.0 5 votes vote down vote up
def test_imds(self):

        testing.assert_allclose(
            tuple(self.gw.imds.get_rhs(i) for i in self.gw.imds.entire_space[0]),
            tuple(self.kgw.imds.get_rhs(i) for i in zip(self.order_k, self.order_p)),
            atol=1e-8
        )
        testing.assert_allclose(
            self.gw.imds.tdm,
            self.kgw.imds.tdm[numpy.ix_(range(len(self.kgw.imds.td_e)), range(2), self.order, self.order)],
            atol=1e-7
        ) 
Example 15
Project: pyscf   Author: pyscf   File: kccsd_rhf.py    License: Apache License 2.0 5 votes vote down vote up
def _get_epq(pindices,qindices,fac=[1.0,1.0],large_num=LARGE_DENOM):
    '''Create a denominator

        fac[0]*e[kp,p0:p1] + fac[1]*e[kq,q0:q1]

    where padded elements have been replaced by a large number.

    Args:
        pindices (5-list of object):
            A list of p0, p1, kp, orbital values, and non-zero indicess for the first
            denominator indices.
        qindices (5-list of object):
            A list of q0, q1, kq, orbital values, and non-zero indicess for the second
            denominator element.
        fac (3-list of float):
            Factors to multiply the first and second denominator elements.
        large_num (float):
            Number to replace the padded elements.
    '''
    def get_idx(x0,x1,kx,n0_p):
        return np.logical_and(n0_p[kx] >= x0, n0_p[kx] < x1)

    assert(all([len(x) == 5 for x in [pindices,qindices]]))
    p0,p1,kp,mo_e_p,nonzero_p = pindices
    q0,q1,kq,mo_e_q,nonzero_q = qindices
    fac_p, fac_q = fac

    epq = large_num * np.ones((p1-p0,q1-q0), dtype=mo_e_p[0].dtype)
    idxp = get_idx(p0,p1,kp,nonzero_p)
    idxq = get_idx(q0,q1,kq,nonzero_q)
    n0_ovp_pq = np.ix_(nonzero_p[kp][idxp], nonzero_q[kq][idxq])
    epq[n0_ovp_pq] = lib.direct_sum('p,q->pq', fac_p*mo_e_p[kp][p0:p1],
                                               fac_q*mo_e_q[kq][q0:q1])[n0_ovp_pq]
    return epq


# TODO: pull these 3 methods to pyscf.util and make tests 
Example 16
Project: pyscf   Author: pyscf   File: kccsd.py    License: Apache License 2.0 5 votes vote down vote up
def init_amps(self, eris):
        time0 = time.clock(), time.time()
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        mo_e_o = [eris.mo_energy[k][:nocc] for k in range(nkpts)]
        mo_e_v = [eris.mo_energy[k][nocc:] for k in range(nkpts)]
        t1 = numpy.zeros((nkpts, nocc, nvir), dtype=numpy.complex128)
        t2 = numpy.zeros((nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir), dtype=numpy.complex128)
        self.emp2 = 0
        eris_oovv = eris.oovv.copy()

        # Get location of padded elements in occupied and virtual space
        nonzero_opadding, nonzero_vpadding = padding_k_idx(self, kind="split")

        kconserv = kpts_helper.get_kconserv(self._scf.cell, self.kpts)
        for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
            kb = kconserv[ki, ka, kj]
            # For LARGE_DENOM, see t1new update above
            eia = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
            n0_ovp_ia = numpy.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
            eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]

            ejb = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
            n0_ovp_jb = numpy.ix_(nonzero_opadding[kj], nonzero_vpadding[kb])
            ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb]
            eijab = eia[:, None, :, None] + ejb[:, None, :]

            t2[ki, kj, ka] = eris_oovv[ki, kj, ka] / eijab

        t2 = numpy.conj(t2)
        self.emp2 = 0.25 * numpy.einsum('pqrijab,pqrijab', t2, eris_oovv).real
        self.emp2 /= nkpts

        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2.real)
        logger.timer(self, 'init mp2', *time0)
        return self.emp2, t1, t2 
Example 17
Project: pyscf   Author: pyscf   File: kccsd_t_rhf.py    License: Apache License 2.0 5 votes vote down vote up
def _get_epqr(pindices,qindices,rindices,fac=[1.0,1.0,1.0],large_num=LARGE_DENOM):
    '''Create a denominator

        fac[0]*e[kp,p0:p1] + fac[1]*e[kq,q0:q1] + fac[2]*e[kr,r0:r1]

    where padded elements have been replaced by a large number.

    Args:
        pindices (5-list of object):
            A list of p0, p1, kp, orbital values, and non-zero indices for the first
            denominator element.
        qindices (5-list of object):
            A list of q0, q1, kq, orbital values, and non-zero indices for the second
            denominator element.
        rindices (5-list of object):
            A list of r0, r1, kr, orbital values, and non-zero indices for the third
            denominator element.
        fac (3-list of float):
            Factors to multiply the first and second denominator elements.
        large_num (float):
            Number to replace the padded elements.
    '''
    def get_idx(x0,x1,kx,n0_p):
        return np.logical_and(n0_p[kx] >= x0, n0_p[kx] < x1)

    assert(all([len(x) == 5 for x in [pindices,qindices]]))
    p0,p1,kp,mo_e_p,nonzero_p = pindices
    q0,q1,kq,mo_e_q,nonzero_q = qindices
    r0,r1,kr,mo_e_r,nonzero_r = rindices
    fac_p, fac_q, fac_r = fac

    epqr = large_num * np.ones((p1-p0,q1-q0,r1-r0), dtype=mo_e_p[0].dtype)
    idxp = get_idx(p0,p1,kp,nonzero_p)
    idxq = get_idx(q0,q1,kq,nonzero_q)
    idxr = get_idx(r0,r1,kr,nonzero_r)
    n0_ovp_pqr = np.ix_(nonzero_p[kp][idxp]-p0, nonzero_q[kq][idxq]-q0, nonzero_r[kr][idxr]-r0)
    epqr[n0_ovp_pqr] = lib.direct_sum('p,q,r->pqr', fac_p*mo_e_p[kp][p0:p1],
                                      fac_q*mo_e_q[kq][q0:q1],
                                      fac_r*mo_e_r[kr][r0:r1])[n0_ovp_pqr]
    #epqr[n0_ovp_pqr] = temp[n0_ovp_pqr]
    return epqr 
Example 18
Project: pyscf   Author: pyscf   File: eom_kccsd_rhf_ip.py    License: Apache License 2.0 5 votes vote down vote up
def amplitudes_to_vector(cc_or_eom, t1, t2, kshift=0, kconserv=None):
    """IP amplitudes to vector."""
    itr = iter_12(cc_or_eom, kshift)
    t1, t2 = np.asarray(t1), np.asarray(t2)

    vc = VectorComposer(t1.dtype)
    vc.put(t1[np.ix_(*next(itr))])
    for slc in itr:
        vc.put(t2[np.ix_(*slc)])
    return vc.flush() 
Example 19
Project: pyscf   Author: pyscf   File: test_krhf_slow.py    License: Apache License 2.0 5 votes vote down vote up
def test_eri(self):
        """Tests all ERI implementations: with and without symmetries (supercell)."""
        o = v = 4
        for eri in (ktd.PhysERI, ktd.PhysERI4, ktd.PhysERI8):
            if eri is not ktd.PhysERI8 or self.test8:
                try:
                    e = eri(self.model_krhf)
                    if eri is ktd.PhysERI:
                        self.assertIn("__full_eri_k__", dir(e))
                    else:
                        self.assertNotIn("__full_eri_k__", dir(e))
                    self.assertNotIn("__full_eri__", dir(e))

                    s = (2 * self.k * self.k, 2 * self.k * self.k, o*v, o*v)
                    m = numpy.zeros(s, dtype=complex)

                    for k in range(self.k):
                        # Prepare indexes
                        r1, r2, c1, c2 = ktd.get_block_k_ix(e, k)

                        r = k2k(r1, r2)
                        c = k2k(c1, c2)

                        # Build matrix
                        _m = e.tdhf_full_form(k)

                        # Assign the submatrix
                        m[numpy.ix_(r, c)] = _m.reshape((2*self.k, o*v, 2*self.k, o*v)).transpose(0, 2, 1, 3)

                    m = m.transpose(0, 2, 1, 3).reshape(self.ref_m_supercell.shape)
                    testing.assert_allclose(self.ref_m_supercell, m, atol=1e-11)
                except Exception:
                    print("When testing {} the following exception occurred:".format(eri))
                    raise 
Example 20
Project: pyscf   Author: pyscf   File: addons.py    License: Apache License 2.0 5 votes vote down vote up
def partial_cholesky_orth_(S, cholthr=1e-9, canthr=1e-7):
    '''Partial Cholesky orthogonalization for curing overcompleteness.

    References:

    Susi Lehtola, Curing basis set overcompleteness with pivoted
    Cholesky decompositions, J. Chem. Phys. 151, 241102 (2019),
    doi:10.1063/1.5139948.

    Susi Lehtola, Accurate reproduction of strongly repulsive
    interatomic potentials, Phys. Rev. A 101, 032504 (2020),
    doi:10.1103/PhysRevA.101.032504.
    '''
    # Ensure the basis functions are normalized
    normlz = numpy.power(numpy.diag(S), -0.5)
    Snorm = numpy.dot(numpy.diag(normlz), numpy.dot(S, numpy.diag(normlz)))

    # Sort the basis functions according to the Gershgorin circle
    # theorem so that the Cholesky routine is well-initialized
    odS = numpy.abs(Snorm)
    numpy.fill_diagonal(odS, 0.0)
    odSs = numpy.sum(odS, axis=0)
    sortidx = numpy.argsort(odSs)

    # Run the pivoted Cholesky decomposition
    Ssort = Snorm[numpy.ix_(sortidx, sortidx)].copy()
    pstrf = scipy.linalg.lapack.get_lapack_funcs('pstrf')
    c, piv, r_c, info = pstrf(Ssort, tol=cholthr)
    # The functions we're going to use are given by the pivot as
    idx = sortidx[piv[:r_c]-1]

    # Get the (un-normalized) sub-basis
    Ssub = S[numpy.ix_(idx, idx)].copy()
    # Orthogonalize sub-basis
    Xsub = canonical_orth_(Ssub, thr=canthr)

    # Full X
    X = numpy.zeros((S.shape[0], Xsub.shape[1]))
    X[idx,:] = Xsub

    return X 
Example 21
Project: pyscf   Author: pyscf   File: test_rhf_slow.py    License: Apache License 2.0 5 votes vote down vote up
def test_eri(self):
        """Tests all ERI implementations: with and without symmetries."""
        for eri in (PhysERI, PhysERI4, PhysERI8):

            # Test plain
            try:
                e = eri(self.model_rhf)
                m = e.tdhf_full_form()

                # Test matrix vs ref
                testing.assert_allclose(m, retrieve_m_hf(e), atol=1e-14)

                # Test matrix vs pyscf
                testing.assert_allclose(self.ref_m, m, atol=1e-14)
                vals, vecs = eig(m, nroots=self.td_model_rhf.nroots)
                testing.assert_allclose(vals, self.td_model_rhf.e, atol=1e-5)

            except Exception:
                print("When testing {} the following exception occurred:".format(eri))
                raise

            # Test frozen
            for frozen in (1, [0, -1]):
                try:
                    e = eri(self.model_rhf, frozen=frozen)
                    m = e.tdhf_full_form()
                    ov_mask = tdhf_frozen_mask(e)
                    ref_m = self.ref_m[numpy.ix_(ov_mask, ov_mask)]
                    testing.assert_allclose(ref_m, m, atol=1e-14)

                except Exception:
                    print("When testing {} with frozen={} the following exception occurred:".format(eri, repr(frozen)))
                    raise 
Example 22
Project: simnibs   Author: simnibs   File: optimization_methods.py    License: GNU General Public License v3.0 5 votes vote down vote up
def _bb_upper_bound(l, Q, target_mean, max_total_current, max_el_current,
                    max_l0, state):
    n = l.shape[1]
    if len(state.active) > max_l0:
        return np.zeros(n), 1e20
    # we can't use the lower bound x because it is biased
    ac = state.active + state.unassigned
    x_ac = optimize_focality(l[:, ac], Q[np.ix_(ac, ac)],
                             target_mean, max_total_current,
                             max_el_current, log_level=10)
    x = np.zeros(n)
    x[ac] = x_ac
    # Select the "l0 - active" largest unassigned electrodes
    order_unasigned = np.argsort(-np.abs(x[state.unassigned]))
    selected_unasigned = \
        [state.unassigned[i] for i in order_unasigned[:max_l0-len(state.active)]]
    # Select the active electrodes plus the largest unassigned electrodes
    s = state.active + selected_unasigned
    x_ = optimize_focality(l[:, s], Q[np.ix_(s, s)],
                           target_mean, max_total_current,
                           max_el_current, log_level=10)
    x_ub = np.zeros(n)
    x_ub[s] = x_

    if np.any(np.abs(l.dot(x_ub) - target_mean) >= np.abs(1e-2 * target_mean)):
        return x_ub, 1e10

    return x_ub, x_ub.dot(Q).dot(x_ub) 
Example 23
Project: simnibs   Author: simnibs   File: test_mesh_io.py    License: GNU General Public License v3.0 5 votes vote down vote up
def test_advanced(self):
        array = np.arange(1, 101, dtype=int).reshape(-1, 5)
        rows = np.array([0, 5])
        cols = np.array([0, 3])
        g = mesh_io._GetitemTester(array)
        assert np.all(g[rows+1, cols] == array[rows, cols])
        assert np.all(g[np.ix_(rows+1, cols)] == array[np.ix_(rows, cols)]) 
Example 24
Project: pyqmc   Author: WagnerGroup   File: jastrowspin.py    License: MIT License 5 votes vote down vote up
def _update_b_partial(self, e, epos, mask):
        r"""
          Calculate b (e-e) partial sum contributions from electron e
        _b_partial_e is the array :math:`B^p_{ils} = \sum_s b_l(r^i_{es}`, with e fixed; :math:`s` indexes over :math:`\uparrow` (:math:`\alpha`) and :math:`\downarrow` (:math:`\beta`) sums, not including electron e. 
          Since :math:`B^p_{ils}` is summed over other electrons, moving electron e will affect other partial sums. This function updates all the necessary partial sums instead of just evaluating the one for electron e.
          :math:`i` is the configuration index.
          Args:
              e: fixed electron index
              epos: configs object for electron e
              mask: mask over configs axis, only return values for configs where mask==True. b_partial_e might have a smaller configs axis than epos, _configscurrent, and _b_partial because of the mask.
        """
        nup = self._mol.nelec[0]
        sep = nup - int(e < nup)
        not_e = np.arange(self._nelec) != e
        edown = int(e >= nup)
        d = epos.dist.dist_i(
            self._configscurrent.configs[mask][:, not_e], epos.configs[mask]
        )
        r = np.linalg.norm(d, axis=-1)
        dold = epos.dist.dist_i(
            self._configscurrent.configs[mask][:, not_e],
            self._configscurrent.configs[mask, e],
        )
        rold = np.linalg.norm(dold, axis=-1)
        b_partial_e = np.zeros((np.sum(mask), *self._b_partial.shape[2:]))
        eind, mind = np.ix_(not_e, mask)
        for l, b in enumerate(self.b_basis):
            bval = b.value(d, r)
            bdiff = bval - b.value(dold, rold)
            self._b_partial[eind, mind, l, edown] += bdiff.transpose((1, 0))
            self._b_partial[e, mask, l, 0] = bval[:, :sep].sum(axis=1)
            self._b_partial[e, mask, l, 1] = bval[:, sep:].sum(axis=1) 
Example 25
Project: recruit   Author: Frank-qlu   File: test_index_tricks.py    License: Apache License 2.0 5 votes vote down vote up
def test_regression_1(self):
        # Test empty inputs create outputs of indexing type, gh-5804
        # Test both lists and arrays
        for func in (range, np.arange):
            a, = np.ix_(func(0))
            assert_equal(a.dtype, np.intp) 
Example 26
Project: recruit   Author: Frank-qlu   File: test_index_tricks.py    License: Apache License 2.0 5 votes vote down vote up
def test_shape_and_dtype(self):
        sizes = (4, 5, 3, 2)
        # Test both lists and arrays
        for func in (range, np.arange):
            arrays = np.ix_(*[func(sz) for sz in sizes])
            for k, (a, sz) in enumerate(zip(arrays, sizes)):
                assert_equal(a.shape[k], sz)
                assert_(all(sh == 1 for j, sh in enumerate(a.shape) if j != k))
                assert_(np.issubdtype(a.dtype, np.integer)) 
Example 27
Project: recruit   Author: Frank-qlu   File: test_index_tricks.py    License: Apache License 2.0 5 votes vote down vote up
def test_bool(self):
        bool_a = [True, False, True, True]
        int_a, = np.nonzero(bool_a)
        assert_equal(np.ix_(bool_a)[0], int_a) 
Example 28
Project: recruit   Author: Frank-qlu   File: test_index_tricks.py    License: Apache License 2.0 5 votes vote down vote up
def test_1d_only(self):
        idx2d = [[1, 2, 3], [4, 5, 6]]
        assert_raises(ValueError, np.ix_, idx2d) 
Example 29
Project: recruit   Author: Frank-qlu   File: test_index_tricks.py    License: Apache License 2.0 5 votes vote down vote up
def test_repeated_input(self):
        length_of_vector = 5
        x = np.arange(length_of_vector)
        out = ix_(x, x)
        assert_equal(out[0].shape, (length_of_vector, 1))
        assert_equal(out[1].shape, (1, length_of_vector))
        # check that input shape is not modified
        assert_equal(x.shape, (length_of_vector,)) 
Example 30
Project: recruit   Author: Frank-qlu   File: indexing.py    License: Apache License 2.0 5 votes vote down vote up
def maybe_convert_ix(*args):
    """
    We likely want to take the cross-product
    """

    ixify = True
    for arg in args:
        if not isinstance(arg, (np.ndarray, list, ABCSeries, Index)):
            ixify = False

    if ixify:
        return np.ix_(*args)
    else:
        return args