<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

#
#  Copyright 2019 The FATE Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import numpy as np

from arch.api.utils import log_utils
from federatedml.linear_model.linear_model_weight import LinearModelWeights
from federatedml.util import consts

LOGGER = log_utils.getLogger()


class _Optimizer(object):
    def __init__(self, learning_rate, alpha, penalty, decay, decay_sqrt, mu):
        self.learning_rate = learning_rate
        self.iters = 0
        self.alpha = alpha
        self.penalty = penalty
        self.decay = decay
        self.decay_sqrt = decay_sqrt
        self.mu = mu

    def decay_learning_rate(self):
        if self.decay_sqrt:
            lr = self.learning_rate / np.sqrt(1 + self.decay * self.iters)
        else:
            lr = self.learning_rate / (1 + self.decay * self.iters)
        return lr

    @property
    def shrinkage_val(self):
        this_step_size = self.learning_rate / np.sqrt(self.iters)
        return self.alpha * this_step_size

    def set_iters(self, iters):
        self.iters = iters

    def apply_gradients(self, grad):
        raise NotImplementedError("Should not call here")

    def _l1_updator(self, model_weights: LinearModelWeights, gradient):
        coef_ = model_weights.coef_
        if model_weights.fit_intercept:
            gradient_without_intercept = gradient[: -1]
        else:
            gradient_without_intercept = gradient

        new_weights = np.sign(coef_ - gradient_without_intercept) * np.maximum(0, np.abs(
            coef_ - gradient_without_intercept) - self.shrinkage_val)

        if model_weights.fit_intercept:
            new_weights = np.append(new_weights, model_weights.intercept_)
            new_weights[-1] -= gradient[-1]
        new_param = LinearModelWeights(new_weights, model_weights.fit_intercept)
        LOGGER.debug("In _l1_updator, original weight: {}, new_weights: {}".format(
            model_weights.unboxed, new_weights
        ))
        return new_param

    def _l2_updator(self, lr_weights: LinearModelWeights, gradient):
        """
        For l2 regularization, the regular term has been added in gradients.
        """

        new_weights = lr_weights.unboxed - gradient
        new_param = LinearModelWeights(new_weights, lr_weights.fit_intercept)

        return new_param

    def add_regular_to_grad(self, grad, lr_weights):

        if self.penalty == consts.L2_PENALTY:
            if lr_weights.fit_intercept:
                gradient_without_intercept = grad[: -1]
                gradient_without_intercept += self.alpha * lr_weights.coef_
                new_grad = np.append(gradient_without_intercept, grad[-1])
            else:
                new_grad = grad + self.alpha * lr_weights.coef_
        else:
            new_grad = grad

        return new_grad

    def regularization_update(self, model_weights: LinearModelWeights, prev_round_weights: LinearModelWeights, grad):
        if self.penalty == consts.L1_PENALTY:
            model_weights = self._l1_updator(model_weights, grad)
        elif self.penalty == consts.L2_PENALTY:
            model_weights = self._l2_updator(model_weights, grad)
        else:
            new_vars = model_weights.unboxed - grad
            model_weights = LinearModelWeights(new_vars, model_weights.fit_intercept)

        if prev_round_weights is not None: # additional proximal term
            coef_ = model_weights.unboxed

            if model_weights.fit_intercept:
                coef_without_intercept = coef_[: -1] 
            else:
                coef_without_intercept = coef_
            
            LOGGER.debug("before applying additional proximal terms, weights {}".format(coef_without_intercept))
            coef_without_intercept -= self.mu * (model_weights.coef_ - prev_round_weights.coef_)
            LOGGER.debug("after applying additional proximal terms, new weights {}, with difference {}".format(coef_without_intercept,model_weights.coef_ - prev_round_weights.coef_))

            if model_weights.fit_intercept:
                new_coef_ = np.append(coef_without_intercept, coef_[-1])
            else:
                new_coef_ = coef_without_intercept
            
            model_weights = LinearModelWeights(new_coef_, model_weights.fit_intercept)
        return model_weights

    def __l1_loss_norm(self, model_weights: LinearModelWeights):
        coef_ = model_weights.coef_
        loss_norm = np.sum(self.alpha * np.abs(coef_))
        return loss_norm

    def __l2_loss_norm(self, model_weights: LinearModelWeights):
        coef_ = model_weights.coef_
        loss_norm = 0.5 * self.alpha * np.dot(coef_, coef_)
        return loss_norm

    def __add_proximal(self, model_weights, prev_round_weights):
        prev_round_coef_ = prev_round_weights.coef_
        coef_ = model_weights.coef_
        diff = coef_ - prev_round_coef_
        loss_norm = self.mu * 0.5 * np.dot(diff, diff)
        return loss_norm

    def loss_norm(self, model_weights: LinearModelWeights, prev_round_weights: LinearModelWeights = None):
        proximal_term = None

        if prev_round_weights is not None:
            proximal_term = self.__add_proximal(model_weights, prev_round_weights)

        if self.penalty == consts.L1_PENALTY:
            loss_norm_value = self.__l1_loss_norm(model_weights)
        elif self.penalty == consts.L2_PENALTY:
            loss_norm_value = self.__l2_loss_norm(model_weights)
        else:
            loss_norm_value = None

        # additional proximal term
        if loss_norm_value is None: 
            loss_norm_value = proximal_term
        elif proximal_term is not None: 
            loss_norm_value += proximal_term
        return loss_norm_value

    def hess_vector_norm(self, delta_s: LinearModelWeights):
        if self.penalty == consts.L1_PENALTY:
            return LinearModelWeights(np.zeros_like(delta_s.unboxed), fit_intercept=delta_s.fit_intercept)
        elif self.penalty == consts.L2_PENALTY:
            return LinearModelWeights(self.alpha * np.array(delta_s.unboxed), fit_intercept=delta_s.fit_intercept)
        else:
            return LinearModelWeights(np.zeros_like(delta_s.unboxed), fit_intercept=delta_s.fit_intercept)

    def update_model(self, model_weights: LinearModelWeights, grad, prev_round_weights: LinearModelWeights = None, has_applied=True):

        if not has_applied:
            grad = self.add_regular_to_grad(grad, model_weights)
            delta_grad = self.apply_gradients(grad)
        else:
            delta_grad = grad
        model_weights = self.regularization_update(model_weights, prev_round_weights, delta_grad)
        return model_weights


