python - 在 Fortran 中实现像快速插值一样的 numpy

Question

我有一个数值例程，我需要运行它来求解某个方程，其中包含几个嵌套的四个循环。我最初将此例程写入 Python，使用 numba.jit 来实现可接受的性能。然而，对于大型系统，这种方法变得相当慢，所以我一直在将例程重写为 Fortran，希望能加快速度。但是我发现我的 Fortran 版本比第一个 Python 版本慢 2-3 倍。

我认为瓶颈是在每个最内层循环调用的线性插值函数。在 Python 实现中，我使用了 numpy.interp，与 numba.jit 结合使用时速度似乎相当快。在 Fortran 中，我编写了自己的插值函数，内容如下:

  complex*16 function interp(x_dat, y_dat, N_dat, x)
    implicit none
    integer, intent(in) :: N_dat
    real*8, dimension(N_dat), intent(in) :: x_dat
    complex*16, dimension(N_dat), intent(in) :: y_dat
    real*8, intent(in) :: x

    complex*16 :: y, y1, y2
    integer :: i, i1, i2, im

    if(x <= x_dat(1)) then
      y = y_dat(1)
    else if(x >= x_dat(N_dat)) then
      y = y_dat(N_dat)
    else
      im = MINLOC(DABS(x_dat - x), DIM=1)
      if(x_dat(im) >=x ) then
        i1 = im
        i2 = im - 1
      else
        i1 = im + 1
        i2 = im
      end if

      y1 = y_dat(i1)
      y2 = y_dat(i2)

      y = y1 + (x-x_dat(i1))*(y2 - y1)/(x_dat(i2) - x_dat(i1))
    end if
    interp = y
    return
  end function interp

请注意，我需要插入复杂的数据。如果我的诊断是正确的，这个函数比 numpy.interp 慢得多，因为必须在每个循环中调用插值，这大大降低了整个程序的速度。

有谁知道是否有办法在 Fortran 中实现类似于 Numpy 的插值速度？或者如果我上面显示的插值函数效率低得可怕？我还没有太多编写 Fortran 代码的经验。

谢谢!

Answer 1

猜测(如果你想让我们做得比猜测更好，请参阅@IanBush 的评论)，就是这条线

im = MINLOC(DABS(x_dat - x), DIM=1)

这占用了你所有的时间，因为这一行是 O(N)，大小为 x_dat，而其他所有内容都是 O(1).

如果 x_dat 是线性间隔的，那么您可以将此行替换为

im = 1 + nint((N_dat-1)*(x-x_dat(1))/(x_dat(N_dat)-x_dat(1)))

或者更好的是，完全跳过 im 并计算 i1 和 i2 为

i1 = 1 + floor((N_dat-1)*(x-x_dat(1))/(x_dat(N_dat)-x_dat(1)))
i2 = i1 + 1

如果 x_dat 不是线性间隔的，但具有其他有用的属性，那么如果可能，您希望使用这些属性来计算 im。