Index: explore/realspace.py
===================================================================
--- explore/realspace.py	(revision 362a66f0d054d3213d4e56147f9378d15b5616f0)
+++ explore/realspace.py	(revision 57782799b2cffcbf9e50fd15e7e764c516ec7e9b)
@@ -9,5 +9,5 @@
 import numpy as np
 from numpy import pi, radians, sin, cos, sqrt
-from numpy.random import poisson, uniform, randn, rand
+from numpy.random import poisson, uniform, randn, rand, randint
 from numpy.polynomial.legendre import leggauss
 from scipy.integrate import simps
@@ -78,8 +78,15 @@
 
 
+I3 = np.matrix([[1., 0, 0], [0, 1, 0], [0, 0, 1]])
+
 class Shape:
-    rotation = np.matrix([[1., 0, 0], [0, 1, 0], [0, 0, 1]])
+    rotation = I3
     center = np.array([0., 0., 0.])[:, None]
     r_max = None
+    lattice_size = np.array((1, 1, 1))
+    lattice_spacing = np.array((1., 1., 1.))
+    lattice_distortion = 0.0
+    lattice_rotation = 0.0
+    lattice_type = ""
 
     def volume(self):
@@ -96,5 +103,6 @@
 
     def rotate(self, theta, phi, psi):
-        self.rotation = rotation(theta, phi, psi) * self.rotation
+        if theta != 0. or phi != 0. or psi != 0.:
+            self.rotation = rotation(theta, phi, psi) * self.rotation
         return self
 
@@ -103,14 +111,79 @@
         return self
 
+    def lattice(self, size=(1, 1, 1), spacing=(2, 2, 2), type="sc",
+                distortion=0.0, rotation=0.0):
+        self.lattice_size = np.asarray(size, 'i')
+        self.lattice_spacing = np.asarray(spacing, 'd')
+        self.lattice_type = type
+        self.lattice_distortion = distortion
+        self.lattice_rotation = rotation
+
     def _adjust(self, points):
-        points = np.asarray(self.rotation * np.matrix(points.T)) + self.center
+        if self.rotation is I3:
+            points = points.T + self.center
+        else:
+            points = np.asarray(self.rotation * np.matrix(points.T)) + self.center
+        if self.lattice_type:
+            points = self._apply_lattice(points)
         return points.T
 
-    def r_bins(self, q, over_sampling=1, r_step=0.):
-        r_max = min(2 * pi / q[0], self.r_max)
+    def r_bins(self, q, over_sampling=10, r_step=0.):
+        if self.lattice_type:
+            r_max = np.sqrt(np.sum(self.lattice_size*self.lattice_spacing*self.dims)**2)/2
+        else:
+            r_max = self.r_max
+        #r_max = min(2 * pi / q[0], r_max)
         if r_step == 0.:
             r_step = 2 * pi / q[-1] / over_sampling
         #r_step = 0.01
         return np.arange(r_step, r_max, r_step)
+
+    def _apply_lattice(self, points):
+        """Spread points to different lattice positions"""
+        size = self.lattice_size
+        spacing = self.lattice_spacing
+        shuffle = self.lattice_distortion
+        rotate = self.lattice_rotation
+        lattice = self.lattice_type
+
+        if rotate != 0:
+            # To vectorize the rotations we will need to unwrap the matrix multiply
+            raise NotImplementedError("don't handle rotations yet")
+
+        # Determine the number of lattice points in the lattice
+        shapes_per_cell = 2 if lattice == "bcc" else 4 if lattice == "fcc" else 1
+        number_of_lattice_points = np.prod(size) * shapes_per_cell
+
+        # For each point in the original shape, figure out which lattice point
+        # to translate it to.  This is both cell index (i*ny*nz + j*nz  + k) as
+        # well as the point in the cell (corner, body center or face center).
+        nsamples = points.shape[1]
+        lattice_point = randint(number_of_lattice_points, size=nsamples)
+
+        # Translate the cell index into the i,j,k coordinates of the senter
+        cell_index = lattice_point // shapes_per_cell
+        center = np.vstack((cell_index//(size[1]*size[2]),
+                            (cell_index%(size[1]*size[2]))//size[2],
+                            cell_index%size[2]))
+        center = np.asarray(center, dtype='d')
+        if lattice == "bcc":
+            center[:, lattice_point % shapes_per_cell == 1] += [[0.5], [0.5], [0.5]]
+        elif lattice == "fcc":
+            center[:, lattice_point % shapes_per_cell == 1] += [[0.0], [0.5], [0.5]]
+            center[:, lattice_point % shapes_per_cell == 2] += [[0.5], [0.0], [0.5]]
+            center[:, lattice_point % shapes_per_cell == 3] += [[0.5], [0.5], [0.0]]
+
+        # Each lattice point has its own displacement from the ideal position.
+        # Not checking that shapes do not overlap if displacement is too large.
+        offset = shuffle*(randn(3, number_of_lattice_points) if shuffle < 0.3
+                          else rand(3, number_of_lattice_points))
+        center += offset[:, cell_index]
+
+        # Each lattice point has its own rotation.  Rotate the point prior to
+        # applying any displacement.
+        # rotation = rotate*(randn(size=(shapes, 3)) if shuffle < 30 else rand(size=(nsamples, 3)))
+        # for k in shapes: points[k] = rotation[k]*points[k]
+        points += center*(np.array([spacing])*np.array(self.dims)).T
+        return points
 
 class Composite(Shape):
