<commit_before>import time
from itertools import product

import emcee
import matplotlib.patches as mpatches
from scipy.cluster.vq import whiten, kmeans2
from scipy.optimize import leastsq, minimize
from tqdm.auto import tqdm
from pymultinest.solve import  solve

from .utils import *
from . import MPLSTYLE


class MaculaModeler(object):
    def __init__(self, t, y, nspots, dy=None, inc=None, Peq=None, k2=None, k4=None, c1=None, c2=None, c3=None, c4=None,
                 d1=None, d2=None, d3=None, d4=None, long=None, lat=None, alpha=None, fspot=None, tmax=None,
                 life=None, ingress=None, egress=None, U=None, B=None, tstart=None, tend=None):
        self.t = t
        self.y = y
        self.nspots = nspots
        self.dy = dy
        self.inc = inc
        self.Peq = Peq
        self.k2 = k2
        self.k4 = k4
        self.c1 = c1
        self.c2 = c2
        self.c3 = c3
        self.c4 = c4
        self.d1 = d1
        self.d2 = d2
        self.d3 = d3
        self.d4 = d4
        self.long = long
        self.lat = lat
        self.alpha = alpha
        self.fspot = fspot
        self.tmax = tmax
        self.life = life
        self.ingress = ingress
        self.egress = egress
        self.U = U
        self.B = B
        self.tstart = tstart
        self.tend = tend

        if self.dy is None:
            self.dy = np.ones_like(y)

        if self.inc is None:
            self.inc = np.array([0, np.pi/2])
        if self.Peq is None:
            self.Peq = np.array([0, 50])
        if self.k2 is None:
            self.k2 = np.array([-1, 1])
        if self.k4 is None:
            self.k4 = np.array([-1, 1])
        if self.c1 is None:
            self.c1 = np.array([-1, 1])
        if self.c2 is None:
            self.c2 = np.array([-1, 1])
        if self.c3 is None:
            self.c3 = np.array([-1, 1])
        if self.c4 is None:
            self.c4 = np.array([-1, 1])
        if self.d1 is None:
            self.d1 = np.array([-1, 1])
        if self.d2 is None:
            self.d2 = np.array([-1, 1])
        if self.d3 is None:
            self.d3 = np.array([-1, 1])
        if self.d4 is None:
            self.d4 = np.array([-1, 1])

        if self.long is None:
            self.long = np.array([[-np.pi, np.pi] for _ in range(nspots)])
        if self.lat is None:
            self.lat = np.array([[-np.pi/2, np.pi/2] for _ in range(nspots)])
        if self.alpha is None:
            self.alpha = np.array([[0, np.pi/4] for _ in range(nspots)])
        if self.fspot is None:
            self.fspot = np.array([[0, 2] for _ in range(nspots)])
        if self.tmax is None:
            self.tmax = np.array([[t[0], t[-1]] for _ in range(nspots)])
        if self.life is None:
            self.life = np.array([[0, t[-1]-t[0]] for _ in range(nspots)])
        if self.ingress is None:
            self.ingress = np.array([[0, t[-1]-t[0]] for _ in range(nspots)])
        if self.egress is None:
            self.egress = np.array([[0, t[-1]-t[0]] for _ in range(nspots)])

        self.mmax = np.size(self.tstart)
        if self.U is None:
            self.U = np.array([[.9, 1.1] for _ in range(self.mmax)])
        if self.B is None:
            self.B = np.array([[.9, 1.1] for _ in range(self.mmax)])

        self.fixed_params = {}
        self.fit_names = []
        for key, val in self.parameters.items():
            if key in self.star_pars.keys():
                if np.ndim(val) < 1:
                    self.fixed_params[key] = val
                else:
                    self.fit_names.append(key)
            else:
                if np.ndim(val) <= 1:
                    self.fixed_params[key] = val
                else:
                    self.fit_names.append(key)
        self.bounds = []
        for key in self.fit_names:
            bounds = getattr(self, key)
            if key in self.star_pars.keys():
                bounds = np.array([bounds])
            for b in bounds:
                self.bounds.append(b)
        self.bounds = np.array(self.bounds)

    @property
    def parameters(self):
        return {**self.star_pars, **self.spot_pars, **self.inst_pars}

    @property
    def star_pars(self):
        return dict(inc=self.inc, Peq=self.Peq, k2=self.k2, k4=self.k4, c1=self.c1, c2=self.c2,
                    c3=self.c3, c4=self.c4, d1=self.d1, d2=self.d2, d3=self.d3, d4=self.d4)

    @property
    def spot_pars(self):
        return dict(long=self.long, lat=self.lat, alpha=self.alpha, fspot=self.fspot,
                    tmax=self.tmax, life=self.life, ingress=self.ingress, egress=self.egress)

    @property
    def inst_pars(self):
        return dict(U=self.U, B=self.B)

    def predict(self, t, theta):
        fit_params = self.get_fit_params(theta)
        theta_full = {**fit_params, **self.fixed_params}
        theta_star = np.array([theta_full[key] for key in self.star_pars.keys()])
        theta_spot = np.array([theta_full[key] for key in self.spot_pars.keys()])
        theta_inst = np.array([theta_full[key] for key in self.inst_pars.keys()])
        yf = macula(t, theta_star, theta_spot, theta_inst, tstart=self.tstart, tend=self.tend)
        return yf

    def get_fit_params(self, theta):
        fit_params = {}
        i = 0
        for key in self.fit_names:
            if key in self.spot_pars.keys():
                fit_params[key] = np.array(theta[i:i + self.nspots])
                i += self.nspots
            elif key in self.inst_pars.keys():
                fit_params[key] = np.array(theta[i:i + self.mmax])
                i += self.mmax
            else:
                fit_params[key] = theta[i]
                i += 1
        return fit_params

    def lnprior(self, theta):
        fit_params = self.get_fit_params(theta)
        for key, val in fit_params.items():
            bounds = getattr(self, key)
            if key in self.star_pars.keys():
                if not (bounds[0] < val < bounds[1]):
                    return -np.inf
            else:
                for bound, v in zip(bounds, val):
                    if not (bound[0] < v < bound[1]):
                        return -np.inf
        return 0.

    def chi(self, theta):
        fit_params = self.get_fit_params(theta)
        theta_full = {**fit_params, **self.fixed_params}
        theta_star = np.array([theta_full[key] for key in self.star_pars.keys()])
        theta_spot = np.array([theta_full[key] for key in self.spot_pars.keys()])
        theta_inst = np.array([theta_full[key] for key in self.inst_pars.keys()])
        yf = macula(self.t, theta_star, theta_spot, theta_inst, tstart=self.tstart, tend=self.tend)
        sse = np.sum(np.square(yf - self.y)) / np.std(self.y) ** 2
        return sse

    def lnprob(self, theta):
        lp = self.lnprior(theta)
        if not np.isfinite(lp):
            return -np.inf
        return lp - self.chi(theta)

    def grad_chi(self, theta):
        fit_params = self.get_fit_params(theta)
        theta_full = {**fit_params, **self.fixed_params}
        theta_star = np.array([theta_full[key] for key in self.star_pars.keys()])
        theta_spot = np.array([theta_full[key] for key in self.spot_pars.keys()])
        theta_inst = np.array([theta_full[key] for key in self.inst_pars.keys()])
        result = macula(self.t, theta_star, theta_spot, theta_inst, tstart=self.tstart, tend=self.tend,
                        full_output=True, derivatives=True)
        yf = result[0]
        dy_star = result[1]
        dy_spot = result[2]
        dy_inst = result[3]<commit_msg>pondered chi^2 for each mth segment<commit_after>import time
