<commit_before>import rfpipe
import pytest
from astropy import time
import numpy as np

# simulate no flag, transient/no flag, transient/flag
inprefs = [{'flaglist': [], 'npix_max': 32, 'sigma_image1': 0, 'spw': [0, 1]},
           {'flaglist': [], 'npix_max': 32, 'sigma_image1': 0, 'spw': [2, 3],
           'maxdm': 100, 'dtarr': [1, 2]}
           ]


@pytest.fixture(scope="module", params=inprefs)
def mockstate(request):
    t0 = time.Time.now().mjd
    meta = rfpipe.metadata.mock_metadata(t0, t0+0.3/(24*3600), 27, 4, 2, 10e3,
                                         datasource='sim')
    return rfpipe.state.State(inmeta=meta, inprefs=request.param)


def mockdata(mockstate, segment):
    data = rfpipe.source.read_segment(mockstate, segment)
    return rfpipe.source.data_prep(mockstate, data)


def test_dataprep(mockstate):
    data_prep = mockdata(mockstate, 0)
    assert data_prep.shape == mockstate.datashape


def test_search(mockstate):

    wisdom = rfpipe.search.set_wisdom(mockstate.npixx, mockstate.npixy)

    featurelist = []
    for segment in range(mockstate.nsegment):
        data_prep = mockdata(mockstate, segment)

        for dmind in range(len(mockstate.dmarr)):
            delay = rfpipe.util.calc_delay(mockstate.freq,
                                           mockstate.freq.max(),
                                           mockstate.dmarr[dmind],
                                           mockstate.inttime)
            data_dm = rfpipe.search.dedisperse(data_prep, delay)

            for dtind in range(len(mockstate.dtarr)):
                data_dmdt = rfpipe.search.resample(data_dm,
                                                   mockstate.dtarr[dtind])

                canddatalist = rfpipe.search.search_thresh(mockstate,
                                                           data_dmdt,
                                                           segment, dmind,
                                                           dtind,
                                                           wisdom=wisdom)

                features = rfpipe.search.calc_features(canddatalist)
                featurelist.append(features)
                print(features.keys())
                assert len(canddatalist) == (mockstate.readints-mockstate.dmshifts[dmind])//mockstate.dtarr[dtind]

    alltimes = np.linspace(mockstate.metadata.starttime_mjd,
                           mockstate.metadata.stoptime_mjd, mockstate.nints)
    integs0 = []
    integs1 = []
    for feature in features:
        for (seg, integ, dmind, dtind, beamnum) in feature:
            if dtind == 0:
                integs0.append(integ)
            elif dtind == 1:
                integs1.append(integ)

    assert len(alltimes) == len(integs0)
    assert len(alltimes)//2 == len(integs1)<commit_msg>conditional assertion for dmind dtind<commit_after>import rfpipe
import pytest
from astropy import time
import numpy as np

# simulate no flag, transient/no flag, transient/flag
inprefs = [{'flaglist': [], 'npix_max': 32, 'sigma_image1': 0, 'spw': [0, 1]},
           {'flaglist': [], 'npix_max': 32, 'sigma_image1': 0, 'spw': [2, 3],
           'maxdm': 100, 'dtarr': [1, 2]}
           ]


@pytest.fixture(scope="module", params=inprefs)
def mockstate(request):
    t0 = time.Time.now().mjd
    meta = rfpipe.metadata.mock_metadata(t0, t0+1./(24*3600), 27, 4, 2, 10e3,
                                         datasource='sim')
    return rfpipe.state.State(inmeta=meta, inprefs=request.param)


def mockdata(mockstate, segment):
    data = rfpipe.source.read_segment(mockstate, segment)
    return rfpipe.source.data_prep(mockstate, data)


def test_dataprep(mockstate):
    data_prep = mockdata(mockstate, 0)
    assert data_prep.shape == mockstate.datashape


def test_search(mockstate):

    wisdom = rfpipe.search.set_wisdom(mockstate.npixx, mockstate.npixy)

    featurelist = []
    for segment in range(mockstate.nsegment):
        data_prep = mockdata(mockstate, segment)

        for dmind in range(len(mockstate.dmarr)):
            delay = rfpipe.util.calc_delay(mockstate.freq,
                                           mockstate.freq.max(),
                                           mockstate.dmarr[dmind],
                                           mockstate.inttime)
            data_dm = rfpipe.search.dedisperse(data_prep, delay)

            for dtind in range(len(mockstate.dtarr)):
                data_dmdt = rfpipe.search.resample(data_dm,
                                                   mockstate.dtarr[dtind])

                canddatalist = rfpipe.search.search_thresh(mockstate,
                                                           data_dmdt,
                                                           segment, dmind,
                                                           dtind,
                                                           wisdom=wisdom)

                features = rfpipe.search.calc_features(canddatalist)
                featurelist.append(features)
                print(features.keys())
                assert len(canddatalist) == (mockstate.readints-mockstate.dmshifts[dmind])//mockstate.dtarr[dtind]

    integs0_0 = []
    integs1_0 = []
    integs0_1 = []
    for features in featurelist:
        for (seg, integ, dmind, dtind, beamnum) in features.keys():
            if dtind == 0 and dmind == 0:
                integs0_0.append(integ)
            elif dtind == 1 and dmind == 0:
                integs1_0.append(integ)
            elif dtind == 0 and dmind == 1:
                integs0_1.append(integ)

    assert mockstate.nints == len(integs0_0)
    if 1 in mockstate.dtarr:
        assert mockstate.nints//2 == len(integs1_0)
    if 1 in mockstate.dmarr:
        assert mockstate.nints-mockstate.dmshifts[1] == len(integs0_1)
{'fraction_non_alphanumeric': 0.10472103004291845, 'fraction_numerical': 0.018540772532188842, 'mean_word_length': 1.702226345083488, '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': 36, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '871199', 'n_tokens_mistral': 1805, 'n_tokens_neox': 1695, 'n_words': 328}