<commit_before>#!/usr/bin/env python

# import modules
import warnings
import sys, re, abc
import numpy.ma as ma
from os import listdir
from netCDF4 import Dataset as nc
from optparse import OptionParser
from os.path import split, splitext, sep
from scipy.stats import scoreatpercentile
from numpy import ceil, double, intersect1d, where, sqrt, polyfit, diff, nan, isnan, arange, zeros, ones

# warnings.filterwarnings('error')

# DETRENDER CLASSES
class Detrender(object):
    @abc.abstractmethod
    def detrend(self, y): return
class NoDetrender(Detrender):
    def detrend(self, y): return y, None # no trend line
class LinDetrender(Detrender):
    def detrend(self, y):
        x = arange(1, 1 + len(y))
        coef = polyfit(x, y, 1)
        line = coef[0] * x + coef[1]
        dy = y - line
        return dy, line
class QuadDetrender(Detrender):
    def detrend(self, y):
        x = arange(1, 1 + len(y))
        coef = polyfit(x, y, 2)
        line = coef[0] * x ** 2 + coef[1] * x + coef[2]
        dy = y - line
        return dy, line
class MovingAveDetrender(Detrender):
    def detrend(self, y):
        dy = y.copy(); dy[:] = nan
        line = y.copy(); line[:] = nan
        line[2 : -2] = self.__winave(y, 5, 7)
        dy[2 : -2] = y[2 : -2] - line[2 : -2]
        dy = ma.masked_where(isnan(dy), dy)
        line = ma.masked_where(isnan(line), line)
        return dy, line
    def __winave(self, d, nstart, nend):
        # perform n-point window average of data d
        nd = len(d)
        dave = d.copy()
        numramp = (nend - nstart) / 2
        n = nstart - 2
        for i in range(nd):
            if i <= numramp:
                n += 2
            elif i >= nd - numramp:
                n -= 2
            idx0 = max(i - n / 2, 0)
            idx1 = min(i + n / 2 + 1, nd)
            dave[i] = d[idx0 : idx1].mean()
        dave = dave[n / 2 : -(n / 2)] # remove left and right edges
        return dave
class FFDDetrender(Detrender):
    def detrend(self, y):
        dy = y.copy()
        dy[:] = nan
        dy[1 :] = diff(y) / y[: -1]
        dy = ma.masked_where(isnan(dy), dy)
        line = None # no trend line
        return dy, line
class FFDRemoveDetrender(Detrender):
    def detrend(self, y):
        dy = y.copy(); dy[:] = nan
        line = y.copy(); line[:] = nan
        x = arange(2, 1 + len(y))
        ffd = diff(y) / y[: -1]
        coef = polyfit(x, ffd, 1)
        line[1 :] = coef[0] * x + coef[1]
        dy[1 :] = ffd - line[1 :]
        dy = ma.masked_where(isnan(dy), dy)
        line = ma.masked_where(isnan(line), line)
        return dy, line
class DetrenderWrapper(Detrender):
    def __init__(self, method):
        if method == 'none':
            self.dt = NoDetrender()
        elif method == 'lin':
            self.dt = LinDetrender()
        elif method == 'ffd':
            self.dt = FFDDetrender()
        else:
            raise Exception('Recognized detrend method')
    def detrend(self, y): return self.dt.detrend(y)

# TRANSFORMER CLASSES
class Transformer(object):
    @abc.abstractmethod
    def transform(p1, p2): return
class VSTransformer(Transformer):
    def transform(self, p1, p2):
        p1 -= p1.mean() # decenter
        p1 *= sqrt(p2.var() / p1.var()) # scale
        p1 += p2.mean() - p1.mean() # shift
        return p1
class QMTransformer(Transformer):
    def transform(self, p, q):
        p2 = p.copy() # copy
        q2 = q.copy()
        
        q2mean = q2.mean() # target mean
        
        p2 -= ma.median(p2) # center
        q2 -= ma.median(q2)

        q5 = self.__percentile(q2, 5)
        q25 = self.__percentile(q2, 25)
        q75 = self.__percentile(q2, 75)
        q95 = self.__percentile(q2, 95)
        
        p25 = self.__percentile(p2, 25) # 25th
        p2[p2 < 0] *= q25 / p25
        
        p75 = self.__percentile(p2, 75) # 75th
        p2[p2 > 0] *= q75 / p75

        p5 = self.__percentile(p2, 5) # 5th
        p2[p2 < q25] *= q5 / p5
        
        p95 = self.__percentile(p2, 95) # 95th
        p2[p2 > q75] *= q95 / p95

        p2 += q2mean - p2.mean() # reshift
        
        return p2
    def __percentile(self, p, per):
        p2 = p[~p.mask] if ma.isMaskedArray(p) else p.copy()
        return scoreatpercentile(p2, per)