@@ -669,6 +742,6 @@
     Iq = 100 * np.ones_like(qx)
     data = Data2D(x=qx, y=qy, z=Iq, dx=None, dy=None, dz=np.sqrt(Iq))
-    data.x_bins = qx[0,:]
-    data.y_bins = qy[:,0]
+    data.x_bins = qx[0, :]
+    data.y_bins = qy[:, 0]
     data.filename = "fake data"
 
@@ -695,5 +768,7 @@
     return shape, fn, fn_xy
 
-def build_sphere(radius=125, rho=2):
+DEFAULT_SPHERE_RADIUS = 125
+DEFAULT_SPHERE_CONTRAST = 2
+def build_sphere(radius=DEFAULT_SPHERE_RADIUS, rho=DEFAULT_SPHERE_CONTRAST):
     shape = TriaxialEllipsoid(radius, radius, radius, rho)
     fn = lambda q: sphere_Iq(q, radius)*rho**2
@@ -751,8 +826,8 @@
     return shape, fn, fn_xy
 
-def build_cubic_lattice(shape, nx=1, ny=1, nz=1, dx=2, dy=2, dz=2,
-                  shuffle=0, rotate=0):
+def build_sc_lattice(shape, nx=1, ny=1, nz=1, dx=2, dy=2, dz=2,
+                        shuffle=0, rotate=0):
     a, b, c = shape.dims
