Illustrate broadcasting over wavenumber grids (or wavenumber PencilArray) in gradient.jl ?

[glwagner]

In the gradient.jl example, a gradient is computed via FFTs by multiplying a transformed field by a wavenumber vector and then transforming back to physical space.

The multiplication between the wavenumbers and the transformed field is by looping over indices

PencilFFTs.jl/docs/examples/gradient.jl

Lines 140 to 154 in 4e0cdfe

	for I in CartesianIndices(θ_glob)
	i, j, k = Tuple(I) # unpack indices
	## Wave number vector associated to current Cartesian index.
	local kx, ky, kz # hide
	kx = kvec[1][i]
	ky = kvec[2][j]
	kz = kvec[3][k]
	## Compute gradient in Fourier space.
	## Note that modifying ∇θ_glob also modifies the original PencilArray ∇θ_hat.
	∇θ_glob[1][I] = im * kx * θ_glob[I]
	∇θ_glob[2][I] = im * ky * θ_glob[I]
	∇θ_glob[3][I] = im * kz * θ_glob[I]
	end

I'm wondering if it might make sense to illustrate how this computation might be done with broadcasting syntax? Broadcasting might preferred in some applications: it's more compact, and generalizes (in principle...) to GPU.

Partly I'm motivated to propose this modification because I'd like to use this syntax but don't understand quite how to achieve it! I see that

PencilFFTs.jl/examples/gradient.jl

Line 258 in 4e0cdfe

grid_fourier_col = localgrid(pencil(θ_hat), collect.(kvec))

is used in

PencilFFTs.jl/examples/gradient.jl

Lines 133 to 138 in 4e0cdfe

	function gradient_grid_broadcast!(∇θ_hat, θ_hat, g::LocalRectilinearGrid)
	@. ∇θ_hat[1] = im * g[1] * θ_hat
	@. ∇θ_hat[2] = im * g[2] * θ_hat
	@. ∇θ_hat[3] = im * g[3] * θ_hat
	∇θ_hat
	end

and might illustrate how localgrid abstracts the notion of three "coordinates" (basically, three arrays with size (Nx, 1, 1), (1, Ny, 1), and (1, 1, Nz)... ?)

An alternative might use broadcasting with wavenumbers wrapped in a PencilArray. Not sure which is preferred (is localgrid just a wrapper for 3 "special" PencilArray with dimensions as above?)

Thanks for the super useful package btw!

[glwagner]

@johnryantaylor @simone-silvestri might be interested too

[jipolanco]

Right, so localgrid was introduced quite recently (jipolanco/PencilArrays.jl#45) precisely to be able to broadcast PencilArrays with coordinates / wavenumbers in an efficient manner. So this is the recommended way of doing this kind of broadcasting.

I'm wondering if it might make sense to illustrate how this computation might be done with broadcasting syntax? Broadcasting might preferred in some applications: it's more compact, and generalizes (in principle...) to GPU.

Maybe I'm misunderstanding the question, but this is precisely what "Method 4" illustrates. Perhaps it should appear first, as it's by far the simplest implementation. It's just that this wasn't possible until quite recently...

This is also illustrated in the PencilArrays docs (but maybe this is not that easy to find?).

and might illustrate how localgrid abstracts the notion of three "coordinates" (basically, three arrays with size (Nx, 1, 1), (1, Ny, 1), and (1, 1, Nz)... ?)

Yeah, that's precisely the way one should think about localgrid.

An alternative might use broadcasting with wavenumbers wrapped in a PencilArray. Not sure which is preferred (is localgrid just a wrapper for 3 "special" PencilArray with dimensions as above?)

The problem is that one cannot construct a 1D PencilArray over a 3D geometry, because that would be incompatible with the domain partitioning scheme (how does one partition a 1D array?). Maybe there is a clever way to somehow allow this, but for now it's not possible. In any case, that's what localgrid is for.

Alternatively, one could create a full 3D array of wavenumbers (with e.g. u[i, j, k] = kx[i] * ky[j] * kz[k]) and broadcast with that, and that would be totally fine, except that it's a bit wasteful as it allocates a whole array (but, once again, that might be totally fine for some use cases).

[glwagner]

Maybe I'm misunderstanding the question, but this is precisely what "Method 4" illustrates.

Bah! I missed that!

The problem is that one cannot construct a 1D PencilArray over a 3D geometry, because that would be incompatible with the domain partitioning scheme (how does one partition a 1D array?).

Ah I see the issue. I think what I'm proposing is to partition three 3D PencilArray that have singleton dimensions, but maybe this doesn't make sense. (I guess this requires behavior where singleton dimensions are "extruded" to anticipate their use in broadcasting...)

This is also illustrated in the PencilArrays docs (but maybe this is not that easy to find?).

Ah, indeed! However, that example doesn't use "provided" coordinates but rather generates x, y, z under the hood, I suppose.

[jipolanco]

I think what I'm proposing is to partition three 3D PencilArray that have singleton dimensions, but maybe this doesn't make sense.

I see, that should be possible, and could be added if there's an interest (but I guess localgrid already solves part of the need for that).

Ah, indeed! However, that example doesn't use "provided" coordinates but rather generates x, y, z under the hood, I suppose.

The coordinates are provided further above :)

It's the same as in the gradient example really, only that here we're using the grid.x, grid.y, grid.z syntax (instead of grid[1], grid[2], grid[3]) to extract the coordinates. Note that the x, y and z names are hardcoded (this was inspired by a recent StaticArrays change). I guess all of this should be clarified in the docs...

[glwagner]

The coordinates are provided further above :)

Wow, you're right, I missed it again...

[glwagner]

I guess all of this should be clarified in the docs...

I think part of it is that the name "grid" carries some baggage / meaning; when I see that I immediately assume that we're talking about a physical grid (until I read further and realized it could be a wavenumber grid in Fourier space). But now I understand that it doesn't have to be a "coordinate" at all --- the restriction I guess is that the three component of the triplet have to vary only with i, j, k respectively, where i, j, k are the indices of the 3D array...

But I'm not sure this means that we need a different name... just musing...

[glwagner]

Anyways since broadcasting is actually in the gradient.jl example I'll close this.