<commit_before>from __future__ import print_function, division

import pytest

import math as m
import numpy as np
import scipy as sc

from sisl import Atom
from sisl.geom import fcc, graphene
from sisl.sparse_geometry import *


@pytest.fixture
def setup():
    class t():
        def __init__(self):
            self.g = fcc(1., Atom(1, R=1.5)) * 2
            self.s1 = SparseAtom(self.g)
            self.s2 = SparseAtom(self.g, 2)
    return t()


@pytest.mark.sparse
@pytest.mark.sparse_geometry
class TestSparseAtom(object):

    @pytest.mark.xfail(raises=ValueError)
    def test_fail_align1(self, setup):
        s = SparseAtom(setup.g * 2)
        print(s)
        setup.s1.spalign(s)

    def test_align1(self, setup):
        s = SparseAtom(setup.g)
        setup.s1.spalign(s)
        setup.s1.spalign(s._csr)

    def test_create1(self, setup):
        setup.s1[0, [1, 2, 3]] = 1
        assert setup.s1.nnz == 3
        assert setup.s1[0, 1] == 1
        setup.s1[2, [1, 2, 3]] = 1
        assert setup.s1.nnz == 6
        setup.s1.empty(keep_nnz=True)
        assert setup.s1.nnz == 6
        setup.s1.empty()
        assert setup.s1.nnz == 0
        setup.s1[0, 0] = np.nan
        assert setup.s1.nnz == 0
        setup.s1.empty()

    def test_create2(self, setup):
        setup.s1[0, [1, 2, 3], (1, 1, 1)] = 1
        assert setup.s1[0, 2, (1, 1, 1)] == 1
        setup.s1.empty()

    def test_create3(self, setup):
        setup.s2[0, [1, 2, 3], (1, 1, 1)] = [1, 2]
        assert np.allclose(setup.s2[0, 2, (1, 1, 1)], [1, 2])
        assert setup.s2[0, 2, 0, (1, 1, 1)] == 1
        assert setup.s2[0, 2, 1, (1, 1, 1)] == 2
        setup.s2.empty()

    def test_in1(self, setup):
        setup.s1[0, [1, 2, 3]] = 1
        assert setup.s1[0, 2] == 1
        assert [0, 4] not in setup.s1
        assert [0, 2] in setup.s1
        setup.s1.empty()

    def test_reset1(self, setup):
        setup.s1[0, [1, 2, 3], (1, 1, 1)] = 1
        assert setup.s1[0, 2, (1, 1, 1)] == 1
        setup.s1.reset()
        assert setup.s1.nnz == 0
        setup.s1.empty()

    def test_eliminate_zeros1(self, setup):
        setup.s1[0, [1, 2, 3]] = 0
        assert setup.s1.nnz == 3
        setup.s1.eliminate_zeros()
        assert setup.s1.nnz == 0
        setup.s1.empty()

    def test_nonzero1(self, setup):
        s2 = setup.s2.copy()
        s2[0, [0, 2, 3]] = [1, 2]
        s2[3, [1, 2, 3]] = [1, 2]

        r, c = s2.nonzero()
        assert np.allclose(r, [0, 0, 0, 3, 3, 3])
        assert np.allclose(c, [0, 2, 3, 1, 2, 3])
        c = s2.nonzero(only_col=True)
        assert np.allclose(c, [0, 2, 3, 1, 2, 3])
        r, c = s2.nonzero(atom=1)
        assert len(r) == 0
        assert len(c) == 0
        r, c = s2.nonzero(atom=0)
        assert np.allclose(r, [0, 0, 0])
        assert np.allclose(c, [0, 2, 3])
        c = s2.nonzero(atom=0, only_col=True)
        assert np.allclose(c, [0, 2, 3])

    def test_cut1(self, setup):
        s1 = SparseAtom(setup.g)
        s1.construct([[0.1, 1.5], [1, 2]])
        s2 = SparseAtom(setup.g * 2)
        s2.construct([[0.1, 1.5], [1, 2]])
        s2 = s2.cut(2, 2).cut(2, 1).cut(2, 0)
        assert s1.spsame(s2)

        s1 = SparseAtom(setup.g)
        s1.construct([[0.1, 1.5], [1, 2]])
        s2 = SparseAtom(setup.g * [2, 1, 1])
        s2.construct([[0.1, 1.5], [1, 2]])
        s2 = s2.cut(2, 0)
        assert s1.spsame(s2)

        s1 = SparseAtom(setup.g)
        s1.construct([[0.1, 1.5], [1, 2]])
        s2 = SparseAtom(setup.g * [1, 2, 1])
        s2.construct([[0.1, 1.5], [1, 2]])
        s2 = s2.cut(2, 1)
        assert s1.spsame(s2)

    def test_iter(self, setup):
        s1 = SparseAtom(setup.g)
        s1.construct([[0.1, 1.5], [1, 2]])
        i = 0
        for r, c in s1:
            i += 1
        assert i == s1.nnz

    @pytest.mark.xfail(raises=ValueError)
    def test_rij_fail1(self, setup):
        s = SparseAtom(setup.g.copy())
        s.construct([[0.1, 1.5], [1, 2]])
        s.rij(what='none')

    @pytest.mark.xfail(raises=NotImplementedError)
    def test_rij_fail2(self, setup):
        s = SparseAtom(setup.g.copy())
        s.construct([[0.1, 1.5], [1, 2]])
        s.rij(what='orbital')

    def test_rij2(self, setup):
        s = SparseAtom(setup.g.copy())
        s.construct([[0.1, 1.5], [1, 2]])
        atom = s.rij()
        assert atom.spsame(s)

    def test_rij3(self, setup):
        sa = SparseAtom(setup.g.copy())
        so = SparseOrbital(setup.g.copy())
        sa.construct([[0.1, 1.5], [1, 2]])
        so.construct([[0.1, 1.5], [1, 2]])
        atom = sa.rij()
        orbital = so.rij()
        assert atom.spsame(orbital)
        atom = sa.rij('atom')
        orbital = so.rij('atom')
        assert atom.spsame(orbital)
        # This only works because there is 1 orbital per atom
        orbital = so.rij()
        assert so.spsame(orbital)
        so.finalize()
        orbital = so.rij()
        assert so.spsame(orbital)

    def test_sp_orb_remove(self, setup):
        so = SparseOrbital(setup.g.copy())
        so2 = so.remove(0)
        assert so.geometry.na - 1 == so2.geometry.na
        so2 = so.remove([0])
        assert so.geometry.na - 1 == so2.geometry.na
        so2 = so.remove([1])
        assert so.geometry.na - 1 == so2.geometry.na

    def test_remove1(self, setup):
        for i in range(len(setup.g)):
            setup.s1.construct([[0.1, 1.5], [1, 2]])
            s1 = setup.s1.remove(i)
            setup.s1.empty()
            s2 = SparseAtom(setup.g.remove(i))
            s2.construct([[0.1, 1.5], [1, 2]])
            assert s1.spsame(s2)

    def test_sub1(self, setup):
        all = range(len(setup.g))
        for i in range(len(setup.g)):
            setup.s1.construct([[0.1, 1.5], [1, 2]])
            # my new sub
            sub = [j for j in all if i != j]
            s1 = setup.s1.sub(sub)
            setup.s1.empty()
            s2 = SparseAtom(setup.g.sub(sub))
            s2.construct([[0.1, 1.5], [1, 2]])
            assert s1.spsame(s2)
            assert len(s2) == len(sub)

    def test_construct_eta(self, setup):
        s = setup.s1.copy()
        s.construct([[0.1, 1.5], [1, 2]], eta=True)<commit_msg>test: added new test for removing an atom and/or atom with orb-index<commit_after>from __future__ import print_function, division

