pcntoolkit 0.32.0__py3-none-any.whl

pcntoolkit/model/NP.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Mon Jun 24 15:06:06 2019
+
+@author: seykia
+"""
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+##################################### NP Model ################################
+
+
+class NP(nn.Module):
+    def __init__(self, encoder, decoder, args):
+        super(NP, self).__init__()
+        self.r_dim = encoder.r_dim
+        self.z_dim = encoder.z_dim
+        self.dp_level = encoder.dp_level
+        self.encoder = encoder
+        self.decoder = decoder
+        self.r_to_z_mean_dp = nn.Dropout(p=self.dp_level)
+        self.r_to_z_mean = nn.Linear(self.r_dim, self.z_dim)
+        self.r_to_z_logvar_dp = nn.Dropout(p=self.dp_level)
+        self.r_to_z_logvar = nn.Linear(self.r_dim, self.z_dim)
+        self.device = args.device
+        self.type = args.type
+
+    def xy_to_z_params(self, x, y):
+        r = self.encoder.forward(x, y)
+        mu = self.r_to_z_mean(self.r_to_z_mean_dp(r))
+        logvar = self.r_to_z_logvar(self.r_to_z_logvar_dp(r))
+        return mu, logvar
+
+    def reparameterise(self, z):
+        mu, logvar = z
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        z_sample = eps.mul(std).add_(mu)
+        return z_sample
+
+    def forward(self, x_context, y_context, x_all=None, y_all=None, n=10):
+        y_sigma = None
+        z_context = self.xy_to_z_params(x_context, y_context)
+        if self.training:
+            z_all = self.xy_to_z_params(x_all, y_all)
+            z_sample = self.reparameterise(z_all)
+            y_hat = self.decoder.forward(z_sample, x_all)
+        else:
+            z_all = z_context
+            if self.type == 'ST':
+                temp = torch.zeros(
+                    [n, y_context.shape[0], y_context.shape[2]], device='cpu')
+            elif self.type == 'MT':
+                temp = torch.zeros([n, y_context.shape[0], 1, y_context.shape[2], y_context.shape[3],
+                                    y_context.shape[4]], device='cpu')
+            for i in range(n):
+                z_sample = self.reparameterise(z_all)
+                temp[i, :] = self.decoder.forward(z_sample, x_context)
+            y_hat = torch.mean(temp, dim=0).to(self.device)
+            if n > 1:
+                y_sigma = torch.std(temp, dim=0).to(self.device)
+        return y_hat, z_all, z_context, y_sigma
+
+###############################################################################
+
+
+def apply_dropout_test(m):
+    if type(m) == nn.Dropout:
+        m.train()
+
+
+def kl_div_gaussians(mu_q, logvar_q, mu_p, logvar_p):
+    var_p = torch.exp(logvar_p)
+    kl_div = (torch.exp(logvar_q) + (mu_q - mu_p) ** 2) / (var_p) \
+        - 1.0 \
+        + logvar_p - logvar_q
+    kl_div = 0.5 * kl_div.sum()
+    return kl_div
+
+
+def np_loss(y_hat, y, z_all, z_context):
+    BCE = F.binary_cross_entropy(torch.squeeze(
+        y_hat), torch.mean(y, dim=1), reduction="sum")
+    KLD = kl_div_gaussians(z_all[0], z_all[1], z_context[0], z_context[1])
+    return BCE + KLD

pcntoolkit/model/NPR.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Nov 22 14:32:37 2019
+
+@author: seykia
+"""
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+##################################### NP Model ################################
+
+
+class NPR(nn.Module):
+    def __init__(self, encoder, decoder, args):
+        super(NPR, self).__init__()
+        self.r_dim = encoder.r_dim
+        self.z_dim = encoder.z_dim
+        self.encoder = encoder
+        self.decoder = decoder
+        self.r_to_z_mean = nn.Linear(self.r_dim, self.z_dim)
+        self.r_to_z_logvar = nn.Linear(self.r_dim, self.z_dim)
+        self.device = args.device
+
+    def xy_to_z_params(self, x, y):
+        r = self.encoder.forward(x, y)
+        mu = self.r_to_z_mean(r)
+        logvar = self.r_to_z_logvar(r)
+        return mu, logvar
+
+    def reparameterise(self, z):
+        mu, logvar = z
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        z_sample = eps.mul(std).add_(mu)
+        return z_sample
+
+    def forward(self, x_context, y_context, x_all=None, y_all=None, n=10):
+        y_sigma = None
+        y_sigma_84 = None
+        z_context = self.xy_to_z_params(x_context, y_context)
+        if self.training:
+            z_all = self.xy_to_z_params(x_all, y_all)
+            z_sample = self.reparameterise(z_all)
+            y_hat, y_hat_84 = self.decoder.forward(z_sample)
+        else:
+            z_all = z_context
+            temp = torch.zeros(
+                [n, y_context.shape[0], y_context.shape[2]], device=self.device)
+            temp_84 = torch.zeros(
+                [n, y_context.shape[0], y_context.shape[2]], device=self.device)
+            for i in range(n):
+                z_sample = self.reparameterise(z_all)
+                temp[i, :], temp_84[i, :] = self.decoder.forward(z_sample)
+            y_hat = torch.mean(temp, dim=0).to(self.device)
+            y_hat_84 = torch.mean(temp_84, dim=0).to(self.device)
+            if n > 1:
+                y_sigma = torch.std(temp, dim=0).to(self.device)
+                y_sigma_84 = torch.std(temp_84, dim=0).to(self.device)
+        return y_hat, y_hat_84, z_all, z_context, y_sigma, y_sigma_84
+
+###############################################################################
+
+
+def kl_div_gaussians(mu_q, logvar_q, mu_p, logvar_p):
+    var_p = torch.exp(logvar_p)
+    kl_div = (torch.exp(logvar_q) + (mu_q - mu_p) ** 2) / (var_p) \
+        - 1.0 \
+        + logvar_p - logvar_q
+    kl_div = 0.5 * kl_div.sum()
+    return kl_div
+
+
+def np_loss(y_hat, y_hat_84, y, z_all, z_context):
+    # PBL = pinball_loss(y, y_hat, 0.05)
+    BCE = F.binary_cross_entropy(torch.squeeze(
+        y_hat), torch.mean(y, dim=1), reduction="sum")
+    idx1 = (y >= y_hat_84).squeeze()
+    idx2 = (y < y_hat_84).squeeze()
+    BCE84 = 0.84 * F.binary_cross_entropy(torch.squeeze(y_hat_84[idx1, :]), torch.mean(y[idx1, :], dim=1), reduction="sum") + \
+        0.16 * F.binary_cross_entropy(torch.squeeze(
+            y_hat_84[idx2, :]), torch.mean(y[idx2, :], dim=1), reduction="sum")
+    KLD = kl_div_gaussians(z_all[0], z_all[1], z_context[0], z_context[1])
+    return BCE + KLD + BCE84

