<commit_before>from __future__ import division, absolute_import, print_function
from past.builtins import xrange

import numpy as np
import esutil
import scipy.integrate
import scipy.interpolate
import copy

from .galaxy import GalaxyCatalog
from .utilities import interpol

class ZredColor(object):
    """
    """

    def __init__(self, zredstr, sigint=0.001, do_correction=True, adaptive=True,
                 use_chisq=True, use_photoerr=True, zrange=None):
        self.zredstr = zredstr

        self.sigint = sigint
        self.do_correction = do_correction
        self.adaptive = adaptive
        self.use_chisq = use_chisq
        self.use_photoerr = use_photoerr
        self.zrange = zrange

        if (zredstr.z[1] - zredstr.z[0]) >= 0.01:
            # Must turn off adaptive if the stepsize is too large
            self.adaptive = False

    def compute_zred(self, galaxy):
        """
        """

        nz = self.zredstr.z.size
        refmagindex = self.zredstr.refmagindex(galaxy.refmag)

        if self.adaptive:
            step = 2
        else:
            step = 1

        notextrap, = np.where(~self.zredstr.extrapolated)

        if self.zrange is None:
            zbinstart = 0
            zbinstop = nz - 1
        else:
            u, = np.where((self.zredstr.z > self.zrange[0]) &
                          (self.zredstr.z < self.zrange[1]))
            zbinstart = u[0]
            zbinstop = u[-1]

        lndist = np.zeros(nz) - 1e12

        zbins = np.arange(zbinstart, zbinstop + step, step)

        # Mark the bins that are completely out of range
        # This last check makes sure we don't hit the overflow bin
        good = ((galaxy.refmag < self.zredstr.maxrefmag[zbins]) &
                (galaxy.refmag > self.zredstr.minrefmag[zbins]) &
                (self.zredstr.z[zbins] < 100.0))

        lndist[zbins[~good]] = -1e11

        if np.nonzero(good)[0].size > 0:
            # we have at least one good bin
            zbins = zbins[good]
            lndist[zbins] = self._calculate_lndist(galaxy, zbins)
        else:
            self._reset_bad_values(galaxy)
            return

        if self.adaptive:
            # only consider a maximum in the non-extrapolated region
            ind_temp = np.argmax(lndist[notextrap])
            ind = notextrap[ind_temp]

            # go over the nearest neighbors
            neighbors = 5

            minindex = ind - neighbors if ind - neighbors >= 0 else 0
            maxindex = ind + neighbors if ind + neighbors <= nz else nz

            if minindex == 0:
                maxindex = 1 + 2*neighbors
            if maxindex == (nz - 1):
                minindex = nz - 2 - 2 * neighbors

            zbins = np.arange(minindex, maxindex + 1)
            # select out the values that have not been run yet
            #  (these are very negative)
            to_run, = np.where(lndist[zbins] < -1e10)

            if to_run.size > 0:
                zbins = zbins[to_run]
                lndist[zbins] = self._calculate_lndist(galaxy, zbins)

        # move from log space to regular space
        maxlndist = np.max(lndist)
        with np.errstate(invalid='ignore'):
            dist = np.exp(lndist - maxlndist)

        # fix infinities and NaNs
        bad, = np.where(~np.isfinite(dist))
        dist[bad] = 0.0

        # did we hit a boundary?
        good, = np.where(dist > 0.0)
        if good.size >= 2:
            if (dist[good[0]] > 1e-5) or (dist[good[-1]] > 1e-5):
                # we did hit a boundary since dist didn't go to zero
                neighbors = 10

                if dist[good[0]] > 1e-5:
                    minindex = good[0] - neighbors if good[0] - neighbors >= 0 else 0
                    maxindex = good[0] - 1 if good[0] - 1 >= 0 else 0
                else:
                    maxindex = good[-1] + neighbors if good[-1] + neighbors <= nz else nz
                    minindex = good[-1] + 1 if good[-1] + 1 <= nz else nz

                zbins = np.arange(minindex, maxindex + 1)
                to_run, = np.where(lndist[zbins] < -1e10)
                if to_run.size > 0:
                    zbins = zbins[to_run]
                    lndist[zbins] = self._calculate_lndist(galaxy, zbins)

                    with np.errstate(invalid='ignore'):
                        dist[zbins] = np.exp(lndist[zbins] - maxlndist)

                    bad, = np.where(~np.isfinite(dist))
                    dist[bad] = 0.0

        # Now estimate zred and zred_e
        good, = np.where(dist > 0.0)
        if good.size < 2:
            self._reset_bad_values(galaxy)<commit_msg>Fix high end index overflow<commit_after>from __future__ import division, absolute_import, print_function
from past.builtins import xrange

