Add patchy spheres primitive

doc/source/primitives.rst

@@ -83,6 +83,9 @@ Base Drawing Module
 .. autoclass:: Mesh
    :members:
 
+.. autoclass:: PatchySpheres
+   :members:
+
 .. autoclass:: SpherePoints
    :members:
 

doc/source/vispy.rst

@@ -48,6 +48,9 @@ Vispy Backend
 .. autoclass:: Mesh
    :members:
 
+.. autoclass:: PatchySpheres
+   :members:
+
 .. autoclass:: SpherePoints
    :members:
 

plato/draw/PatchySpheres.py

@@ -0,0 +1,53 @@
+import itertools
+
+import numpy as np
+
+from .internal import ShapeDecorator, ShapeAttribute
+from .Spheres import Spheres
+
+@ShapeDecorator
+class PatchySpheres(Spheres):
+    """A collection of patchy spheres in 3D.
+
+    Each sphere can have a different position, orientation, base
+    color, and diameter. Spheres also have a set of patches, specified
+    by plane equations (a normal vector (x, y, z) and offset w) and
+    associated colors (RGBA). Plane equations are specified for a
+    diameter 2 sphere and each shape is scaled by its diameter.
+    """
+
+    _ATTRIBUTES = Spheres._ATTRIBUTES + list(itertools.starmap(ShapeAttribute, [
+        ('orientations', np.float32, (1, 0, 0, 0), 2, True,
+         'Orientation of each particle'),
+        ('patch_planes', np.float32, (0, 0, 0, 0), 2, False,
+         'Plane equations (x, y, z, w) for patches, specified for a sphere of diameter 2'),
+        ('patch_colors', np.float32, (0.5, 0.5, 0.5, 1), 2, False,
+         'Colors (RGBA) for each patch'),
+        ('shape_color_fraction', np.float32, 0, 0, False,
+         'Fraction of a patch\'s color that should be assigned based on colors'),
+        ]))
+
+    @property
+    def patch_unit_angles(self):
+        """Unit+angle specification for patch directions and sizes.
+
+        Patches are specified in this way as an array of (x, y, z,
+        theta) values, with x, y, and z specifying a unit vector and
+        theta specifying the angle swept out by the patch.
+        """
+        result = self.patch_planes.copy()
+        n, h = result[:, :3], result[:, 3]
+        length = np.linalg.norm(n, axis=-1)
+        result[:, :3] /= length[:, np.newaxis]
+        result[:, 3] = 2*np.arccos(h/length)
+        result[np.any(np.logical_not(np.isfinite(result)), axis=-1)] = (1, 0, 0, 0)
+        return result
+
+    @patch_unit_angles.setter
+    def patch_unit_angles(self, value):
+        value = np.atleast_2d(value).copy()
+        length = np.linalg.norm(value[:, :3], -1, keepdims=True)
+        value[:, :3] /= length
+        value[:, 3] = np.cos(value[:, 3]*.5)
+        value[np.any(np.logical_not(np.isfinite(value)), axis=-1)] = (1, 0, 0, 0)
+        self.patch_planes = value

plato/draw/__init__.py

@@ -9,6 +9,7 @@
 from .Spheropolygons import Spheropolygons
 from .Voronoi import Voronoi
 from .Lines import Lines
+from .PatchySpheres import PatchySpheres
 from .Spheres import Spheres
 from .SpherePoints import SpherePoints
 from .SphereUnions import SphereUnions

plato/draw/vispy/PatchySpheres.py