-    shapes = [copy(shape)
+    corners= [copy(shape)
               .shift((ix+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
                      (iy+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
@@ -762,7 +837,66 @@
               for iy in range(ny)
               for iz in range(nz)]
-    lattice = Composite(shapes)
+    lattice = Composite(corners)
     return lattice
 
+def build_bcc_lattice(shape, nx=1, ny=1, nz=1, dx=2, dy=2, dz=2,
+                      shuffle=0, rotate=0):
+    a, b, c = shape.dims
+    corners = [copy(shape)
+               .shift((ix+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
+                      (iy+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
+                      (iz+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
+               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
+               for ix in range(nx)
+               for iy in range(ny)
+               for iz in range(nz)]
+    centers = [copy(shape)
+               .shift((ix+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
+                      (iy+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
+                      (iz+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
+               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
+               for ix in range(nx)
+               for iy in range(ny)
+               for iz in range(nz)]
+    lattice = Composite(corners + centers)
+    return lattice
+
+def build_fcc_lattice(shape, nx=1, ny=1, nz=1, dx=2, dy=2, dz=2,
+                      shuffle=0, rotate=0):
+    a, b, c = shape.dims
+    corners = [copy(shape)
+               .shift((ix+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
+                      (iy+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
+                      (iz+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
+               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
+               for ix in range(nx)
+               for iy in range(ny)
+               for iz in range(nz)]
+    faces_a = [copy(shape)
+               .shift((ix+0.0+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
+                      (iy+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
+                      (iz+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
+               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
+               for ix in range(nx)
+               for iy in range(ny)
+               for iz in range(nz)]
+    faces_b = [copy(shape)
+               .shift((ix+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
+                      (iy+0.0+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
+                      (iz+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
+               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
+               for ix in range(nx)
+               for iy in range(ny)
+               for iz in range(nz)]
+    faces_c = [copy(shape)
+               .shift((ix+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
+                      (iy+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
+                      (iz+0.0+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
+               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
+               for ix in range(nx)
+               for iy in range(ny)
+               for iz in range(nz)]
+    lattice = Composite(corners + faces_a + faces_b + faces_c)
+    return lattice
 
 SHAPE_FUNCTIONS = OrderedDict([
@@ -775,4 +909,10 @@
 ])
 SHAPES = list(SHAPE_FUNCTIONS.keys())
+LATTICE_FUNCTIONS = OrderedDict([
+    ("sc", build_sc_lattice),
+    ("bcc", build_bcc_lattice),
+    ("fcc", build_fcc_lattice),
+])
+LATTICE_TYPES = list(LATTICE_FUNCTIONS.keys())
 
 def check_shape(title, shape, fn=None, show_points=False,
@@ -783,5 +923,7 @@
     r = shape.r_bins(q, r_step=r_step)
     sampling_density = samples / shape.volume
+    print("sampling points")
     rho, points = shape.sample(sampling_density)
+    print("calculating Pr")
     t0 = time.time()
     Pr = calc_Pr(r, rho-rho_solvent, points)
@@ -792,5 +934,5 @@
     import pylab
     if show_points:
-         plot_points(rho, points); pylab.figure()
+        plot_points(rho, points); pylab.figure()
     plot_calc(r, Pr, q, Iq, theory=theory, title=title)
     pylab.gcf().canvas.set_window_title(title)
@@ -806,4 +948,5 @@
     Qx, Qy = np.meshgrid(qx, qy)
     sampling_density = samples / shape.volume
+    print("sampling points")
     t0 = time.time()
     rho, points = shape.sample(sampling_density)
@@ -844,5 +987,8 @@
                         help='lattice size')
     parser.add_argument('-z', '--spacing', type=str, default='2,2,2',
-                        help='lattice spacing')
+                        help='lattice spacing (relative to shape)')
+    parser.add_argument('-t', '--type', choices=LATTICE_TYPES,
+                        default=LATTICE_TYPES[0],
+                        help='lattice type')
     parser.add_argument('-r', '--rotate', type=float, default=0.,
                         help="rotation relative to lattice, gaussian < 30 degrees, uniform otherwise")
@@ -858,8 +1004,36 @@
     nx, ny, nz = [int(v) for v in opts.lattice.split(',')]
     dx, dy, dz = [float(v) for v in opts.spacing.split(',')]
-    shuffle, rotate = opts.shuffle, opts.rotate
+    distortion, rotation = opts.shuffle, opts.rotate
     shape, fn, fn_xy = SHAPE_FUNCTIONS[opts.shape](**pars)
-    if nx > 1 or ny > 1 or nz > 1:
-        shape = build_cubic_lattice(shape, nx, ny, nz, dx, dy, dz, shuffle, rotate)
+    view = tuple(float(v) for v in opts.view.split(','))
+    # If comparing a sphere in a cubic lattice, compare against the
+    # corresponding paracrystalline model.
+    if opts.shape == "sphere" and dx == dy == dz and nx*ny*nz > 1:
+        radius = pars.get('radius', DEFAULT_SPHERE_RADIUS)
+        model_name = opts.type + "_paracrystal"
+        model_pars = {
+            "scale": 1.,
+            "background": 0.,
+            "lattice_spacing": 2*radius*dx,
+            "lattice_distortion": distortion,
+            "radius": radius,
+            "sld": pars.get('rho', DEFAULT_SPHERE_CONTRAST),
+            "sld_solvent": 0.,
+            "theta": view[0],
+            "phi": view[1],
+            "psi": view[2],
+        }
+        fn, fn_xy = wrap_sasmodel(model_name, **model_pars)
+    if nx*ny*nz > 1:
+        if rotation != 0:
+            print("building %s lattice"%opts.type)
+            build_lattice = LATTICE_FUNCTIONS[opts.type]
+            shape = build_lattice(shape, nx, ny, nz, dx, dy, dz,
+                                  distortion, rotation)
+        else:
+            shape.lattice(size=(nx, ny, nz), spacing=(dx, dy, dz),
+                          type=opts.type,
+                          rotation=rotation, distortion=distortion)
+
     title = "%s(%s)" % (opts.shape, " ".join(opts.pars))
     if opts.dim == 1:
@@ -867,5 +1041,4 @@
                     mesh=opts.mesh, qmax=opts.qmax, samples=opts.samples)
     else:
-        view = tuple(float(v) for v in opts.view.split(','))
         check_shape_2d(title, shape, fn_xy, view=view, show_points=opts.plot,
                        mesh=opts.mesh, qmax=opts.qmax, samples=opts.samples)
