<commit_before>import os
import glob
from contextlib import contextmanager

import pytest

from numpy.testing import assert_allclose
import numpy as np
from scipy import ndimage as ndi
from sklearn.linear_model import LogisticRegression as LR
import subprocess as sp
from gala import imio, features, agglo, evaluate as ev


@contextmanager
def tar_extract(fn):
    sp.call(['tar', '-xzf', fn + '.tar.gz'])
    ext_fn = os.path.basename(fn)
    yield ext_fn
    os.remove(ext_fn)
    for sub_fn in glob.glob(ext_fn + '_*'):
        os.remove(sub_fn)


rundir = os.path.dirname(__file__)

### fixtures

@pytest.fixture
def dummy_data():
    frag = np.arange(1, 17, dtype=int).reshape((4, 4))
    gt = np.array([[1, 1, 2, 2], [1, 1, 2, 2], [3] * 4, [3] * 4], dtype=int)
    fman = features.base.Mock(frag, gt)
    g = agglo.Rag(frag, feature_manager=fman, use_slow=True)
    return frag, gt, g, fman


@pytest.fixture
def dummy_data_fast():
    frag, gt, _, fman = dummy_data()
    frag = ndi.zoom(frag, 2, order=0)
    gt = ndi.zoom(gt, 2, order=0)
    g = agglo.Rag(frag, feature_manager=fman)
    return frag, gt, g, fman


### tests

def test_generate_flat_learning_edges(dummy_data):
    """Run a flat epoch and ensure all edges are correctly represented."""
    frag, gt, g, fman = dummy_data
    feat, target, weights, edges = g.learn_flat(gt, fman)
    assert feat.shape == (24, 2)
    assert tuple(edges[0]) == (1, 2)
    assert tuple(edges[-1]) == (15, 16)
    assert np.sum(target[:, 0] == 1) == 6  # number of non-merge edges


def test_generate_flat_learning_edges_fast(dummy_data_fast):
    """Run a flat epoch and ensure all edges are correctly represented."""
    frag, gt, g, fman = dummy_data_fast
    feat, target, weights, edges = g.learn_flat(gt, fman)
    assert feat.shape == (24, 2)
    assert tuple(edges[0]) == (1, 2)
    assert tuple(edges[-1]) == (15, 16)
    assert np.sum(target[:, 0] == 1) == 6  # number of non-merge edges


def test_generate_lash_examples(dummy_data):
    """Run a flat epoch and an active epoch of learning, compare learned sets.

    The mock feature manager places all merge examples at (0, 0) in feature
    space, and all non-merge examples at (1, 0), *in flat learning*. During
    agglomeration, non-merge examples go to (0, 1), which confuses the flat
    classifier (which has only learned the difference along the first feature
    dimension).

    This test checks for those differences in learning using a simple
    logistic regression.
    """
    frag, gt, g, fman = dummy_data
    np.random.seed(99)
    summary, allepochs = g.learn_agglomerate(gt, fman,
                                             learning_mode='permissive',
                                             classifier='logistic regression')
    feat, target, weights, edges = summary
    ffeat, ftarget, fweights, fedges = allepochs[0]  # flat
    lr = LR().fit(feat, target[:, 0])
    flr = LR().fit(ffeat, ftarget[:, 0])
    def pred(v):
        return lr.predict_proba([v])[0, 1]
    def fpred(v):
        return flr.predict_proba([v])[0, 1]
    assert len(allepochs[1][0]) == 15  # number of merges is |nodes| - 1

    # approx. same learning results at (0., 0.) and (1., 0.)
    print([(fpred(i), pred(i)) for i in [[0, 0], [1, 0], [0, 1]]])
    assert_allclose(fpred([0, 0]), 0.2, atol=0.025)
    assert_allclose(pred([0, 0]), 0.2, atol=0.025)
    assert_allclose(fpred([1, 0]), 0.65, atol=0.025)
    assert_allclose(pred([1, 0]), 0.65, atol=0.025)

    # difference between agglomerative and flat learning in point (0., 1.)
    assert_allclose(fpred([0, 1]), 0.2, atol=0.025)
    assert_allclose(pred([0, 1]), 0.5, atol=0.025)


def test_generate_gala_examples(dummy_data):
    """As `test_generate_lash_examples`, but using strict learning. """<commit_msg>Add test for fast RAG with LASH<commit_after>import os
import glob
from contextlib import contextmanager

import pytest

from numpy.testing import assert_allclose
import numpy as np
from scipy import ndimage as ndi
from sklearn.linear_model import LogisticRegression as LR
import subprocess as sp
from gala import imio, features, agglo, evaluate as ev


@contextmanager
def tar_extract(fn):
    sp.call(['tar', '-xzf', fn + '.tar.gz'])
    ext_fn = os.path.basename(fn)
    yield ext_fn
    os.remove(ext_fn)
    for sub_fn in glob.glob(ext_fn + '_*'):
        os.remove(sub_fn)


rundir = os.path.dirname(__file__)

### fixtures

@pytest.fixture
def dummy_data():
    frag = np.arange(1, 17, dtype=int).reshape((4, 4))
    gt = np.array([[1, 1, 2, 2], [1, 1, 2, 2], [3] * 4, [3] * 4], dtype=int)
    fman = features.base.Mock(frag, gt)
    g = agglo.Rag(frag, feature_manager=fman, use_slow=True)
    return frag, gt, g, fman


