Джулия: Применение одномерной функции Джулии к многомерному массиву - PullRequest
Новая
       5

Джулия: Применение одномерной функции Джулии к многомерному массиву

0 голосов
/ 29 мая 2020

Я из тех людей, которые «пишут Fortran на всех языках», пытаюсь изучить современные методы программирования. У меня есть одномерная функция ft(lx)=HT(x,f(x),lx), где x и f(x) - одномерные массивы размера nx, а lx - размер выходного массива ft. Я хочу применить HT к многомерному массиву f(x,y,z).

В основном я хочу применить HT ко всем трем измерениям к go от f(x,y,z), определенного на размерной сетке (nx,ny,nz), до ft(lx,ly,lz), определенного на размерной сетке (lx,ly,lz): 

ft(lx,y,z)   = HT(x,f(x,y,z)   ,lx)
ft(lx,ly,z)  = HT(y,ft(lx,y,z) ,ly)
ft(lx,ly,lz) = HT(z,ft(lx,ly,z),lz)

В стиле f95 я бы предпочел написать что-то вроде: 

FTx=zeros((lx,ny,nz))
for k=1:nz
for j=1:ny
    FTx[:,j,k]=HT(x,f[:,j,k],lx)
end
end

FTxy=zeros((lx,ly,nz))
for k=1:nz
for i=1:lx
    FTxy[i,:,k]=HT(y,FTx[i,:,k],ly)
end
end

FTxyz=zeros((lx,ly,lz))
for j=1:ly
for i=1:lx
    FTxyz[i,j,:]=HT(z,FTxy[i,j,:],lz)
end
end

Я знаю идиомати c Джулии потребуется что-то вроде mapslices. Я не мог понять, как это сделать go из документации mapslices.

Итак, мой вопрос: каков будет идиоматический c код Джулии вместе с правильными объявлениями типов, эквивалентный версии в стиле Фортран?

Последующий подвопрос: Можно ли написать функцию 

FT = HTnD((Tuple of x,y,z etc.),f(x,y,z), (Tuple of lx,ly,lz etc.))

, которая работает с произвольными размерами? Т.е. он будет автоматически корректировать вычисления для измерений 1,2,3 в зависимости от размеров входных кортежей и функции?

Ответы [ 2 ]

2 голосов
/ 05 июня 2020

Следуя коду @gTcV, ваша функция будет выглядеть так: 

using Base.Cartesian

ht(x,F,d) = mapslices(f -> HT(x, f, d), F, dims = 1)

@generated function HTnD(
    xx::NTuple{N,Any},
    F::AbstractArray{<:Any,N},
    newdims::NTuple{N,Int}
) where {N}
    quote
        F_0 = F
        Base.Cartesian.@nexprs $N k->begin
            tmp_k = reshape(F_{k-1},(size(F_{k-1},1),prod(Base.tail(size(F_{k-1})))))
            tmp_k = ht(xx[k], tmp_k, newdims[k])
            F_k = Array(reshape(permutedims(tmp_k),(Base.tail(size(F_{k-1}))...,size(tmp_k,1))))
            # https://github.com/JuliaLang/julia/issues/30988
        end
        return $(Symbol("F_",N))
    end
end

Более простая версия, которая показывает использование картографических файлов, будет выглядеть так: 

function simpleHTnD(
    xx::NTuple{N,Any},
    F::AbstractArray{<:Any,N},
    newdims::NTuple{N,Int}
) where {N}
    for k = 1:N
        F = mapslices(f -> HT(xx[k], f, newdims[k]), F, dims = k)
    end
    return F
end

вы даже можете использовать foldl если вы друг однострочников; -) 

fold_HTnD(xx, F, newdims) = foldl((F, k) -> mapslices(f -> HT(xx[k], f, newdims[k]), F, dims = k), 1:length(xx), init = F)
2 голосов
/ 04 июня 2020

У меня есть фрагмент кода здесь , который довольно близок к тому, что вы хотите. Ключевой инструмент - Base.Cartesian.@nexprs, о котором вы можете прочитать в связанной документации.

Три основные строки в моем коде - это строки с 30 по 32. Вот словесное описание того, что они делают.

  • Строка 30: преобразовать массив размером n1 x n2 x ... nN C_{k-1} в n1 x prod(n2,...,nN) матрицу tmp_k.
  • Строка 31: Применить функцию B[k] к каждому столбец tmp_k. В моем коде есть несколько косвенных указаний, поскольку я хочу разрешить B[k] быть матрицей или функцией, но основная идея c такая, как описано выше. Это та часть, куда вы хотели бы добавить свою функцию HT.
  • Строка 32: Измените форму tmp_k обратно в N -мерный массив и циклически переставьте измерения так, чтобы второе измерение tmp_k получилось как первое измерение C_k. Это гарантирует, что следующая итерация «l oop», подразумеваемая @nexprs, работает со вторым измерением исходного массива и так далее.

Как видите, мой код избегает формирования срезов по произвольным измерениям путем перестановки, так что нам всегда нужно срезать только по первому измерению. Это значительно упрощает программирование, а также может иметь некоторые преимущества в производительности. Например, вычисление произведений матрица-вектор B * C[i1,:,i3] для всех i1, i3 может быть выполнено легко и очень эффективно, переместив второе измерение C в первую позицию tmp и используя gemm для вычисления B * tmp. Сделать то же самое эффективно без перестановки было бы намного сложнее.

...