@@ -0,0 +1,265 @@
+import itertools
+
+import numpy as np
+
+from ... import draw
+from ... import mesh
+from .internal import GLPrimitive, GLShapeDecorator
+from ..internal import ShapeAttribute
+from ..Scene import DEFAULT_DIRECTIONAL_LIGHTS
+
+@GLShapeDecorator
+class PatchySpheres(draw.PatchySpheres, GLPrimitive):
+    __doc__ = draw.PatchySpheres.__doc__
+
+    shaders = {}
+
+    shaders['vertex'] = """
+       uniform mat4 camera;
+       uniform vec4 rotation;
+       uniform vec3 translation;
+       uniform int transparency_mode;
+
+       attribute vec4 color;
+       attribute vec3 position;
+       attribute vec4 orientation;
+       attribute vec2 image;
+       attribute float radius;
+
+       varying vec4 v_color;
+       varying vec3 v_position;
+       varying vec4 v_orientation;
+       varying vec2 v_image;
+       varying float v_radius;
+       varying float v_depth;
+
+       vec3 rotate(vec3 point, vec4 quat)
+       {
+           vec3 result = (quat.x*quat.x - dot(quat.yzw, quat.yzw))*point;
+           result += 2.0*quat.x*cross(quat.yzw, point);
+           result += 2.0*dot(quat.yzw, point)*quat.yzw;
+           return result;
+       }
+
+       vec4 quatquat(vec4 a, vec4 b)
+       {
+           float real = a.x*b.x - dot(a.yzw, b.yzw);
+           vec3 imag = a.x*b.yzw + b.x*a.yzw + cross(a.yzw, b.yzw);
+           return vec4(real, imag);
+       }
+
+       vec4 conj(vec4 quat)
+       {
+           return vec4(-quat.x, quat.yzw);
+       }
+
+       void main()
+       {
+           vec3 vertexPos = position;
+           vec2 localPos = image*radius;
+           vertexPos = rotate(vertexPos, rotation) + vec3(localPos, 0.0) + translation;
+           vec4 screenPosition = camera * vec4(vertexPos, 1.0);
+
+           int should_discard = 0;
+           should_discard += int(transparency_mode < 0 && color.a < 1.0);
+           should_discard += int(transparency_mode > 0 && color.a >= 1.0);
+           if(should_discard > 0)
+               screenPosition = vec4(2.0, 2.0, 2.0, 2.0);
+
+           // transform to screen coordinates
+           gl_Position = screenPosition;
+           v_color = color;
+           v_image = localPos;
+           v_radius = radius;
+           v_position = vertexPos;
+           v_depth = vertexPos.z;
+           v_orientation = conj(quatquat(rotation, orientation));
+       }
+       """
+
+    shaders['fragment'] = """
+#ifdef GL_ES
+       precision highp float;
+#endif
+
+       // base light level
+       uniform float ambientLight;
+       // (x, y, z) direction*intensity
+       uniform vec3 diffuseLight[NUM_DIFFUSELIGHT];
+       uniform vec4 patch_planes[NUM_PATCH_PLANES];
+       uniform vec4 patch_colors[NUM_PATCH_COLORS];
+       uniform int transparency_mode;
+       uniform mat4 camera;
+       uniform float light_levels;
+       uniform float outline;
+       uniform float shape_color_fraction;
+
+       varying vec4 v_color;
+       varying vec4 v_orientation;
+       varying vec2 v_image;
+       varying float v_radius;
+       varying float v_depth;
+
+       vec3 rotate(vec3 point, vec4 quat)
+       {
+           vec3 result = (quat.x*quat.x - dot(quat.yzw, quat.yzw))*point;
+           result += 2.0*quat.x*cross(quat.yzw, point);
+           result += 2.0*dot(quat.yzw, point)*quat.yzw;
+           return result;
+       }
+
+       void main()
+       {
+           float rsq = dot(v_image, v_image);
+           float Rsq = v_radius*v_radius;
+
+           float r = sqrt(rsq);
+
+           float lambda1 = 1.0;
+           if(outline > 1e-6)
+           {
+               lambda1 = (v_radius - r)/outline;
+               lambda1 *= lambda1;
+               lambda1 *= lambda1;
+               lambda1 *= lambda1;
+               lambda1 *= lambda1;
+               lambda1 = min(lambda1, 1.0);
+           }
+
+           if(r > v_radius) discard;
+
+           vec3 r_local = vec3(v_image.xy, sqrt(Rsq - rsq));
+           vec3 normal = normalize(r_local);
+           float light = ambientLight;
+           for(int i = 0; i < NUM_DIFFUSELIGHT; ++i)
+               light += max(0.0, -dot(normal, diffuseLight[i]));
+
+           if(light_levels > 0.0)
+           {
+               light *= light_levels;
+               light = floor(light);
+               light /= light_levels;
+           }
+
+           vec4 color = v_color;
+           vec3 rotated_normal = rotate(normal, v_orientation);
+
+           for(int i = 0; i < NUM_PATCH_PLANES; i++)
+           {
+               vec4 plane = patch_planes[i];
+               if(dot(rotated_normal, plane.xyz) > plane.w)
+                   color = mix(patch_colors[i], v_color, shape_color_fraction);
+           }
+
+           color = vec4(color.xyz*lambda1*light, color.w);
+
+           #ifndef WEBGL
+           float depth = v_depth + r_local.z;
+           gl_FragDepth = 0.5*(camera[2][2]*depth + camera[3][2] +
+               camera[2][3]*depth + camera[3][3])/(camera[2][3]*depth + camera[3][3]);
+           #endif
+
+           float z = abs(v_depth);
+           float alpha = color.a;
+           float weight = alpha*max(3e3*pow(
+               (1.0 - gl_FragCoord.z), 3.0), 1e-2);
+
+           if(transparency_mode < 1)
+               gl_FragColor = color;
+           else if(transparency_mode == 1)
+               gl_FragColor = vec4(color.rgb*alpha, alpha)*weight;
+           else
+               gl_FragColor = vec4(alpha);
+       }
+       """
+
+    shaders['fragment_plane'] = """
+       // base light level
+       uniform mat4 camera;
+       uniform float render_positions = 0.0;
+
+       varying vec3 v_position;
+       varying vec2 v_image;
+       varying float v_radius;
+       varying float v_depth;
+
+       void main()
+       {
+           float rsq = dot(v_image, v_image);
+           float Rsq = v_radius*v_radius;
+
+           if(rsq > Rsq)
+               discard;
+
+           vec3 r_local = vec3(v_image.xy, sqrt(Rsq - rsq));
+           vec3 normal = normalize(r_local);
+           #ifndef WEBGL
+           float depth = v_depth + r_local.z;
+           gl_FragDepth = 0.5*(camera[2][2]*depth + camera[3][2] +
+               camera[2][3]*depth + camera[3][3])/(camera[2][3]*depth + camera[3][3]);
+           #endif
+
+           if(render_positions > 0.5)
+               gl_FragColor = vec4(gl_FragCoord.xyz, 1.0);
+           else if(render_positions < -0.5)
+               gl_FragColor = 0.5 + 0.5*vec4(normal.xyz, 1.0);
+           else // Store the plane equation as a color
+               gl_FragColor = vec4(normal, dot(normal, v_position.xyz));
+       }
+       """
+
+    _vertex_attribute_names = ['position', 'orientation', 'color', 'radius', 'image']
+
+    _GL_UNIFORMS = list(itertools.starmap(ShapeAttribute, [
+        ('camera', np.float32, np.eye(4), 2, False,
+         'Internal: 4x4 Camera matrix for world projection'),
+        ('ambientLight', np.float32, .25, 0, False,
+         'Internal: Ambient (minimum) light level for all surfaces'),
+        ('diffuseLight[]', np.float32, DEFAULT_DIRECTIONAL_LIGHTS, 2, False,
+         'Internal: Diffuse light direction*magnitude'),
+        ('rotation', np.float32, (1, 0, 0, 0), 1, False,
+         'Internal: Rotation to be applied to each scene as a quaternion'),
+        ('translation', np.float32, (0, 0, 0), 1, False,
+         'Internal: Translation to be applied to the scene'),
+        ('transparency_mode', np.int32, 0, 0, False,
+         'Internal: Transparency stage (<0: opaque, 0: all, 1: '
+         'translucency stage 1, 2: translucency stage 2)'),
+        ('light_levels', np.float32, 0, 0, False,
+         'Number of light levels to quantize to (0: disable)'),
+        ('outline', np.float32, 0, 0, False,
+         'Outline for all particles'),
+        ('patch_planes[]', np.float32, (0, 0, 0, 0), 2, False,
+         'Internal: Plane equations (x, y, z, w) for patches, specified for a sphere of diameter 2'),
+        ('patch_colors[]', np.float32, (0.5, 0.5, 0.5, 1), 2, False,
+         'Internal: Colors (RGBA) for each patch'),
+        ('shape_color_fraction', np.float32, 0, 0, False,
+         'Internal: Fraction of a patch\'s color that should be assigned based on colors')
+        ]))
+
+    def __init__(self, *args, **kwargs):
+        GLPrimitive.__init__(self)
+        draw.PatchySpheres.__init__(self, *args, **kwargs)
+
+    def update_arrays(self):
+        try:
+            for name in self._dirty_attributes:
+                self._gl_vertex_arrays[name][:] = self._attributes[name]
+                self._dirty_vertex_attribs.add(name)
+        except (ValueError, KeyError):
+            # vertices for an equilateral triangle
+            triangle = np.array([[2, 0],
+                                 [-1, np.sqrt(3)],
+                                 [-1, -np.sqrt(3)]], dtype=np.float32)*1.01
+
+            vertex_arrays = mesh.unfoldProperties(
+                [self.positions, self.orientations, self.colors, self.radii.reshape((-1, 1))],
+                [triangle])
+
+            unfolded_shape = vertex_arrays[0].shape[:-1]
+            indices = (np.arange(unfolded_shape[0])[:, np.newaxis, np.newaxis]*unfolded_shape[1] +
+                       np.array([[0, 1, 2]], dtype=np.uint32))
+            indices = indices.reshape((-1, 3))
+
+            self._finalize_array_updates(indices, vertex_arrays)
+
+        self._dirty_attributes.clear()