# RETRENDER CLASSES
class Retrender(object):
    @abc.abstractmethod
    def retrend(self, T, R, o): return
class SumRetrender(Retrender):
    def retrend(self, T, R, o): return T + R
class FFDLeftRetrender(Retrender):
    def retrend(self, T, R, o):
        r = zeros((len(T),))
        r[0] = o[0] # initial condition = leftmost point
        # print r[0]
        # print T
        # print R
        for i in range(1, len(r)):
            # try:
                # print i, r[i-1], T[i], R[i]
                r[i] = r[i - 1] * (1. + T[i] + R[i])
            # except:
            #     print T
            #     print R
        return r
class FFDRightRetrender(Retrender):
    def retrend(self, T, R, o):
        r = zeros((len(T),))
        r[-1] = o[-1] # initial condition = rightmost point
        for i in range(-1, -len(r), -1):
            r[i - 1] = r[i] / (1. + T[i] + R[i])
        return r

# WRAPPER CLASS
class BiasCorrecter(object):
    def __init__(self, method):
        detrend_method, transform_method = method.split(' ')
        if transform_method == 'VS':
            self.transformer = VSTransformer()
        elif transform_method == 'QM':
            self.transformer = QMTransformer()
        else:
            raise Exception('Unrecognized transform method')
        if detrend_method in ['ffd-left', 'ffd-right']:
            self.detrenderObs = FFDRemoveDetrender()
            self.detrenderSim = FFDRemoveDetrender()
            if detrend_method == 'ffd-left':
                self.retrender = FFDLeftRetrender() 
            else:
                self.retrender = FFDRightRetrender()
        else:
            self.retrender = SumRetrender()
            self.detrenderSim = LinDetrender()
            if detrend_method == 'lin':
                self.detrenderObs = LinDetrender()
            elif detrend_method == 'quad':
                self.detrenderObs = QuadDetrender()
            elif detrend_method == 'moving-ave':
                self.detrenderObs = MovingAveDetrender()
            else:
                raise Exception('Unrecognized detrend method')
    
    def correct(self, sim, obs):
        """ Corrects simulated sequence sim using observations obs """
        # detrend observations
        odT, T = self.detrenderObs.detrend(obs)
        
        # detrend simualtions
        sdT = self.detrenderSim.detrend(sim)[0]

        if isnan(odT).sum() == len(odT):
            print odT
            print obs
            raise Exception('blah')
        if isnan(sdT).sum() == len(sdT):
            print sdT
            print sim
            raise Exception('blah2')

        # residuals
        # print 'sdT =', sdT
        # print 'odT =', odT
        if sdT.mask.all() or odT.mask.all(): # if either is totally masked
            R = ma.masked_array(zeros((len(T),)), mask = ones((len(T),)))
            yhat = R.copy()
        else:
            # transform distribution
            R = self.transformer.transform(sdT, odT)
            
            # combine
            yhat = self.retrender.retrend(T, R, obs)
        
        return yhat, T, R

# AGGREGATED BIAS-CORRECTED FILE CLASS
class AggBiasCorrectedFile(object):
    def __init__(self, filename, gadm, time, scen, detrend_methods, corr_methods):
        # create file
        self.filename = filename
        f = nc(filename, 'w', format = 'NETCDF4_CLASSIC')
        
        # create gadm
        f.createDimension('gadm0_index', len(gadm))
        gadmvar = f.createVariable('gadm0_index', 'i4', ('gadm0_index',))
        gadmvar[:] = gadm
        gadmvar.units = 'GADM L0 index'
        gadmvar.long_name = '253 countries'
        
        # create time
        f.createDimension('time', len(time))
        timevar = f.createVariable('time', 'i4', ('time',))
        timevar[:] = time
        timevar.units = 'years since {:d}-01-01'.format(int(time[0]))
        timevar.long_name = 'time'
        
        # create scen<commit_msg>Small change to biascorrect.py<commit_after>#!/usr/bin/env python

