Better implementation of Fourier filtering?

Question

my general idea is to

1. Fourier transform some data
2. Select out a few modes which meet certain criteria, by putting all other modes to zero
3. Do an inverse Fourier transformation

The physics scenario could be to select only the modes with some given wave number.

Here is my implementation. I'm using a two-dimensional example but ultimately my data is 3D.

(* generate data *)
n = 16;
l = RandomReal[{-1, 1}, {n, n}];

(* define the wave vectors *)
kx = {Range[0, n/2 - 1], Range[-n/2, -1]} // Flatten // N;
kvec = Outer[List, kx, kx];

(* compute the norm of each wave vector *)
knorm = Map[Round@*Norm, kvec, {2}];

(* let's say we want to filter out the wave number = 3 modes *)
p = 3;

(* find out where the modes are *)
pos = Position[knorm, p];

(* using the factor g to make all other modes 0 *)
g = ConstantArray[0, {n, n}] // ReplacePart[#, pos -> 1] &;
Re@InverseFourier[g*Fourier[l]]

Altogether I can write some function like

Clear[fourierFilter]
fourierFilter[l_, p_] := Module[{n, kx, knorm, pos, g},
  n = Dimensions[l][[1]];
  kx = {Range[0, n/2 - 1], Range[-n/2, -1]} // Flatten // N;
  knorm = Map[Round@*Norm, Outer[List, kx, kx], {2}];
  
  pos = Position[knorm, p];
  g = ConstantArray[0, {n, n}] // ReplacePart[#, pos -> 1] &;
  Re@InverseFourier[g*Fourier[l]]
  ]

Is there a better way to implement the idea, either in terms of time or memory efficiency? I know Fourier can take a second argument which specifies which mode I want, but I'm not sure if there is a straightforward way to inverse transform with these selected modes. I tried to do that by hand but it turns out to be slower than the implementation above.

Answer 1

Note: I am guessing that in the 3D situation that OP has, a Fourier-based approach will probably be best. I provide another approach, just for comparison.

Here is an approach that does not use Fourier, except in a preliminary step. This will be competitive if only very few modes meet the criterium, i.e. if pos is short. Note that I implicitly make assumptions about what kind of filter one has. It works in the example that OP provides, I did not check beyond that.

Here is the code, and I assume the first code block provided by OP has already been evaluated:

(* construct auxiliary matrix, this is a preliminary step
   I made not effort at all to make this fast
   one can check that P.Transpose[P] is the identity matrix *)
P=With[{aux=Map[Flatten[Fourier[ReplacePart[ConstantArray[0.,{n,n}],#->1]]]&,pos]},
    Re[MatrixPower[aux.Transpose[aux],-1/2].aux]];

(* two filters *)
filterOP[x_]:=Re[InverseFourier[g*Fourier[x]]];
filterNew[x_]:=ArrayReshape[(P.Flatten[x]).P,{n,n}];

(* check that they give the same result *)
Chop[Norm[filterOP[l]-filterNew[l]]]
(* 0 *)

(* timing *)
RepeatedTiming[filterOP[l];]
(* about 0.000012 *)
RepeatedTiming[filterNew[l];]
(* about 0.000005 *)

It is slightly faster in this case.

Note. I have assumed that the construction of g in OPs code, and the construction of P in my code, is not part of the timing. Whether they should be part of the timing depends on the situation.