from itertools import product

import emcee
import matplotlib.patches as mpatches
from scipy.cluster.vq import whiten, kmeans2
from scipy.optimize import leastsq, minimize
from tqdm.auto import tqdm
from pymultinest.solve import  solve

from .utils import *
from . import MPLSTYLE


class MaculaModeler(object):
    def __init__(self, t, y, nspots, dy=None, inc=None, Peq=None, k2=None, k4=None, c1=None, c2=None, c3=None, c4=None,
                 d1=None, d2=None, d3=None, d4=None, long=None, lat=None, alpha=None, fspot=None, tmax=None,
                 life=None, ingress=None, egress=None, U=None, B=None, tstart=None, tend=None):
        self.t = t
        self.y = y
        self.nspots = nspots
        self.dy = dy
        self.inc = inc
        self.Peq = Peq
        self.k2 = k2
        self.k4 = k4
        self.c1 = c1
        self.c2 = c2
        self.c3 = c3
        self.c4 = c4
        self.d1 = d1
        self.d2 = d2
        self.d3 = d3
        self.d4 = d4
        self.long = long
        self.lat = lat
        self.alpha = alpha
        self.fspot = fspot
        self.tmax = tmax
        self.life = life
        self.ingress = ingress
        self.egress = egress
        self.U = U
        self.B = B
        self.tstart = tstart
        self.tend = tend

        if self.dy is None:
            self.dy = np.ones_like(y)

        if self.inc is None:
            self.inc = np.array([0, np.pi/2])
        if self.Peq is None:
            self.Peq = np.array([0, 50])
        if self.k2 is None:
            self.k2 = np.array([-1, 1])
        if self.k4 is None:
            self.k4 = np.array([-1, 1])
        if self.c1 is None:
            self.c1 = np.array([-1, 1])
        if self.c2 is None:
            self.c2 = np.array([-1, 1])
        if self.c3 is None:
            self.c3 = np.array([-1, 1])
        if self.c4 is None:
            self.c4 = np.array([-1, 1])
        if self.d1 is None:
            self.d1 = np.array([-1, 1])
        if self.d2 is None:
            self.d2 = np.array([-1, 1])
        if self.d3 is None:
            self.d3 = np.array([-1, 1])
        if self.d4 is None:
            self.d4 = np.array([-1, 1])

        if self.long is None:
            self.long = np.array([[-np.pi, np.pi] for _ in range(nspots)])
        if self.lat is None:
            self.lat = np.array([[-np.pi/2, np.pi/2] for _ in range(nspots)])
        if self.alpha is None:
            self.alpha = np.array([[0, np.pi/4] for _ in range(nspots)])
        if self.fspot is None:
            self.fspot = np.array([[0, 2] for _ in range(nspots)])
        if self.tmax is None:
            self.tmax = np.array([[t[0], t[-1]] for _ in range(nspots)])
        if self.life is None:
            self.life = np.array([[0, t[-1]-t[0]] for _ in range(nspots)])
        if self.ingress is None:
            self.ingress = np.array([[0, t[-1]-t[0]] for _ in range(nspots)])
        if self.egress is None:
            self.egress = np.array([[0, t[-1]-t[0]] for _ in range(nspots)])

        self.mmax = np.size(self.tstart)
        if self.U is None:
            self.U = np.array([[.9, 1.1] for _ in range(self.mmax)])
        if self.B is None:
            self.B = np.array([[.9, 1.1] for _ in range(self.mmax)])

        self.fixed_params = {}
        self.fit_names = []
        for key, val in self.parameters.items():
            if key in self.star_pars.keys():
                if np.ndim(val) < 1:
                    self.fixed_params[key] = val
                else:
                    self.fit_names.append(key)
            else:
                if np.ndim(val) <= 1:
                    self.fixed_params[key] = val
                else:
                    self.fit_names.append(key)
        self.bounds = []
        for key in self.fit_names:
            bounds = getattr(self, key)
            if key in self.star_pars.keys():
                bounds = np.array([bounds])
            for b in bounds:
                self.bounds.append(b)
        self.bounds = np.array(self.bounds)

    @property
    def parameters(self):
        return {**self.star_pars, **self.spot_pars, **self.inst_pars}

    @property
    def star_pars(self):
        return dict(inc=self.inc, Peq=self.Peq, k2=self.k2, k4=self.k4, c1=self.c1, c2=self.c2,
                    c3=self.c3, c4=self.c4, d1=self.d1, d2=self.d2, d3=self.d3, d4=self.d4)

    @property
    def spot_pars(self):
        return dict(long=self.long, lat=self.lat, alpha=self.alpha, fspot=self.fspot,
                    tmax=self.tmax, life=self.life, ingress=self.ingress, egress=self.egress)

    @property
    def inst_pars(self):
        return dict(U=self.U, B=self.B)

    def predict(self, t, theta):
        fit_params = self.get_fit_params(theta)
        theta_full = {**fit_params, **self.fixed_params}
        theta_star = np.array([theta_full[key] for key in self.star_pars.keys()])
        theta_spot = np.array([theta_full[key] for key in self.spot_pars.keys()])
        theta_inst = np.array([theta_full[key] for key in self.inst_pars.keys()])
        yf = macula(t, theta_star, theta_spot, theta_inst, tstart=self.tstart, tend=self.tend)
        return yf

    def get_fit_params(self, theta):
        fit_params = {}
        i = 0
        for key in self.fit_names:
            if key in self.spot_pars.keys():
                fit_params[key] = np.array(theta[i:i + self.nspots])
                i += self.nspots
            elif key in self.inst_pars.keys():
                fit_params[key] = np.array(theta[i:i + self.mmax])
                i += self.mmax
            else:
                fit_params[key] = theta[i]
                i += 1
        return fit_params

    def lnprior(self, theta):
        fit_params = self.get_fit_params(theta)
        for key, val in fit_params.items():
            bounds = getattr(self, key)
            if key in self.star_pars.keys():
                if not (bounds[0] < val < bounds[1]):
                    return -np.inf
            else:
                for bound, v in zip(bounds, val):
                    if not (bound[0] < v < bound[1]):
                        return -np.inf
        return 0.

    def chi(self, theta):
        fit_params = self.get_fit_params(theta)
        theta_full = {**fit_params, **self.fixed_params}
        theta_star = np.array([theta_full[key] for key in self.star_pars.keys()])
        theta_spot = np.array([theta_full[key] for key in self.spot_pars.keys()])
        theta_inst = np.array([theta_full[key] for key in self.inst_pars.keys()])
        yf = macula(self.t, theta_star, theta_spot, theta_inst, tstart=self.tstart, tend=self.tend)
        eps = []
        for t1, t2 in zip(self.tstart, self.tend):
            mask = np.logical_and(t1 <= self.t, self.t <= t2)
            eps = np.append(eps, np.square(yf[mask] - self.y[mask]) / np.std(self.y[mask]) ** 2)
        sse = np.sum(eps)
        return sse

    def lnprob(self, theta):
        lp = self.lnprior(theta)
        if not np.isfinite(lp):
            return -np.inf
        return lp - self.chi(theta)

    def grad_chi(self, theta):
        fit_params = self.get_fit_params(theta)
        theta_full = {**fit_params, **self.fixed_params}
        theta_star = np.array([theta_full[key] for key in self.star_pars.keys()])
        theta_spot = np.array([theta_full[key] for key in self.spot_pars.keys()])
        theta_inst = np.array([theta_full[key] for key in self.inst_pars.keys()])
        result = macula(self.t, theta_star, theta_spot, theta_inst, tstart=self.tstart, tend=self.tend,
                        full_output=True, derivatives=True)
        yf = result[0]
        dy_star = result[1]
        dy_spot = result[2]
        dy_inst = result[3]
{'fraction_non_alphanumeric': 0.13063004425930747, 'fraction_numerical': 0.01763863577193439, 'mean_word_length': 1.7669728074914461, 'pattern_counts': {'":': 0, '<': 17, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2827968', 'n_tokens_mistral': 5537, 'n_tokens_neox': 5312, 'n_words': 1241}