grogu/tests/test_utilities.py

# Copyright (c) [2024] []
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np
import pytest
from hypothesis import given
from hypothesis import strategies as st
from numpy.testing import assert_allclose, assert_array_almost_equal

from grogupy.constants import TAU_0, TAU_X, TAU_Y, TAU_Z
from grogupy.utilities import (
    RotM,
    RotMa2b,
    crossM,
    hsk,
    int_de_ke,
    make_contour,
    make_kset,
    read_siesta_emin,
    tau_u,
)


# Test basic Pauli matrix properties
def test_pauli_matrices_properties():
    """Test fundamental properties of Pauli matrices"""
    # Test anticommutation relations
    assert_array_almost_equal(TAU_X @ TAU_Y + TAU_Y @ TAU_X, np.zeros((2, 2)))
    assert_array_almost_equal(TAU_Y @ TAU_Z + TAU_Z @ TAU_Y, np.zeros((2, 2)))
    assert_array_almost_equal(TAU_Z @ TAU_X + TAU_X @ TAU_Z, np.zeros((2, 2)))

    # Test square of Pauli matrices equals identity
    assert_array_almost_equal(TAU_X @ TAU_X, TAU_0)
    assert_array_almost_equal(TAU_Y @ TAU_Y, TAU_0)
    assert_array_almost_equal(TAU_Z @ TAU_Z, TAU_0)

    # Test Hermiticity
    assert_array_almost_equal(TAU_X, TAU_X.conj().T)
    assert_array_almost_equal(TAU_Y, TAU_Y.conj().T)
    assert_array_almost_equal(TAU_Z, TAU_Z.conj().T)


# Test hsk function
def test_hsk_basics():
    """Test basic properties of the hsk function"""
    size = 4
    H = np.random.random((3, size, size)) + 1j * np.random.random((3, size, size))
    ss = np.random.random((3, size, size)) + 1j * np.random.random((3, size, size))
    sc_off = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0]])
    k = np.array([0.1, 0.2, 0.3])

    HK, SK = hsk(H, ss, sc_off, k)

    # Check shapes
    assert HK.shape == (size, size)
    assert SK.shape == (size, size)

    # Test k=0 case gives the central cell contribution
    HK0, SK0 = hsk(H, ss, sc_off, np.zeros(3))
    assert_array_almost_equal(HK0, H[0])
    assert_array_almost_equal(SK0, ss[0])


@given(st.integers(min_value=1, max_value=20))
def test_make_kset_dimensions(num_k):
    """Test k-point set generation with different dimensions"""
    # Test 1D
    kset_x = make_kset("x", num_k)
    assert kset_x.shape[1] == 3
    assert len(np.unique(kset_x[:, 1])) == 1  # y coordinates should be same
    assert len(np.unique(kset_x[:, 2])) == 1  # z coordinates should be same

    # Test 2D
    kset_xy = make_kset("xy", num_k)
    assert kset_xy.shape[1] == 3
    assert len(np.unique(kset_xy[:, 2])) == 1  # z coordinates should be same

    # Test 3D
    kset_xyz = make_kset("xyz", num_k)
    assert kset_xyz.shape[1] == 3
    assert len(kset_xyz) == num_k**3


def test_make_contour():
    """Test contour generation for energy integration"""
    emin = -20
    emax = 0.0
    enum = 42
    p = 150

    cont = make_contour(emin, emax, enum, p)

    # Check attributes
    assert hasattr(cont, "R")
    assert hasattr(cont, "z0")
    assert hasattr(cont, "ze")
    assert hasattr(cont, "we")
    assert hasattr(cont, "enum")

    # Check dimensions
    assert len(cont.ze) == enum
    assert len(cont.we) == enum

    # Check basic properties
    assert cont.z0 == (emax + emin) / 2
    assert cont.R == (emax - emin) / 2
    assert cont.enum == enum


