Adjoint solver performance bottleneck

[CharlesDove]

Hi folks, thanks for the great package! While trying out the adjoint solver, I noticed what appears to be a significant performance bottleneck. Specifically, filter_source.py, starting on line 40

        self.bf = [
            lambda t, i=i: 0
            if t > self.T
            else (
                self.nuttall(t, self.center_frequencies)
                / (self.dt / np.sqrt(2 * np.pi))
            )[i]
            for i in range(len(self.center_frequencies))
        ]

produces a list of python lambdas representing sources. This list is then eventually passed into the c++ engine, which calls back to the python lambda at each step. Removing this call seems to yield a ~10x speed improvement. Perhaps this functionality should be pushed into the c++ engine or otherwise compiled?

[smartalecH]

Hi @CharlesDove, thanks for reaching out! Handling the adjoint sources is rather tricky (as we talk about in our paper). And I think the behavior you're observing is a red herring. Let me try to walk through things:

So this piece of the code simply computes the different basis functions needed to properly compute a broadband source. Evaluating this (that for loop) is actually relatively cheap and only happens once up front. Then, all the sources' spatial profiles are cached in C++, and we call the time profile like any other custom time profile. In other words, this isn't in a "hot loop."

I suspect your speedup is because you are getting rid of the long run times that comes with these narrow basis functions (we talk a lot about this in the paper). But moving this to C++ shouldn't really change anything (unless your forward and adjoint simulations are really really fast, and the overhead is dominated by this?)

[CharlesDove]

Hi @smartalecH, thanks for the reply and info!

Right, I saw that this line is only run once, before the actual time stepping loop. My concern is that self.bf is a list of python lambdas, each of which is eventually associated with an adjoint source. These sources are then passed into the c++ engine, which applies them during the time stepping loop. So if there are 50 adjoint sources, we seem to call back into these python lambdas, and the nuttal function, 50 times for every timestep.

I noticed this by using a python debugger while running an adjoint simulation. Manually pausing the debugger during the c++ run should only pause once python code is referenced. For the forward run, the whole sim is finished before pausing (what I'd expect). For the adjoint run, the sim pauses immediately in the nuttal function, making me think that that function is being run frequently during the actual time stepping loop.

This might be expected behavior for a custom source, however, as you mentioned. I was mainly thinking that this functionality might be compiled using something like Cython, then the function pointer passed into the C++ loop, which would prevent this call back into python.