Hi there,
Apologies for the delay in response.
The function I was performing inference on is simply
where M is a diagonal 2 x 2 matrix, comprising unkowns m_1 and m_2 .
Code thus far is as follows:
import pytensor
import numpy
import pytensor.tensor as pt
from pytensor import function
import numpy as np
import pymc as pm
import matplotlib.pyplot as plt
import arviz as az
from pymc import pytensorf
plt.close('all')
def generate_synthetic_data():
#function for observed data
# Define the shared variable for the frequencies 'w'
f = np.linspace(0.1, 1.0, 100)
w_values = 2*np.pi*f
# Define the 2x2 matrix M as a shared variable, the non-zero values are to be inferred by this routine.
M_shared = np.array([[8,0],[0,2]])
#Define Z function (true values)
Z_true = np.zeros((len(f),2,2),dtype = complex)
for i in range(len(f)):
Z_true[i,:,:] = 1j*w_values[i]*M_shared
# Add noise to create the observed data
n_std = 0.01
Z_observed_values = np.abs(Z_true) + np.random.normal(0, n_std, Z_true.shape)
return f, w_values, np.abs(Z_true), Z_observed_values # Return w and Z_obs for later use
def create_pymc_model(w, Z_obs):
# Create a PyMC model to sample masses and sigma
with pm.Model() as model:
# Priors for unknown model parameters
mass1 = pm.Uniform("mass1", 0.1, 10)
mass2 = pm.Uniform("mass2", 0.1, 10)
sigma = pm.HalfNormal("sigma", 1)
mass_mat = pt.stacklists([[mass1, 0], [0, mass2]])
# loop step function over w to calculate/predict Z with priors
def step(w, mass_mat):
Z_pred = 1j*w*mass_mat
return np.abs(Z_pred)
# store results using scan instead of for looping
results, _ = pytensor.scan(fn=step, sequences=[w], outputs_info=None, non_sequences=[mass_mat])
# Likelihood term
likelihood = pm.Normal("likelihood", mu=results.flatten(), sigma=sigma, observed=Z_obs.flatten())
# Run the MCMC sampling
trace = pm.sample(1000, tune=1000,step=pm.Metropolis(),chains = 4, cores=4)
return trace
def main():
plt.close('all')
#main section and plotting
f, w, Z_true, Z_obs = generate_synthetic_data()
trace = create_pymc_model(w, Z_obs)
az.plot_trace(trace);
if __name__ == "__main__":
main()
Based on that, I hope I am right in saying that
From here, would it be best to take the absolute value and leave it in matrix form? or take the absolute value and flatten/vectorise, so that there is a single scalar value at each point in the flattened vector (though there will be gradients of zero for half the points)?
Also, as an aside (and sorry for squeezing this in here 
) - there will be a need for me in the future to add more lines to the scan element of the code so that I can use more complicated functions involving matrix inversions. An example extending the scan function in my previous code block is provided below…
def step(w, mass_mat):
Z_pred = 1j*w*mass_mat
Y_pred = pt.linalg.inv(Z_pred)
return pt.abs(Y_pred)
results, _ = pytensor.scan(fn=step, sequences=[w], outputs_info=None, non_sequences=[mass_mat])
It would appear that the inference goes out the window when I define the scan in this way. If I do it all in one line, then the inference works ok. NOTE - I realise this is applicable for MCMC-MH type algorithms and not gradient-based ones.
def step(w, mass_mat):
Y_pred = pt.linalg.inv(1j*w*mass_mat)
return pt.abs(Y_pred)
I could code something up like this super easy via ‘for’ loops but trying to stick to the rules of using scan. I’m at a loss as to why this added step ceases to work but executing in one line does not. Is there a way to use multiple lines as intended? I hope it is a simple error on my part.
Many thanks once again.
M.