import numpy
from pylearn.preprocessing.pca import *

def test_max_energy_fraction():
    rng = numpy.random.RandomState(234)
    a = rng.randn(10,5)
    a[:,0] *=2
    a+= 5
    (eigvals, eigvecs), centered_a = pca_from_examples(a)
    assert len(eigvals)==5
    vartot = eigvals.sum()

    def f(frac, n):
        (evals, evecs), c_a = pca_from_examples(a, max_energy_fraction=frac)
        assert len(evals)==n
        assert evals.sum() <= (frac * vartot + 1e-8)

    f(.87,1)
    f(.88,2)
    f(.99999999, 4)
    f(1.0, 5)
    f(.5, 0)