pcntoolkit 0.32.0__py3-none-any.whl

pcntoolkit/model/rfa.py

@@ -0,0 +1,245 @@
+from __future__ import print_function
+from __future__ import division
+
+import numpy as np
+import torch
+
+
+class GPRRFA:
+    """Random Feature Approximation for Gaussian Process Regression
+
+    Estimation and prediction of Bayesian linear regression models
+
+    Basic usage::
+
+        R = GPRRFA()
+        hyp = R.estimate(hyp0, X, y)
+        ys,s2 = R.predict(hyp, X, y, Xs)
+
+    where the variables are
+
+    :param hyp: vector of hyperparmaters.
+    :param X: N x D data array
+    :param y: 1D Array of targets (length N)
+    :param Xs: Nte x D array of test cases
+    :param hyp0: starting estimates for hyperparameter optimisation
+
+    :returns: * ys - predictive mean
+              * s2 - predictive variance
+
+    The hyperparameters are::
+
+        hyp = [ log(sn), log(ell), log(sf) ]  # hyp is a numpy array
+
+    where sn^2 is the noise variance, ell are lengthscale parameters and 
+    sf^2 is the signal variance. This provides an approximation to the
+    covariance function::
+
+        k(x,z) = x'*z + sn2*exp(0.5*(x-z)'*Lambda*(x-z))
+
+    where Lambda = diag((ell_1^2, ... ell_D^2))
+
+    Written by A. Marquand
+    """
+
+    def __init__(self, hyp=None, X=None, y=None, n_feat=None,
+                 n_iter=100, tol=1e-3, verbose=False):
+
+        self.hyp = np.nan
+        self.nlZ = np.nan
+        self.tol = tol          # not used at present
+        self.Nf = n_feat
+        self.n_iter = n_iter
+        self.verbose = verbose
+        self._n_restarts = 5
+
+        if (hyp is not None) and (X is not None) and (y is not None):
+            self.post(hyp, X, y)
+
+    def _numpy2torch(self, X, y=None, hyp=None):
+
+        if type(X) is torch.Tensor:
+            pass
+        elif type(X) is np.ndarray:
+            X = torch.from_numpy(X)
+        else:
+            ValueError('Unknown data type (X)')
+        X = X.double()
+
+        if y is not None:
+            if type(y) is torch.Tensor:
+                pass
+            elif type(y) is np.ndarray:
+                y = torch.from_numpy(y)
+            else:
+                raise ValueError('Unknown data type (y)')
+
+            if len(y.shape) == 1:
+                y.resize_(y.shape[0], 1)
+            y = y.double()
+
+        if hyp is not None:
+            if type(hyp) is torch.Tensor:
+                pass
+            else:
+                hyp = torch.tensor(hyp, requires_grad=True)
+
+        return X, y, hyp
+
+    def get_n_params(self, X):
+
+        return X.shape[1] + 2
+
+    def post(self, hyp, X, y):
+        """ Generic function to compute posterior distribution.
+
+            This function will save the posterior mean and precision matrix as
+            self.m and self.A and will also update internal parameters (e.g.
+            N, D and the prior covariance (Sigma) and precision (iSigma).
+        """
+
+        # make sure all variables are the right type
+        X, y, hyp = self._numpy2torch(X, y, hyp)
+
+        self.N, self.Dx = X.shape
+
+        # ensure the number of features is specified (use 75% as a default)
+        if self.Nf is None:
+            self.Nf = int(0.75 * self.N)
+
+        self.Omega = torch.zeros((self.Dx, self.Nf), dtype=torch.double)
+        for f in range(self.Nf):
+            self.Omega[:, f] = torch.exp(hyp[1:-1]) * \
+                torch.randn((self.Dx, 1), dtype=torch.double).squeeze()
+
+        XO = torch.mm(X, self.Omega)
+        self.Phi = torch.exp(hyp[-1])/np.sqrt(self.Nf) *  \
+            torch.cat((torch.cos(XO), torch.sin(XO)), 1)
+
+        # concatenate linear weights
+        self.Phi = torch.cat((self.Phi, X), 1)
+        self.D = self.Phi.shape[1]
+
+        if self.verbose:
+            print("estimating posterior ... | hyp=", hyp)
+
+        self.A = torch.mm(torch.t(self.Phi), self.Phi) / torch.exp(2*hyp[0]) + \
+            torch.eye(self.D, dtype=torch.double)
+        self.m = torch.mm(torch.solve(torch.t(self.Phi), self.A)[0], y) / \
+            torch.exp(2*hyp[0])
+
+        # save hyperparameters
+        self.hyp = hyp
+
+        # update optimizer iteration count
+        if hasattr(self, '_iterations'):
+            self._iterations += 1
+
+    def loglik(self, hyp, X, y):
+        """ Function to compute compute log (marginal) likelihood """
+        X, y, hyp = self._numpy2torch(X, y, hyp)
+
+        # always recompute the posterior
+        self.post(hyp, X, y)
+
+        # logdetA = 2*torch.sum(torch.log(torch.diag(torch.cholesky(self.A))))
+        try:
+            # compute the log determinants in a numerically stable way
+            logdetA = 2 * \
+                torch.sum(torch.log(torch.diag(torch.cholesky(self.A))))
+        except Exception as e:
+            print("Warning: Estimation of posterior distribution failed")
+            print(e)
+            # nlZ = torch.tensor(1/np.finfo(float).eps)
+            nlZ = torch.tensor(np.nan)
+            self._optim_failed = True
+            return nlZ
+
+        # compute negative marginal log likelihood
+        nlZ = -0.5 * (self.N*torch.log(1/torch.exp(2*hyp[0])) -
+                      self.N*np.log(2*np.pi) -
+                      torch.mm(torch.t(y - torch.mm(self.Phi, self.m)),
+                               (y - torch.mm(self.Phi, self.m))) /
+                      torch.exp(2*hyp[0]) -
+                      torch.mm(torch.t(self.m), self.m) - logdetA)
+
+        if self.verbose:
+            print("nlZ= ", nlZ, " | hyp=", hyp)
+
+        # save marginal likelihood
+        self.nlZ = nlZ
+        return nlZ
+
+    def dloglik(self, hyp, X, y):
+        """ Function to compute derivatives """
+
+        print("derivatives not available")
+
+        return
+
+    def estimate(self, hyp0, X, y, optimizer='lbfgs'):
+        """ Function to estimate the model """
+
+        if type(hyp0) is torch.Tensor:
+            hyp = hyp0
+            hyp0.requires_grad_()
+        else:
+            hyp = torch.tensor(hyp0, requires_grad=True)
+        # save the starting values
+        self.hyp0 = hyp
+
+        if optimizer.lower() == 'lbfgs':
+            opt = torch.optim.LBFGS([hyp])
+        else:
+            raise ValueError("Optimizer " + " not implemented")
+        self._iterations = 0
+
+        def closure():
+            opt.zero_grad()
+            nlZ = self.loglik(hyp, X, y)
+            if not torch.isnan(nlZ):
+                nlZ.backward()
+            return nlZ
+
+        for r in range(self._n_restarts):
+            self._optim_failed = False
+
+            nlZ = opt.step(closure)
+
+            if self._optim_failed:
+                print("optimization failed. retrying (", r+1, "of",
+                      self._n_restarts, ")")
+                hyp = torch.randn_like(hyp, requires_grad=True)
+                self.hyp0 = hyp
+            else:
+                print("Optimzation complete after", self._iterations,
+                      "evaluations. Function value =",
+                      nlZ.detach().numpy().squeeze())
+                break
+
+        return self.hyp.detach().numpy()
+
+    def predict(self, hyp, X, y, Xs):
+        """ Function to make predictions from the model """
+
+        X, y, hyp = self._numpy2torch(X, y, hyp)
+        Xs, *_ = self._numpy2torch(Xs)
+
+        if (hyp != self.hyp).all() or not (hasattr(self, 'A')):
+            self.post(hyp, X, y)
+
+        # generate prediction tensors
+        XsO = torch.mm(Xs, self.Omega)
+        Phis = torch.exp(hyp[-1])/np.sqrt(self.Nf) * \
+            torch.cat((torch.cos(XsO), torch.sin(XsO)), 1)
+        # add linear component
+        Phis = torch.cat((Phis, Xs), 1)
+
+        ys = torch.mm(Phis, self.m)
+
+        # compute diag(Phis*(Phis'\A)) avoiding computing off-diagonal entries
+        s2 = torch.exp(2*hyp[0]) + \
+            torch.sum(Phis * torch.t(torch.solve(torch.t(Phis), self.A)[0]), 1)
+
+        # return output as numpy arrays
+        return ys.detach().numpy().squeeze(), s2.detach().numpy().squeeze()