<commit_before>"""Tests for the general many-body problem."""

import pytest
from sympy import IndexedBase

from drudge import GenMBDrudge, CR, AN


@pytest.fixture(scope='module')
def genmb(spark_ctx):
    """Initialize the environment for a free algebra."""

    dr = GenMBDrudge(spark_ctx)
    return dr


def test_genmb_has_basic_properties(genmb):
    """Test the general many-body model has basic properties."""
    dr = genmb

    assert len(dr.orb_ranges) == 1
    assert len(dr.spin_vals) == 0

    assert dr.one_body == dr.names.t == IndexedBase('t')
    assert dr.two_body == dr.names.u == IndexedBase('u')

    # The Hamiltonian should already be simplified for this simple model.
    assert dr.ham.n_terms == 2
    assert dr.ham == dr.orig_ham
    # The details of the Hamiltonian will be tested in other ways.


@pytest.mark.parametrize('par_level', [0, 1, 2])
@pytest.mark.parametrize('full_simplify', [True, False])
@pytest.mark.parametrize('simple_merge', [True, False])
def test_genmb_simplify_simple_expressions(
        genmb, par_level, full_simplify, simple_merge
):
    """Test the basic Wick expansion facility on a single Fermion expression."""

    dr = genmb  # type: GenMBDrudge

    c_ = dr.op[AN]
    c_dag = dr.op[CR]
    r = dr.names.L
    a, b, c, d = dr.names.L_dumms[:4]

    t = IndexedBase('t')
    u = IndexedBase('u')

    inp = dr.sum(
        (a, r), (b, r), (c, r), (d, r),
        t[a, b] * u[c, d] * c_dag[a] * c_[b] * c_dag[c] * c_[d]
    )

    dr.wick_parallel = par_level
    assert dr.wick_parallel == par_level
    dr.full_simplify = full_simplify
    assert dr.full_simplify == full_simplify
    dr.simple_merge = simple_merge
    assert dr.simple_merge == simple_merge

    res = inp.simplify()

    dr.wick_parallel = 0
    assert dr.wick_parallel == 0
    dr.full_simplify = True
    assert dr.full_simplify
    dr.simple_merge = False
    assert not dr.simple_merge

    assert res.n_terms == 2

    expected = dr.einst(
        t[a, c] * u[b, d] * c_dag[a] * c_dag[b] * c_[d] * c_[c] +
        t[a, c] * u[c, b] * c_dag[a] * c_[b]
    ).simplify()

    assert res == expected


def test_genmb_derives_spin_orbit_hartree_fock(genmb):
    """Test general many-body model can derive HF theory in spin-orbital basis.
    """

    dr = genmb
    p = genmb.names
    c_ = p.c_
    c_dag = p.c_dag
    r = p.L
    a, b, c, d = p.L_dumms[:4]

    rot = c_dag[a] * c_[b]
    comm = (dr.ham | rot).simplify()
    assert comm.n_terms == 4

    rho = IndexedBase('rho')
    # Following Ring and Schuck, here all creation comes before the
    # annihilation.
    res = comm.eval_vev(lambda op1, op2, _: (
        rho[op2.indices[1], op1.indices[1]]
        if op1.indices[0] == CR and op2.indices[0] == AN
        else 0
    )).simplify()
    assert res.n_terms == 2

    # The correct result: [\rho, f]^b_a

    f = IndexedBase('f')
    expected = dr.sum((c, r), rho[b, c] * f[c, a] - f[b, c] * rho[c, a])
    expected = expected.subst(f[a, b], p.t[a, b] + dr.sum(
        (c, r), (d, r),
        p.u[a, c, b, d] * rho[d, c] - p.u[a, c, d, b] * rho[d, c]
    ))
    expected = expected.simplify()

    assert res == expected<commit_msg>Add test for LaTeX form of field operators<commit_after>"""Tests for the general many-body problem."""

import pytest
from sympy import IndexedBase

from drudge import GenMBDrudge, CR, AN


@pytest.fixture(scope='module')
def genmb(spark_ctx):
    """Initialize the environment for a free algebra."""

    dr = GenMBDrudge(spark_ctx)
    return dr


