stein_ksd/kernels.py at master · singhalrk/stein_ksd · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
import autograd.numpy as np
from autograd import grad, jacobian

"""
Add some commong kernels
- rbf
- imq
- polynomial
 ------ find some more  ---------
"""


######## make this vectorized ##########

class rbf_kernel:
    def __init__(self, params = dict(beta=1e-2)):
        self.beta = params['beta']
        assert self.beta > 0

    def value(self, x, y):
        r = ((x - y)**2).sum()
        return np.exp(-self.beta * r)

    def grad_x(self, x, y):
        r = ((x - y)**2).sum()
        return - 2 * self.beta * (x - y) * np.exp(-self.beta * r)

    def grad_y(self, x, y):
        r = ((x - y)**2).sum()
        return 2 * self.beta * (x - y) * np.exp(-self.beta * r)

    def grad_xy(self, x, y):
        assert len(x) == len(y)
        n = len(x)
        r = ((x - y)**2).sum()
        _y = 2 * self.beta * np.exp(-self.beta * r) * np.ones(n)
        _xy = 4 * self.beta**2 * (x - y)**2 * np.exp(-self.beta * r)
        return _y + _xy

class imq_kernel:
    def __init__(self, params = dict(beta = 1, c=1)):
        self.beta = params['beta']
        self.c = params['c']
        assert self.beta > 0

    def value(self, x, y):
        r = ((x - y)**2).sum()
        return 1./(self.c**2 + r)**(self.beta)

    def grad_x(self, x, y):
        r = ((x - y)**2).sum()
        return -2 * self.beta * (x - y) / (self.c**2 + r)**(self.beta + 1.)

    def grad_y(self, x, y):
        r = ((x - y)**2).sum()
        return 2 * self.beta * (x - y) / (self.c**2 + r)**(self.beta + 1.)

    def grad_xy(self, x, y):
        r = ((x - y)**2).sum()
        n = len(x)
        _y = 2 * self.beta * np.ones(n) / (self.c**2 + r)**(self.beta + 1.)
        _xy = -4 * self.beta* (self.beta + 1) * (x - y)**2 / (self.c**2 + r)**(self.beta + 2.)
        return _y + _xy

class poly_kernel:
    def __init__(self, params = dict(c=1, degree=1)):
        self.degree = params['degree']
        self.c = params['c']
        assert self.degree > 0

    def value(self, x, y):
        r = np.dot(x,y)
        return (self.c + r) ** self.degree

    def grad_x(self, x, y):
        r = np.dot(x, y)
        return self.degree * (self.c + r) ** (self.degree - 1) * y

    def grad_y(self, x, y):
        r = np.dot(x, y)
        return self.degree * (self.c + r) ** (self.degree - 1) * x

    def grad_xy(self, x, y):
        n = len(x)
        r = np.dot(x, y)
        return self.degree * (self.c + r) ** (self.degree - 1) * np.ones(n)

class kernels_1:
    def __init__(self, name='rbf', beta=None, c=None, degree=None):
        self.name = name
        self.beta = beta
        self.c = c
        self.degree = degree
        self.params = dict(beta=self.beta, degree=self.degree, c=self.c)

##  write kernels gradients manually
class kernels:
    def __init__(self, name='rbf', beta=None, c=None, degree=None):
        self.name = name
        self.beta = beta
        self.c = c
        self.degree = degree

    def rbf(self, x, y, beta=0.01):
        return np.exp(- ((x - y)**2).sum() * beta)

    def imq(self, x, y, c=1, beta=-1):
        return (c**2 + ((x - y)**2).sum())**(-beta)

    def polynomial(self, x, y, c=1, degree=2):
        return (c + np.dot(x,y))**degree

    # here we fix y and take a derivative wrt x
    def grad_kx(self,x,y):
        k = self.get_kernel(self.name)
        k_x = lambda x_: k(x_,y)
        return grad(k_x)(x)

    # here we fix x and take a derivative wrt y
    def grad_ky(self,x, y):
        k = self.get_kernel(self.name)
        k_y = lambda y_: k(x,y_)
        return grad(k_y)(y)

    # here we take the derivative of k wrt x and y
    def grad_kxy(self,x,y):
        kx_y = lambda y_: self.grad_kx(x, y_)
        return jacobian(kx_y)(y)

    def get_kernel(self,name=None):
        if name == None:
            name = self.name

        if name == 'rbf':
            return self.rbf

        if name == 'imq':
            return self.imq

        if name == 'polynomial':
            return self.polynomial