How can I optimize sampling rate in a matrix based Bayesian Model?

Questions version agnostic
1

Hello Everybody, I need a bit of help regarding a project that employs Bayesian modeling to determine parameters for fluorescence decay of excited fluorophores for use in FRET modeling.

(I am running PyMC 5.16.2 with numpy 1.26 on a MacBook with Apple Silicon if this serves any value)

My current model code is as follows:

def fluorescenceModel(r, t, R_0, f_D, data, numGauss, numComponents) -> pm.Model:

    deltaR = r[1] - r[0]

    forsterAcceleration = 1 + np.power(R_0 / r, 6)

    T = t[:, np.newaxis]

    with pm.Model() as fluorescentModel:

        tauFluorescence = pm.TruncatedNormal('rateFluorescence', mu=70, sigma=8, lower=0.0, shape=numComponents)
        alphaFluorescence = pm.Dirichlet('alphaF', a=np.ones(numComponents), shape=numComponents)

        alphaDistance = pm.Dirichlet('alphaD', a=np.ones(numGauss), shape=numGauss)
        muDistance = pm.TruncatedNormal('r_0D', mu=50, sigma=60, lower=0.0, upper=120, shape=numGauss)
        sigmaDistance = pm.TruncatedNormal('sigmaD', mu=10, sigma=13, lower=0.0, upper=30, shape=numGauss)

        K_Transpose = (1 / tauFluorescence)[np.newaxis, :] # 1 x m vector/matrix

        Alpha = alphaFluorescence[:, np.newaxis]

        lambdaArg = - T @ K_Transpose # n x m

        Lambda = np.exp(lambdaArg)

        I_D = Lambda @ Alpha # n x 1, where n = t.size

        AcceleratedTimeMatrix = T @ forsterAcceleration[np.newaxis, :]

        Epsilon = pt.zeros((t.size, r.size)) # n x k

        for j in range(r.size):
            Epsilon = pt.set_subtensor(Epsilon[j, :], expData(AcceleratedTimeMatrix[j, :], tauFluorescence, alphaFluorescence))

        
        P = pt.as_tensor_variable(distDistribution(r, alphaDistance, muDistance, sigmaDistance, numGauss) * deltaR)
        
        I_DA = Epsilon @ P
        
        I_F = f_D * I_D + (1 - f_D) * I_DA
        
        likelihood = pm.Poisson('likelihood', mu=I_F, observed=data[:, np.newaxis])
        
        return fluorescentModel

with some helper functions I wrote being

def distDistribution(r, alpha, mu, sigma, numComp):
    P = 0
    for i in range(numComp):
        P += (alpha[i] / sigma[i] ) * np.exp(- 0.5 * np.square( (r[:, np.newaxis] - mu[i]) / sigma[i]))

    return P / np.sqrt(2 * np.pi)

and

def expData(
        domain: np.ndarray,
        tau: np.ndarray,
        alphas: np.ndarray) -> np.ndarray:

    arg = domain[:, np.newaxis] / tau[np.newaxis, :] # n x m array


    arg = np.exp(-arg) # n x m array

    ans = arg @ alphas[:, np.newaxis] # n x m @ m x 1 = n x 1 2D array

   return ans.flatten()

The model works for miniscule ranges of inputs, such as when the r and t parameters are np.ndarrays of size 5 and 6, respectively. However, when I try increasing the size of the parameters (I obtain them using np.linspace), the time required blows up massively to a point where I haven’t yet managed to obtain samples even after waiting for ~1 hour. I have verified through the use of rudimentary print statements that the graph computations for the matrices are done fairly quickly ( < 0.01 sec) after the method is called with the appropriate arguments.

Does anyone have any recommendations on how I can possibly optimize my model? In addition, if anything seems glaringly wrong with the implementation, please do let me know.
(Briefly edited as a function definition was cutoff in OG)

2

For recursive computations that are long enough you’ll probably need to use Scan: scan – Looping in PyTensor — PyTensor dev documentation