import pytest

import math as m
import numpy as np
import scipy as sc

from sisl import Geometry, Atom
from sisl.geom import fcc, graphene
from sisl.sparse_geometry import *


@pytest.fixture
def setup():
    class t():
        def __init__(self):
            self.g = fcc(1., Atom(1, R=1.5)) * 2
            self.s1 = SparseAtom(self.g)
            self.s2 = SparseAtom(self.g, 2)
    return t()


@pytest.mark.sparse
@pytest.mark.sparse_geometry
class TestSparseAtom(object):

    @pytest.mark.xfail(raises=ValueError)
    def test_fail_align1(self, setup):
        s = SparseAtom(setup.g * 2)
        print(s)
        setup.s1.spalign(s)

    def test_align1(self, setup):
        s = SparseAtom(setup.g)
        setup.s1.spalign(s)
        setup.s1.spalign(s._csr)

    def test_create1(self, setup):
        setup.s1[0, [1, 2, 3]] = 1
        assert setup.s1.nnz == 3
        assert setup.s1[0, 1] == 1
        setup.s1[2, [1, 2, 3]] = 1
        assert setup.s1.nnz == 6
        setup.s1.empty(keep_nnz=True)
        assert setup.s1.nnz == 6
        setup.s1.empty()
        assert setup.s1.nnz == 0
        setup.s1[0, 0] = np.nan
        assert setup.s1.nnz == 0
        setup.s1.empty()

    def test_create2(self, setup):
        setup.s1[0, [1, 2, 3], (1, 1, 1)] = 1
        assert setup.s1[0, 2, (1, 1, 1)] == 1
        setup.s1.empty()

    def test_create3(self, setup):
        setup.s2[0, [1, 2, 3], (1, 1, 1)] = [1, 2]
        assert np.allclose(setup.s2[0, 2, (1, 1, 1)], [1, 2])
        assert setup.s2[0, 2, 0, (1, 1, 1)] == 1
        assert setup.s2[0, 2, 1, (1, 1, 1)] == 2
        setup.s2.empty()

    def test_in1(self, setup):
        setup.s1[0, [1, 2, 3]] = 1
        assert setup.s1[0, 2] == 1
        assert [0, 4] not in setup.s1
        assert [0, 2] in setup.s1
        setup.s1.empty()

    def test_reset1(self, setup):
        setup.s1[0, [1, 2, 3], (1, 1, 1)] = 1
        assert setup.s1[0, 2, (1, 1, 1)] == 1
        setup.s1.reset()
        assert setup.s1.nnz == 0
        setup.s1.empty()

    def test_eliminate_zeros1(self, setup):
        setup.s1[0, [1, 2, 3]] = 0
        assert setup.s1.nnz == 3
        setup.s1.eliminate_zeros()
        assert setup.s1.nnz == 0
        setup.s1.empty()

    def test_nonzero1(self, setup):
        s2 = setup.s2.copy()
        s2[0, [0, 2, 3]] = [1, 2]
        s2[3, [1, 2, 3]] = [1, 2]

        r, c = s2.nonzero()
        assert np.allclose(r, [0, 0, 0, 3, 3, 3])
        assert np.allclose(c, [0, 2, 3, 1, 2, 3])
        c = s2.nonzero(only_col=True)
        assert np.allclose(c, [0, 2, 3, 1, 2, 3])
        r, c = s2.nonzero(atom=1)
        assert len(r) == 0
        assert len(c) == 0
        r, c = s2.nonzero(atom=0)
        assert np.allclose(r, [0, 0, 0])
        assert np.allclose(c, [0, 2, 3])
        c = s2.nonzero(atom=0, only_col=True)
        assert np.allclose(c, [0, 2, 3])

    def test_cut1(self, setup):
        s1 = SparseAtom(setup.g)
        s1.construct([[0.1, 1.5], [1, 2]])
        s2 = SparseAtom(setup.g * 2)
        s2.construct([[0.1, 1.5], [1, 2]])
        s2 = s2.cut(2, 2).cut(2, 1).cut(2, 0)
        assert s1.spsame(s2)

        s1 = SparseAtom(setup.g)
        s1.construct([[0.1, 1.5], [1, 2]])
        s2 = SparseAtom(setup.g * [2, 1, 1])
        s2.construct([[0.1, 1.5], [1, 2]])
        s2 = s2.cut(2, 0)
        assert s1.spsame(s2)

        s1 = SparseAtom(setup.g)
        s1.construct([[0.1, 1.5], [1, 2]])
        s2 = SparseAtom(setup.g * [1, 2, 1])
        s2.construct([[0.1, 1.5], [1, 2]])
        s2 = s2.cut(2, 1)
        assert s1.spsame(s2)

    def test_iter(self, setup):
        s1 = SparseAtom(setup.g)
        s1.construct([[0.1, 1.5], [1, 2]])
        i = 0
        for r, c in s1:
            i += 1
        assert i == s1.nnz

    @pytest.mark.xfail(raises=ValueError)
    def test_rij_fail1(self, setup):
        s = SparseAtom(setup.g.copy())
        s.construct([[0.1, 1.5], [1, 2]])
        s.rij(what='none')

    @pytest.mark.xfail(raises=NotImplementedError)
    def test_rij_fail2(self, setup):
        s = SparseAtom(setup.g.copy())
        s.construct([[0.1, 1.5], [1, 2]])
        s.rij(what='orbital')

    def test_rij2(self, setup):
        s = SparseAtom(setup.g.copy())
        s.construct([[0.1, 1.5], [1, 2]])
        atom = s.rij()
        assert atom.spsame(s)

    def test_rij3(self, setup):
        sa = SparseAtom(setup.g.copy())
        so = SparseOrbital(setup.g.copy())
        sa.construct([[0.1, 1.5], [1, 2]])
        so.construct([[0.1, 1.5], [1, 2]])
        atom = sa.rij()
        orbital = so.rij()
        assert atom.spsame(orbital)
        atom = sa.rij('atom')
        orbital = so.rij('atom')
        assert atom.spsame(orbital)
        # This only works because there is 1 orbital per atom
        orbital = so.rij()
        assert so.spsame(orbital)
        so.finalize()
        orbital = so.rij()
        assert so.spsame(orbital)

    def test_sp_orb_remove(self, setup):
        so = SparseOrbital(setup.g.copy())
        so2 = so.remove(0)
        assert so.geometry.na - 1 == so2.geometry.na
        so2 = so.remove([0])
        assert so.geometry.na - 1 == so2.geometry.na
        so2 = so.remove([1])
        assert so.geometry.na - 1 == so2.geometry.na

    def test_sp_orb_remove_atom(self):
        so = SparseOrbital(Geometry([[0] *3, [1]* 3], [Atom[1], Atom[2]], 2))
        so2 = so.remove(Atom[1])
        assert so.geometry.na - 1 == so2.geometry.na
        so = SparseOrbital(Geometry([[0] *3, [1]* 3], [Atom(1, [1, 2]), Atom(2, [1, 2])], 2))
        so2 = so.remove(so.geometry.atoms[0], [0])
        assert so.geometry.na == so2.geometry.na
        assert so.geometry.no -1 == so2.geometry.no

    def test_remove1(self, setup):
        for i in range(len(setup.g)):
            setup.s1.construct([[0.1, 1.5], [1, 2]])
            s1 = setup.s1.remove(i)
            setup.s1.empty()
            s2 = SparseAtom(setup.g.remove(i))
            s2.construct([[0.1, 1.5], [1, 2]])
            assert s1.spsame(s2)

    def test_sub1(self, setup):
        all = range(len(setup.g))
        for i in range(len(setup.g)):
            setup.s1.construct([[0.1, 1.5], [1, 2]])
            # my new sub
            sub = [j for j in all if i != j]
            s1 = setup.s1.sub(sub)
            setup.s1.empty()
            s2 = SparseAtom(setup.g.sub(sub))
            s2.construct([[0.1, 1.5], [1, 2]])
            assert s1.spsame(s2)
            assert len(s2) == len(sub)

    def test_construct_eta(self, setup):
        s = setup.s1.copy()
        s.construct([[0.1, 1.5], [1, 2]], eta=True)
{'fraction_non_alphanumeric': 0.17288528389339514, 'fraction_numerical': 0.06303592120509849, 'mean_word_length': 1.870509977827051, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '228555', 'n_tokens_mistral': 5546, 'n_tokens_neox': 5227, 'n_words': 1173}