Python numpy.ix_() Examples

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) 

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 

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 

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 

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 

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 

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 

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) 

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) 

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 

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 

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 

def test_regression_1(self):
        # Test empty inputs create ouputs 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) 

def subinv_k(xi,k,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 k: the column and row to leave out
    :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
    N = xi.shape[0]
    rng = np.arange(N)
    k_rng = rng[rng != k]
    out = 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



# ---------------------------------------------
# single sub-matrix
# --------------------------------------------- 

def single_grad(
    N0, N1, in_controls, splits, b_i, w_pen, treated_units, Y_treated, Y_control
):
    """
    wrapper for the real function 
    """

    in_controls2 = [np.ix_(i, i) for i in in_controls]

    def inner(A, weights, dA_dV_ki_k, dB_dV_ki_k):
        """
        Calculate a single component of the gradient
        """
        dPI_dV = np.zeros((N0, N1))  # stupid notation: PI = W.T
        for i, (_, (_, test)) in enumerate(zip(in_controls, splits)):
            dA = dA_dV_ki_k[i]
            dB = dB_dV_ki_k[i]
            try:
                b = np.linalg.solve(A[in_controls2[i]], dB - dA.dot(b_i[i]))
            except np.linalg.LinAlgError as exc:
                print("Unique weights not possible.")
                if w_pen == 0:
                    print("Try specifying a very small w_pen rather than 0.")
                raise exc
            dPI_dV[np.ix_(in_controls[i], treated_units[test])] = b
        # einsum is faster than the equivalent (Ey * Y_control.T.dot(dPI_dV).T.getA()).sum()
        return 2 * np.einsum(
            "ij,kj,ki->", (weights.T.dot(Y_control) - Y_treated), Y_control, dPI_dV
        )

    return inner 

def get_weights(self, include_trivial_donors=True):
        """
        getter for the sc_weights. By default, the trivial
        """
        if include_trivial_donors or not self.trivial_units.any():
            return self._sc_weights

        if self.model_type != "full":
            trivial_donors = self.trivial_units[self.control_units]
        else:
            trivial_donors = self.trivial_units

        __weights = self._sc_weights.copy()
        __weights[np.ix_(np.logical_not(self.trivial_units), trivial_donors)] = 0
        return __weights 

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, int)) 

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)) 

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) 

def statcon(K,f,cd):
    """
    Condensation of static FE-equations according to the vector cd.

    Parameters:
    
        K                       global stiffness matrix, dim(K) = nd x nd
        f                       global load vector, dim(f)= nd x 1

        cd                      vector containing dof's to be eliminated
                                dim(cd)= nc x 1, nc: number of condensed dof's
    Returns:
    
        K1                      condensed stiffness matrix,
                                dim(K1)= (nd-nc) x (nd-nc)
        f1                      condensed load vector, dim(f1)= (nd-nc) x 1
    """
    nd,nd = np.shape(K)
    cd = (cd-1).flatten()
  
    aindx = np.arange(nd)
    aindx = np.delete(aindx,cd,0)
    bindx = cd

    Kaa = np.mat(K[np.ix_(aindx,aindx)])
    Kab = np.mat(K[np.ix_(aindx,bindx)])
    Kbb = np.mat(K[np.ix_(bindx,bindx)])

    fa = np.mat(f[aindx])
    fb = np.mat(f[bindx])
    
    K1 = Kaa-Kab*Kbb.I*Kab.T
    f1 = fa-Kab*Kbb.I*fb
    
    return K1,f1 

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,)) 

def test_1d_only(self):
        idx2d = [[1, 2, 3], [4, 5, 6]]
        assert_raises(ValueError, np.ix_, idx2d) 

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) 

def test_regression_1(self):
        # Test empty inputs create ouputs 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) 

def get_args(argspec, n):
    if isinstance(argspec, np.ndarray):
        args = argspec.copy()
    else:
        nargs = len(argspec)
        ms = np.asarray([1.5 if isinstance(spec, ComplexArg) else 1.0 for spec in argspec])
        ms = (n**(ms/sum(ms))).astype(int) + 1

        args = []
        for spec, m in zip(argspec, ms):
            args.append(spec.values(m))
        args = np.array(np.broadcast_arrays(*np.ix_(*args))).reshape(nargs, -1).T

    return args 