class _SgdOptimizer(_Optimizer):
    def apply_gradients(self, grad):
        learning_rate = self.decay_learning_rate()

        delta_grad = learning_rate * grad
        LOGGER.debug("In sgd optimizer, learning_rate: {}, delta_grad: {}".format(learning_rate, delta_grad))

        return delta_grad


class _RMSPropOptimizer(_Optimizer):
    def __init__(self, learning_rate, alpha, penalty, decay, decay_sqrt, mu):
        super().__init__(learning_rate, alpha, penalty, decay, decay_sqrt)<commit_msg>fix mu<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

#
#  Copyright 2019 The FATE Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import numpy as np

from arch.api.utils import log_utils
from federatedml.linear_model.linear_model_weight import LinearModelWeights
from federatedml.util import consts

LOGGER = log_utils.getLogger()


class _Optimizer(object):
    def __init__(self, learning_rate, alpha, penalty, decay, decay_sqrt, mu = 0):
        self.learning_rate = learning_rate
        self.iters = 0
        self.alpha = alpha
        self.penalty = penalty
        self.decay = decay
        self.decay_sqrt = decay_sqrt
        self.mu = mu

    def decay_learning_rate(self):
        if self.decay_sqrt:
            lr = self.learning_rate / np.sqrt(1 + self.decay * self.iters)
        else:
            lr = self.learning_rate / (1 + self.decay * self.iters)
        return lr

    @property
    def shrinkage_val(self):
        this_step_size = self.learning_rate / np.sqrt(self.iters)
        return self.alpha * this_step_size

    def set_iters(self, iters):
        self.iters = iters

    def apply_gradients(self, grad):
        raise NotImplementedError("Should not call here")

    def _l1_updator(self, model_weights: LinearModelWeights, gradient):
        coef_ = model_weights.coef_
        if model_weights.fit_intercept:
            gradient_without_intercept = gradient[: -1]
        else:
            gradient_without_intercept = gradient

        new_weights = np.sign(coef_ - gradient_without_intercept) * np.maximum(0, np.abs(
            coef_ - gradient_without_intercept) - self.shrinkage_val)

        if model_weights.fit_intercept:
            new_weights = np.append(new_weights, model_weights.intercept_)
            new_weights[-1] -= gradient[-1]
        new_param = LinearModelWeights(new_weights, model_weights.fit_intercept)
        LOGGER.debug("In _l1_updator, original weight: {}, new_weights: {}".format(
            model_weights.unboxed, new_weights
        ))
        return new_param

    def _l2_updator(self, lr_weights: LinearModelWeights, gradient):
        """
        For l2 regularization, the regular term has been added in gradients.
        """

        new_weights = lr_weights.unboxed - gradient
        new_param = LinearModelWeights(new_weights, lr_weights.fit_intercept)

        return new_param

    def add_regular_to_grad(self, grad, lr_weights):

        if self.penalty == consts.L2_PENALTY:
            if lr_weights.fit_intercept:
                gradient_without_intercept = grad[: -1]
                gradient_without_intercept += self.alpha * lr_weights.coef_
                new_grad = np.append(gradient_without_intercept, grad[-1])
            else:
                new_grad = grad + self.alpha * lr_weights.coef_
        else:
            new_grad = grad

        return new_grad

    def regularization_update(self, model_weights: LinearModelWeights, grad, prev_round_weights: LinearModelWeights = None):
        if self.penalty == consts.L1_PENALTY:
            model_weights = self._l1_updator(model_weights, grad)
        elif self.penalty == consts.L2_PENALTY:
            model_weights = self._l2_updator(model_weights, grad)
        else:
            new_vars = model_weights.unboxed - grad
            model_weights = LinearModelWeights(new_vars, model_weights.fit_intercept)

        if prev_round_weights is not None: # additional proximal term
            coef_ = model_weights.unboxed

            if model_weights.fit_intercept:
                coef_without_intercept = coef_[: -1] 
            else:
                coef_without_intercept = coef_
            
            LOGGER.debug("before applying additional proximal terms, weights {}".format(coef_without_intercept))
            coef_without_intercept -= self.mu * (model_weights.coef_ - prev_round_weights.coef_)
            LOGGER.debug("after applying additional proximal terms, new weights {}, with difference {}".format(coef_without_intercept,model_weights.coef_ - prev_round_weights.coef_))

            if model_weights.fit_intercept:
                new_coef_ = np.append(coef_without_intercept, coef_[-1])
            else:
                new_coef_ = coef_without_intercept
            
            model_weights = LinearModelWeights(new_coef_, model_weights.fit_intercept)
        return model_weights

    def __l1_loss_norm(self, model_weights: LinearModelWeights):
        coef_ = model_weights.coef_
        loss_norm = np.sum(self.alpha * np.abs(coef_))
        return loss_norm

    def __l2_loss_norm(self, model_weights: LinearModelWeights):
        coef_ = model_weights.coef_
        loss_norm = 0.5 * self.alpha * np.dot(coef_, coef_)
        return loss_norm

    def __add_proximal(self, model_weights, prev_round_weights):
        prev_round_coef_ = prev_round_weights.coef_
        coef_ = model_weights.coef_
        diff = coef_ - prev_round_coef_
        loss_norm = self.mu * 0.5 * np.dot(diff, diff)
        return loss_norm

    def loss_norm(self, model_weights: LinearModelWeights, prev_round_weights: LinearModelWeights = None):
        proximal_term = None

        if prev_round_weights is not None:
            proximal_term = self.__add_proximal(model_weights, prev_round_weights)

        if self.penalty == consts.L1_PENALTY:
            loss_norm_value = self.__l1_loss_norm(model_weights)
        elif self.penalty == consts.L2_PENALTY:
            loss_norm_value = self.__l2_loss_norm(model_weights)
        else:
            loss_norm_value = None

        # additional proximal term
        if loss_norm_value is None: 
            loss_norm_value = proximal_term
        elif proximal_term is not None: 
            loss_norm_value += proximal_term
        return loss_norm_value

    def hess_vector_norm(self, delta_s: LinearModelWeights):
        if self.penalty == consts.L1_PENALTY:
            return LinearModelWeights(np.zeros_like(delta_s.unboxed), fit_intercept=delta_s.fit_intercept)
        elif self.penalty == consts.L2_PENALTY:
            return LinearModelWeights(self.alpha * np.array(delta_s.unboxed), fit_intercept=delta_s.fit_intercept)
        else:
            return LinearModelWeights(np.zeros_like(delta_s.unboxed), fit_intercept=delta_s.fit_intercept)

    def update_model(self, model_weights: LinearModelWeights, grad, prev_round_weights: LinearModelWeights = None, has_applied=True):

        if not has_applied:
            grad = self.add_regular_to_grad(grad, model_weights)
            delta_grad = self.apply_gradients(grad)
        else:
            delta_grad = grad
        model_weights = self.regularization_update(model_weights, delta_grad, prev_round_weights)
        return model_weights


class _SgdOptimizer(_Optimizer):
    def apply_gradients(self, grad):
        learning_rate = self.decay_learning_rate()

        delta_grad = learning_rate * grad
        LOGGER.debug("In sgd optimizer, learning_rate: {}, delta_grad: {}".format(learning_rate, delta_grad))

        return delta_grad


class _RMSPropOptimizer(_Optimizer):
    def __init__(self, learning_rate, alpha, penalty, decay, decay_sqrt, mu):
        super().__init__(learning_rate, alpha, penalty, decay, decay_sqrt)
{'fraction_non_alphanumeric': 0.07717041800643087, 'fraction_numerical': 0.005337620578778135, 'mean_word_length': 2.7763477416221467, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '4772287', 'n_tokens_mistral': 4891, 'n_tokens_neox': 4694, 'n_words': 1034}