Commit 346c7ca3 authored by Chris Kerr's avatar Chris Kerr

Merge branch 'master' into ci-playground

parents 3e1a8c6e 0d37044f
Pipeline #3831 failed with stage
in 2 minutes and 20 seconds
.cache/*
.tox/*
.hypothesis/*
*.egg-info
......@@ -24,4 +24,4 @@ tox_tests:
- python2 pyopencl/demo.py
- python3 pyopencl/demo.py
script:
- tox -e py27,py35
- tox -e py27,py35 -- -sv
recursive-include galumph/cl-src *.cl.c
include COPYING
include README.md
......@@ -9,16 +9,19 @@ Example code:
import numpy as np
import Bio.PDB
import periodictable
import pyopencl
import galumph
ctx = pyopencl.create_some_context()
NS = 4096 # Number of S values at which to calculate the scattering
smax = 1.0 # Maximum S value
LMAX = 63 # Maximum harmonic order to use for the calculations
## Initialise the S array and allocate the ALM storage on the GPU
s = np.linspace(0, smax, NS)
atomalm = galumph.AtomAlm(LMAX, s)
kernel = galumph.AtomicScattering(LMAX, s, ctx=ctx)
## Use Bio.PDB to read the structure and periodictable to calculate the atomic form factors
pdb = Bio.PDB.PDBParser().get_structure("6lyz", "6lyz.pdb")
......@@ -26,6 +29,7 @@ xyz = np.array([aa.get_vector().get_array() for aa in pdb.get_atoms()])
ff = np.array([periodictable.elements.symbol(aa.element).xray.f0(s) for aa in pdb.get_atoms()])
## Run the GPU calculation
atomalm.add_many_atoms(xyz, ff)
Icalc = atomalm.sum_intensity()
alm = kernel.zeros()
kernel.add_many_atoms(alm, xyz, ff)
Icalc = kernel.sum_intensity(alm)
```
\ No newline at end of file
......@@ -3,9 +3,7 @@
__copyright__ = "European Molecular Biology Laboratory (EMBL)"
__license__ = "LGPLv3+"
from .almarray import AlmArray
from .atomicscattering import AtomicScattering
from .sphericalbessel import SphericalBessel
from .sphericalharmonic import SphericalHarmonic
from .atomalm import AtomAlm
__all__ = ["SphericalBessel", "SphericalHarmonic", "AtomAlm"]
__all__ = ["AlmArray", "AtomicScattering"]
# -*- coding: utf-8 -*-
from __future__ import absolute_import, division
__copyright__ = "Christopher Kerr"
__license__ = "LGPLv3+"
import numpy as np
import pyopencl as cl
from pyopencl import array as cl_array
from .util import check_shape, PackedLMMixin
class AlmArray(PackedLMMixin, cl_array.Array):
"""Array holding A(L,M,s) representation of scattering amplitudes."""
def __init__(self, cq, LMAX, NS, dtype='c8', **kwargs):
"""Work out the shape from LMAX and call the superclass constructor."""
assert LMAX >= 0
assert NS > 0
self.LMAX = LMAX
self.NS = NS
shape = (self.NLM, self.NS)
super(AlmArray, self).__init__(cq, shape, dtype=dtype, **kwargs)
def get(self, queue=None, ary=None, unpack=False, **kwargs):
"""Copy from the device into a numpy array.
If unpack is True, unpack the array into a lower triangular array
with dimensions [L,M,S].
"""
if not unpack:
return super(AlmArray, self).get(queue=queue, ary=ary, **kwargs)
unpacked_shape = (self.LMAX1, self.LMAX1, self.NS)
if ary is None:
ary = np.zeros(unpacked_shape, dtype=self.dtype)
else:
check_shape(ary, unpacked_shape)
for L in range(self.LMAX1):
packed_range = slice(self.indexLM(L, 0), self.indexLM(L+1, 0))
ary[L,0:L+1,:] = self[packed_range,:].get(queue=queue, **kwargs)
return ary
def empty(cq, LMAX, NS, dtype='c8', **kwargs):
"""Create an AlmArray without initializing the memory."""
return AlmArray(cq, LMAX, NS, dtype=dtype, **kwargs)
def zeros(cq, LMAX, NS, dtype='c8', **kwargs):
"""Create an AlmArray initialized to zero."""
alm = AlmArray(cq, LMAX, NS, dtype=dtype, **kwargs)
if isinstance(cq, cl.CommandQueue):
alm.fill(0, queue=cq)
else:
with cl.CommandQueue(cq) as queue:
alm.fill(0, queue)
return alm
# -*- coding: utf-8 -*-
from __future__ import absolute_import, division
__copyright__ = "European Molecular Biology Laboratory (EMBL)"
__license__ = "LGPLv3+"
......@@ -9,8 +11,10 @@ import pyopencl as cl
from pyopencl import array as cl_array
import pytest
from . import almarray
from .sphericalbessel import SphericalBessel
from .sphericalharmonic import SphericalHarmonic
from .util import check_shape, PackedLMMixin
from pkg_resources import resource_string
program_src = resource_string(__name__, "cl-src/alm-add-atom.cl.c")
......@@ -45,20 +49,27 @@ def sicopolar_array(xyz):
z = xyz[:,2]
rho = np.linalg.norm(xyz[:,:2], axis=1)
rho_gt0 = (rho > 0)
eiphi = x + 1j * y
eiphi[rho>0] /= rho
eiphi[rho_gt0] /= rho[rho_gt0]
r = np.linalg.norm(xyz, axis=1)
r_gt0 = (r > 0)
costheta = np.zeros_like(r)
sintheta = np.zeros_like(r)
costheta[r>0] = z[r>0]/r[r>0]
sintheta[r>0] = rho[r>0]/r[r>0]
costheta[r_gt0] = z[r_gt0]/r[r_gt0]
sintheta[r_gt0] = rho[r_gt0]/r[r_gt0]
return r, costheta, sintheta, eiphi
class AtomAlm:
class AtomicScattering(PackedLMMixin):
"""Kernel for calculating the ALM scattering from an atomic structure."""
def __init__(self, ctx, LMAX, s, worksize=64, natwork=16):
def __init__(self, LMAX, s,
worksize=64,
natwork=16,
ctx=None,
queue=None):
assert LMAX >= 0
self.LMAX = LMAX
s = np.asanyarray(s)
......@@ -73,13 +84,17 @@ class AtomAlm:
assert natwork > 0
self.natwork = natwork
if ctx is None:
ctx = cl.create_some_context()
self.ctx = ctx
if queue is None:
queue = cl.CommandQueue(ctx)
self.queue = queue
# For some reason `queue.context is ctx` is false
assert queue.context == ctx
self._sphbes = SphericalBessel(ctx, LMAX, s, worksize, natwork)
self._sphharm = SphericalHarmonic(ctx, LMAX, natwork)
NLM = self._sphharm.NLM
self.queue = cl.CommandQueue(ctx)
with self.queue:
self.alm_dev = cl_array.empty(self.queue, (NLM, NS), np.complex64)
self.alm_dev.fill(0.0)
......@@ -88,13 +103,27 @@ class AtomAlm:
self._program.build(["-DNATWORK=%d" % natwork,
"-DWORKSIZE=%d" % worksize])
def add_atom(self, xyz, ff):
LMAX1 = self.LMAX + 1
def zeros(self):
"""Return an AlmArray with the correct shape initialized to zero."""
return almarray.zeros(self.queue, self.LMAX, self.NS)
def add_atom(self, alm, xyz, ff):
"""Add the scattering from a single atom to the ALM.
Args:
alm: AlmArray
xyz: Cartesian coordinates of the atom
ff: Atomic form factor: single value or an array of length NS
"""
assert alm.context == self.ctx
check_shape(alm, (self.NLM, self.NS))
LMAX1 = self.LMAX1
NS = self.NS
worksize = self.worksize
ff = np.asanyarray(ff)
if np.isscalar(ff):
if np.isscalar(ff) or (np.shape(ff) == ()):
ff = np.ones_like(self.s) * ff
else:
assert ff.shape == self.s.shape
......@@ -113,12 +142,22 @@ class AtomAlm:
ylm_dev.data,
jl_dev.data,
ff_dev.data,
self.alm_dev.data,
alm.data,
wait_for=(evY, evJ))
ev.wait()
def add_many_atoms(self, xyz, ff):
LMAX1 = self.LMAX + 1
def add_many_atoms(self, alm, xyz, ff):
"""Add the scattering from an array of atoms to the ALM.
Args:
alm: AlmArray
xyz: Cartesian coordinates of the atoms
ff: Atomic form factors: shape either (natoms) or (natoms, NS)
"""
assert alm.context == self.ctx
check_shape(alm, (self.NLM, self.NS))
LMAX1 = self.LMAX1
NS = self.NS
natwork = self.natwork
worksize = self.worksize
......@@ -130,7 +169,7 @@ class AtomAlm:
ff = np.asanyarray(ff)
if np.ndim(ff) == 1:
ff = np.ones_like(self.s) * ff[:, np.newaxis]
assert np.shape(ff) == (natoms, NS)
check_shape(ff, (natoms, NS))
ylm_local = cl.LocalMemory(8*natwork*LMAX1)
......@@ -152,13 +191,29 @@ class AtomAlm:
ylm_dev.data,
jl_dev.data,
ff_dev.data,
self.alm_dev.data,
alm.data,
ylm_local,
wait_for=(evY, evJ))
ev.wait()
def sum_intensity(self):
def __call__(self, xyz, ff):
"""Calculate the atomic scattering from the given atoms.
Allocates an AtomAlm array, zeros it, then adds the scattering from
the atoms given the (x,y.z) positions and the atomic form factors.
"""
alm = self.zeros()
self.add_many_atoms(alm, xyz, ff)
return alm
def sum_intensity(self, alm):
"""Calculate the scattering intensity from the given ALM array.
The intensity at each S point is the sum of the squared magnitude of
all [L,M] indices at that S point.
"""
assert alm.context == self.ctx
check_shape(alm, (self.NLM, self.NS))
NS = self.NS
worksize = self.worksize
......@@ -168,7 +223,7 @@ class AtomAlm:
(NS,),
(worksize,),
np.uint32(self.LMAX),
self.alm_dev.data,
alm.data,
int_dev.data)
ev.wait()
int_host = int_dev.get()
......
......@@ -25,7 +25,7 @@
* to give the final answer.
*
* Because the values can go outside the range of a 32-bit float,
* the exponent is returned in a separate integer array
* the exponent is returned in a separate integer array.
*/
void ylm_Mpart_local(const unsigned int L,
const float sintheta,
......@@ -49,7 +49,7 @@ void ylm_Mpart_local(const unsigned int L,
// sin(theta) part of the Legendre polynomial
termF = ( sintheta * (1 - 2*(signed)M)
// pre-factor sqrt( (L-M)! / (L+M)! )
* rsqrt((float)((L+M)*(L+1-M))) );
* rsqrt(convert_float((L+M)*(L+1-M))) );
}
MpartF[M] = frexp(termF, &termE);
MpartE[M] = termE;
......@@ -77,6 +77,7 @@ void ylm_Mpart_local(const unsigned int L,
barrier(CLK_LOCAL_MEM_FENCE);
}
// Raise a complex number to a non-negative integer power
cfloat_t cfloat_pown(cfloat_t Z, uint N)
{
cfloat_t res = cfloat_new(1.0f, 0.0f);
......@@ -136,13 +137,24 @@ __kernel void ylm(__constant float *costheta,
ylm_Mpart_local(L, sintheta[iAt], MpartF, MpartE);
// This calculates the cos(theta) part of the Legendre polynomial in term0
// (the sin(theta) part is calculated in ylm_Mpart_local)
if (L == M) {
term0 = 1;
} else if (L > M) {
// Recurrence property
// (L-M+1)*P[L+1,M](x) = (2*L+1)*x*P[L,M](x) - (L+M)*P[L-1,M](x)
// i.e. (L-M)*P[L,M](x) = (2*L-1)*x*P[L-1,M](x) - (L+M-1)*P[L-2,M](x)
// For L == M + 1, P[L-2,M](x) is taken to be zero
term0 = ( costheta[iAt] * (2*L - 1) * term1
- (L+M-1) * term2) / (L-M);
}
// The cos(theta) part of the Legendre polynomial is now given by:
// ldexp(term0, sumexp)
// and the sin(theta) part with the sqrt((L-M)!/(L+M)!) prefactor by:
// ldexp(mpartF[M], mpartE[M])
// Now calculate Ylm = (cospart) * (sinpart) * exp(i*M*phi)
if (L >= M) {
size_t ilm = (L*(L+1))/2 + M;
size_t iAtlm = iAt * NLM + ilm;
......@@ -150,12 +162,19 @@ __kernel void ylm(__constant float *costheta,
ylm[iAtlm] = cfloat_mulr(eiMphi, magnitude);
}
// Rotate term0 to term1 and term1 to term2
//
// To avoid over- or underflow of the floating point exponent, store the
// exponent of the cos(theta) part in sumexp
termE = ilogb(term0);
// Avoid scaling when a term is very close to zero
if (termE >= -32) {
// Avoid scaling when the term is very close to zero
if (termE >= -24) {
term2 = ldexp(term1, -termE);
term1 = ldexp(term0, -termE);
sumexp += termE;
} else {
term2 = term1;
term1 = term0;
}
}
}
# -*- coding: utf-8 -*-
from __future__ import division
__copyright__ = "Christopher Kerr"
__license__ = "LGPLv3+"
def check_shape(array, shape, exact=True):
"""Check that the array has the correct shape.
If exact is False, allow the leading dimension to be larger than the
expected shape.
"""
if array.shape == shape:
return
elif ((not exact) and
(array.shape[1:] == shape[1:]) and
(array.shape[0] >= shape[0])):
return
else:
raise ValueError('Array shape %r does not match expected shape %s' %
(array.shape, shape))
class PackedLMMixin:
"""Mixin class with methods for arrays with packed [L,M] dimensions.
This version is for arrays only containing non-negative M values.
If the LMAX attribute is set on the class, LMAX1 and NLM are available
as properties.
"""
@staticmethod
def indexLM(L, M=0):
"""Index into a packed [L,M] array."""
assert M >= 0
assert L >= M
return L * (L + 1) // 2 + M
@property
def LMAX1(self):
return self.LMAX + 1
@property
def NLM(self):
return self.indexLM(self.LMAX1, 0)
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from setuptools import setup
setup(
name='galumph',
version='0.0.1',
description='Calculate ALM using GPU acceleration',
url='https://git.embl.de/grp-svergun/galumph',
classifiers=[
......@@ -13,7 +15,10 @@ setup(
'Intended Audience :: Science/Research',
'License :: OSI Approved :: GNU Lesser General Public License v3 or later (LGPLv3+)',
'Natural Language :: English',
'Programming Language :: Python :: 2',
'Programming Language :: Python :: 2.7',
'Programming Language :: Python :: 3',
'Programming Language :: Python :: 3.5',
'Programming Language :: Python :: 3.6',
],
install_requires = ["pyopencl"],
packages=[
......
"""Shared pytest fixtures."""
import numpy as np
import pyopencl as cl
import pytest
@pytest.fixture(scope='session')
def cl_context():
"""Get a pyopencl context."""
return cl.create_some_context()
@pytest.fixture()
def cl_queue(cl_context):
"""Get a pyopencl Queue."""
with cl.CommandQueue(cl_context) as queue:
yield queue
@pytest.fixture(scope='module')
def rng():
# N.B. numpy.random.randint(a, b) excludes b,
# random.randint(a, b) includes b as a possible value
return np.random.RandomState(1234)
# -*- coding: utf-8 -*-
from __future__ import division
from hypothesis import given
import hypothesis.strategies as st
import numpy as np
import numpy.testing as npt
from galumph import almarray
@given(
cq=st.sampled_from(('context', 'queue')),
LMAX=st.integers(min_value=0, max_value=63),
NS=st.integers(min_value=1, max_value=1024),
)
def test_alm_constructor(cl_context, cl_queue, cq, LMAX, NS):
"""Test allocating an Alm array."""
if cq == 'context':
cq = cl_context
else:
cq = cl_queue
alm = almarray.AlmArray(cq, LMAX, NS)
NLM = (LMAX + 1) * (LMAX + 2) // 2
assert alm.NLM == NLM
assert alm.shape == (NLM, NS)
@given(
LMAX=st.integers(min_value=0, max_value=7),
NS=st.integers(min_value=1, max_value=64),
)
def test_alm_unpack(cl_queue, LMAX, NS):
"""Test the get() function on an AlmArray."""
alm = almarray.zeros(cl_queue, LMAX, NS)
# TODO fill with non-uniform values to test that unpacked values are in
# the correct order
alm.fill(np.complex64(1-2j))
npalm = alm.get(unpack=True)
assert npalm.shape == (LMAX+1, LMAX+1, NS)
for L in range(LMAX+1):
for M in range(LMAX+1):
if M > L:
npt.assert_array_equal(npalm[L,M,:], 0)
else:
npt.assert_array_equal(npalm[L,M,:], 1-2j)
# -*- coding: utf-8 -*-
import unittest
import numpy as np
from numpy.testing import assert_approx_equal
import pyopencl as cl
from scipy.stats import linregress
from galumph import AtomAlm
## N.B. numpy.random.randint(a, b) excludes b,
## random.randint(a, b) includes b as a possible value
rng = np.random.RandomState(123)
def structure_Guinier(xyz, ff):
I0 = np.sum(ff)**2
centre = np.dot(xyz.T, ff) / np.sum(ff)
r2 = np.sum((xyz - centre[np.newaxis, :])**2, axis=1)
rG = np.dot(r2, ff) / np.sum(ff)
return rG, I0
def scattering_Guinier(s, I, rGest=None):
if rGest:
fitrange = (s <= (1/rGest))
s = s[fitrange]
I = I[fitrange]
result = linregress(s**2, np.log(I))
rG = -3 * result[0]
I0 = np.exp(result[1])
return rG, I0
class TestAtomAlm(unittest.TestCase):
def setUp(self):
self.ctx = cl.create_some_context()
self.worksize = 8
self.LMAX = 15
self.NS = 32
self.s = np.linspace(0, 3, self.NS)
self.natwork = 4
self.alm = AtomAlm(self.ctx, self.LMAX, self.s, self.worksize, self.natwork)
def test_randxyz(self):
"""Generates a random structure and checks the scattering has the correct rG and I0"""
NATOMS = 1000
xyz = rng.randn(NATOMS, 3) * 3
ff = rng.rand(NATOMS) + 1
self.alm.add_many_atoms(xyz, ff)
Icalc = self.alm.sum_intensity()
rGxyz, I0xyz = structure_Guinier(xyz, ff)
rGalm, I0alm = scattering_Guinier(self.s, Icalc, rGxyz)
assert_approx_equal(I0alm, I0xyz, significant=3)
assert_approx_equal(rGalm, rGxyz, significant=2)
# -*- coding: utf-8 -*-
from hypothesis import assume, given, settings
import hypothesis.strategies as st
import hypothesis.extra.numpy as npst
import numpy as np
import numpy.testing as npt
from scipy.stats import linregress
from galumph.atomicscattering import AtomicScattering
@st.composite
def alm_parameters(draw,
LMAX=st.integers(min_value=0, max_value=20),
smax=st.floats(min_value=0.1, max_value=10),
worksize=st.integers(min_value=8, max_value=64),
natwork=st.integers(min_value=4, max_value=16),
):
"""Hypothesis strategy to get parameters for an Alm array."""
LMAX = draw(LMAX)
smax = draw(smax)
nchunks = draw(st.integers(min_value=4, max_value=16))
worksize = draw(worksize