1D Convolutions in pm.Model

1

Hi
I am new to pymc and to wanted to implement an easy example of histogram data
with an exponential shape that is convolved with a gaussian kernel. For the convolution I’m using the function shown here How to apply signal filtering in PyMC? but when I run my code

import pytensor
import pytensor.tensor as pt

# Generate toy data
xlim = (0, 3000)
bin_width = 2
bins = np.arange(xlim[0], xlim[1], bin_width)
bin_midpoints = (bins[1:] + bins[:-1]) / 2
A_true = 1e4
tau_true = 1e-3
counts = A_true * np.exp(-tau_true * bin_midpoints)
counts += np.random.randn(counts.shape[0])*100

# Create Gaussian kernel
kernel_gaussian_range = 5
kernel_sigma_mean = 3.3
kernel_sigma_std = 0.1
kernel_range = np.arange(-kernel_gaussian_range*kernel_sigma_mean, kernel_gaussian_range*kernel_sigma_mean, bin_width)
kernel_values = np.exp((-kernel_range ** 2) / (2 * kernel_sigma_mean ** 2))
kernel_values /= np.sum(kernel_values)  

observed_counts = np.convolve(counts, kernel_values, mode='same')

def conv_1d(idx, data, kernel):
    return (data[idx:idx+len(kernel_range)] * kernel).sum()

with pm.Model() as model:
    # Priors for the parameters
    A_log = pm.Normal('A_log', mu=np.log(A_true)/2, sigma=0.1)  
    tau_log = pm.Normal('tau_log', mu=np.log(tau_true)/2, sigma=0.1)  
    sigma = pm.Normal('sigma', mu=kernel_sigma_mean, sigma=kernel_sigma_std)  # Sigma of the Gaussian kernel

    # Transform parameters back to original scale
    A = pm.Deterministic('A', pm.math.exp(A_log))
    tau = pm.Deterministic('tau', pm.math.exp(tau_log))

    # Generate Gaussian kernel based on sampled sigma
    kernel = pm.Deterministic('kernel', pm.math.exp(-pm.math.sqr(kernel_range) / (2 * pm.math.sqr(sigma))))
    kernel /= pm.math.sum(kernel)  # Normalize the kernel

    expected_counts = A * pm.math.exp(-tau * bin_midpoints)

    convolved_counts = pt.reshape(
            pt.conv.conv2d(
                expected_counts.reshape((1, 1, 1, expected_counts.shape[0])),
                kernel.reshape((1, 1, 1, kernel.shape[0])),
                border_mode="valid"
            ), (expected_counts.shape[0] - (kernel.shape[0] - 1),)
        ).eval()

    # convolved_counts, _ = pytensor.scan(conv_1d,
    #                                     sequences = pytensor.tensor.arange(len(observed_counts) - len(kernel_range) + 1),
    #                                     non_sequences=[expected_counts, kernel])

    # Likelihood (sampling distribution) of the observed data
    counts_likelihood = pm.Poisson('counts_likelihood', mu=convolved_counts, observed=observed_counts)

    # Run the MCMC sampling
    trace = pm.sample(2000, tune=2000, cores=4)

I get the error ValueError: Incompatible Elemwise input shapes [(1499,), (1483,)]. I also tried using the scan function as you can see in the code but I get the same error, which is even more confusing to me since this functions works just fine if I compile it with pymc_convolution = pytensor.function(inputs=[data, kernel], outputs=rolling_average) and pass it two numpy arrays.

Any help would be greatly appreciated :slight_smile:

2

You are using border_mode="valid". Checking the docstring for conv2d:

    ``'valid'``: apply filter wherever it completely overlaps with the
        input. Generates output of shape: input shape - filter shape + 1

So you end up dropping a number of observations equal to the width of your kernel. If you check the shape error, you see that 1499 - 1483 = 16, which is one less than the kernel width of 17, as promised.

Since your model cannot produce estimates for the first 16 datapoints (they are consumed to make the prediction for the 17th), you need to drop these from your observed data. Alternatively, you could pad your data with a length 16 initial state vector, whose values you would need to estimate.

Your final choice is to fiddle with the handling of the boundaries. This gist might be helpful if you choose to go that route.