# import modules
import warnings
import sys, re, abc
import numpy.ma as ma
from os import listdir
from netCDF4 import Dataset as nc
from optparse import OptionParser
from os.path import split, splitext, sep
from scipy.stats import scoreatpercentile
from numpy import ceil, double, intersect1d, where, sqrt, polyfit, diff, nan, isnan, arange, zeros, ones

# warnings.filterwarnings('error')

# DETRENDER CLASSES
class Detrender(object):
    @abc.abstractmethod
    def detrend(self, y): return
class NoDetrender(Detrender):
    def detrend(self, y): return y, None # no trend line
class LinDetrender(Detrender):
    def detrend(self, y):
        x = arange(1, 1 + len(y))
        coef = polyfit(x, y, 1)
        line = coef[0] * x + coef[1]
        dy = y - line
        return dy, line
class QuadDetrender(Detrender):
    def detrend(self, y):
        x = arange(1, 1 + len(y))
        coef = polyfit(x, y, 2)
        line = coef[0] * x ** 2 + coef[1] * x + coef[2]
        dy = y - line
        return dy, line
class MovingAveDetrender(Detrender):
    def detrend(self, y):
        dy = y.copy(); dy[:] = nan
        line = y.copy(); line[:] = nan
        line[2 : -2] = self.__winave(y, 5, 7)
        dy[2 : -2] = y[2 : -2] - line[2 : -2]
        dy = ma.masked_where(isnan(dy), dy)
        line = ma.masked_where(isnan(line), line)
        return dy, line
    def __winave(self, d, nstart, nend):
        # perform n-point window average of data d
        nd = len(d)
        dave = d.copy()
        numramp = (nend - nstart) / 2
        n = nstart - 2
        for i in range(nd):
            if i <= numramp:
                n += 2
            elif i >= nd - numramp:
                n -= 2
            idx0 = max(i - n / 2, 0)
            idx1 = min(i + n / 2 + 1, nd)
            dave[i] = d[idx0 : idx1].mean()
        dave = dave[n / 2 : -(n / 2)] # remove left and right edges
        return dave
class FFDDetrender(Detrender):
    def detrend(self, y):
        dy = y.copy()
        dy[:] = nan
        dy[1 :] = diff(y) / y[: -1]
        dy = ma.masked_where(isnan(dy), dy)
        line = None # no trend line
        return dy, line
class FFDRemoveDetrender(Detrender):
    def detrend(self, y):
        dy = y.copy(); dy[:] = nan
        line = y.copy(); line[:] = nan
        x = arange(2, 1 + len(y))
        ffd = diff(y) / y[: -1]
        coef = polyfit(x, ffd, 1)
        line[1 :] = coef[0] * x + coef[1]
        dy[1 :] = ffd - line[1 :]
        dy = ma.masked_where(isnan(dy), dy)
        line = ma.masked_where(isnan(line), line)
        return dy, line
class DetrenderWrapper(Detrender):
    def __init__(self, method):
        if method == 'none':
            self.dt = NoDetrender()
        elif method == 'lin':
            self.dt = LinDetrender()
        elif method == 'ffd':
            self.dt = FFDDetrender()
        else:
            raise Exception('Recognized detrend method')
    def detrend(self, y): return self.dt.detrend(y)

# TRANSFORMER CLASSES
class Transformer(object):
    @abc.abstractmethod
    def transform(p1, p2): return
class VSTransformer(Transformer):
    def transform(self, p1, p2):
        p1 -= p1.mean() # decenter
        p1 *= sqrt(p2.var() / p1.var()) # scale
        p1 += p2.mean() - p1.mean() # shift
        return p1
class QMTransformer(Transformer):
    def transform(self, p, q):
        p2 = p.copy() # copy
        q2 = q.copy()
        
        q2mean = q2.mean() # target mean
        
        p2 -= ma.median(p2) # center
        q2 -= ma.median(q2)

        q5 = self.__percentile(q2, 5)
        q25 = self.__percentile(q2, 25)
        q75 = self.__percentile(q2, 75)
        q95 = self.__percentile(q2, 95)
        
        p25 = self.__percentile(p2, 25) # 25th
        p2[p2 < 0] *= q25 / p25
        
        p75 = self.__percentile(p2, 75) # 75th
        p2[p2 > 0] *= q75 / p75

        p5 = self.__percentile(p2, 5) # 5th
        p2[p2 < q25] *= q5 / p5
        
        p95 = self.__percentile(p2, 95) # 95th
        p2[p2 > q75] *= q95 / p95

        p2 += q2mean - p2.mean() # reshift
        
        return p2
    def __percentile(self, p, per):
        p2 = p[~p.mask] if ma.isMaskedArray(p) else p.copy()
        return scoreatpercentile(p2, per)

