<commit_before>from __future__ import division, print_function

import os
import numpy as np
import scipy.sparse

from . import rebin, Spectrum, Target
import redrock

def calc_zchi2(redshifts, spectra, template):
    '''
    TODO: document
    '''
    targets = [Target(0, spectra), ]
    zchi2, zcoeff = calc_zchi2_targets(redshifts, targets, template)
    return zchi2[0], zcoeff[0]

def parallel_calc_zchi2_targets(redshifts, targets, template, verbose=False, \
    ncpu=None, numthreads=-1):
    '''
    Parallel version of calc_zchi2_targets; see that docstring for details

    TODO: document ncpu and numthreads
    '''
    import multiprocessing as mp
    if ncpu is None:
        ncpu = mp.cpu_count()

    if numthreads == 0:
        n = max(mp.cpu_count() // ncpu, 1)
        os.environ['OMP_NUM_THREADS'] = str(n)
        os.environ['MKL_NUM_THREADS'] = str(n)
        if verbose:
            print('DEBUG: setting OMP_NUM_THREADS={}'.format(n))
            print('DEBUG: setting MKL_NUM_THREADS={}'.format(n))
    elif numthreads > 0:
        os.environ['OMP_NUM_THREADS'] = str(numthreads)
        os.environ['MKL_NUM_THREADS'] = str(numthreads)

    if verbose:
        print('DEBUG: $OMP_NUM_THREADS={}'.format(os.getenv('OMP_NUM_THREADS')))
        print('DEBUG: $MKL_NUM_THREADS={}'.format(os.getenv('MKL_NUM_THREADS')))

    def foo(i, qin, qout):
        zz = qin.get()
        if verbose:
            print('Process {}: {} redshifts from {:.4f} to {:.4f}'.format(i, len(zz), zz[0], zz[-1]))
        qout.put((zz[0], redrock.zscan.calc_zchi2_targets(zz, targets, template, verbose=verbose)))

    ii = np.linspace(0, len(redshifts), ncpu+1).astype(int)
    qin = mp.Queue()
    qout = mp.Queue()
    for i in range(ncpu):
        qin.put(redshifts[ii[i]:ii[i+1]])<commit_msg>zscan with ncpu > nz<commit_after>from __future__ import division, print_function

import os
import numpy as np
import scipy.sparse

from . import rebin, Spectrum, Target
import redrock

def calc_zchi2(redshifts, spectra, template):
    '''
    TODO: document
    '''
    targets = [Target(0, spectra), ]
    zchi2, zcoeff = calc_zchi2_targets(redshifts, targets, template)
    return zchi2[0], zcoeff[0]

def parallel_calc_zchi2_targets(redshifts, targets, template, verbose=False, \
    ncpu=None, numthreads=-1):
    '''
    Parallel version of calc_zchi2_targets; see that docstring for details

    TODO: document ncpu and numthreads
    '''
    import multiprocessing as mp
    if ncpu is None:
        ncpu = mp.cpu_count()

    if ncpu > len(redshifts):
        ncpu = len(redshifts)
        print('WARNING: Using {} cores for {} redshifts'.format(ncpu, ncpu))

    if numthreads == 0:
        n = max(mp.cpu_count() // ncpu, 1)
        os.environ['OMP_NUM_THREADS'] = str(n)
        os.environ['MKL_NUM_THREADS'] = str(n)
        if verbose:
            print('DEBUG: setting OMP_NUM_THREADS={}'.format(n))
            print('DEBUG: setting MKL_NUM_THREADS={}'.format(n))
    elif numthreads > 0:
        os.environ['OMP_NUM_THREADS'] = str(numthreads)
        os.environ['MKL_NUM_THREADS'] = str(numthreads)

    if verbose:
        print('DEBUG: $OMP_NUM_THREADS={}'.format(os.getenv('OMP_NUM_THREADS')))
        print('DEBUG: $MKL_NUM_THREADS={}'.format(os.getenv('MKL_NUM_THREADS')))

    def foo(i, qin, qout):
        zz = qin.get()
        if verbose:
            print('Process {}: {} redshifts from {:.4f} to {:.4f}'.format(i, len(zz), zz[0], zz[-1]))
        qout.put((zz[0], redrock.zscan.calc_zchi2_targets(zz, targets, template, verbose=verbose)))

    ii = np.linspace(0, len(redshifts), ncpu+1).astype(int)
    qin = mp.Queue()
    qout = mp.Queue()
    for i in range(ncpu):
        tmpz = redshifts[ii[i]:ii[i+1]]
        assert len(tmpz) > 0
        qin.put(tmpz)
{'fraction_non_alphanumeric': 0.14289487585842578, 'fraction_numerical': 0.011885895404120444, 'mean_word_length': 3.1661166116611663, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1414813', 'n_tokens_mistral': 1413, 'n_tokens_neox': 1353, 'n_words': 288}