google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Definitions of kernels for Gaussian Process models for UQ experiments."""
17

18
from __future__ import absolute_import
19
from __future__ import division
20
from __future__ import print_function
21

22
import tensorflow.compat.v2 as tf
23
import tensorflow_probability as tfp
24

25

26
class RBFKernelFn(tf.keras.layers.Layer):
27
  """ExponentiatedQuadratic kernel provider."""
28

29
  def __init__(self,
30
               num_classes,
31
               per_class_kernel,
32
               feature_size,
33
               initial_amplitude,
34
               initial_length_scale,
35
               initial_linear_bias,
36
               initial_linear_slope,
37
               add_linear=False,
38
               name='vgp_kernel',
39
               **kwargs):
40
    super(RBFKernelFn, self).__init__(**kwargs)
41
    self._per_class_kernel = per_class_kernel
42
    self._initial_linear_bias = initial_linear_bias
43
    self._initial_linear_slope = initial_linear_slope
44
    self._add_linear = add_linear
45

46
    with tf.compat.v1.variable_scope(name):
47
      if self._per_class_kernel and num_classes > 1:
48
        amplitude_shape = (num_classes,)
49
        length_scale_shape = (num_classes, feature_size)
50
      else:
51
        amplitude_shape = ()
52
        length_scale_shape = (feature_size,)
53

54
      self._amplitude = self.add_variable(
55
          initializer=tf.constant_initializer(initial_amplitude),
56
          shape=amplitude_shape,
57
          name='amplitude')
58

59
      self._length_scale = self.add_variable(
60
          initializer=tf.constant_initializer(initial_length_scale),
61
          shape=length_scale_shape,
62
          name='length_scale')
63

64
      if self._add_linear:
65
        self._linear_bias = self.add_variable(
66
            initializer=tf.constant_initializer(self._initial_linear_bias),
67
            shape=amplitude_shape,
68
            name='linear_bias')
69
        self._linear_slope = self.add_variable(
70
            initializer=tf.constant_initializer(self._initial_linear_slope),
71
            shape=amplitude_shape,
72
            name='linear_slope')
73

74
  def call(self, x):
75
    # Never called -- this is just a layer so it can hold variables
76
    # in a way Keras understands.
77
    return x
78

79
  @property
80
  def kernel(self):
81
    k = tfp.math.psd_kernels.FeatureScaled(
82
        tfp.math.psd_kernels.ExponentiatedQuadratic(
83
            amplitude=tf.nn.softplus(self._amplitude)),
84
        scale_diag=tf.math.sqrt(tf.nn.softplus(self._length_scale)))
85
    if self._add_linear:
86
      k += tfp.math.psd_kernels.Linear(
87
          bias_amplitude=self._linear_bias,
88
          slope_amplitude=self._linear_slope)
89
    return k
90

91

92
class MaternKernelFn(tf.keras.layers.Layer):
93
  """Matern kernel provider."""
94

95
  def __init__(self,
96
               num_classes,
97
               degree,
98
               per_class_kernel,
99
               feature_size,
100
               initial_amplitude,
101
               initial_length_scale,
102
               initial_linear_bias,
103
               initial_linear_slope,
104
               add_linear=False,
105
               name='vgp_kernel',
106
               **kwargs):
107
    super(MaternKernelFn, self).__init__(**kwargs)
108
    self._per_class_kernel = per_class_kernel
109
    self._initial_linear_bias = initial_linear_bias
110
    self._initial_linear_slope = initial_linear_slope
111
    self._add_linear = add_linear
112

113
    if degree not in [1, 3, 5]:
114
      raise ValueError(
115
          'Matern degree must be one of [1, 3, 5]: {}'.format(degree))
116

117
    self._degree = degree
118

119
    with tf.compat.v1.variable_scope(name):
120
      if self._per_class_kernel and num_classes > 1:
121
        amplitude_shape = (num_classes,)
122
        length_scale_shape = (num_classes, feature_size)
123
      else:
124
        amplitude_shape = ()
125
        length_scale_shape = (feature_size,)
126

127
      self._amplitude = self.add_variable(
128
          initializer=tf.constant_initializer(initial_amplitude),
129
          shape=amplitude_shape,
130
          name='amplitude')
131

132
      self._length_scale = self.add_variable(
133
          initializer=tf.constant_initializer(initial_length_scale),
134
          shape=length_scale_shape,
135
          name='length_scale')
136

137
      if self._add_linear:
138
        self._linear_bias = self.add_variable(
139
            initializer=tf.constant_initializer(self._initial_linear_bias),
140
            shape=amplitude_shape,
141
            name='linear_bias')
142
        self._linear_slope = self.add_variable(
143
            initializer=tf.constant_initializer(self._initial_linear_slope),
144
            shape=amplitude_shape,
145
            name='linear_slope')
146

147
  def call(self, x):
148
    # Never called -- this is just a layer so it can hold variables
149
    # in a way Keras understands.
150
    return x
151

152
  @property
153
  def kernel(self):
154
    if self._degree == 1:
155
      kernel_class = tfp.math.psd_kernels.MaternOneHalf
156
    if self._degree == 3:
157
      kernel_class = tfp.math.psd_kernels.MaternThreeHalves
158
    if self._degree == 5:
159
      kernel_class = tfp.math.psd_kernels.MaternFiveHalves
160

161
    k = tfp.math.psd_kernels.FeatureScaled(
162
        kernel_class(amplitude=tf.nn.softplus(self._amplitude)),
163
        scale_diag=tf.math.sqrt(tf.nn.softplus(self._length_scale)))
164
    if self._add_linear:
165
      k += tfp.math.psd_kernels.Linear(
166
          bias_amplitude=self._linear_bias,
167
          slope_amplitude=self._linear_slope)
168
    return k
169

170

171
class LinearKernelFn(tf.keras.layers.Layer):
172
  """Matern kernel provider."""
173

174
  def __init__(self,
175
               num_classes,
176
               per_class_kernel,
177
               initial_linear_bias,
178
               initial_linear_slope,
179
               name='vgp_kernel',
180
               **kwargs):
181
    super(LinearKernelFn, self).__init__(**kwargs)
182
    self._per_class_kernel = per_class_kernel
183
    self._initial_linear_bias = initial_linear_bias
184
    self._initial_linear_slope = initial_linear_slope
185

186
    with tf.compat.v1.variable_scope(name):
187
      if self._per_class_kernel and num_classes > 1:
188
        shape = (num_classes,)
189
      else:
190
        shape = ()
191

192
      self._linear_bias = self.add_variable(
193
          initializer=tf.constant_initializer(self._initial_linear_bias),
194
          shape=shape,
195
          name='linear_bias')
196
      self._linear_slope = self.add_variable(
197
          initializer=tf.constant_initializer(self._initial_linear_slope),
198
          shape=shape,
199
          name='linear_slope')
200

201
  def call(self, x):
202
    # Never called -- this is just a layer so it can hold variables
203
    # in a way Keras understands.
204
    return x
205

206
  @property
207
  def kernel(self):
208
    return tfp.math.psd_kernels.Linear(
209
        bias_amplitude=self._linear_bias,
210
        slope_amplitude=self._linear_slope)
211