def test_genmb_has_basic_properties(genmb):
    """Test the general many-body model has basic properties."""
    dr = genmb

    assert len(dr.orb_ranges) == 1
    assert len(dr.spin_vals) == 0

    assert dr.one_body == dr.names.t == IndexedBase('t')
    assert dr.two_body == dr.names.u == IndexedBase('u')

    # The Hamiltonian should already be simplified for this simple model.
    assert dr.ham.n_terms == 2
    assert dr.ham == dr.orig_ham
    # The details of the Hamiltonian will be tested in other ways.


@pytest.mark.parametrize('par_level', [0, 1, 2])
@pytest.mark.parametrize('full_simplify', [True, False])
@pytest.mark.parametrize('simple_merge', [True, False])
def test_genmb_simplify_simple_expressions(
        genmb, par_level, full_simplify, simple_merge
):
    """Test the basic Wick expansion facility on a single Fermion expression."""

    dr = genmb  # type: GenMBDrudge

    c_ = dr.op[AN]
    c_dag = dr.op[CR]
    r = dr.names.L
    a, b, c, d = dr.names.L_dumms[:4]

    t = IndexedBase('t')
    u = IndexedBase('u')

    inp = dr.sum(
        (a, r), (b, r), (c, r), (d, r),
        t[a, b] * u[c, d] * c_dag[a] * c_[b] * c_dag[c] * c_[d]
    )

    dr.wick_parallel = par_level
    assert dr.wick_parallel == par_level
    dr.full_simplify = full_simplify
    assert dr.full_simplify == full_simplify
    dr.simple_merge = simple_merge
    assert dr.simple_merge == simple_merge

    res = inp.simplify()

    dr.wick_parallel = 0
    assert dr.wick_parallel == 0
    dr.full_simplify = True
    assert dr.full_simplify
    dr.simple_merge = False
    assert not dr.simple_merge

    assert res.n_terms == 2

    expected = dr.einst(
        t[a, c] * u[b, d] * c_dag[a] * c_dag[b] * c_[d] * c_[c] +
        t[a, c] * u[c, b] * c_dag[a] * c_[b]
    ).simplify()

    assert res == expected


def test_genmb_derives_spin_orbit_hartree_fock(genmb):
    """Test general many-body model can derive HF theory in spin-orbital basis.
    """

    dr = genmb
    p = genmb.names
    c_ = p.c_
    c_dag = p.c_dag
    r = p.L
    a, b, c, d = p.L_dumms[:4]

    rot = c_dag[a] * c_[b]
    comm = (dr.ham | rot).simplify()
    assert comm.n_terms == 4

    rho = IndexedBase('rho')
    # Following Ring and Schuck, here all creation comes before the
    # annihilation.
    res = comm.eval_vev(lambda op1, op2, _: (
        rho[op2.indices[1], op1.indices[1]]
        if op1.indices[0] == CR and op2.indices[0] == AN
        else 0
    )).simplify()
    assert res.n_terms == 2

    # The correct result: [\rho, f]^b_a

    f = IndexedBase('f')
    expected = dr.sum((c, r), rho[b, c] * f[c, a] - f[b, c] * rho[c, a])
    expected = expected.subst(f[a, b], p.t[a, b] + dr.sum(
        (c, r), (d, r),
        p.u[a, c, b, d] * rho[d, c] - p.u[a, c, d, b] * rho[d, c]
    ))
    expected = expected.simplify()

    assert res == expected


def test_fock_drudge_prints_operators(genmb):
    """Test the LaTeX printing by Fock drudge.

    Things like term linkage should be tested for the base class.  Here we
    concentrate on the vector part, which is turned for field operators.
    """

    dr = genmb
    p = dr.names

    x = IndexedBase('x')
    a, b = p.L_dumms[:2]

    tensor = dr.einst(- x[a, b] * p.c_dag[a] * p.c_[b])
    assert tensor.latex() == (
        r'- \sum_{a \in L} \sum_{b \in L} x_{a,b} c^{\dagger}_{a} c^{}_{b}'
    )
{'fraction_non_alphanumeric': 0.14372733865119652, 'fraction_numerical': 0.007106598984771574, 'mean_word_length': 2.7316017316017316, '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': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3785024', 'n_tokens_mistral': 2738, 'n_tokens_neox': 2604, 'n_words': 698}