# RETRENDER CLASSES
class Retrender(object):
    @abc.abstractmethod
    def retrend(self, T, R, o): return
class SumRetrender(Retrender):
    def retrend(self, T, R, o): return T + R
class FFDLeftRetrender(Retrender):
    def retrend(self, T, R, o):
        r = zeros((len(T),))
        r[0] = o[0] # initial condition = leftmost point
        # print r[0]
        # print T
        # print R
        for i in range(1, len(r)):
            # try:
                # print i, r[i-1], T[i], R[i]
                r[i] = r[i - 1] * (1. + T[i] + R[i])
            # except:
            #     print T
            #     print R
        return r
class FFDRightRetrender(Retrender):
    def retrend(self, T, R, o):
        r = zeros((len(T),))
        r[-1] = o[-1] # initial condition = rightmost point
        for i in range(-1, -len(r), -1):
            r[i - 1] = r[i] / (1. + T[i] + R[i])
        return r

# WRAPPER CLASS
class BiasCorrecter(object):
    def __init__(self, method):
        detrend_method, transform_method = method.split(' ')
        if transform_method == 'VS':
            self.transformer = VSTransformer()
        elif transform_method == 'QM':
            self.transformer = QMTransformer()
        else:
            raise Exception('Unrecognized transform method')
        if detrend_method in ['ffd-left', 'ffd-right']:
            self.detrenderObs = FFDRemoveDetrender()
            self.detrenderSim = FFDRemoveDetrender()
            if detrend_method == 'ffd-left':
                self.retrender = FFDLeftRetrender() 
            else:
                self.retrender = FFDRightRetrender()
        else:
            self.retrender = SumRetrender()
            self.detrenderSim = LinDetrender()
            if detrend_method == 'lin':
                self.detrenderObs = LinDetrender()
            elif detrend_method == 'quad':
                self.detrenderObs = QuadDetrender()
            elif detrend_method == 'moving-ave':
                self.detrenderObs = MovingAveDetrender()
            else:
                raise Exception('Unrecognized detrend method')
    
    def correct(self, sim, obs):
        """ Corrects simulated sequence sim using observations obs """
        # detrend observations
        odT, T = self.detrenderObs.detrend(obs)
        
        # detrend simualtions
        sdT = self.detrenderSim.detrend(sim)[0]

        if isnan(odT).sum() == len(odT):
            print odT
            print obs
            raise Exception('blah')
        if isnan(sdT).sum() == len(sdT):
            print sdT
            print sim
            raise Exception('blah2')

        # residuals
        # print 'sdT =', sdT
        # print 'odT =', odT
        if sdT.mask.all() or odT.mask.all(): # if either is totally masked
            R = ma.masked_array(zeros((len(T),)), mask = ones((len(T),)))
            yhat = R.copy()
        else:
            # transform distribution
            R = self.transformer.transform(sdT, odT)
            
            # combine
            yhat = self.retrender.retrend(T, R, obs)
        
        return yhat, T, R

# AGGREGATED BIAS-CORRECTED FILE CLASS
class AggBiasCorrectedFile(object):
    def __init__(self, filename, gadm, time, scen, detrend_methods, corr_methods):
        # create file
        self.filename = filename
        f = nc(filename, 'w', format = 'NETCDF4_CLASSIC')
        
        # create gadm
        f.createDimension('gadm0_index', len(gadm))
        gadmvar = f.createVariable('gadm0_index', 'i4', ('gadm0_index',))
        gadmvar[:] = gadm
        gadmvar.units = 'GADM L0 index'
        gadmvar.long_name = '253 countries'
        
        # create time
        f.createDimension('time', len(time))
        timevar = f.createVariable('time', 'i4', ('time',))
        timevar[:] = time - time[0]
        timevar.units = 'years since {:d}-01-01'.format(int(time[0]))
        timevar.long_name = 'time'
        
        # create scen
{'fraction_non_alphanumeric': 0.12969822089987249, 'fraction_numerical': 0.021555650009108022, 'mean_word_length': 1.859375, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, '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': '6392562', 'n_tokens_mistral': 5764, 'n_tokens_neox': 5487, 'n_words': 1445}