Using params class with sympy covariance. Adding conditional statements for presence of weave.

This commit is contained in:
Neil Lawrence 2014-02-18 19:37:53 -05:00
parent d1b6d18ddf
commit f6484bcbd0

View file

@ -1,17 +1,31 @@
import numpy as np try:
import sympy as sp import sympy as sp
from sympy.utilities.codegen import codegen sympy_available=True
from sympy.core.cache import clear_cache except ImportError:
sympy_available=False
exit()
from sympy.core.cache import clear_cache
from sympy.utilities.codegen import codegen
try:
from scipy import weave
weave_available = True
except ImportError:
weave_available = False
from scipy import weave
import re
import os import os
import sys
current_dir = os.path.dirname(os.path.abspath(os.path.dirname(__file__))) current_dir = os.path.dirname(os.path.abspath(os.path.dirname(__file__)))
import sys
import numpy as np
import re
import tempfile import tempfile
import pdb import pdb
import ast import ast
from kernpart import Kernpart from kernpart import Kernpart
from ...core.parameterization import Param
from ...core.parameterization.transformations import Logexp
class spkern(Kernpart): class spkern(Kernpart):
""" """
@ -75,17 +89,20 @@ class spkern(Kernpart):
self.num_split_params = len(self._sp_theta_i) self.num_split_params = len(self._sp_theta_i)
self._split_theta_names = ["%s"%theta.name[:-2] for theta in self._sp_theta_i] self._split_theta_names = ["%s"%theta.name[:-2] for theta in self._sp_theta_i]
for theta in self._split_theta_names: for theta in self._split_theta_names:
setattr(self, theta, np.ones(self.output_dim)) setattr(self, theta, Param(theta, np.ones(self.output_dim), None))
self.add_parameters(getattr(self, theta))
#setattr(self, theta, np.ones(self.output_dim))
self.num_shared_params = len(self._sp_theta) self.num_shared_params = len(self._sp_theta)
self.num_params = self.num_shared_params+self.num_split_params*self.output_dim #self.num_params = self.num_shared_params+self.num_split_params*self.output_dim
else: else:
self.num_split_params = 0 self.num_split_params = 0
self._split_theta_names = [] self._split_theta_names = []
self._sp_theta = thetas self._sp_theta = thetas
self.num_shared_params = len(self._sp_theta) self.num_shared_params = len(self._sp_theta)
self.num_params = self.num_shared_params #self.num_params = self.num_shared_params
# Add parameters to the model. # Add parameters to the model.
for theta in self._sp_theta: for theta in self._sp_theta:
@ -93,9 +110,12 @@ class spkern(Kernpart):
if param is not None: if param is not None:
if param.has_key(theta): if param.has_key(theta):
val = param[theta] val = param[theta]
setattr(self, theta.name, val) #setattr(self, theta.name, val)
setattr(self, theta.name, Param(theta.name, val, None))
self.add_parameters(getattr(self, theta.name))
self.parameters_changed() # initializes cache
#deal with param #deal with param
self._set_params(self._get_params()) #self._set_params(self._get_params())
# Differentiate with respect to parameters. # Differentiate with respect to parameters.
self._sp_dk_dtheta = [sp.diff(k,theta).simplify() for theta in self._sp_theta] self._sp_dk_dtheta = [sp.diff(k,theta).simplify() for theta in self._sp_theta]
@ -112,7 +132,7 @@ class spkern(Kernpart):
# generate the code for the covariance functions # generate the code for the covariance functions
self._gen_code() self._gen_code()
if weave if weave_available:
if False: if False:
extra_compile_args = ['-ftree-vectorize', '-mssse3', '-ftree-vectorizer-verbose=5'] extra_compile_args = ['-ftree-vectorize', '-mssse3', '-ftree-vectorizer-verbose=5']
else: else:
@ -131,7 +151,7 @@ class spkern(Kernpart):
return spkern(self._sp_k+other._sp_k) return spkern(self._sp_k+other._sp_k)
def _gen_code(self): def _gen_code(self):
#generate c functions from sympy objects
argument_sequence = self._sp_x+self._sp_z+self._sp_theta argument_sequence = self._sp_x+self._sp_z+self._sp_theta
code_list = [('k',self._sp_k)] code_list = [('k',self._sp_k)]
# gradients with respect to covariance input # gradients with respect to covariance input
@ -142,31 +162,47 @@ class spkern(Kernpart):
if self.output_dim > 1: if self.output_dim > 1:
argument_sequence += self._sp_theta_i + self._sp_theta_j argument_sequence += self._sp_theta_i + self._sp_theta_j
code_list += [('dk_d%s'%theta.name,dtheta) for theta,dtheta in zip(self._sp_theta_i,self._sp_dk_dtheta_i)] code_list += [('dk_d%s'%theta.name,dtheta) for theta,dtheta in zip(self._sp_theta_i,self._sp_dk_dtheta_i)]
# generate c functions from sympy objects
if weave_available:
code_type = "C"
else:
code_type = "PYTHON"
(foo_c,self._function_code), (foo_h,self._function_header) = \ (foo_c,self._function_code), (foo_h,self._function_header) = \
codegen(code_list, "C",'foobar',argument_sequence=argument_sequence) codegen(code_list,
#put the header file where we can find it code_type,
f = file(os.path.join(tempfile.gettempdir(),'foobar.h'),'w') self.name,
argument_sequence=argument_sequence)
# Use weave to compute the underlying functions.
if weave_available:
# put the header file where we can find it
f = file(os.path.join(tempfile.gettempdir(), self.name + '.h'),'w')
f.write(self._function_header) f.write(self._function_header)
f.close() f.close()
# Substitute any known derivatives which sympy doesn't compute # Substitute any known derivatives which sympy doesn't compute
self._function_code = re.sub('DiracDelta\(.+?,.+?\)','0.0',self._function_code) self._function_code = re.sub('DiracDelta\(.+?,.+?\)','0.0',self._function_code)
# This is the basic argument construction for the C code. if weave_available:
#arg_list = (["X[i*input_dim+%s]"%x.name[2:] for x in self._sp_x] # arg_list will store the arguments required for the C code.
# + ["Z[j*input_dim+%s]"%z.name[2:] for z in self._sp_z])
arg_list = (["X2(i, %s)"%x.name[2:] for x in self._sp_x] arg_list = (["X2(i, %s)"%x.name[2:] for x in self._sp_x]
+ ["Z2(j, %s)"%z.name[2:] for z in self._sp_z]) + ["Z2(j, %s)"%z.name[2:] for z in self._sp_z])
# for multiple outputs reverse argument list is also required
if self.output_dim>1: if self.output_dim>1:
reverse_arg_list = list(arg_list) reverse_arg_list = list(arg_list)
reverse_arg_list.reverse() reverse_arg_list.reverse()
# This gives the parameters for the arg list.
param_arg_list = [shared_params.name for shared_params in self._sp_theta] param_arg_list = [shared_params.name for shared_params in self._sp_theta]
arg_list += param_arg_list arg_list += param_arg_list
precompute_list=[] precompute_list=[]
if self.output_dim > 1: if self.output_dim > 1:
reverse_arg_list+=list(param_arg_list) reverse_arg_list+=list(param_arg_list)
# For multiple outputs, also need the split parameters.
split_param_arg_list = ["%s1(%s)"%(theta.name[:-2].upper(),index) for index in ['ii', 'jj'] for theta in self._sp_theta_i] split_param_arg_list = ["%s1(%s)"%(theta.name[:-2].upper(),index) for index in ['ii', 'jj'] for theta in self._sp_theta_i]
split_param_reverse_arg_list = ["%s1(%s)"%(theta.name[:-2].upper(),index) for index in ['jj', 'ii'] for theta in self._sp_theta_i] split_param_reverse_arg_list = ["%s1(%s)"%(theta.name[:-2].upper(),index) for index in ['jj', 'ii'] for theta in self._sp_theta_i]
arg_list += split_param_arg_list arg_list += split_param_arg_list
@ -177,6 +213,15 @@ class spkern(Kernpart):
reverse_arg_string = ", ".join(reverse_arg_list) reverse_arg_string = ", ".join(reverse_arg_list)
arg_string = ", ".join(arg_list) arg_string = ", ".join(arg_list)
precompute_string = "\n".join(precompute_list) precompute_string = "\n".join(precompute_list)
# Now we use the arguments in code that computes the separate parts.
# Any precomputations will be done here eventually.
self._precompute = \
"""
// Precompute code would go here. It will be called when parameters are updated.
"""
# Here's the code to do the looping for K # Here's the code to do the looping for K
self._K_code =\ self._K_code =\
""" """
@ -190,13 +235,15 @@ class spkern(Kernpart):
//#pragma omp parallel for private(j) //#pragma omp parallel for private(j)
for (i=0;i<N;i++){ for (i=0;i<N;i++){
for (j=0;j<num_inducing;j++){ for (j=0;j<num_inducing;j++){
%s %s
//target[i*num_inducing+j] = //target[i*num_inducing+j] =
TARGET2(i, j) += k(%s); TARGET2(i, j) += k(%s);
} }
} }
%s %s
"""%(precompute_string,arg_string,"/*"+str(self._sp_k)+"*/") #adding a string representation forces recompile when needed """%(precompute_string,arg_string,"/*"+str(self._sp_k)+"*/")
# adding a string representation of the function in the
# comment forces recompile when needed
# Code to compute diagonal of covariance. # Code to compute diagonal of covariance.
@ -244,7 +291,7 @@ class spkern(Kernpart):
//#pragma omp parallel for private(j) //#pragma omp parallel for private(j)
for (i=0;i<N;i++){ for (i=0;i<N;i++){
for (j=0;j<num_inducing;j++){ for (j=0;j<num_inducing;j++){
%s %s
} }
} }
%s %s
@ -328,7 +375,11 @@ class spkern(Kernpart):
#TODO: insert multiple functions here via string manipulation #TODO: insert multiple functions here via string manipulation
#TODO: similar functions for psi_stats #TODO: similar functions for psi_stats
#TODO: similar functions when cython available.
#TODO: similar functions when only python available.
def _get_arg_names(self, Z=None, partial=None): def _get_arg_names(self, Z=None, partial=None):
arg_names = ['target','X'] arg_names = ['target','X']
for shared_params in self._sp_theta: for shared_params in self._sp_theta:
arg_names += [shared_params.name] arg_names += [shared_params.name]
@ -341,7 +392,7 @@ class spkern(Kernpart):
arg_names += ['output_dim'] arg_names += ['output_dim']
return arg_names return arg_names
def _weave_inline(self, code, X, target, Z=None, partial=None): def _generate_inline(self, code, X, target, Z=None, partial=None):
output_dim = self.output_dim output_dim = self.output_dim
for shared_params in self._sp_theta: for shared_params in self._sp_theta:
locals()[shared_params.name] = getattr(self, shared_params.name) locals()[shared_params.name] = getattr(self, shared_params.name)
@ -350,35 +401,38 @@ class spkern(Kernpart):
for split_params in self._split_theta_names: for split_params in self._split_theta_names:
locals()[split_params] = getattr(self, split_params) locals()[split_params] = getattr(self, split_params)
arg_names = self._get_arg_names(Z, partial) arg_names = self._get_arg_names(Z, partial)
if weave_available:
weave.inline(code=code, arg_names=arg_names,**self.weave_kwargs) weave.inline(code=code, arg_names=arg_names,**self.weave_kwargs)
else:
raise RuntimeError('Weave not available and other variants of sympy covariance not yet implemented')
def K(self,X,Z,target): def K(self,X,Z,target):
if Z is None: if Z is None:
self._weave_inline(self._K_code_X, X, target) self._generate_inline(self._K_code_X, X, target)
else: else:
self._weave_inline(self._K_code, X, target, Z) self._generate_inline(self._K_code, X, target, Z)
def Kdiag(self,X,target): def Kdiag(self,X,target):
self._weave_inline(self._Kdiag_code, X, target) self._generate_inline(self._Kdiag_code, X, target)
def _param_grad_helper(self,partial,X,Z,target): def _param_grad_helper(self,partial,X,Z,target):
if Z is None: if Z is None:
self._weave_inline(self._dK_dtheta_code_X, X, target, Z, partial) self._generate_inline(self._dK_dtheta_code_X, X, target, Z, partial)
else: else:
self._weave_inline(self._dK_dtheta_code, X, target, Z, partial) self._generate_inline(self._dK_dtheta_code, X, target, Z, partial)
def dKdiag_dtheta(self,partial,X,target): def dKdiag_dtheta(self,partial,X,target):
self._weave_inline(self._dKdiag_dtheta_code, X, target, Z=None, partial=partial) self._generate_inline(self._dKdiag_dtheta_code, X, target, Z=None, partial=partial)
def gradients_X(self,partial,X,Z,target): def gradients_X(self,partial,X,Z,target):
if Z is None: if Z is None:
self._weave_inline(self._dK_dX_code_X, X, target, Z, partial) self._generate_inline(self._dK_dX_code_X, X, target, Z, partial)
else: else:
self._weave_inline(self._dK_dX_code, X, target, Z, partial) self._generate_inline(self._dK_dX_code, X, target, Z, partial)
def dKdiag_dX(self,partial,X,target): def dKdiag_dX(self,partial,X,target):
self._weave.inline(self._dKdiag_dX_code, X, target, Z, partial) self._generate_inline(self._dKdiag_dX_code, X, target, Z, partial)
def compute_psi_stats(self): def compute_psi_stats(self):
#define some normal distributions #define some normal distributions
@ -407,31 +461,34 @@ class spkern(Kernpart):
self._sp_psi2 = sp.integrate(self._sp_psi2,(self._sp_x[i],-sp.oo,sp.oo)) self._sp_psi2 = sp.integrate(self._sp_psi2,(self._sp_x[i],-sp.oo,sp.oo))
clear_cache() clear_cache()
self._sp_psi2 = self._sp_psi2.simplify() self._sp_psi2 = self._sp_psi2.simplify()
def parameters_changed(self):
# Do anything here that needs to happen when parameters change, like precompute.
self._generate_inline(self._precompute, X, target, Z, partial)
def _set_params(self,param): # def _set_params(self,param):
assert param.size == (self.num_params) # assert param.size == (self.num_params)
for i, shared_params in enumerate(self._sp_theta): # for i, shared_params in enumerate(self._sp_theta):
setattr(self, shared_params.name, param[i]) # setattr(self, shared_params.name, param[i])
if self.output_dim>1: # if self.output_dim>1:
for i, split_params in enumerate(self._split_theta_names): # for i, split_params in enumerate(self._split_theta_names):
start = self.num_shared_params + i*self.output_dim # start = self.num_shared_params + i*self.output_dim
end = self.num_shared_params + (i+1)*self.output_dim # end = self.num_shared_params + (i+1)*self.output_dim
setattr(self, split_params, param[start:end]) # setattr(self, split_params, param[start:end])
def _get_params(self): # def _get_params(self):
params = np.zeros(0) # params = np.zeros(0)
for shared_params in self._sp_theta: # for shared_params in self._sp_theta:
params = np.hstack((params, getattr(self, shared_params.name))) # params = np.hstack((params, getattr(self, shared_params.name)))
if self.output_dim>1: # if self.output_dim>1:
for split_params in self._split_theta_names: # for split_params in self._split_theta_names:
params = np.hstack((params, getattr(self, split_params).flatten())) # params = np.hstack((params, getattr(self, split_params).flatten()))
return params # return params
def _get_param_names(self): # def _get_param_names(self):
if self.output_dim>1: # if self.output_dim>1:
return [x.name for x in self._sp_theta] + [x.name[:-2] + str(i) for x in self._sp_theta_i for i in range(self.output_dim)] # return [x.name for x in self._sp_theta] + [x.name[:-2] + str(i) for x in self._sp_theta_i for i in range(self.output_dim)]
else: # else:
return [x.name for x in self._sp_theta] # return [x.name for x in self._sp_theta]