vectorization - 在 Julia 中结合 Repmat 和 Transpose

Question

我正在将一些代码从 R 移植到 julia 以熟悉该语言，但我发现一些模式无法顺利翻译。考虑以下函数，

# Ricatti-Bessel and derivatives up to nmax, vectorised over x
function rb(x, nmax)

  n = 1:nmax
  nu = 0.5 + [0, n]

  bj = hcat([besselj(nu, _x) for _x in x]...).'
  # ^ first question ^ 
  sq = repmat(sqrt(pi/2*x), 1, nmax+1)

  bj .*= sq
  xm = repmat(x, 1, nmax)
  nm = repmat(n', length(x), 1)
  # ^ second question ^ 

  dpsi = bj[:,n] - nm .* bj[:,n+1] ./ xm 
  psi = bj[:,n+1]
  return psi, dpsi # it'd be nice to return a "named list" instead

end
# rb(1:5,3)

• 第一个问题:这是使用 besselj() 获取具有 nmax 列和 length(x) 行的矩阵的最佳方法吗？在找到有效的模式之前，我不得不挠头好一阵子。

• 第二个问题:我发现自己经常需要在repmat 中和/或从repmat 中转置对象，是否有其他方法可以指定输出大小和填充方向(按行或按列)？

也许我对整个事情采取了错误的方法:我习惯于使用矢量化函数(在 R 中，以及 Matlab 的旧内存)，因为它们通常是线性代数快速例程的最短路线。始终保持 x 为标量并仅在最高级别循环是否更有意义？我担心这样做，我将无法利用 BLAS 等的快速矩阵/向量函数，并且基本上无法在 Julia 中重写它们，更不用说明显的可读性损失了。我应该强调，我对最佳性能感兴趣，因为对于许多 x 值，该函数将在内部被多次调用。

Answer 1

对于你的第一个问题，我将用以下矩阵理解替换它:

nu = (0:nmax)+0.5
bj = [besselj(i,j) for j in x, i in nu]

对于你的第二个问题，我认为在 Julia 中编写高性能代码的一个好原则是避免不必要的分配(当然，还要阅读 performance tips !)Julia 在可能的情况下生成非常快的指令 - 这就是为什么 for循环完全没问题，除了线性代数(例如矩阵乘法)之外，向量化任何东西并不重要。它做得不好的是避免分配不必要的内存(像您的 sq 这样的临时矩阵)。我更换了bj/sq仅包含以下内容的行

nu = (0:nmax)+0.5
bj = [besselj(i,j)*sqrt(pi/2*j) for j in x, i in nu]

这很好，因为它只是一次分配，而且是在我们之前分配的(假设我们从我对第一个问题的回答开始):

1. 分配bj
2. 分配sq
3. 分配bj.*sq并重新绑定(bind)bj到这个新的内存

(请注意，.*= 不是就地操作!)

您对“命名列表”的请求现在可能最好通过创建 type 来满足。用于返回此函数(这根本不是一个昂贵的操作，并且在 Julia 的矩阵分解代码中很常见，其中需要返回多个值)。或者，您可以返回 Dict ，但这感觉不太惯用。

对于dpsi行，我给你两个选择。第一个是另一种矩阵理解:

dpsi = [ bj[i,j] - j * bj[i,j+1] / i 
           for i in 1:length(x), j in 1:nmax]

另一个是for循环-y:

dpsi = zeros(length(x),nmax)
for i in 1:length(x), j in 1:nmax
  dpsi[i,j] = bj[i,j] - j * bj[i,j+1] / i
end

在这两种情况下，我都避免了临时分配。同样，您的原始设备具有以下分配:

1. 对于xm
2. 对于nm
3. 对于bj[:,n] (这将在 0.4 中更改为 View )
4. 对于bj[:,n+1] (同上)
5. 对于nm .* bj[:,n+1]
6. 对于nm .* bj[:,n+1] ./ xm
7. 对于整个结果

我提出的两个版本都只有一个分配，并且可能更接近问题的原始数学陈述

我的最终版本是

function myrb(x, nmax)
  bj   = [ besselj(i,j)*sqrt(pi/2*j)
              for j in x, i in (0:nmax)+0.5]
  dpsi = [ bj[i,j] - j * bj[i,j+1] / i 
              for i in 1:length(x), j in 1:nmax]
  psi  = bj[:,2:nmax+1]
  return psi, dpsi
end

我对 besselj 不太了解，但我猜这是整个事情中最慢的部分，所以在这种特殊情况下，所有这些在速度方面可能并不重要。对这个微观案例进行基准测试表明:

# original
elapsed time: 9.7578e-5 seconds (7176 bytes allocated)
elapsed time: 7.2644e-5 seconds (7176 bytes allocated)
elapsed time: 7.5709e-5 seconds (7176 bytes allocated)
# revised
elapsed time: 2.7536e-5 seconds (728 bytes allocated)
elapsed time: 2.7097e-5 seconds (728 bytes allocated)
elapsed time: 1.6601e-5 seconds (728 bytes allocated)

您可以使用分析器确认这一点(尽管我必须使用更大的输入:)

@profile myrb(1:500,300)
Profile.print()

在我的机器上，该函数收集了 429 个样本，其中 426 个位于 bessel.jl 中。 Julia 内的文件，2 个用于 dpsi ，1 代表 psi .