@given(st.lists(st.floats(min_value=-1, max_value=1), min_size=3, max_size=3))
def test_tau_u(u):
    """Test generation of Pauli matrix in arbitrary direction"""
    u = np.array(u)
    if np.all(u == 0):
        u = np.array([1, 0, 0])  # Handle zero vector case
    u = u / np.linalg.norm(u)

    result = tau_u(u)

    # Check result is 2x2
    assert result.shape == (2, 2)

    # Check Hermiticity
    assert_array_almost_equal(result, result.conj().T)

    # Check trace is zero
    assert abs(np.trace(result)) < 1e-10

    # Check square equals identity
    assert_array_almost_equal(result @ result, np.eye(2))


def test_crossM():
    """Test cross product matrix generation"""
    u = np.array([1.0, 2.0, 3.0])
    v = np.array([4.0, 5.0, 6.0])

    # Test that crossM(u) @ v equals np.cross(u, v)
    assert_array_almost_equal(crossM(u) @ v, np.cross(u, v))

    # Test antisymmetry
    assert_array_almost_equal(crossM(u), -crossM(u).T)


@pytest.mark.parametrize("theta", [0, np.pi / 4, np.pi / 2, np.pi])
def test_RotM(theta):
    """Test rotation matrix generation"""
    # Test rotation around z-axis
    u = np.array([0, 0, 1])
    R = RotM(theta, u)

    # Check orthogonality
    assert_array_almost_equal(R @ R.T, np.eye(3))

    # Check determinant is 1
    assert_allclose(np.linalg.det(R), 1.0)

    # Test rotation of vector
    v = np.array([1, 0, 0])
    rotated_v = R @ v
    expected_v = np.array([np.cos(theta), np.sin(theta), 0])
    assert_array_almost_equal(rotated_v, expected_v)


def test_RotMa2b():
    """Test rotation matrix between two vectors"""
    # Test simple cases
    a = np.array([1, 0, 0])
    b = np.array([0, 1, 0])
    R = RotMa2b(a, b)

    # Check that R rotates a to b
    assert_array_almost_equal(R @ a, b)

    # Check orthogonality
    assert_array_almost_equal(R @ R.T, np.eye(3))

    # Check determinant is 1
    assert_allclose(np.linalg.det(R), 1.0)


def test_read_siesta_emin(tmpdir):
    """Test reading minimum energy from SIESTA .EIG file"""
    # Create mock .EIG file
    eig_file = tmpdir / "test.EIG"
    with open(eig_file, "w") as f:
        # Write mock eigenvalues
        f.write("1.0 2.0 3.0\n-1.0 0.0 1.0\n")

    # Test reading minimum energy
    emin = read_siesta_emin(str(eig_file))
    assert emin == -1.0


def test_int_de_ke():
    """Test energy integral calculation"""
    # Create test data
    energy = np.linspace(-10, 10, 1000)
    traced = np.exp(-(energy**2))  # Gaussian function
    we = np.ones_like(energy)  # Unit weights

    result = int_de_ke(traced, we)

    # Result should be real for real input
    assert np.imag(result) < 1e-10

    # Test with complex input
    traced_complex = traced + 1j * traced
    result_complex = int_de_ke(traced_complex, we)
    assert isinstance(result_complex, complex)


# Property-based tests
@given(st.integers(min_value=2, max_value=10))
def test_make_kset_dimensions_property(num_k):
    """Property-based test for k-point set generation"""
    kset = make_kset("xyz", num_k)

    # All k-points should be in [0,1)
    assert np.all(kset >= 0)
    assert np.all(kset < 1)

    # Shape should be (num_k^3, 3)
    assert kset.shape == (num_k**3, 3)


@given(
    st.floats(min_value=-20, max_value=-1),
    st.floats(min_value=0, max_value=1),
    st.integers(min_value=10, max_value=100),
)
def test_make_contour_properties(emin, emax, enum):
    """Property-based test for contour generation"""
    cont = make_contour(emin, emax, enum, 150)

    # Check basic properties
    assert len(cont.ze) == enum
    assert len(cont.we) == enum
    assert cont.R > 0
    assert cont.z0 == (emax + emin) / 2


if __name__ == "__main__":
    pytest.main([__file__])