Indexing an array with a theano vector

I am trying to address array elements by index with a Theano vector created in pymc3. Although I have cast the vector as int64, using theano.tensor.cast(), I still get an indexing error. In the code below I am modifying a beta prior distribution with an interpolating function I have written, but I need to use a similar interpolating function again in another part of the code, and will get a similar error. (So, I’m not just looking for a different way to modify my prior, but a way to properly index the array. Thanks for any help!

Here is the code:

# minimal indexing example

import numpy as np
import pymc3 as pm
import theano.tensor as tht

def cubic_spline(x, y):
  """
  x  : list of floats
  y  : list of floats

  Returns list of list of floats
  """
  n = len(x)
  b = [0] * (n)
  c = [0] * (n)
  d = [0] * (n)
  return [y, b, c, d]

def pri(x,xs,a):
    """
    Returns interpolated y value given x input, x values use to create spline (xs),
    and first x (x0) and Δx values for xs.
    """
    i=tht.cast(np.floor((x-xs[0])/(xs[1]-xs[0])),'int64')
    return a[0][i]+a[1][i]*(x-xs[i])+a[2][i]*(x-xs[i])**2+a[3][i]*(x-xs[i])**3  

def fx(x,p,ra):
    """
    """
    return(np.sin(x/p)+ra*np.random.random())

#Dummy Data
Npts=10000
xs=np.linspace(0,20,Npts)
ys=fx(xs,5.,0)

#Interpolation function to change range and distribution of prior
NSpts=100
sxs=np.linspace(0,1,NSpts)
sys=np.linspace(0,1,NSpts)
a=cubic_spline(sxs, sys)

#Call pymc3
with pm.Model() as model:
    pr = pri(pm.Beta('pr', alpha=1,beta=1),sxs,a)  
    eps = .1
    mu=fx(xs,pr,0.)
    y_pred = pm.Normal('y_pred', mu=mu, sd=eps, observed=ys)
    step = pm.NUTS()
    trace = pm.sample(1000)

The Error:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-9-7c02edbe0924> in <module>()
     12 #Call pymc3
     13 with pm.Model() as model:
---> 14     pr = pri(pm.Beta('pr', alpha=1,beta=1),sxs,a)
     15     eps = .1
     16     mu=fx(xs,pr,0.)

<ipython-input-8-4beabf6a93a6> in pri(x, xs, a)
     26     """
     27     i=tht.cast(np.floor((x-xs[0])/(xs[1]-xs[0])),'int64')
---> 28     return a[0][i]+a[1][i]*(x-xs[i])+a[2][i]*(x-xs[i])**2+a[3][i]*(x-xs[i])**3
     29 
     30 def fx(x,p,ra):

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

You cannot index a numpy array or a list using a theano tensor. You can either cast the input into theano tensor:

import theano
def pri(x,xs,a):
    """
    Returns interpolated y value given x input, x values use to create spline (xs),
    and first x (x0) and Δx values for xs.
    """
    i=tht.cast(np.floor((x-xs[0])/(xs[1]-xs[0])),'int64')
    y = theano.shared(np.asarray(a[0]))
    b = theano.shared(np.asarray(a[1]))
    c = theano.shared(np.asarray(a[2]))
    d = theano.shared(np.asarray(a[3]))
    xs_ = theano.shared(np.asarray(xs))
    return y[i]+b[i]*(x-xs_[i])+c[i]*(x-xs_[i])**2+d[i]*(x-xs_[i])**3  

or wrap your function with theano.as_op

Also, there is room for improvement of the implementation, you might find the Interpolated class in pymc3 useful as a point of reference.

That works. Thank you! And thanks for pointing me to the Interpolated class.

1 Like