A different nuts sampler like numpyro may be needed to stay performant: Faster Sampling with JAX and Numba — PyMC example gallery

1 Like
3

Hello, thank you for suggesting the use of scan. I had previously tried to use it and had failed to make it work, so I tried it again. I tried to use the following to recursively construct the array I need, but an issue that is popping up is that PyTensor is not liking that I am passing in a (Random)Variable instead of a numeric array, with an error saying “Expected an array-like object, but found a Variable: maybe you are trying to call a function on a (possibly shared) variable instead of a numeric array”. Do you have any advice on how I can get PyTensor to operate on the RandomVariables rather than arrays?

def monogaussianHelper2(r, fluorescenceModel, t, forsterAcceleration):
    deltaR = r[1] - r[0]
    # multiply final output by deltaR

    rF = fluorescenceModel['rateFluorescence']
    #rF = pm.TruncatedNormal('rateF', mu=70, sigma=80, lower=0.0)
    aD = fluorescenceModel['alphaDD']
    mD = fluorescenceModel['r_0DD']
    sD = fluorescenceModel['sigmaDD']

    R = pt.matrix("r") # m x 1 dimensions
    F = pt.matrix("forsterAcceleration") # same dimensions as R (m x 1 dimensions)
    T = pt.matrix("t") # n x 1 dimensions



    def accumulatorFunction(r, forsterAcceleration, prevArray, t, rateF, alphaD, muD, sigmaD):
        
        P_r = distDistribution(r, alphaD, muD, sigmaD, sigmaD.eval().size)

        donorAcceptorTerm = pt.exp(- t * rateF * forsterAcceleration) #check if this causes issues
        
        return pt.inc_subtensor(prevArray[:,:], P_r * donorAcceptorTerm * deltaR)
    
    initialCondition = pt.zeros((t.size, 1))

    results, updates = pytensor.scan(
        fn=accumulatorFunction,
        sequences=[R, F],
        outputs_info=initialCondition, # this may be the issue, as there are random variables involved whereas a matrix is likely determined already
        non_sequences=[T, rF, aD, mD, sD],
        n_steps=R.shape[0])
    
    answer = results[-1]

    compiledFunction = pytensor.function(inputs=[R, F,T, rF, aD, mD, sD], outputs=answer, updates=updates)
    
    return compiledFunction(r[:, None], forsterAcceleration[:, None], t[:, None], rF, aD, mD, sD)[:, None]

def fluorescenceModelPyTensor2(r, t, data, numGauss, numComponents):


    R_0 = 67
    f_Donor = 0.07 #generally below .10


    forsterAcceleration = np.power(R_0 / r, 6)



    with pm.Model() as fluorescentModel:
        rateFluorescence = pm.TruncatedNormal('rateFluorescence', mu=70, sigma=80, lower=0.0, shape=numComponents) #can put shape = numComponents when normal works

        alphaDistance = pm.Dirichlet('alphaDD', a=np.ones(numGauss), shape=numGauss)
        muDistance = pm.TruncatedNormal('r_0DD', mu=50, sigma=60, lower=0.0, upper=120, shape=numGauss)
        sigmaDistance = pm.TruncatedNormal('sigmaDD', mu=10, sigma=13, lower=0.0, upper=30, shape=numGauss)

        I_D = np.exp(- t[:, np.newaxis] * rateFluorescence)

        I_DA = monogaussianHelper2(r, fluorescentModel, t, forsterAcceleration)

        I_F = f_Donor * I_D + (1 - f_Donor) * I_DA

        likelihood = pm.Poisson("likelihood", mu=I_F, observed=data[:, None])

        with fluorescentModel:
            trace = pm.sample(1000)

I will also try using a different NUTS sampler to try and make sure that PyMC can keep up sampling rates (once I have remediated the main issue of course).

4

The error you mention is likely coming from numpy or scipy, not from pytensor. You cannot use numpy functions on symbolic pytensor variables (or random variables, which are the same thing).