Calculation of grid shape

[fvoliom]

Hi Matthew,

in the propagation.direct_integration_method.py in the calculate_meshgrid function there is the line

grid_shape = [field_shape[i] + propagation_plane.shape[i] - 1 for i in range(2)]

Where does this -1 come from? For planes of size (1024,1024) I get a resulting grid of (2047, 2047), which creates problems with the following FFT.

Just wondering whether this is a bug or intended behaviour.

[fvoliom]

So looking more into the codebase, it seems like this intentional in how it interacts with other functions like getting the grid boundaries and similar. However, it is still computationally somewhat inefficient as it leads to FFTs that are not a multiple of 2.

I am currently considering changing conv2d_fft to perform padding to the next multiple of 2

def next_power_of_two(x: int) -> int:
        return 1 << (x - 1).bit_length()

def conv2d_fft(input: Tensor, weight: Tensor) -> Tensor:
    """
   ...
    Original code:
    input_fr = fft2(input)
    output_size = (input_fr.size(-2) - weight.size(-2) + 1, input_fr.size(-1) - weight.size(-1) + 1)
    weight_fr = fft2(weight.flip(-1, -2).conj(), s=(input_fr.size(-2), input_fr.size(-1)))
    output_fr = input_fr * weight_fr.conj()
    output = ifft2(output_fr)[..., : output_size[0], : output_size[1]]
    return output
    """
    iH, iW = input.shape[-2:]
    kH, kW = weight.shape[-2:]

    # size for valid convolution
    outH = iH - kH + 1
    outW = iW - kW + 1

    # pad to next power of two for FFT
    fft_H = next_power_of_two(iH)
    fft_W = next_power_of_two(iW)

    # pad input and weight to fft shape
    input_padded = pad(input, (0, fft_W - iW, 0, fft_H - iH))
    weight_flipped = weight.flip(-1, -2).conj()
    weight_padded = pad(weight_flipped, (0, fft_W - kW, 0, fft_H - kH))

   #same logic as before
    input_fr = fft2(input_padded)
    weight_fr = fft2(weight_padded)
    output_fr = input_fr * weight_fr.conj()
    output = ifft2(output_fr)[..., :outH, :outW]
    return output

This still passes all testcases and leads to

# original
DIM propagation: (1024, 2048):  86.926 ms (avg over 100 runs)​
DIM propagation: (1023, 2047):  143.347 ms (avg over 100 runs)​
DIM propagation: (1025, 2049):  251.682 ms (avg over 100 runs)​

# modified
DIM propagation: (1024, 2048):  52.849 ms (avg over 100 runs)​
DIM propagation: (1023, 2047):  53.318 ms (avg over 100 runs)​
DIM propagation: (1025, 2049):  161.516 ms (avg over 100 runs)

on a simple propagation benchmark

def benchmark_dim_propagation(shape:tuple[int, int], device="cuda", runs:int=20):
    field  = Field(
        gaussian(shape, waist_radius=30, wavelength = 0.800, waist_z = 0, spacing=0.5),
        spacing=0.5,
        offset = (0,0),
        wavelength=0.8,
        )
    field.to(device)
    for _ in range(5):
        field.propagate_to_z(
            z=500,
            propagation_method="dim",
            )
    if device != "cpu":
        torch.cuda.synchronize()
    start = time.time()
    for _ in range(runs):
        field.propagate_to_z(
            z=500,
            propagation_method="dim",
            )
    if device != "cpu":
        torch.cuda.synchronize()
    end = time.time()
    avg_time = (end - start) / runs * 1000  # ms
    print(f"DIM propagation: {shape}:  {avg_time:.3f} ms (avg over {runs} runs)")

[MatthewFilipovich]

Hi Florian, apologies for the delayed response. Yes the -1 is intentional, it ensures the convolution produces the correct output size. I added a padding of +1 in the FFT call (#92), which improves performance for power-of-two grid shapes. I didn't add next_power_of_two by default as it can lead to significantly larger padding (e.g., 1025->2048), which would hurt performance.