Complicated source functions

Hello together,

I’m new to Fenics and I haven’t seen by now any post about really complicated source functions as e.g.

  • numpy code

  • pytorch code

My ultimate goal is to solve Poisson’s problem with a neural network (in pytorch) as the source function. In order to do so I’ve started with a minimal example to get pytorch/numpy source functions work: As mentioned below I reduced the working example to the following



import ufl
from mpi4py import MPI
import dolfinx
from dolfinx.fem.petsc import LinearProblem
from dolfinx import mesh
from dolfinx.fem import FunctionSpace
from dolfinx import fem

import torch as tr
import numpy as np

def my_f(x):
    # pytorch stuff
    #array_np = np.array([x[0], x[1]])
    #array_tr = tr.tensor(array_np, dtype=tr.double)
    #value    = tr.sum(array_tr).to_numpy()

    # numpy stuff
    array_np = np.array([x[0], x[1]], dtype=np.float64)
    value    = np.linalg.norm(array_np)

    return value



domain = mesh.create_rectangle(MPI.COMM_WORLD, [[-1.0, -1.0], [1.0, 1.0]], [50, 50])
V      = FunctionSpace(domain, ("Lagrange", 1))


# ## Defining the source term
f = fem.Function(V)
f.interpolate(my_f)

The Ouput for this code is

Traceback (most recent call last):
  File "/usr/lib/petsc/lib/python3/dist-packages/dolfinx/fem/function.py", line 397, in interpolate
    _interpolate(u, cells)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/lib/petsc/lib/python3/dist-packages/dolfinx/fem/function.py", line 373, in _interpolate
    self._cpp_object.interpolate(u, cells, nmm_interpolation_data)
TypeError: interpolate(): incompatible function arguments. The following argument types are supported:
    1. (self: dolfinx.cpp.fem.Function_float64, f: numpy.ndarray[numpy.float64], cells: numpy.ndarray[numpy.int32]) -> None
    2. (self: dolfinx.cpp.fem.Function_float64, u: dolfinx.cpp.fem.Function_float64, cells: numpy.ndarray[numpy.int32], nmm_interpolation_data: Tuple[List[int], List[int], List[float], List[int]]) -> None
    3. (self: dolfinx.cpp.fem.Function_float64, expr: dolfinx::fem::Expression<double, double>, cells: numpy.ndarray[numpy.int32]) -> None

Invoked with: <dolfinx.cpp.fem.Function_float64 object at 0x7f0721adccb0>, <function my_f at 0x7f08008cd750>, array([   0,    1,    2, ..., 4997, 4998, 4999], dtype=int32), ((), (), (), ())

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "~/08_test_interpolate.py", line 33, in <module>
    f.interpolate(my_f)
  File "/usr/lib/petsc/lib/python3/dist-packages/dolfinx/fem/function.py", line 402, in interpolate
    self._cpp_object.interpolate(np.asarray(u(x), dtype=self.dtype), cells)
IndexError: invalid axis: 0 (ndim = 0)

I really don’t have a clue why this shouldn’t be possible. The dolfinx version is:
DOLFINx version: 0.7.3

I appreciate any help, thanks in advance.

This is the right way to go. Just use that same syntax for f, replacing the lambda function with my_f (without providing any value to the input argument, especially not the ufl spatial coordinates).

Hi francesco-ballarin,

thanks for your fast response. Of course, simple functions can be represented using lambda expressions. But if the source function is represented by, lets say 150 lines of numpy/pytorch code, then this way won’t work.
Best, LeErmst

To clarify what I previously stated, it is not mandatory to use lambda functions in combination with interpolate. Just use

f.interpolate(my_f)

Ah yes, it should work like that but it don’t. Or at least I haven’t been able to make it work. I think the problem is kind of related to this question I’ve found: numerical-values-from-ufl-spatialcoordinate.

We can’t help much further until you clarify what the issue is. Please simplify your code (defining the problem and solving it is unnecessary here, just interpolate an expression which is representative of the one you have), ideally without using pytorch. If interpolation works when the output is a plain numpy array but doesn’t work when you pass it through pytorch, then the issue is likely in how you pass x to pytorch, and you should ask their community for help.

You’re right, I reduced the working example to the absolute minimum. I’ve looked at the code in dolfinx/fem/function.py", line 402, in interpolate
but I could not figure out why it isn’t working. It doesn’t matter whether I use pytorch or numpy, I get the same error for both.

Adding some prints to hopefully clarify what is going on.

def my_f(x):
    print("x shape is", x.shape)
    value = np.linalg.norm(x, axis=0)
    print("value shape is", value.shape)
    return value

Cool, it worked with that. The output is:

x shape is (3, 15000)
value shape is (15000,)

I think that x[0] and x[1] are the spatial coordinates and x[2] is the third one in case the problem is three dimensional, because np.sum(x[2]) returns zero.
Many thanks to you francesco, I’m pretty confident now that I can do what I’m wanting.

Yes, internally DOLFINx stores coordinates in 3 dimensions, too speed up accessibility in integration kernels (and various other places).

1 Like