tests/test_mhm.py

"""
Sanity tests for the open BatteryMHM method. Run: pytest -q
These verify the harmonic algebra, the Chi matrix, the feature library, and
the ensemble all work — so a reviewer can confirm the release is functional.
"""

import numpy as np

from batterymhm import (
    CHI_MATRIX,
    MHMEnsemble,
    chi_score,
    cmr,
    compute_metrics,
    creative_add,
    element_hin,
    f9,
    hin,
    mhm_full_features,
    miller_sequence,
    miller_sub,
    seq_to_harmonics,
)


def test_fold_map_range():
    # HIN(k) = 1 + ((k-1) mod 9) always lands in {1..9}
    for k in range(-50, 200):
        assert 1 <= f9(k) <= 9
    assert f9(1) == 1 and f9(9) == 9 and f9(10) == 1


def test_hin_matches_fold_map():
    # hin(Z) is the fold map on atomic number
    for z in range(1, 119):
        assert hin(z) == 1 + ((z - 1) % 9)
    assert element_hin("Li") == hin(3)      # Li, Z=3
    assert element_hin("O") == hin(8)       # O,  Z=8


def test_operations_stay_in_harmonic_space():
    # Both binary operations are closed on {1..9} and commutative.
    for a in range(1, 10):
        for b in range(1, 10):
            assert 1 <= creative_add(a, b) <= 9
            assert 1 <= miller_sub(a, b) <= 9
            assert creative_add(a, b) == creative_add(b, a)   # commutative


def test_chi_matrix_symmetric_unit_diagonal():
    assert CHI_MATRIX.shape == (9, 9)
    assert np.allclose(np.diag(CHI_MATRIX), 1.0)
    assert np.allclose(CHI_MATRIX, CHI_MATRIX.T)
    assert 0.0 < CHI_MATRIX.min() <= CHI_MATRIX.max() <= 1.0
    assert chi_score(3, 9) == CHI_MATRIX[2, 8]


def test_miller_sequence_seed():
    assert miller_sequence(6) == [1, 1, 3, 4, 7, 11]
    assert cmr([3, 2, 2]) >= 1   # CMR returns a harmonic digit


def test_feature_vector_is_finite_and_sized():
    hins = seq_to_harmonics(list(np.linspace(1.0, 0.8, 60)), bins=9)
    feats = mhm_full_features(hins)
    assert len(feats) == 557
    assert all(np.isfinite(v) for v in feats.values())


def test_ensemble_learns_a_trivial_signal():
    # A model that consumes MHM features should fit a simple monotone target
    # better than predicting the training mean.
    rng = np.random.default_rng(0)
    curves = [np.linspace(1.0, 1.0 - r, 60) for r in rng.uniform(0.05, 0.4, 80)]
    y = np.array([c[-1] for c in curves])
    X = np.array([[d[k] for k in sorted(d)]
                  for d in (mhm_full_features(seq_to_harmonics(list(c), 9))
                            for c in curves)])
    m = MHMEnsemble(use_xgb=False, et_params={"n_estimators": 80},
                    use_stacking=False).fit(X[:60], y[:60])
    pred = m.predict(X[60:])
    metrics = compute_metrics(y[60:], pred)
    baseline = float(np.mean(np.abs(y[60:] - y[:60].mean())))
    assert metrics.mae < baseline