# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_lagrangians.ipynb.

# %% auto 0
__all__ = ['lagr_identity', 'lagr_repu', 'lagr_relu', 'lagr_softmax', 'lagr_exp', 'lagr_rexp', 'lagr_tanh', 'lagr_sigmoid',
           'lagr_layernorm', 'lagr_spherical_norm', 'LIdentity', 'LRepu', 'LRelu', 'LSigmoid', 'LSoftmax', 'LExp',
           'LRexp', 'LTanh', 'LLayerNorm', 'LSphericalNorm']

# %% ../nbs/00_lagrangians.ipynb 3
import jax.numpy as jnp
import jax
import numpy as np
from fastcore.test import *
import functools as ft
from typing import *
import treex as tx
from dataclasses import dataclass
from typing import *

# %% ../nbs/00_lagrangians.ipynb 7
def lagr_identity(x): 
    """The Lagrangian whose activation function is simply the identity."""
    return 1 / 2 * jnp.power(x, 2)

# %% ../nbs/00_lagrangians.ipynb 9
def lagr_repu(x, 
              n): # Degree of the polynomial in the power unit
    """Rectified Power Unit of degree `n`"""
    return 1 / n * jnp.power(jnp.maximum(x, 0), n)

# %% ../nbs/00_lagrangians.ipynb 12
def lagr_relu(x):
    """Rectified Linear Unit. Same as repu of degree 2"""
    return lagr_repu(x, 2)

# %% ../nbs/00_lagrangians.ipynb 15
def lagr_softmax(x,
                 beta:float=1.0, # Inverse temperature
                 axis:int=-1): # Dimension over which to apply logsumexp
    """The lagrangian of the softmax -- the logsumexp"""
    return (1/beta * jax.nn.logsumexp(beta * x, axis=axis, keepdims=True))

# %% ../nbs/00_lagrangians.ipynb 18
def lagr_exp(x, 
             beta:float=1.0): # Inverse temperature
    """Exponential activation function, as in [Demicirgil et al.](https://arxiv.org/abs/1702.01929). Operates elementwise"""
    return 1 / beta * jnp.exp(beta * x)

# %% ../nbs/00_lagrangians.ipynb 21
def lagr_rexp(x, 
             beta:float=1.0): # Inverse temperature
    """Rectified exponential activation function"""
    xclipped = jnp.maximum(x, 0)
    return (1 / beta * jnp.exp(beta * xclipped)-xclipped)

# %% ../nbs/00_lagrangians.ipynb 25
@jax.custom_jvp
def _lagr_tanh(x, beta=1.0):
    return 1 / beta * jnp.log(jnp.cosh(beta * x))

@_lagr_tanh.defjvp
def _lagr_tanh_defjvp(primals, tangents):
    x, beta = primals
    x_dot, beta_dot = tangents
    primal_out = _lagr_tanh(x, beta)
    tangent_out = jnp.tanh(beta * x) * x_dot
    return primal_out, tangent_out

def lagr_tanh(x, 
              beta=1.0): # Inverse temperature
    """Lagrangian of the tanh activation function"""
    return _lagr_tanh(x, beta)

# %% ../nbs/00_lagrangians.ipynb 29
@jax.custom_jvp
def _lagr_sigmoid(x, 
                  beta=1.0, # Inverse temperature
                  scale=1.0): # Amount to stretch the range of the sigmoid's lagrangian
    """The lagrangian of a sigmoid that we can define custom JVPs of"""
    return scale / beta * jnp.log(jnp.exp(beta * x) + 1)

def _tempered_sigmoid(x, 
                     beta=1.0, # Inverse temperature
                     scale=1.0): # Amount to stretch the range of the sigmoid
    """The basic sigmoid, but with a scaling factor"""
    return scale / (1 + jnp.exp(-beta * x))

@_lagr_sigmoid.defjvp
def _lagr_sigmoid_jvp(primals, tangents):
    x, beta, scale = primals
    x_dot, beta_dot, scale_dot = tangents
    primal_out = _lagr_sigmoid(x, beta, scale)
    tangent_out = _tempered_sigmoid(x, beta=beta, scale=scale) * x_dot # Manually defined sigmoid
    return primal_out, tangent_out

