def test_tocoo():
    syst = kwant.Builder()
    lat1 = kwant.lattice.chain(norbs=1)
    syst[lat1(0)] = syst[lat1(1)] = 0
    syst = syst.finalized()

    op = ops.Density(syst)
    assert isinstance(op.tocoo(), coo_matrix)

    # Constant and non-constant values.
    assert np.all(op.tocoo().todense() == np.eye(2))
    op = ops.Density(syst, lambda site: 1)
    assert np.all(op.tocoo().todense() == np.eye(2))

    # Correct treatment of where
    op = ops.Density(syst, where=[lat1(0)])
    assert np.all(op.tocoo().todense() == [[1, 0], [0, 0]])

    # No accidental transpose.
    syst = kwant.Builder()
    lat2 = kwant.lattice.chain(norbs=2)
    syst[lat2(0)] = lambda site, paramerer: np.eye(2)
    syst = syst.finalized()
    op = ops.Density(syst, [[1, 1], [0, 1]], check_hermiticity=False)
    assert np.all(op.tocoo().todense() == [[1, 1], [0, 1]])

    op = ops.Density(syst,
                     lambda site, p: [[1, 1], [0, 1]],
                     check_hermiticity=False)
    op = op.bind(args=(1, ))
    raises(ValueError, op.tocoo, [1])
def test_opservables_finite():
    lat, syst = _random_square_system(3)
    fsyst = syst.finalized()
    ev, wfs = la.eigh(fsyst.hamiltonian_submatrix())

    Q = ops.Density(fsyst)
    Qtot = ops.Density(fsyst, sum=True)
    J = ops.Current(fsyst)
    K = ops.Source(fsyst)

    for i, wf in enumerate(wfs.T):  # wfs[:, i] is i'th eigenvector
        assert np.allclose(Q.act(wf), wf)  # this operation is identity
        _test(Q, wf, reduced_val=1)  # eigenvectors are normalized
        _test(Qtot, wf, per_el_val=1)  # eigenvectors are normalized
        _test(J, wf, per_el_val=0)  # time-reversal symmetry: no current
        _test(K, wf, per_el_val=0)  # onsite commutes with hamiltonian

    # check that we get correct (complex) output
    for bra, ket in zip(wfs.T, wfs.T):
        _test(Q, bra, ket, per_el_val=(bra * ket))

    # check with get_hermiticity=False
    Qi = ops.Density(fsyst, 1j, check_hermiticity=False)
    for wf in wfs.T:  # wfs[:, i] is i'th eigenvector
        assert np.allclose(Qi.act(wf), 1j * wf)

    # test with different numbers of orbitals
    lat2 = kwant.lattice.chain(norbs=2)
    extra_sites = [lat2(i) for i in range(len(fsyst.sites))]
    syst[extra_sites] = np.eye(2)
    syst[zip(fsyst.sites, extra_sites)] = ta.matrix([1, 1])
    fsyst = syst.finalized()
    ev, wfs = la.eigh(fsyst.hamiltonian_submatrix())

    Q = ops.Density(fsyst)
    Qtot = ops.Density(fsyst, sum=True)
    J = ops.Current(fsyst)
    K = ops.Source(fsyst)

    for wf in wfs.T:  # wfs[:, i] is i'th eigenvector
        assert np.allclose(Q.act(wf), wf)  # this operation is identity
        _test(Q, wf, reduced_val=1)  # eigenvectors are normalized
        _test(Qtot, wf, per_el_val=1)  # eigenvectors are normalized
        _test(J, wf, per_el_val=0)  # time-reversal symmetry: no current
        _test(K, wf, per_el_val=0)  # onsite commutes with hamiltonian
