<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

"""
Poisson point process penalised likelihood regression.
"""

try:
    import autograd
    import autograd.numpy as np
    # import autograd.scipy as sp
    have_autograd = True
except ImportError as e:
    import numpy as np
    # import scipy as sp
    have_autograd = False

from scipy.optimize import minimize

from . import influence
from . import background
from . import model


class ContinuousExact(object):
    def __init__(
            self,
            debug=False
            ):
        self.debug = debug

    def _debug_print(self, *args, **kwargs):
        if self.debug:
            print(*args, **kwargs)

    def _pack(
            self,
            mu=1.0,
            kappa=0.0,
            tau=None,
            omega=None
            ):
        """
        pack all params into one vector for blackbox optimisation;
        currently unused
        """

        packed = np.zeros(
            1 +
            self.n_phi_bases +
            self.n_tau +
            self.n_omega
        )
        packed[0] = mu
        packed[
            1:self.n_phi_bases + 1
        ] = kappa
        packed[
            self.n_phi_bases + 1:
            self.n_phi_bases + self.n_tau + 1
        ] = tau if tau is not None else 1.0
        packed[
            self.n_phi_bases + self.n_tau + 1:
        ] = omega if omega is not None else 0.0
        return packed

    def _unpack(
            self,
            packed):
        """
        unpack all params from one vector for blackbox optimisation;
        currently unused.
        """
        unpacked = dict(
            mu=packed[0],
            kappa=packed[1:self.n_phi_bases+1],
        )
        # are we fitting time parameters?
        if self.n_tau > 0:
            unpacked = packed[
                self.n_phi_bases + 1:
                self.n_phi_bases + self.n_tau + 1
            ]
        # omega parameters - exogenous variation
        if self.n_omega > 0:
            unpacked = packed[
                self.n_phi_bases + self.n_tau + 1:
            ]
        return unpacked

    def negloglik(
            self,
            **kwargs):

        return -model.loglik(
            ts=self._ts,
            eval_ts=self._eval_ts,
            phi_kernel=self.phi_kernel,
            mu_kernel=self.mu_kernel,
            **kwargs
        )

    def _penalty_weight_packed(
            self,
            pi_kappa=0.0,
            pi_omega=1e-8,
            **kwargs
            ):
        return self._pack(
            mu=0.0,
            kappa=pi_kappa,
            omega=pi_omega
        )  # / self.penalty_scale

    def penalty(
            self,
            mu=1.0,
            kappa=0.0,
            tau=1.0,
            omega=0.0,
            pi_kappa=0.0,
            pi_omega=0.0,
            **kwargs
            ):
        # 2 different l_1 penalties
        return np.sum(
            np.abs(kappa) * pi_kappa
        ) + np.sum(
            np.abs(omega) * pi_omega
        )

    def objective(
            self,
            **kwargs):

        g_negloglik = self.negloglik(
            **kwargs
        )
        g_penalty = self.penalty(
            **kwargs
        )
        penalty_dominated = np.abs(
            g_negloglik
        ) < (
            g_penalty
        )
        return (g_negloglik + g_penalty) * penalty_dominated

    def _setup_graphs(
            self,
            ts,
            phi_kernel=None,
            mu_kernel=None,
            n_phi_bases=1,
            n_mu_bases=0,
            fit_tau=False,
            fit_omega=False,
            t_start=0.0,
            t_end=None,
            param_vector=None,
            **kwargs):

        self._t_start = t_start
        self._t_end = t_end or ts[-1]
        self._ts = ts
        self._n_ts = np.sum(ts[
            np.logical_and(
                ts >= self._t_start,
                ts < self._t_end
            )
        ])
        # Full data likelihood must evaluate at the t_end also
        if ts[-1] < self._t_end:
            _eval_ts = np.append(
                ts[ts > self._t_start],
                [self._t_end]
            )
        else:
            _eval_ts = ts[
                np.logical_and(
                    ts >= self._t_start,
                    ts < self._t_end
                )
            ]

        self._eval_ts = _eval_ts

        if phi_kernel is None:
            phi_kernel = influence.MaxwellKernel(n_bases=5)
        else:
            phi_kernel = influence.as_influence_kernel(
                phi_kernel, n_bases=n_phi_bases
            )
        self.phi_kernel = phi_kernel
        self.n_phi_bases = phi_kernel.n_bases
        if mu_kernel is None:
            if n_mu_bases > 0:
                mu_kernel = background.LinearStepKernel(
                    t_start=self._t_start,
                    t_end=self._t_end,
                    n_bases=n_mu_bases)
            else:
                mu_kernel = background.ConstKernel()
        self.mu_kernel = mu_kernel
        self.n_mu_bases = n_mu_bases

        self._fit_tau = fit_tau
        self._fit_omega = fit_omega
        self.n_tau = self.n_phi_bases * self._fit_tau
        self.n_omega = self.n_mu_bases * self._fit_omega

    def _teardown_graphs(self):

        del(self._ts)
        del(self._eval_ts)

    def fit(
            self,
            ts,
            phi_kernel=None,
            mu_kernel=None,
            n_phi_bases=1,
            n_mu_bases=0,
            fit_tau=False,
            fit_omega=False,
            t_start=0.0,
            t_end=None,
            **kwargs
            ):
        self._setup_graphs(
            ts=ts,
            phi_kernel=phi_kernel,
            mu_kernel=mu_kernel,
            n_phi_bases=n_phi_bases,
            n_mu_bases=n_mu_bases,
            fit_tau=fit_tau,
            fit_omega=fit_omega,
            **kwargs
        )
        return self._fit(
            **kwargs
        )

    def _fit(
            self,
            mu=None,  # ignored?
            kappa=None,
            tau=None,
            omega=None,
            pi_kappa=0.0,
            pi_omega=1e-8,
            max_steps=10,
            step_iter=10,
            eps=1e-8,
            **kwargs
            ):
        """
        fit by Truncated Newton in each coordinate group
        """
        if tau is None:
            tau = self.phi_kernel.get_param('tau')
        if omega is None:
            omega = self.mu_kernel.get_param('kappa', 0.0)
        if kappa is None:
            kappa = self.phi_kernel.get_param('kappa')
        if mu is None:
            mu = self.mu_kernel.get_param(
                'kappa',
                self._n_ts/(self._t_end - self._t_start)
            )

        # def obj_mu(mu):
        #     return self.objective(
        #         mu=mu,
        #         kappa=kappa,
        #         tau=tau,
        #         omega=omega,
        #         pi_kappa=pi_kappa,
        #         pi_omega=pi_omega)
        #
        # grad_mu = autograd.value_and_grad(
        #     obj_mu, 0)

        def obj_kappa(kappa):
            return self.objective(
                mu=mu,
                kappa=kappa,
                tau=tau,
                omega=omega,
                pi_kappa=pi_kappa,
                pi_omega=pi_omega)

        grad_kappa = autograd.value_and_grad(
            obj_kappa, 0)
        kappa_bounds = [(0, 1)] * self.n_phi_bases

        def obj_tau(tau):
            return self.objective(
                mu=mu,
                kappa=kappa,
                tau=tau,
                omega=omega,
                pi_kappa=pi_kappa,
                pi_omega=pi_omega)

        grad_tau = autograd.value_and_grad(
            obj_tau, 0)
        tau_bounds = [(0, None)] * self.n_tau

        def obj_omega(mu):
            return self.objective(
                mu=mu,
                kappa=kappa,
                tau=tau,
                omega=omega,
                pi_kappa=pi_kappa,
                pi_omega=pi_omega)

        grad_omega = autograd.value_and_grad(
            obj_omega, 0)

        # self._hessian_diag_negloglik = autograd.elementwise_grad(
        #     self._grad_negloglik, 0)

        for i in range(max_steps):
            fit = {}
            if self._fit_tau:
                res = minimize(
                    grad_tau,
                    tau,
                    method='L-BFGS-B',  # ?
                    jac=True,
                    bounds=tau_bounds,
                    callback=lambda x: self._debug_print('tau', x),
                    options=dict(
                        maxiter=step_iter,
                        disp=self.debug
                    )
                )
                tau = res.x
                fit['tau'] = tau
            res = minimize(
                grad_kappa,
                kappa,
                method='TNC',
                jac=True,
                bounds=kappa_bounds,
                callback=lambda x: self._debug_print('kappa', x),
                options=dict(
                    maxiter=step_iter,
                    disp=self.debug
                )
            )
            kappa = res.x
            fit['kappa'] = kappa

            if self._fit_omega:
                omega_bounds = [(0, None)] * self.n_omega
                res = minimize(
                    grad_omega,
                    omega,
                    method='TNC',
                    jac=True,
                    bounds=omega_bounds,
                    callback=lambda x: self._debug_print('omega', x),
                    options=dict(
                        maxiter=step_iter,
                        disp=self.debug
                    )
                )
                omega = res.x
                fit['omega'] = omega

            # one-step mu update
            big_lam = model.big_lam_hawkes(
                mu=0.0,
                ts=self._ts,
                eval_ts=self._eval_ts,
                phi_kernel=self.phi_kernel,
                mu_kernel=self.mu_kernel,
                **fit,
                **kwargs
            )
            mu = (self._n_ts - big_lam)/(self._t_end-self._t_start)
            self._debug_print('mu', mu)
            fit['mu'] = mu

        return fit

    # def _fit_path(
    #         self,
    #         param_vector=None,
    #         pi_kappa=0.0,
    #         pi_omega=1e-8,
    #         max_steps=50,
    #         step_iter=20,
    #         lr=0.1,
    #         gamma=0.9,
    #         eps=1e-8,
    #         backoff=0.75,
    #         warm_start=False,
    #         **kwargs
    #         ):
    #
    #     n_params = param_vector.size
    #     param_path = np.zeros((n_params, max_steps))
    #     pi_kappa_path = np.zeros(max_steps)
    #     pi_omega_path = np.zeros(max_steps)
    #     loglik_path = np.zeros(max_steps)
    #     dof_path = np.zeros(max_steps)
    #     aic_path = np.zeros(max_steps)<commit_msg>fix one-step mu fit<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

