Skip to content

Abs2 gpu #26

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kif wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Abs2 gpu #26

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
481056c
fix GPU code
kif Sep 13, 2022
8b9ec14
performance optimization of the function abs2 on GPU
kif Sep 14, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 16 additions & 5 deletions saxstats/saxstats.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,6 @@
#it works, but often results in nans randomly
PYFFTW = False


def myfftn(x, DENSS_GPU=False):
if DENSS_GPU:
return cp.fft.fftn(x)
Expand Down Expand Up @@ -118,10 +117,16 @@ def myabs(x, out=None,DENSS_GPU=False):
else:
return np.abs(x,out=out)

# @numba.vectorize([numba.float64(numba.complex128),numba.float32(numba.complex64)])
def abs2(x):
if CUPY_LOADED:
gpu_abs2 = cp.ElementwiseKernel(
'complex128 z', 'float64 d', 'double re=real(z), im=imag(z); d=re*re+im*im;', "gpu_abs2")

def abs2(x, DENSS_GPU=False):
#a faster way to calculate abs(x)**2, for calculating intensities
return x.real**2 + x.imag**2
if DENSS_GPU:
return gpu_abs2(x)
else:
return x.real**2 + x.imag**2

# @numba.jit(nopython=True)
# def mybincount(x, weights):
Expand All @@ -131,9 +136,13 @@ def abs2(x):

def mybinmean(xravel,binsravel,xcount=None,DENSS_GPU=False):
if DENSS_GPU:
if isinstance(binsravel, np.ndarray):
binsravel = cp.asarray(binsravel)
xsum = cp.bincount(binsravel, xravel)
if xcount is None:
xcount = cp.bincount(binsravel)
elif isinstance(xcount, np.ndarray):
xcount = cp.asarray(xcount)
return xsum/xcount
else:
xsum = np.bincount(binsravel, xravel)
Expand Down Expand Up @@ -1208,7 +1217,7 @@ def denss(q, I, sigq, dmax, ne=None, voxel=5., oversampling=3., recenter=True, r
#APPLY RECIPROCAL SPACE RESTRAINTS
#calculate spherical average of intensities from 3D Fs
# I3D = myabs(F, out=I3D, DENSS_GPU=DENSS_GPU)**2
I3D = abs2(F)
I3D = abs2(F, DENSS_GPU=DENSS_GPU)
Imean = mybinmean(I3D.ravel(), qblravel, xcount=xcount, DENSS_GPU=DENSS_GPU)

#scale Fs to match data
Expand All @@ -1230,6 +1239,8 @@ def denss(q, I, sigq, dmax, ne=None, voxel=5., oversampling=3., recenter=True, r

#Error Reduction
newrho *= 0
if DENSS_GPU:
newrho = cp.asarray(newrho)
newrho[support] = rhoprime[support]

# enforce positivity by making all negative density points zero.
Expand Down