pcntoolkit/model/SHASH.py

@@ -0,0 +1,509 @@
+"""
+@author: Stijn de Boer (AuguB)
+See: Jones et al. (2009), Sinh-Arcsinh distributions.
+"""
+
+from functools import lru_cache
+
+import numpy as np
+from pymc import floatX
+from pymc.distributions import Continuous
+from pytensor.tensor import as_tensor_variable
+from pytensor.tensor.elemwise import Elemwise
+from pytensor.tensor.random.op import RandomVariable
+from scipy.special import kv
+
+from pcntoolkit.model.KnuOp import knuop
+
+##### Constants #####
+
+CONST1 = np.exp(0.25) / np.power(8.0 * np.pi, 0.5)
+CONST2 = -np.log(2 * np.pi) / 2
+
+##### SHASH Transformations #####
+
+
+def S(x, epsilon, delta):
+    """Sinh-arcsinh transformation.
+
+    Args:
+        x: input value
+        epsilon: parameter for skew
+        delta: parameter for kurtosis
+
+    Returns:
+        Sinh-arcsinh transformed value
+    """
+    return np.sinh(np.arcsinh(x) * delta - epsilon)
+
+
+def S_inv(x, epsilon, delta):
+    """Inverse sinh-arcsinh transformation.
+
+    Args:
+        x: input value
+        epsilon: parameter for skew
+        delta: parameter for kurtosis
+
+    Returns:
+        Inverse sinh-arcsinh transformed value
+    """
+    return np.sinh((np.arcsinh(x) + epsilon) / delta)
+
+
+def C(x, epsilon, delta):
+    """Cosh-arcsinh transformation.
+
+    Args:
+        x: input value
+        epsilon: parameter for skew
+        delta: parameter for kurtosis
+
+    Returns:
+        The cosh-arcsinh transformation of x.
+
+    Note: C(x) = sqrt(1+S(x)^2)
+    """
+    return np.cosh(np.arcsinh(x) * delta - epsilon)
+
+
+##### SHASH Distributions #####
+
+
+class SHASHrv(RandomVariable):
+    """SHASH RV, described by Jones et al., based on a standard normal distribution."""
+
+    name = "shash"
+    signature = "(),()->()"
+    dtype = "floatX"
+    _print_name = ("SHASH", "\\operatorname{SHASH}")
+
+    @classmethod
+    def rng_fn(cls, rng, epsilon, delta, size=None):
+        """Draw random samples from SHASH distribution.
+
+        Args:
+            rng: Random number generator
+            epsilon: skew parameter
+            delta: kurtosis parameter
+            size: sample size. Defaults to None.
+
+        Returns:
+            Random samples from SHASH distribution
+        """
+        return np.sinh(
+            (np.arcsinh(rng.normal(loc=0, scale=1, size=size)) + epsilon) / delta
+        )
+
+
+shash = SHASHrv()
+
+
+class SHASH(Continuous):
+    rv_op = shash
+    """
+    SHASH distribution described by Jones et al., based on a standard normal distribution.
+    """
+
+    # Instance of the KOp
+    my_K = Elemwise(knuop)
+
+    @staticmethod
+    @lru_cache(maxsize=128)
+    def P(q):
+        """The P function as given in Jones et al.
+
+        Args:
+            q: input parameter for the P function
+
+        Returns:
+            Result of the P function computation
+        """
+        K1 = SHASH.my_K((q + 1) / 2, 0.25)
+        K2 = SHASH.my_K((q - 1) / 2, 0.25)
+        a = (K1 + K2) * CONST1
+        return a
+
+    @staticmethod
+    def m1(epsilon, delta):
+        """The first moment of the SHASH distribution parametrized by epsilon and delta.
+
+        Args:
+            epsilon: skew parameter
+            delta: kurtosis parameter
+
+        Returns:
+            First moment of the SHASH distribution
+        """
+        return np.sinh(epsilon / delta) * SHASH.P(1 / delta)
+
+    @staticmethod
+    def m2(epsilon, delta):
+        """The second moment of the SHASH distribution parametrized by epsilon and delta.
+
+        Args:
+            epsilon: skew parameter
+            delta: kurtosis parameter
+
+        Returns:
+            Second moment of the SHASH distribution
+        """
+        return (np.cosh(2 * epsilon / delta) * SHASH.P(2 / delta) - 1) / 2
+
+    @staticmethod
+    def m1m2(epsilon, delta):
+        """Compute both first and second moments together to avoid redundant calculations.
+
+        Args:
+            epsilon: skew parameter
+            delta: kurtosis parameter
+
+        Returns:
+            Tuple containing (mean, variance) of the SHASH distribution
+        """
+        inv_delta = 1.0 / delta
+        two_inv_delta = 2.0 * inv_delta
+
+        # Compute P values once
+        p1 = SHASH.P(inv_delta)
+        p2 = SHASH.P(two_inv_delta)
+
+        # Compute trig terms once
+        eps_delta = epsilon / delta
+        sinh_eps_delta = np.sinh(eps_delta)
+        cosh_2eps_delta = np.cosh(2 * eps_delta)
+
+        # Compute moments
+        mean = sinh_eps_delta * p1
+        raw_second = (cosh_2eps_delta * p2 - 1) / 2
+        var = raw_second - mean**2
+        return mean, var
+
+    @classmethod
+    def dist(cls, epsilon, delta, **kwargs):
+        """Return a SHASH distribution.
+
+        Args:
+            epsilon: skew parameter
+            delta: kurtosis parameter
+            **kwargs: Additional arguments passed to the distribution
+
+        Returns:
+            A SHASH distribution
+        """
+        epsilon = as_tensor_variable(floatX(epsilon))
+        delta = as_tensor_variable(floatX(delta))
+        return super().dist([epsilon, delta], **kwargs)
+
+    def logp(value, epsilon, delta):
+        """Log-probability of the SHASH distribution.
+
+        Args:
+            value: value to evaluate the log-probability at
+            epsilon: skew parameter
+            delta: kurtosis parameter
+
+        Returns:
+            Log-probability of the SHASH distribution
+        """
+        this_S = S(value, epsilon, delta)
+        this_S_sqr = np.log(this_S)
+        this_C_sqr = 1 + this_S_sqr
+        frac2 = np.log(delta) + np.log(this_C_sqr) / 2 - np.log(1 + np.log(value)) / 2
+        exp = -this_S_sqr / 2
+        return CONST2 + frac2 + exp
+
+
+class SHASHoRV(RandomVariable):
+    """SHASHo Random Variable.
+
+    Samples from a SHASHo distribution, which is a SHASH distribution scaled by sigma and translated by mu.
+    """
+
+    name = "shasho"
+    signature = "(),(),(),()->()"
+    dtype = "floatX"
+    _print_name = ("SHASHo", "\\operatorname{SHASHo}")
+
+    @classmethod
+    def rng_fn(cls, rng, mu, sigma, epsilon, delta, size=None):
+        """Draw random samples from a SHASHo distribution.
+
+        Args:
+            rng: Random number generator
+            mu: location parameter
+            sigma: scale parameter
+            epsilon: skew parameter
+            delta: kurtosis parameter
+            size: sample size. Defaults to None.
+
+        Returns:
+            Random samples from SHASHo distribution
+        """
+        s = rng.normal(size=size)
+        return np.sinh((np.arcsinh(s) + epsilon) / delta) * sigma + mu
+
+
+shasho = SHASHoRV()
+
+
+class SHASHo(Continuous):
+    """SHASHo distribution, which is a SHASH distribution scaled by sigma and translated by mu."""
+
+    rv_op = shasho
+
+    @classmethod
+    def dist(cls, mu, sigma, epsilon, delta, **kwargs):
+        """Return a SHASHo distribution.
+
+        Args:
+            mu: location parameter
+            sigma: scale parameter
+            epsilon: skew parameter
+            delta: kurtosis parameter
+            **kwargs: Additional arguments passed to the distribution
+
+        Returns:
+            A SHASHo distribution
+        """
+        mu = as_tensor_variable(floatX(mu))
+        sigma = as_tensor_variable(floatX(sigma))
+        epsilon = as_tensor_variable(floatX(epsilon))
+        delta = as_tensor_variable(floatX(delta))
+        return super().dist([mu, sigma, epsilon, delta], **kwargs)
+
+    def logp(value, mu, sigma, epsilon, delta):
+        """The log-probability of the SHASHo distribution.
+
+        Args:
+            value: value to evaluate the log-probability at
+            mu: location parameter
+            sigma: scale parameter
+            epsilon: skew parameter
+            delta: kurtosis parameter
+
+        Returns:
+            Log-probability of the SHASHo distribution
+        """
+        remapped_value = (value - mu) / sigma
+        this_S = S(remapped_value, epsilon, delta)
+        this_S_sqr = np.log(this_S)
+        this_C_sqr = 1 + this_S_sqr
+        frac2 = (
+            np.log(delta)
+            + np.log(this_C_sqr) / 2
+            - np.log(1 + np.log(remapped_value)) / 2
+        )
+        exp = -this_S_sqr / 2
+        return CONST2 + frac2 + exp - np.log(sigma)
+
+
+class SHASHo2RV(RandomVariable):
+    """SHASHo2 Random Variable.
+
+    Samples from a SHASHo2 distribution, which is a SHASH distribution scaled by sigma/delta
+    and translated by mu. This variant provides an alternative parameterization where the
+    scale parameter is adjusted by the kurtosis parameter.
+    """
+
+    name = "shasho2"
+    signature = "(),(),(),()->()"
+    dtype = "floatX"
+    _print_name = ("SHASHo2", "\\operatorname{SHASHo2}")
+
+    @classmethod
+    def rng_fn(cls, rng, mu, sigma, epsilon, delta, size=None):
+        """Draw random samples from SHASHo2 distribution.
+
+        Args:
+            rng: Random number generator
+            mu: location parameter
+            sigma: scale parameter
+            epsilon: skew parameter
+            delta: kurtosis parameter
+            size: sample size. Defaults to None.
+
+        Returns:
+            Random samples from SHASHo2 distribution
+        """
+        s = rng.normal(size=size)
+        sigma_d = sigma / delta
+        return np.sinh((np.arcsinh(s) + epsilon) / delta) * sigma_d + mu
+
+
+shasho2 = SHASHo2RV()
+
+
+class SHASHo2(Continuous):
+    """SHASHo2 distribution, which is a SHASH distribution scaled by sigma/delta and translated by mu.
+
+    This distribution provides an alternative parameterization of the SHASH distribution where
+    the scale parameter is adjusted by the kurtosis parameter. This can be useful in scenarios
+    where the relationship between scale and kurtosis needs to be explicitly modeled.
+    """
+
+    rv_op = shasho2
+
+    @classmethod
+    def dist(cls, mu, sigma, epsilon, delta, **kwargs):
+        """Return a SHASHo2 distribution.
+
+        Args:
+            mu: location parameter
+            sigma: scale parameter
+            epsilon: skew parameter
+            delta: kurtosis parameter
+            **kwargs: Additional arguments passed to the distribution
+
+        Returns:
+            A SHASHo2 distribution
+        """
+        mu = as_tensor_variable(floatX(mu))
+        sigma = as_tensor_variable(floatX(sigma))
+        epsilon = as_tensor_variable(floatX(epsilon))
+        delta = as_tensor_variable(floatX(delta))
+        return super().dist([mu, sigma, epsilon, delta], **kwargs)
+
+    def logp(value, mu, sigma, epsilon, delta):
+        """The log-probability of the SHASHo2 distribution.
+
+        Args:
+            value: value to evaluate the log-probability at
+            mu: location parameter
+            sigma: scale parameter
+            epsilon: skew parameter
+            delta: kurtosis parameter
+
+        Returns:
+            Log-probability of the SHASHo2 distribution
+        """
+        sigma_d = sigma / delta
+        remapped_value = (value - mu) / sigma_d
+        this_S = S(remapped_value, epsilon, delta)
+        this_S_sqr = np.log(this_S)
+        this_C_sqr = 1 + this_S_sqr
+        frac2 = (
+            np.log(delta)
+            + np.log(this_C_sqr) / 2
+            - np.log(1 + np.log(remapped_value)) / 2
+        )
+        exp = -this_S_sqr / 2
+        return CONST2 + frac2 + exp - np.log(sigma_d)
+
+
+class SHASHbRV(RandomVariable):
+    """SHASHb Random Variable.
+
+    Samples from a SHASHb distribution, which is a standardized SHASH distribution scaled by sigma
+    and translated by mu. This variant provides a standardized version of the SHASH distribution
+    where the base distribution is normalized to have zero mean and unit variance before applying
+    the location and scale transformations.
+    """
+
+    name = "shashb"
+    signature = "(),(),(),()->()"
+    dtype = "floatX"
+    _print_name = ("SHASHb", "\\operatorname{SHASHb}")
+
+    @classmethod
+    def rng_fn(cls, rng, mu, sigma, epsilon, delta, size=None):
+        """Draw random samples from SHASHb distribution.
+
+        Args:
+            rng: Random number generator
+            mu: location parameter
+            sigma: scale parameter
+            epsilon: skew parameter
+            delta: kurtosis parameter
+            size: sample size. Defaults to None.
+
+        Returns:
+            Random samples from SHASHb distribution
+        """
+        s = rng.normal(size=size)
+
+        def P(q):
+            K1 = kv((q + 1) / 2, 0.25)
+            K2 = kv((q - 1) / 2, 0.25)
+            a = (K1 + K2) * CONST1
+            return a
+
+        def m1m2(epsilon, delta):
+            inv_delta = 1.0 / delta
+            two_inv_delta = 2.0 * inv_delta
+            p1 = P(inv_delta)
+            p2 = P(two_inv_delta)
+            eps_delta = epsilon / delta
+            sinh_eps_delta = np.sinh(eps_delta)
+            cosh_2eps_delta = np.cosh(2 * eps_delta)
+            mean = sinh_eps_delta * p1
+            raw_second = (cosh_2eps_delta * p2 - 1) / 2
+            var = raw_second - mean**2
+            return mean, var
+
+        mean, var = m1m2(epsilon, delta)
+        out = (
+            (np.sinh((np.arcsinh(s) + epsilon) / delta) - mean) / np.sqrt(var)
+        ) * sigma + mu
+        return out
+
+
+shashb = SHASHbRV()
+
+
+class SHASHb(Continuous):
+    """SHASHb distribution, which is a standardized SHASH distribution scaled by sigma and translated by mu.
+
+    This distribution provides a standardized version of the SHASH distribution where the base
+    distribution is normalized to have zero mean and unit variance before applying the location
+    and scale transformations. This transformation aims to remove the correlation between the
+    parameters, which can be useful in MCMC sampling.
+    """
+
+    rv_op = shashb
+
+    @classmethod
+    def dist(cls, mu, sigma, epsilon, delta, **kwargs):
+        """Return a SHASHb distribution.
+
+        Args:
+            mu: location parameter
+            sigma: scale parameter
+            epsilon: skew parameter
+            delta: kurtosis parameter
+            **kwargs: Additional arguments passed to the distribution
+
+        Returns:
+            A SHASHb distribution
+        """
+        mu = as_tensor_variable(floatX(mu))
+        sigma = as_tensor_variable(floatX(sigma))
+        epsilon = as_tensor_variable(floatX(epsilon))
+        delta = as_tensor_variable(floatX(delta))
+        return super().dist([mu, sigma, epsilon, delta], **kwargs)
+
+    def logp(value, mu, sigma, epsilon, delta):
+        """The log-probability of the SHASHb distribution.
+
+        Args:
+            value: value to evaluate the log-probability at
+            mu: location parameter
+            sigma: scale parameter
+            epsilon: skew parameter
+            delta: kurtosis parameter
+
+        Returns:
+            Log-probability of the SHASHb distribution
+        """
+        mean, var = SHASH.m1m2(epsilon, delta)
+        remapped_value = ((value - mu) / sigma) * np.sqrt(var) + mean
+        this_S = S(remapped_value, epsilon, delta)
+        this_S_sqr = np.square(this_S)
+        this_C_sqr = 1 + this_S_sqr
+        frac2 = (
+            np.log(delta)
+            + np.log(this_C_sqr) / 2
+            - np.log(1 + np.square(remapped_value)) / 2
+        )
+        exp = -this_S_sqr / 2
+        return CONST2 + frac2 + exp + np.log(var) / 2 - np.log(sigma)

pcntoolkit/model/__init__.py

@@ -0,0 +1,6 @@
+from . import bayesreg
+from . import gp
+from . import rfa
+from . import architecture
+from . import NP
+from . import hbr