def lagr_sigmoid(x, 
                 beta=1.0, # Inverse temperature
                 scale=1.0): # Amount to stretch the range of the sigmoid's lagrangian
    """The lagrangian of the sigmoid activation function"""
    return _lagr_sigmoid(x, beta=beta, scale=scale)

# %% ../nbs/00_lagrangians.ipynb 32
def _simple_layernorm(x:jnp.ndarray, 
                   gamma:float=1.0, # Scale the stdev
                   delta:Union[float, jnp.ndarray]=0., # Shift the mean
                   axis=-1, # Which axis to normalize
                   eps=1e-5, # Prevent division by 0
                  ): 
    """Layer norm activation function"""
    xmean = x.mean(axis, keepdims=True)
    xmeaned = x - xmean
    denominator = jnp.sqrt(jnp.power(xmeaned, 2).mean(axis, keepdims=True) + eps)
    return gamma * xmeaned / denominator + delta

def lagr_layernorm(x:jnp.ndarray, 
                   gamma:float=1.0, # Scale the stdev
                   delta:Union[float, jnp.ndarray]=0., # Shift the mean
                   axis=-1, # Which axis to normalize
                   eps=1e-5, # Prevent division by 0
                  ): 
    """Lagrangian of the layer norm activation function"""
    D = x.shape[axis] if axis is not None else x.size
    xmean = x.mean(axis, keepdims=True)
    xmeaned = x - xmean
    y = jnp.sqrt(jnp.power(xmeaned, 2).mean(axis, keepdims=True) + eps)
    return D * gamma * y + (delta * x).sum()


# %% ../nbs/00_lagrangians.ipynb 34
def _simple_spherical_norm(x:jnp.ndarray, 
                   axis=-1, # Which axis to normalize
                  ): 
    """Spherical norm activation function"""
    xmean = x.mean(axis, keepdims=True)
    xmeaned = x - xmean
    denominator = jnp.sqrt(jnp.power(xmeaned, 2).mean(axis, keepdims=True) + eps)
    return gamma * xmeaned / denominator + delta

def lagr_spherical_norm(x:jnp.ndarray, 
                   gamma:float=1.0, # Scale the stdev
                   delta:Union[float, jnp.ndarray]=0., # Shift the mean
                   axis=-1, # Which axis to normalize
                   eps=1e-5, # Prevent division by 0
                  ): 
    """Lagrangian of the spherical norm activation function"""
    y = jnp.sqrt(jnp.power(x, 2).sum(axis, keepdims=True) + eps)
    return gamma * y + (delta * x).sum()


# %% ../nbs/00_lagrangians.ipynb 37
class LIdentity(tx.Module):
    """Reduced Lagrangian whose activation function is the identity function"""
    def __init__(self): pass
    def __call__(self, x):
        return lagr_identity(x).sum()

class LRepu(tx.Module):
    """Reduced Lagrangian whose activation function is the rectified polynomial unit of specified degree `n`"""
    n: float = 2
    
    # Need a default for `n` to work with layer creation
    def __init__(self,
                 n=2.): # The degree of the RePU. By default, set to the ReLU configuration
        self.n = n

    def __call__(self, x):
        return lagr_repu(x, self.n).sum()
    
class LRelu(tx.Module):
    """Reduced Lagrangian whose activation function is the rectified linear unit"""
    def __init__(self): pass
    def __call__(self, x):
        return lagr_relu(x).sum()
    
class LSigmoid(tx.Module):
    """Reduced Lagrangian whose activation function is the sigmoid
    
    Parameterized by (beta)
    """
    beta: Union[float, jnp.ndarray] = tx.Parameter.node(default=1.0)
    scale: float = tx.Parameter.node(default=1.0)
    min_beta: float = 1e-6
    
    def __init__(self, 
                 beta=1., # Inverse temperature
                 scale=1., # Amount to stretch the sigmoid.
                 min_beta=1e-6): # Minimal accepted value of beta. For energy dynamics, it is important that beta be positive.
        self.beta = beta
        self.scale = scale
        self.min_beta = min_beta

    def __call__(self, x):
        return lagr_sigmoid(x, beta=jnp.clip(self.beta, self.min_beta), scale=self.scale).sum()
    