"""
Poisson point process penalised likelihood regression.
"""

try:
    import autograd
    import autograd.numpy as np
    # import autograd.scipy as sp
    have_autograd = True
except ImportError as e:
    import numpy as np
    # import scipy as sp
    have_autograd = False

from scipy.optimize import minimize

from . import influence
from . import background
from . import model


class ContinuousExact(object):
    def __init__(
            self,
            debug=False
            ):
        self.debug = debug

    def _debug_print(self, *args, **kwargs):
        if self.debug:
            print(*args, **kwargs)

    def _debug_tee(self, label, obj):
        if self.debug:
            print(label, obj)
        return obj

    def _pack(
            self,
            mu=1.0,
            kappa=0.0,
            tau=None,
            omega=None
            ):
        """
        pack all params into one vector for blackbox optimisation;
        currently unused
        """

        packed = np.zeros(
            1 +
            self.n_phi_bases +
            self.n_tau +
            self.n_omega
        )
        packed[0] = mu
        packed[
            1:self.n_phi_bases + 1
        ] = kappa
        packed[
            self.n_phi_bases + 1:
            self.n_phi_bases + self.n_tau + 1
        ] = tau if tau is not None else 1.0
        packed[
            self.n_phi_bases + self.n_tau + 1:
        ] = omega if omega is not None else 0.0
        return packed

    def _unpack(
            self,
            packed):
        """
        unpack all params from one vector for blackbox optimisation;
        currently unused.
        """
        unpacked = dict(
            mu=packed[0],
            kappa=packed[1:self.n_phi_bases+1],
        )
        # are we fitting time parameters?
        if self.n_tau > 0:
            unpacked = packed[
                self.n_phi_bases + 1:
                self.n_phi_bases + self.n_tau + 1
            ]
        # omega parameters - exogenous variation
        if self.n_omega > 0:
            unpacked = packed[
                self.n_phi_bases + self.n_tau + 1:
            ]
        return unpacked

    def negloglik(
            self,
            **kwargs):

        return -model.loglik(
            ts=self._ts,
            eval_ts=self._eval_ts,
            phi_kernel=self.phi_kernel,
            mu_kernel=self.mu_kernel,
            **kwargs
        )

    def _penalty_weight_packed(
            self,
            pi_kappa=0.0,
            pi_omega=1e-8,
            **kwargs
            ):
        return self._pack(
            mu=0.0,
            kappa=pi_kappa,
            omega=pi_omega
        )  # / self.penalty_scale

    def penalty(
            self,
            mu=1.0,
            kappa=0.0,
            tau=1.0,
            omega=0.0,
            pi_kappa=0.0,
            pi_omega=0.0,
            **kwargs
            ):
        # 2 different l_1 penalties
        return np.sum(
            np.abs(kappa) * pi_kappa
        ) + np.sum(
            np.abs(omega) * pi_omega
        )

    def objective(
            self,
            **kwargs):

        g_negloglik = self.negloglik(
            **kwargs
        )
        g_penalty = self.penalty(
            **kwargs
        )
        penalty_dominated = np.abs(
            g_negloglik
        ) < (
            g_penalty
        )
        return (g_negloglik + g_penalty) * penalty_dominated

    def _setup_graphs(
            self,
            ts,
            phi_kernel=None,
            mu_kernel=None,
            n_phi_bases=1,
            n_mu_bases=0,
            fit_tau=False,
            fit_omega=False,
            t_start=0.0,
            t_end=None,
            param_vector=None,
            **kwargs):

        self._t_start = t_start
        self._t_end = t_end or ts[-1]
        self._ts = ts
        self._n_ts = np.sum(ts[
            np.logical_and(
                ts >= self._t_start,
                ts < self._t_end
            )
        ])
        # Full data likelihood must evaluate at the t_end also
        if ts[-1] < self._t_end:
            _eval_ts = np.append(
                ts[ts > self._t_start],
                [self._t_end]
            )
        else:
            _eval_ts = ts[
                np.logical_and(
                    ts >= self._t_start,
                    ts < self._t_end
                )
            ]

        self._eval_ts = _eval_ts

        if phi_kernel is None:
            phi_kernel = influence.MaxwellKernel(n_bases=5)
        else:
            phi_kernel = influence.as_influence_kernel(
                phi_kernel, n_bases=n_phi_bases
            )
        self.phi_kernel = phi_kernel
        self.n_phi_bases = phi_kernel.n_bases
        if mu_kernel is None:
            if n_mu_bases > 0:
                mu_kernel = background.AdditiveStepKernel(
                    t_start=self._t_start,
                    t_end=self._t_end,
                    n_bases=n_mu_bases)
            else:
                mu_kernel = background.ConstKernel()
        self.mu_kernel = mu_kernel
        self.n_mu_bases = n_mu_bases

        self._fit_tau = fit_tau
        self._fit_omega = fit_omega
        self.n_tau = self.n_phi_bases * self._fit_tau
        self.n_omega = self.n_mu_bases * self._fit_omega

    def _teardown_graphs(self):

        del(self._ts)
        del(self._eval_ts)

    def fit(
            self,
            ts,
            phi_kernel=None,
            mu_kernel=None,
            n_phi_bases=1,
            n_mu_bases=0,
            fit_tau=False,
            fit_omega=False,
            t_start=0.0,
            t_end=None,
            **kwargs
            ):
        self._setup_graphs(
            ts=ts,
            phi_kernel=phi_kernel,
            mu_kernel=mu_kernel,
            n_phi_bases=n_phi_bases,
            n_mu_bases=n_mu_bases,
            fit_tau=fit_tau,
            fit_omega=fit_omega,
            **kwargs
        )
        return self._fit(
            **kwargs
        )

    def _fit(
            self,
            mu=None,  # ignored?
            kappa=None,
            tau=None,
            omega=None,
            pi_kappa=0.0,
            pi_omega=1e-8,
            max_steps=3,
            step_iter=15,
            eps=1e-8,
            **kwargs
            ):
        """
        fit by truncated Newton in each coordinate group
        """
        if tau is None:
            tau = self.phi_kernel.get_param('tau')
        if omega is None:
            omega = self.mu_kernel.get_param('kappa', 0.0)
        if kappa is None:
            kappa = self.phi_kernel.get_param('kappa')
        if mu is None:
            mu = self.mu_kernel.get_param(
                'kappa',
                self._n_ts/(self._t_end - self._t_start)
            )

        # def obj_mu(mu):
        #     return self.objective(
        #         mu=mu,
        #         kappa=kappa,
        #         tau=tau,
        #         omega=omega,
        #         pi_kappa=pi_kappa,
        #         pi_omega=pi_omega)
        #
        # grad_mu = autograd.value_and_grad(
        #     obj_mu, 0)

        def obj_kappa(kappa):
            self._debug_tee('kappa', kappa)
            return self._debug_tee('kappa_obj', self.objective(
                mu=mu,
                kappa=kappa,
                tau=tau,
                omega=omega,
                pi_kappa=pi_kappa,
                pi_omega=pi_omega))

        grad_kappa = autograd.value_and_grad(
            obj_kappa, 0)
        kappa_bounds = [(0, 1)] * self.n_phi_bases

        def obj_tau(tau):
            self._debug_tee('tau', tau)
            return self._debug_tee('tau_obj', self.objective(
                mu=mu,
                kappa=kappa,
                tau=tau,
                omega=omega,
                pi_kappa=pi_kappa,
                pi_omega=pi_omega))

        grad_tau = autograd.value_and_grad(
            obj_tau, 0)
        tau_bounds = [(0, None)] * self.n_tau

        def obj_omega(omega):
            self._debug_tee('omega', omega)
            return self._debug_tee('omega_obj', self.objective(
                mu=mu,
                kappa=kappa,
                tau=tau,
                omega=omega,
                pi_kappa=pi_kappa,
                pi_omega=pi_omega))

        grad_omega = autograd.value_and_grad(
            obj_omega, 0)

        # self._hessian_diag_negloglik = autograd.elementwise_grad(
        #     self._grad_negloglik, 0)

        for i in range(max_steps):
            fit = {}
            if self._fit_tau:
                res = minimize(
                    grad_tau,
                    tau,
                    method='L-BFGS-B',  # ?
                    jac=True,
                    bounds=tau_bounds,
                    callback=lambda x: self._debug_print('tau_fit', x),
                    options=dict(
                        maxiter=step_iter,
                        disp=self.debug
                    )
                )
                tau = res.x
                fit['tau'] = tau
            res = minimize(
                grad_kappa,
                kappa,
                method='TNC',
                jac=True,
                bounds=kappa_bounds,
                callback=lambda x: self._debug_print('kappa_fit', x),
                options=dict(
                    maxiter=step_iter,
                    disp=self.debug
                )
            )
            kappa = res.x
            fit['kappa'] = kappa

            if self._fit_omega:
                omega_bounds = [(0, None)] * self.n_omega
                res = minimize(
                    grad_omega,
                    omega,
                    method='TNC',
                    jac=True,
                    bounds=omega_bounds,
                    callback=lambda x: self._debug_print('omega_fit', x),
                    options=dict(
                        maxiter=step_iter,
                        disp=self.debug
                    )
                )
                omega = res.x
                fit['omega'] = omega

            # one-step mu update is possible for known noise structure
            # e.g. additive but not in general
            big_lam = model.big_lam_hawkes(
                mu=0.0,
                ts=self._ts,
                eval_ts=np.array([self._t_end, self._t_end]),
                phi_kernel=self.phi_kernel,
                mu_kernel=self.mu_kernel,
                **fit,
                **kwargs
            )
            big_lam_inc = big_lam[1] - big_lam[0]

            mu = (self._n_ts - big_lam_inc)/(self._t_end-self._t_start)
            self._debug_print('mu_fit', mu)
            fit['mu'] = mu

        return fit

    # def _fit_path(
    #         self,
    #         param_vector=None,
    #         pi_kappa=0.0,
    #         pi_omega=1e-8,
    #         max_steps=50,
    #         step_iter=20,
    #         lr=0.1,
    #         gamma=0.9,
    #         eps=1e-8,
    #         backoff=0.75,
    #         warm_start=False,
    #         **kwargs
    #         ):
    #
    #     n_params = param_vector.size
    #     param_path = np.zeros((n_params, max_steps))
    #     pi_kappa_path = np.zeros(max_steps)
    #     pi_omega_path = np.zeros(max_steps)
    #     loglik_path = np.zeros(max_steps)
    #     dof_path = np.zeros(max_steps)
    #     aic_path = np.zeros(max_steps)
{'fraction_non_alphanumeric': 0.08618409451595839, 'fraction_numerical': 0.008508199612061364, 'mean_word_length': 1.1167304282914994, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 15, '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': '1333983', 'n_tokens_mistral': 7133, 'n_tokens_neox': 6678, 'n_words': 1255}