SBC-Utrecht / PyTom

PyTom software for cryo-tomogram processing
GNU General Public License v2.0
33 stars 8 forks source link

suggested update for pytom agnostic libraries without gpu environment variable #41

Open McHaillet opened 1 year ago

McHaillet commented 1 year ago

I am actually not a big fan of this PYTOM_GPU environment variable. As this complicated test execution shows, you can only run GPU code after setting this variable at the start--which is know taken care of in the pytom executable. It is pretty cumbersome.

I think it would be better to use the cupy adviced method (https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code). You can use cp.get_array_module() to load backend from an ndarray. Its very easy for a simple image filter function, e.g.

import cupy as cp

def low_pass_filter(image):
    xp = cp.get_array_module(image)
    # create_low_pass outputs np/cp array depending on 
    low_pass = create_low_pass_filter(image.shape, backend=xp)
    return xp.fft.ifftn(xp.fft.fftn(image) * low_pass)

Originally posted by @McHaillet in https://github.com/SBC-Utrecht/PyTomPrivate/issues/116#issuecomment-1434584115