class LSoftmax(tx.Module):
    """Reduced Lagrangian whose activation function is the softmax
    
    Parameterized by (beta)
    """
    beta: Union[float, jnp.ndarray] = tx.Parameter.node(default=1.0)
    axis: int = -1
    min_beta: float = 1e-6

    def __init__(self, 
         beta=1., # Inverse temperature
         axis=-1, # Axis over which to apply the softmax
         min_beta=1e-6): # Minimal accepted value of beta. For energy dynamics, it is important that beta be positive.
        self.beta = beta
        self.axis = axis
        self.min_beta = min_beta

    def __call__(self, x):
        return lagr_softmax(x, beta=jnp.clip(self.beta, self.min_beta), axis=self.axis).sum()
    

class LExp(tx.Module):
    """Reduced Lagrangian whose activation function is the exponential function
    
    Parameterized by (beta)
    """
    beta: Union[float, jnp.ndarray] = tx.Parameter.node(default=1.0)
    min_beta: float = 1e-6

    def __init__(self, 
                 beta=1., # Inverse temperature, for the sharpness of the exponent
                 min_beta=1e-6): # Minimal accepted value of beta. For energy dynamics, it is important that beta be positive.
        self.beta = beta
        self.min_beta = min_beta
        
    def __call__(self, x):
        return lagr_exp(x, beta=jnp.clip(self.beta, self.min_beta)).sum()
    
class LRexp(tx.Module):
    """Reduced Lagrangian whose activation function is the rectified exponential function
    
    Parameterized by (beta)
    """
    beta: Union[float, jnp.ndarray] = tx.Parameter.node(default=1.0)
    min_beta: float = 1e-6

    def __init__(self, 
                 beta=1., # Inverse temperature, for the sharpness of the exponent
                 min_beta=1e-6): # Minimal accepted value of beta. For energy dynamics, it is important that beta be positive.
        self.beta = beta
        self.min_beta = min_beta
        
    def __call__(self, x):
        return lagr_rexp(x, beta=jnp.clip(self.beta, self.min_beta)).sum()
    
    
class LTanh(tx.Module):
    """Reduced Lagrangian whose activation function is the tanh
    
    Parameterized by (beta)
    """
    beta: Union[float, jnp.ndarray] = tx.Parameter.node(default=1.0)
    min_beta: float = 1e-6

    def __init__(self, 
                 beta=1., # Inverse temperature, for the sharpness of the exponent
                 min_beta=1e-6): # Minimal accepted value of beta. For energy dynamics, it is important that beta be positive.
        self.beta = beta
        self.min_beta = min_beta

    def __call__(self, x):
        return lagr_tanh(x, beta=jnp.clip(self.beta,self.min_beta)).sum()
    
class LLayerNorm(tx.Module):
    """Reduced Lagrangian whose activation function is the layer norm
    
    Parameterized by (gamma, delta), a scale and a shift
    """
    gamma: Union[float, jnp.ndarray] = tx.Parameter.node(default=1.0)
    delta: Union[float, jnp.ndarray] = tx.Parameter.node(default=0.)
    eps: float = 1e-5

    def __init__(self, 
                 gamma=1., # Inverse temperature, for the sharpness of the exponent
                 delta=0.,
                 eps = 1e-5
                ): # Minimal accepted value of beta. For energy dynamics, it is important that beta be positive.
        self.gamma = gamma
        self.delta = delta
        self.eps = eps

    def __call__(self, x):
        return lagr_layernorm(x, gamma=self.gamma, delta=self.delta, eps=self.eps).sum()
    
class LSphericalNorm(tx.Module):
    """Reduced Lagrangian whose activation function is the spherical L2 norm
    
    Parameterized by (gamma, delta), a scale and a shift
    """
    # gamma: Union[float, jnp.ndarray] = tx.Parameter.node(default=1.0)
    # delta: Union[float, jnp.ndarray] = tx.Parameter.node(default=0.)
    eps: float = 1e-5

    def __init__(self, 
                 gamma=1., # Inverse temperature, for the sharpness of the exponent
                 delta=0.,
                 eps = 1e-5
                ): # Minimal accepted value of beta. For energy dynamics, it is important that beta be positive.
        self.gamma = gamma
        self.delta = delta
        self.eps = eps

    def __call__(self, x):
        return lagr_spherical_norm(x, gamma=self.gamma, delta=self.delta, eps=self.eps).sum()
