Package: mlkernels
MLKernelsKernels, the machine learning ones (also, mean functions)Contents:InstallationUsageAutoGrad, TensorFlow, PyTorch, or JAX? Your Choice!Available KernelsAvailable MeansCompositional DesignDisplaying Kernels and MeansProperties of Kernels and MeansStructured Matrix TypesImplementing Your Own KernelTLDR:>>>frommlkernelsimportEQ,Linear>>>k1=2*Linear()+1>>>k12*Linear()+1>>>k1(np.linspace(0,1,3))# Structured matrices enable efficiency.<low-rankmatrix:shape=3x3,dtype=float64,rank=2left=[[0.1.][0.51.][1.1.]]middle=[[2.0.][0.1.]]right=[[0.1.][0.51.][1.1.]]>>>>importlabasB>>>B.dense(k1(np.linspace(0,1,3)))# Discard structure: get a regular NumPy array.array([[1.,1.,1.],[1.,1.5,2.],[1.,2.,3.]])>>>k2=2+EQ()*Linear()>>>k22*1+EQ()*Linear()>>>k2(np.linspace(0,1,3))<densematrix:shape=3x3,dtype=float64mat=[[2.2.2.][2.2.252.441][2.2.4413.]]>>>>B.dense(k1(np.linspace(0,1,3)))array([[2.,2.,2.],[2.,2.25,2.44124845],[2.,2.44124845,3.]])Installationpip install mlkernelsSee alsothe instructions here.UsageLetkbe a kernel, e.g.k = EQ().k(x, y)constructs thekernel matrixfor all pairs of points betweenxandy.k(x)is shorthand fork(x, x).k.elwise(x, y)constructs thekernel vectorpairing the points inxandyelement-wise, which will be arank-2 column vector.Example:>>>k=EQ()>>>k(np.linspace(0,1,3))<densematrix:shape=3x3,dtype=float64mat=[[1.0.8820.607][0.8821.0.882][0.6070.8821.]]>>>>k.elwise(np.linspace(0,1,3),0)array([[1.],[0.8824969],[0.60653066]])Letmbe a mean, e.g.m = TensorProductMean(lambda x: x ** 2).m(x)constructs themean vectorfor the points inx, which will be arank-2 column vector.Example:>>>m=TensorProductMean(lambdax:x**2)>>>m(np.linspace(0,1,3))array([[0.],[0.25],[1.]])Inputs to kernels and means are interpreted in the following way:If the inputxis arank-0 tensor, i.e. a scalar, thenxrefers to a
single input location.
For example,0simply refers to the sole input location0.If the inputxis arank-1 tensor, i.e. a vector, then every element ofxis
interpreted as a separate input location.
For example,np.linspace(0, 1, 10)generates 10 different input locations
ranging from0to1.If the inputxis arank-2 tensor, i.e. a matrix, then everyrowofxis
interpreted as a separate input location. In this case inputs are
multi-dimensional, and the columns correspond to the various input dimensions.If the inputxis atensor of rank 3 or higher, then the input is
interpreted as a batch of matrices where the matrix dimensions are the
two outermost dimensions.Example:>>>k=EQ()>>>k(0)# One scalar input<densematrix:batch=(),shape=(1,1),dtype=float64mat=[[1.]]>>>>k(np.linspace(0,1,3))# Three single-dimensional inputs<densematrix:batch=(),shape=(3,3),dtype=float64mat=[[1.0.8820.607][0.8821.0.882][0.6070.8821.]]>>>>k(np.random.randn(3,2))# Three two-dimensional inputs<densematrix:batch=(),shape=(3,3),dtype=float64mat=[[1.0.6060.399][0.6061.0.931][0.3990.9311.]]>>>>k(np.random.randn(2,3,2))# A batch of two times three two-dimensional inputs<densematrix:batch=(2,),shape=(3,3),dtype=float64mat=[[[1.0.150.559][0.151.0.678][0.5590.6781.]][[1.0.8910.627][0.8911.0.638][0.6270.6381.]]]>Finally, if you are simultaneously computing means and kernel matrices, then speed-ups
may be possible.
To obtain these speed-ups, usemean_varinstead of first callingm(x)and thenk(x);
and usemean_var_diaginstead of first callingm(x)and thenk.elwise(x).Example:>>>frommlkernelsimportmean_var,mean_var_diag>>>m=TensorProductMean(lambdax:x**2)>>>k=EQ()>>>x=np.linspace(0,1,3)>>>m(x),k(x)# Less efficient(array([[0.],[0.25],[1.]]),<densematrix:batch=(),shape=(3,3),dtype=float64mat=[[1.0.8820.607][0.8821.0.882][0.6070.8821.]]>)>>>mean_var(m,k,x)# More efficient(array([[0.],[0.25],[1.]]),<densematrix:batch=(),shape=(3,3),dtype=float64mat=[[1.0.8820.607][0.8821.0.882][0.6070.8821.]]>)>>>m(x),k.elwise(x)# Less efficient(array([[0.],[0.25],[1.]]),array([[1.],[1.],[1.]]))>>>mean_var_diag(m,k,x)# More efficient(array([[0.],[0.25],[1.]]),array([[1.],[1.],[1.]]))AutoGrad, TensorFlow, PyTorch, or JAX? Your Choice!frommlkernels.autogradimportEQ,Linearfrommlkernels.tensorflowimportEQ,Linearfrommlkernels.torchimportEQ,Linearfrommlkernels.jaximportEQ,LinearAvailable KernelsSeeherefor a nicely rendered version of the list below.Constants function as constant kernels.
Besides that, the following kernels are available:EQ(), the exponentiated quadratic:$$ k(x, y) = \exp\left( -\frac{1}{2}|x - y|^2 \right); $$CEQ(alpha), the causal exponentiated quadratic:$$ k(x, y) = \left(1 - \operatorname{erf}\left( \frac{\alpha}{4}|x - y| \right)\right)\exp\left( -\frac{1}{2}|x - y|^2 \right); $$RQ(alpha), the rational quadratic:$$ k(x, y) = \left( 1 + \frac{|x - y|^2}{2 \alpha} \right)^{-\alpha}; $$Matern12()orExp(), the Matern–1/2 kernel:$$ k(x, y) = \exp\left( -|x - y| \right); $$Matern32(), the Matern–3/2 kernel:$$ k(x, y) = \left(
1 + \sqrt{3}|x - y|
\right)\exp\left(-\sqrt{3}|x - y|\right); $$Matern52(), the Matern–5/2 kernel:$$ k(x, y) = \left(
1 + \sqrt{5}|x - y| + \frac{5}{3} |x - y|^2
\right)\exp\left(-\sqrt{3}|x - y|\right); $$Linear(), the linear kernel:$$ k(x, y) = \langle x, y \rangle; $$Delta(epsilon=1e-6), the Kronecker delta kernel:$$ k(x, y) = \begin{cases}
1 & \text{if } |x - y| < \varepsilon, \
0 & \text{otherwise};
\end{cases} $$DecayingKernel(alpha, beta):$$ k(x, y) = \frac{|\beta|^\alpha}{|x + y + \beta|^\alpha}; $$LogKernel():$$ k(x, y) = \frac{\log(1 + |x - y|)}{|x - y|}; $$PosteriorKernel(k_ij, k_zi, k_zj, z, K_z):$$ k(x, y) = k_{ij}(x, y) - k_{iz}(x, z) K_z^{-1} k_{zj}(x, y); $$SubspaceKernel(k_zi, k_zj, z, A):$$ k(x, y) = k_{iz}(x, z) A^{-1} k_{zj}(x, y); $$TensorProductKernel(f):$$ k(x, y) = f(x)f(y). $$Adding or multiplying aFunctionTypefto or with a kernel will
automatically translateftoTensorProductKernel(f). For example,f * kwill translate toTensorProductKernel(f) * k, andf + kwill
translate toTensorProductKernel(f) + k.Available MeansConstants function as constant means.
Besides that, the following means are available:TensorProductMean(f):$$ m(x) = f(x). $$Adding or multiplying aFunctionTypefto or with a mean will
automatically translateftoTensorProductMean(f). For example,f * mwill translate toTensorProductMean(f) * m, andf + mwill
translate toTensorProductMean(f) + m.Compositional DesignAdd and subtract kernels and means.Example:>>>EQ()+Matern12()EQ()+Matern12()>>>EQ()+EQ()2*EQ()>>>EQ()+1EQ()+1>>>EQ()+0EQ()>>>EQ()-Matern12()EQ()-Matern12()>>>EQ()-EQ()0Multiply kernels and means.Example:>>>EQ()*Matern12()EQ()*Matern12()>>>2*EQ()2*EQ()>>>0*EQ()0Shift kernels and means.Definition:k.shift(c)(x,y)==k(x-c,y-c)k.shift(c1,c2)(x,y)==k(x-c1,y-c2)Example:>>>Linear().shift(1)Linear()shift1>>>EQ().shift(1,2)EQ()shift(1,2)Stretch kernels and means.Definition:k.stretch(c)(x,y)==k(x/c,y/c)k.stretch(c1,c2)(x,y)==k(x/c1,y/c2)Example:>>>EQ().stretch(2)EQ()>2>>>EQ().stretch(2,3)EQ()>(2,3)Select particular input dimensions of kernels and means.Definition:k.select([0])(x,y)==k(x[:,0],y[:,0])k.select([0,1])(x,y)==k(x[:,[0,1]],y[:,[0,1]])k.select([0],[1])(x,y)==k(x[:,0],y[:,1])k.select(None,[1])(x,y)==k(x,y[:,1])Example:>>>EQ().select([0])EQ():[0]>>>EQ().select([0,1])EQ():[0,1]>>>EQ().select([0],[1])EQ():([0],[1])>>>EQ().select(None,[1])EQ():(None,[1])Transform the inputs of kernels and means.Definition:k.transform(f)(x,y)==k(f(x),f(y))k.transform(f1,f2)(x,y)==k(f1(x),f2(y))k.transform(None,f)(x,y)==k(x,f(y))Example:>>>EQ().transform(f)EQ()transformf>>>EQ().transform(f1,f2)EQ()transform(f1,f2)>>>EQ().transform(None,f)EQ()transform(None,f)Numerically, but efficiently, take derivatives of kernels and means.
This currently only works in TensorFlow and does not yet support batched
inputs.Definition:k.diff(0)(x,y)==d/d(x[:,0])d/d(y[:,0])k(x,y)k.diff(0,1)(x,y)==d/d(x[:,0])d/d(y[:,1])k(x,y)k.diff(None,1)(x,y)==d/d(y[:,1])k(x,y)Example:>>>EQ().diff(0)d(0)EQ()>>>EQ().diff(0,1)d(0,1)EQ()>>>EQ().diff(None,1)d(None,1)EQ()Make kernels periodic. This is not implemented for means.Definition:k.periodic(2pi/w)(x,y)==k((sin(w*x),cos(w*x)),(sin(w*y),cos(w*y)))Example:>>>EQ().periodic(1)EQ()per1Reverse the arguments of kernels. This does not apply to means.Definition:reversed(k)(x,y)==k(y,x)Example:>>>reversed(Linear())Reversed(Linear())Extract terms and factors from sums and products respectively.Example:>>>(EQ()+RQ(0.1)+Linear()).term(1)RQ(0.1)>>>(2*EQ()*Linear).factor(0)2Kernels and means "wrapping" others can be "unwrapped" by indexingk[0]orm[0].Example:>>>reversed(Linear())Reversed(Linear())>>>reversed(Linear())[0]Linear()>>>EQ().periodic(1)EQ()per1>>>EQ().periodic(1)[0]EQ()Displaying Kernels and MeansKernels and means have adisplaymethod.
Thedisplaymethod accepts a callable formatter that will be applied before any value
is printed.
This comes in handy when pretty printing kernels.Example:>>>print((2.12345*EQ()).display(lambdax:f"{x:.2f}"))2.12*EQ()Properties of Kernels and MeansKernels and means can be equated to check for equality.
This will attempt basic algebraic manipulations.
If the kernels and means are not equalorequality cannot be proved,
thenFalseis returned.Example of equating kernels:>>>2*EQ()==EQ()+EQ()True>>>EQ()+Matern12()==Matern12()+EQ()True>>>2*Matern12()==EQ()+Matern12()False>>>EQ()+Matern12()+Linear()==Linear()+Matern12()+EQ()# Too hard: cannot prove equality!FalseThe stationarity of a kernelkcan always be determined by queryingk.stationary.Example of querying the stationarity:>>>EQ().stationaryTrue>>>(EQ()+Linear()).stationaryFalseStructured Matrix TypesMLKernels usesan extension of LABto
accelerate linear algebra with structured linear algebra primitives.Example:>>>importlabasB>>>k=2*Delta()>>>x=np.linspace(0,5,10)>>>frommlkernelsimportpairwise>>>k(x)# Preserve structure.<diagonalmatrix:shape=10x10,dtype=float64diag=[2.2.2.2.2.2.2.2.2.2.]>>>>B.dense(k(x))# Do not preserve structure.array([[2.,0.,0.,0.,0.,0.,0.,0.,0.,0.],[0.,2.,0.,0.,0.,0.,0.,0.,0.,0.],[0.,0.,2.,0.,0.,0.,0.,0.,0.,0.],[0.,0.,0.,2.,0.,0.,0.,0.,0.,0.],[0.,0.,0.,0.,2.,0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.,2.,0.,0.,0.,0.],[0.,0.,0.,0.,0.,0.,2.,0.,0.,0.],[0.,0.,0.,0.,0.,0.,0.,2.,0.,0.],[0.,0.,0.,0.,0.,0.,0.,0.,2.,0.],[0.,0.,0.,0.,0.,0.,0.,0.,0.,2.]])These structured matrices are compatible withLAB:
they know how to efficiently add, multiply, and do other linear algebra operations.>>>importlabasB>>>B.matmul(pairwise(k,x),pairwise(k,x))<diagonalmatrix:shape=10x10,dtype=float64diag=[4.4.4.4.4.4.4.4.4.4.]>As in the above example, you can convert structured primitives to regular
NumPy/TensorFlow/PyTorch/JAX arrays by callingB.dense:>>>importlabasB>>>B.dense(B.matmul(pairwise(k,x),pairwise(k,x)))array([[4.,0.,0.,0.,0.,0.,0.,0.,0.,0.],[0.,4.,0.,0.,0.,0.,0.,0.,0.,0.],[0.,0.,4.,0.,0.,0.,0.,0.,0.,0.],[0.,0.,0.,4.,0.,0.,0.,0.,0.,0.],[0.,0.,0.,0.,4.,0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.,4.,0.,0.,0.,0.],[0.,0.,0.,0.,0.,0.,4.,0.,0.,0.],[0.,0.,0.,0.,0.,0.,0.,4.,0.,0.],[0.,0.,0.,0.,0.,0.,0.,0.,4.,0.],[0.,0.,0.,0.,0.,0.,0.,0.,0.,4.]])Implementing Your Own KernelAn example is most helpful:importlabasBfromalgebra.utilimportidenticalfrommatriximportDensefromplumimportdispatchfrommlkernelsimportKernel,pairwise,elwiseclassEQWithLengthScale(Kernel):"""Exponentiated quadratic kernel with a length scale.Args:scale (scalar): Length scale of the kernel."""def__init__(self,scale):self.scale=scaledef_compute(self,dists2):# This computes the kernel given squared distances. We use `B` to provide a# backend-agnostic implementation.returnB.exp(-0.5*dists2/self.scale**2)defrender(self,formatter):# This method determines how the kernel is displayed.returnf"EQWithLengthScale({formatter(self.scale)})"@propertydef_stationary(self):# This method can be defined to return `True` to indicate that the kernel is# stationary. By default, kernels are assumed to not be stationary.returnTrue@dispatchdef__eq__(self,other:"EQWithLengthScale"):# If `other` is also a `EQWithLengthScale`, then this method checks whether# `self` and `other` can be treated as identical for the purpose of# algebraic simplifications. In this case, `self` and `other` are identical# for the purpose of algebraic simplification if `self.scale` and# `other.scale` are. We use `algebra.util.identical` to check this condition.returnidentical(self.scale,other.scale)# It remains to implement pairwise and element-wise computation of the kernel.@pairwise.dispatchdefpairwise(k:EQWithLengthScale,x:B.Numeric,y:B.Numeric):returnDense(k._compute(B.pw_dists2(x,y)))@elwise.dispatchdefelwise(k:EQWithLengthScale,x:B.Numeric,y:B.Numeric):returnk._compute(B.ew_dists2(x,y))>>>k=EQWithLengthScale(2)>>>kEQWithLengthScale(2)>>>k==EQWithLengthScale(2)True>>>2*k==k+EQWithLengthScale(2)True>>>k==Linear()False>>>k_composite=(2*k+Linear())*RQ(2.0)>>>k_composite(2*EQWithLengthScale(2)+Linear())*RQ(2.0)>>>k_composite(np.linspace(0,1,3))<densematrix:shape=3x3,dtype=float64mat=[[2.1.7171.13][1.7172.252.16][1.132.163.]]>Of course, in practice we do not need to implement variants of kernels which include
length scales, because we always adjust the length scale by stretching a base kernel:>>>EQ().stretch(2)(np.linspace(0,1,3))<densematrix:shape=3x3,dtype=float64mat=[[1.0.9690.882][0.9691.0.969][0.8820.9691.]]>>>>EQWithLengthScale(2)(np.linspace(0,1,3))<densematrix:shape=3x3,dtype=float64mat=[[1.0.9690.882][0.9691.0.969][0.8820.9691.]]>