plato/draw/vispy/__init__.py

@@ -46,6 +46,7 @@
 from .Spheropolygons import Spheropolygons
 from .Voronoi import Voronoi
 from .Lines import Lines
+from .PatchySpheres import PatchySpheres
 from .Spheres import Spheres
 from .SpherePoints import SpherePoints
 from .SphereUnions import SphereUnions

test/test_scenes.py

@@ -480,6 +480,19 @@ def field_scene(N=10, use='ellipsoids'):
         prim = draw.Ellipsoids(
             positions=positions, orientations=orientations, colors=colors,
             a=.5, b=.125, c=.125)
+    elif use == 'patchy_spheres':
+        orientations = rowan.vector_vector_rotation([(1, 0, 0)], units)
+        unit_angles = [
+            (1, 0, 0, np.pi/8),
+            (0, 1, 0, np.pi/8),
+            (0, 0, 1, np.pi/8),
+        ]
+        patch_colors = np.array(unit_angles)
+        patch_colors[:, 3] = 1
+        prim = draw.PatchySpheres(
+            positions=positions, orientations=orientations, colors=colors,
+            patch_unit_angles=unit_angles, patch_colors=patch_colors,
+            shape_color_fraction=.5)
     elif use == 'lines':
         features['ambient_light'] = 1
         starts = positions - units/2
@@ -504,6 +517,10 @@ def field_lines(N=10):
 def field_ellipsoids(N=10):
     return field_scene(N, 'ellipsoids')
 
+@register_scene
+def field_patchy_spheres(N=10):
+    return field_scene(N, 'patchy_spheres')
+
 @register_scene
 def simple_cubes_octahedra(N=4):
     xs = np.linspace(-N/2, N/2, N)