def test_operator_construction():
    lat, syst = _random_square_system(3)
    fsyst = syst.finalized()
    N = len(fsyst.sites)

    # test construction failure if norbs not given
    latnone = kwant.lattice.chain()
    syst[latnone(0)] = 1
    for A in opservables:
        raises(ValueError, A, syst.finalized())
    del syst[latnone(0)]

    # test construction failure when dimensions of onsite do not match
    for A in opservables:
        raises(ValueError, A, fsyst, onsite=np.eye(2))

    # test that error is raised when input array is the wrong size
    for A in opservables:
        a = A(fsyst)
        kets = list(map(np.ones, [(0, ), (N - 1, ), (N + 1, ), (N, 1)]))
        for ket in kets:
            raises(ValueError, a, ket)
            raises(ValueError, a, ket, ket)
            raises(ValueError, a.act, ket)
            raises(ValueError, a.act, ket, ket)

    # Test failure on non-hermitian
    for A in (ops.Density, ops.Current, ops.Source):
        raises(ValueError, A, fsyst, 1j)

    # Test output dtype
    ket = np.ones(len(fsyst.sites))
    for A in (ops.Density, ops.Current, ops.Source):
        a = A(fsyst)
        a_nonherm = A(fsyst, check_hermiticity=False)
        assert a(ket, ket).dtype == np.complex128
        assert a(ket).dtype == np.float64
        assert a_nonherm(ket, ket).dtype == np.complex128
        assert a_nonherm(ket).dtype == np.complex128

    # test construction with different numbers of orbitals
    lat2 = kwant.lattice.chain(norbs=2)
    extra_sites = [lat2(i) for i in range(N)]
    syst[extra_sites] = np.eye(2)
    syst[zip(fsyst.sites, extra_sites)] = ta.matrix([1, 1])
    for A in opservables:
        raises(ValueError, A, syst.finalized(), onsite=np.eye(2))
        A(syst.finalized())
        A.onsite == np.eye(2)
    del syst[extra_sites]

    check = [(ops.Density, np.arange(N).reshape(-1, 1), ta.identity(1)),
             (ops.Current, np.array(list(fsyst.graph)), ta.identity(1)),
             (ops.Source, np.arange(N).reshape(-1, 1), ta.identity(1))]

    # test basic construction
    for A, where, onsite in check:
        a = A(fsyst)
        assert np.all(np.asarray(a.where) == where)
        assert all(a.onsite == onsite for i in range(N))

    # test construction with dict `onsite`
    for A in opservables:
        B = A(fsyst, {lat: 1})
        assert all(B.onsite(i) == 1 for i in range(N))

    # test construction with a functional onsite
    for A in opservables:
        B = A(fsyst, lambda site: site.pos[0])  # x-position operator
        assert all(B.onsite(i) == fsyst.sites[i].pos[0] for i in range(N))

    # test construction with `where` given by a sequence
    where = [lat(2, 2), lat(1, 1)]
    fwhere = tuple(fsyst.id_by_site[s] for s in where)
    A = ops.Density(fsyst, where=where)
    assert np.all(np.asarray(A.where).reshape(-1) == fwhere)

    where = [(lat(2, 2), lat(1, 2)), (lat(0, 0), lat(0, 1))]
    fwhere = np.asarray([(fsyst.id_by_site[a], fsyst.id_by_site[b])
                         for a, b in where])
    A = ops.Current(fsyst, where=where)
    assert np.all(np.asarray(A.where) == fwhere)

    # test construction with `where` given by a function
    tag_list = [(1, 0), (1, 1), (1, 2)]

    def where(site):
        return site.tag in tag_list

    A = ops.Density(fsyst, where=where)
    assert all(fsyst.sites[A.where[w, 0]].tag in tag_list
               for w in range(A.where.shape[0]))

    where_list = set(kwant.HoppingKind((1, 0), lat)(syst))
    fwhere_list = set(
        (fsyst.id_by_site[a], fsyst.id_by_site[b]) for a, b in where_list)

    def where(a, b):
        return (a, b) in where_list

    A = ops.Current(fsyst, where=where)
    assert all((a, b) in fwhere_list for a, b in A.where)

    # test that `sum` is passed to constructors correctly
    for A in opservables:
        A(fsyst, sum=True).sum == True