Merge branch 'master' into ci-playground

.gitignore

 .cache/*
 .tox/*
+.hypothesis/*
 *.egg-info

.gitlab-ci.yml

@@ -24,4 +24,4 @@ tox_tests:
     - python2 pyopencl/demo.py
     - python3 pyopencl/demo.py
   script:
-    - tox -e py27,py35
+    - tox -e py27,py35 -- -sv

MANIFEST.in

+recursive-include galumph/cl-src *.cl.c
+include COPYING
+include README.md

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

galumph/__init__.py

@@ -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"]

galumph/almarray.py

+# -*- 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

galumph/atomalm.py → galumph/atomicscattering.py

 # -*- 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()

galumph/cl-src/spherical-harmonic.cl.c

@@ -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;
     }
   }
 }

galumph/util.py

+# -*- 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)

setup.py

+#!/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=[

test/conftest.py

+"""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)

test/test_almarray.py

+# -*- 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)

test/test_atomalm.py

-# -*- 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)

test/test_atomic_scattering.py

+# -*- 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