def _get_kept_col_data(self) -> Dict[str, List[ndarray]]:
    data_dict: Dict[str, List[ndarray]] = defaultdict(list)
    for dtype, locs in self._group_dtype_loc.items():
        ix = np.ix_(self._group_position, locs)
        arr = self._df._data[dtype][ix]
        if arr.ndim == 1:
            arr = arr[:, np.newaxis]
        data_dict[dtype].append(arr)
    return data_dict 

def _get_group_col_data(self) -> Dict[str, List[ndarray]]:
        data_dict: Dict[str, List[ndarray]] = defaultdict(list)
        for dtype, locs in self._group_dtype_loc.items():
            ix = np.ix_(self._group_position, locs)
            arr = self._df._data[dtype][ix]
            if arr.ndim == 1:
                arr = arr[:, np.newaxis]
            data_dict[dtype].append(arr)
        return data_dict 

def test_site():
    chinfo = npc.ChargeInfo([1, 3])
    leg = gen_random_legcharge(chinfo, 8)
    op1 = npc.Array.from_func(np.random.random, [leg, leg.conj()], shape_kw='size')
    op2 = npc.Array.from_func(np.random.random, [leg, leg.conj()], shape_kw='size')
    labels = ['up'] + [None] * (leg.ind_len - 2) + ['down']
    s = site.Site(leg, labels, silly_op=op1)
    assert s.state_index('up') == 0
    assert s.state_index('down') == leg.ind_len - 1
    assert s.opnames == set(['silly_op', 'Id', 'JW'])
    assert s.silly_op is op1
    s.add_op('op2', op2)
    assert s.op2 is op2
    assert s.get_op('op2') is op2
    assert s.get_op('silly_op') is op1
    npt.assert_equal(
        s.get_op('silly_op op2').to_ndarray(),
        npc.tensordot(op1, op2, [1, 0]).to_ndarray())
    leg2 = npc.LegCharge.from_drop_charge(leg, 1)
    leg2 = npc.LegCharge.from_change_charge(leg2, 0, 2, 'changed')
    s2 = copy.deepcopy(s)
    s2.change_charge(leg2)
    perm_qind, leg2s = leg2.sort()
    perm_flat = leg2.perm_flat_from_perm_qind(perm_qind)
    s2s = copy.deepcopy(s2)
    s2s.change_charge(leg2s, perm_flat)
    for site_check in [s2, s2s]:
        print("site_check.leg = ", site_check.leg)
        for opn in site_check.opnames:
            op1 = s.get_op(opn).to_ndarray()
            op2 = site_check.get_op(opn).to_ndarray()
            perm = site_check.perm
            npt.assert_equal(op1[np.ix_(perm, perm)], op2)
    # done 

def get_site_op_flat(site, op):
    """Like ``site.get_op(op)``, but return a flat numpy array and revert permutation from charges.

    site.perm should store the permutation compared to "conserve=None", so we can use that to
    convert to the "standard" flat form with conserve=None.
    """
    op = site.get_op(op).to_ndarray()
    iperm = inverse_permutation(site.perm)
    return op[np.ix_(iperm, iperm)] 

def test_npc_Array_project():
    a = npc.Array.from_ndarray(arr, [lc, lc.conj()])
    p1 = np.array([True, True, False, True, True])
    p2 = np.array([0, 1, 3])

    b = a.copy(True)
    b.iproject([p1, p2], (0, 1))
    b.test_sanity()
    bflat = a.to_ndarray()[np.ix_(p1, p2)]
    npt.assert_equal(b.to_ndarray(), bflat)
    # and again for a being blocked before: can we split the blocks
    print("for blocked")
    _, a = a.sort_legcharge()
    b = a.copy(True)
    b.iproject([p1, p2], (0, 1))
    b.test_sanity()
    bflat = a.to_ndarray()[np.ix_(p1, p2)]
    npt.assert_equal(b.to_ndarray(), bflat)

    print("for trivial charge")
    a = npc.Array.from_func(np.random.random, [lcTr, lcTr.conj()], shape_kw='size')
    p1 = (np.arange(lcTr.ind_len) % 3 == 0)
    b = a.copy(True)
    b.iproject([p1, p2], (0, 1))
    b.test_sanity()
    bflat = a.to_ndarray()[np.ix_(p1, p2)]
    npt.assert_equal(b.to_ndarray(), bflat)