import numpy as np
import esutil
import scipy.integrate
import scipy.interpolate
import copy

from .galaxy import GalaxyCatalog
from .utilities import interpol

class ZredColor(object):
    """
    """

    def __init__(self, zredstr, sigint=0.001, do_correction=True, adaptive=True,
                 use_chisq=True, use_photoerr=True, zrange=None):
        self.zredstr = zredstr

        self.sigint = sigint
        self.do_correction = do_correction
        self.adaptive = adaptive
        self.use_chisq = use_chisq
        self.use_photoerr = use_photoerr
        self.zrange = zrange

        if (zredstr.z[1] - zredstr.z[0]) >= 0.01:
            # Must turn off adaptive if the stepsize is too large
            self.adaptive = False

    def compute_zred(self, galaxy):
        """
        """

        nz = self.zredstr.z.size
        refmagindex = self.zredstr.refmagindex(galaxy.refmag)

        if self.adaptive:
            step = 2
        else:
            step = 1

        notextrap, = np.where(~self.zredstr.extrapolated)

        if self.zrange is None:
            zbinstart = 0
            zbinstop = nz - 1
        else:
            u, = np.where((self.zredstr.z > self.zrange[0]) &
                          (self.zredstr.z < self.zrange[1]))
            zbinstart = u[0]
            zbinstop = u[-1]

        lndist = np.zeros(nz) - 1e12

        zbins = np.arange(zbinstart, zbinstop + step, step)

        # Mark the bins that are completely out of range
        # This last check makes sure we don't hit the overflow bin
        good = ((galaxy.refmag < self.zredstr.maxrefmag[zbins]) &
                (galaxy.refmag > self.zredstr.minrefmag[zbins]) &
                (self.zredstr.z[zbins] < 100.0))

        lndist[zbins[~good]] = -1e11

        if np.nonzero(good)[0].size > 0:
            # we have at least one good bin
            zbins = zbins[good]
            lndist[zbins] = self._calculate_lndist(galaxy, zbins)
        else:
            self._reset_bad_values(galaxy)
            return

        if self.adaptive:
            # only consider a maximum in the non-extrapolated region
            ind_temp = np.argmax(lndist[notextrap])
            ind = notextrap[ind_temp]

            # go over the nearest neighbors
            neighbors = 5

            minindex = ind - neighbors if ind - neighbors >= 0 else 0
            maxindex = ind + neighbors if ind + neighbors <= nz else nz

            if minindex == 0:
                maxindex = 1 + 2*neighbors
            if maxindex == (nz - 1):
                minindex = nz - 2 - 2 * neighbors

            zbins = np.arange(minindex, maxindex + 1)
            # select out the values that have not been run yet
            #  (these are very negative)
            to_run, = np.where(lndist[zbins] < -1e10)

            if to_run.size > 0:
                zbins = zbins[to_run]
                lndist[zbins] = self._calculate_lndist(galaxy, zbins)

        # move from log space to regular space
        maxlndist = np.max(lndist)
        with np.errstate(invalid='ignore'):
            dist = np.exp(lndist - maxlndist)

        # fix infinities and NaNs
        bad, = np.where(~np.isfinite(dist))
        dist[bad] = 0.0

        # did we hit a boundary?
        good, = np.where(dist > 0.0)
        if good.size >= 2:
            if (dist[good[0]] > 1e-5) or (dist[good[-1]] > 1e-5):
                # we did hit a boundary since dist didn't go to zero
                neighbors = 10

                if dist[good[0]] > 1e-5:
                    minindex = good[0] - neighbors if good[0] - neighbors >= 0 else 0
                    maxindex = good[0] - 1 if good[0] - 1 >= 0 else 0
                else:
                    maxindex = good[-1] + neighbors if good[-1] + neighbors < nz else nz-1
                    minindex = good[-1] + 1 if good[-1] + 1 < nz else nz-1

                zbins = np.arange(minindex, maxindex + 1)
                to_run, = np.where(lndist[zbins] < -1e10)
                if to_run.size > 0:
                    zbins = zbins[to_run]
                    lndist[zbins] = self._calculate_lndist(galaxy, zbins)

                    with np.errstate(invalid='ignore'):
                        dist[zbins] = np.exp(lndist[zbins] - maxlndist)

                    bad, = np.where(~np.isfinite(dist))
                    dist[bad] = 0.0

        # Now estimate zred and zred_e
        good, = np.where(dist > 0.0)
        if good.size < 2:
            self._reset_bad_values(galaxy)
{'fraction_non_alphanumeric': 0.10025873221216042, 'fraction_numerical': 0.018542475204829666, 'mean_word_length': 1.6037047431939377, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 33, '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': '1424815', 'n_tokens_mistral': 3147, 'n_tokens_neox': 2893, 'n_words': 747}