@pytest.fixture
def dummy_data_fast():
    frag, gt, _, fman = dummy_data()
    frag = ndi.zoom(frag, 2, order=0)
    gt = ndi.zoom(gt, 2, order=0)
    g = agglo.Rag(frag, feature_manager=fman)
    return frag, gt, g, fman


### tests

def test_generate_flat_learning_edges(dummy_data):
    """Run a flat epoch and ensure all edges are correctly represented."""
    frag, gt, g, fman = dummy_data
    feat, target, weights, edges = g.learn_flat(gt, fman)
    assert feat.shape == (24, 2)
    assert tuple(edges[0]) == (1, 2)
    assert tuple(edges[-1]) == (15, 16)
    assert np.sum(target[:, 0] == 1) == 6  # number of non-merge edges


def test_generate_flat_learning_edges_fast(dummy_data_fast):
    """Run a flat epoch and ensure all edges are correctly represented."""
    frag, gt, g, fman = dummy_data_fast
    feat, target, weights, edges = g.learn_flat(gt, fman)
    assert feat.shape == (24, 2)
    assert tuple(edges[0]) == (1, 2)
    assert tuple(edges[-1]) == (15, 16)
    assert np.sum(target[:, 0] == 1) == 6  # number of non-merge edges


def test_generate_lash_examples(dummy_data):
    """Run a flat epoch and an active epoch of learning, compare learned sets.

    The mock feature manager places all merge examples at (0, 0) in feature
    space, and all non-merge examples at (1, 0), *in flat learning*. During
    agglomeration, non-merge examples go to (0, 1), which confuses the flat
    classifier (which has only learned the difference along the first feature
    dimension).

    This test checks for those differences in learning using a simple
    logistic regression.
    """
    frag, gt, g, fman = dummy_data
    np.random.seed(99)
    summary, allepochs = g.learn_agglomerate(gt, fman,
                                             learning_mode='permissive',
                                             classifier='logistic regression')
    feat, target, weights, edges = summary
    ffeat, ftarget, fweights, fedges = allepochs[0]  # flat
    lr = LR().fit(feat, target[:, 0])
    flr = LR().fit(ffeat, ftarget[:, 0])
    def pred(v):
        return lr.predict_proba([v])[0, 1]
    def fpred(v):
        return flr.predict_proba([v])[0, 1]
    assert len(allepochs[1][0]) == 15  # number of merges is |nodes| - 1

    # approx. same learning results at (0., 0.) and (1., 0.)
    print([(fpred(i), pred(i)) for i in [[0, 0], [1, 0], [0, 1]]])
    assert_allclose(fpred([0, 0]), 0.2, atol=0.025)
    assert_allclose(pred([0, 0]), 0.2, atol=0.025)
    assert_allclose(fpred([1, 0]), 0.65, atol=0.025)
    assert_allclose(pred([1, 0]), 0.65, atol=0.025)

    # difference between agglomerative and flat learning in point (0., 1.)
    assert_allclose(fpred([0, 1]), 0.2, atol=0.025)
    assert_allclose(pred([0, 1]), 0.5, atol=0.025)


def test_generate_lash_examples_fast(dummy_data_fast):
    """Run a flat epoch and an active epoch of learning, compare learned sets.

    The mock feature manager places all merge examples at (0, 0) in feature
    space, and all non-merge examples at (1, 0), *in flat learning*. During
    agglomeration, non-merge examples go to (0, 1), which confuses the flat
    classifier (which has only learned the difference along the first feature
    dimension).

    This test checks for those differences in learning using a simple
    logistic regression.
    """
    frag, gt, g, fman = dummy_data_fast
    np.random.seed(99)
    summary, allepochs = g.learn_agglomerate(gt, fman,
                                             learning_mode='permissive',
                                             classifier='logistic regression')
    feat, target, weights, edges = summary
    ffeat, ftarget, fweights, fedges = allepochs[0]  # flat
    lr = LR().fit(feat, target[:, 0])
    flr = LR().fit(ffeat, ftarget[:, 0])
    def pred(v):
        return lr.predict_proba([v])[0, 1]
    def fpred(v):
        return flr.predict_proba([v])[0, 1]
    assert len(allepochs[1][0]) == 15  # number of merges is |nodes| - 1

    # approx. same learning results at (0., 0.) and (1., 0.)
    print([(fpred(i), pred(i)) for i in [[0, 0], [1, 0], [0, 1]]])
    assert_allclose(fpred([0, 0]), 0.2, atol=0.2)
    assert_allclose(pred([0, 0]), 0.2, atol=0.2)
    assert_allclose(fpred([1, 0]), 0.65, atol=0.15)
    assert_allclose(pred([1, 0]), 0.65, atol=0.15)

    # difference between agglomerative and flat learning in point (0., 1.)
    assert_allclose(fpred([0, 1]), 0.2, atol=0.2)  # < 0.4
    assert_allclose(pred([0, 1]), 0.65, atol=0.2)  # > 0.45


def test_generate_gala_examples(dummy_data):
    """As `test_generate_lash_examples`, but using strict learning. """
{'fraction_non_alphanumeric': 0.1381156316916488, 'fraction_numerical': 0.03640256959314775, 'mean_word_length': 3.1460275188637374, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 15, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2597265', 'n_tokens_mistral': 3454, 'n_tokens_neox': 3219, 'n_words': 992}