Making pymc wrapper for finite element

Questions v5
1

Hi,

I tried to wrap another python software with pymc. My goal is that to use pymc sampling compatible for a finite element analysis input and output. The input for finite element is three parameters (a0,a1,a2) from normal sampling in pymc that needs to be converted into a float (to be used for each separate analysis) and the output is 1x20 displacement vector in numpy array. These outputs would act as mu for likelihood factor from multivariate normal with observation y is also a 1x20 displacement vector in numpy array.

I tried to use the op to wrap things but so far my implementation resulting in “NotImplementedError: Cannot convert <function FEmcOperator at 0x28c676160> to a tensor variable.” Help is appreciated as I am not understand on how the priors sampling would act in this function. Thanks.

def fenicsFEmc(a0,a1,a2):
    g0 = float(a0)
    g1 = float(a1)
    g2 = float(a2)
    dim = domain.topology.dim
    degree = dim
    shape = (dim,)  
    V = fem.functionspace(domain, ("P", degree, shape))
    x = ufl.SpatialCoordinate(domain)
    u_sol = fem.Function(V, name="Displacement")
    E = (g0 + g1 * ufl.classes.Sin(np.pi * x[0]) + g2 * ufl.classes.Sin(2. * np.pi * x[0]))
    nu = 0.3
    lmbd = E * nu / (1 + nu) / (1 - 2 * nu)
    mu_ = E / 2 / (1 + nu)

    W = fem.functionspace(domain, ("Lagrange", 1))

    lmbda_exp = fem.Expression(lmbd, W.element.interpolation_points())
    mu_exp = fem.Expression(mu_, W.element.interpolation_points())

    lmbda = fem.Function(W)
    mu = fem.Function(W)

    lmbda.interpolate(lmbda_exp)
    mu.interpolate(mu_exp)
    
    u = TrialFunction(V)
    v = TestFunction(V)

    f = fem.Constant(domain, np.array([1.]))

    dx = Measure("dx", domain=domain)
    a = inner(sigma(u, lmbda, dim, mu), epsilon(v)) * dx
    L = inner(f, v) * dx
    
    left_dofs = fem.locate_dofs_geometrical(V, left)
    right_dofs = fem.locate_dofs_geometrical(V, right)

    bcs = [fem.dirichletbc(np.zeros((1,)), left_dofs, V), fem.dirichletbc(np.zeros((1,)), right_dofs, V),]
    problem = fem.petsc.LinearProblem(a, L, u=u_sol, bcs=bcs, petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
    problem.solve()

    u = u_sol.x.array
    return u

#Create the op function to wrap things from fenics
class FiniteElementOp(Op):
    def make_node(self,a0,a1,a2) -> Apply:
        # Convert inputs to tensor variables
        a0 = pt.as_tensor_variable(a0)
        a1 = pt.as_tensor_variable(a1)
        a2 = pt.as_tensor_variable(a2)
        inputs = [a0,a1,a2]

        outputs = [pt.tensor()]

        return Apply(self, inputs, outputs)

    def perform(self, node: Apply, inputs: list[np.ndarray], outputs: list[list[None]]) -> None:
        a0,a1,a2 = inputs  # this will contain my variables

        # call fenics function
        ufe = fenicsFEmc(a0,a1,a2)

        outputs[0][0] = np.asarray(ufe, dtype=node.outputs[0].dtype)

def FEmcOperator(a0,a1,a2):
    return FiniteElement(a0,a1,a2)

# Create covariance matrix
covariance = np.eye(len(x_cent)) * perturb

#building pymc model
with pm.Model() as model:
    # Priors
    a0 = pm.Normal('a0', mu=g0, sigma=1.)
    a1 = pm.Normal('a1', mu=g1, sigma=0.5)
    a2 = pm.Normal('a2', mu=g2, sigma=0.2)

    # Likelihood
    likelihood = pm.MvNormal('likelihood', mu=FEmcOperator, cov = covariance, observed=y)

    # MCMC sampling
    trace = pm.sample(draws=1000, tune=1000)
2

You are passing a function object to mu in your likelihood term. This is what is producing the error. You need to actually call the function on inputs a0, a1, a2

3

I called the function with “FEmcOperator(h0,h1,h2)” but now they seem like in an infinite loop.Or should I recreate all the pm.model with custom dist? I am sampling the function FEmcOperator as mean (mu) for multivariate normal sampling with covariance, observation “y”, and inputs “a0,a1,a2”.

I also tried to testing code using the guide in this link and it gives me the result that is expected (1x20 numpy array with displacement vector u).

Testing code :

finiteelement_op = fenicsFEmc()
test_out = finiteelement_op(g0,g1,g2)
test_out.eval()

Sampling code :

# Create covariance matrix
covariance = np.eye(len(x_cent)) * perturb

#building pymc model
with pm.Model() as model:
    # Priors
    a0 = pm.Normal('a0', mu=g0, sigma=1.)
    a1 = pm.Normal('a1', mu=g1, sigma=0.5)
    a2 = pm.Normal('a2', mu=g2, sigma=0.2)
    
    h0 = pt.as_tensor_variable(a0)
    h1 = pt.as_tensor_variable(a1)
    h2 = pt.as_tensor_variable(a2)

    # Likelihood
    fe= FEmcOperator(h0,h1,h2)
    likelihood = pm.MvNormal('likelihood', mu=fe, cov = covariance, observed=y)

with model:
    # Use custom number of draws to replace the HMC based defaults
    trace = pm.sample(1000, tune=1000)

# plot the traces
az.plot_trace(trace);
4

I’m not sure what you mean by infinite loop. Can you call pm.draw on your likelihood variable?

Whether to use a CustomDist or not is a modeling question. Right now you are modeling your data as a system of linear equations with correlated errors, centered on fe. Is that your intent? If so, what you have is fine.

You also do not need to use as_tensor_variable in your model. All PyMC objects are already tensor variables.

5

I can call pm.draw on the “likelihood” function and resulting a 1x20 array similar to the displacement that I was hoping, taking quite a while per draw (±20s). Seems like the sampling on using “pm.sample” is taking a long time based on this drawing time on my device.

That’s the intent to have system of equation sampled as mu, centered on “fe” function and compare with covariance “cov” a set of constant matrix and observation “y” which is basically a comparison set of data.

Thanks for the help, Jesse.