def axes_scene(rotation=[1, 0, 0, 0]):
x = np.array([[6, 0, 0],
[6, -1, -1],
[6, -1, 1],
[6, 1, -1],
[6, 1, 1],
[6, -0.5, -0.5],
[6, -0.5, 0.5],
[6, 0.5, -0.5],
[6, 0.5, 0.5]])
x_sphere_prim = draw.Spheres(positions=x,
radii=0.25*np.ones(len(x)),
colors=[[1, 0, 0, 1]]*len(x))
x_lines_prim = draw.Lines(start_points=np.array([[0, 0, 0], *x[1:5]]),
end_points=np.array([x[0]]*5),
widths=0.5*np.ones(5),
colors=[[1, 0, 0, 1]]*5)
y = np.array([[0, 6, 0],
[1, 6, 1],
[-1, 6, 1],
[0, 6, -1],
[0.5, 6, 0.5],
[-0.5, 6, 0.5],
[0, 6, -0.5]])
y_sphere_prim = draw.Spheres(positions=y,
radii=0.25*np.ones(len(y)),
colors=[[0, 1, 0, 1]]*len(y))
y_lines_prim = draw.Lines(start_points=np.array([[0, 0, 0], *y[1:4]]),
end_points=np.array([y[0]]*4),
widths=0.5*np.ones(4),
colors=[[0, 1, 0, 1]]*4)
z = np.array([[0, 0, 6],
[-1, 1, 6],
[1, 1, 6],
[-1, -1, 6],
[1, -1, 6],
[-0.5, -1, 6],
[0, -1, 6],
[0.5, -1, 6],
[-0.5, 1, 6],
[0, 1, 6],
[0.5, 1, 6],
[-0.5, -0.5, 6],
[0.5, 0.5, 6]])
z_sphere_prim = draw.Spheres(positions=z,
radii=0.25*np.ones(len(z)),
colors=[[0, 0, 1, 1]]*len(z))
z_lines_prim = draw.Lines(start_points=np.array([[0, 0, 0], *z[1:4]]),
end_points=np.array([z[0], *z[2:5]]),
widths=0.5*np.ones(4),
colors=[[0, 0, 1, 1]]*4)
scene = draw.Scene([x_sphere_prim, x_lines_prim,
y_sphere_prim, y_lines_prim,
z_sphere_prim, z_lines_prim],
rotation=rotation)
return scene
def test_Spheres(self):
for N in range(1, 3):
kwargs = dict(positions=N * [(0, 0, 0)],
colors=N * [(1, 1, 1, 1)],
radii=N * [1])
prim = draw.Spheres(**kwargs)
self.assertEqual(len(prim), N)
def test_indexing(self):
prim0 = draw.Spheres()
prim1 = draw.ConvexPolyhedra()
scene = draw.Scene([prim0, prim1])
self.assertIs(scene[0], prim0)
self.assertIs(scene[-1], prim1)
self.assertEqual(scene[:2], [prim0, prim1])
self.assertEqual(scene[:3], [prim0, prim1])
with self.assertRaises(IndexError):
scene[5]
def draw_plato(self):
min_length = min(*self.box[:3])
pixel_scale = 600 / min_length
size = (min_length * 4. / 3, min_length)
box_args = dict(zip(['Lx', 'Ly', 'Lz', 'xy', 'xz', 'yz'], self.box))
box_prim = draw.Box(**box_args)
box_prim.width = min_length * .05
if self.arguments['density']:
box_prim.color = (1, 1, 1, 1)
color_scale = (self.arguments['color_scale'] *
len(self.positions) / len(self.projected_positions))
prim = draw.Spheres(positions=self.projected_positions,
diameters=.05)
colors = np.zeros((len(prim.positions), 4))
colors[:, :3] = plato.cmap.cubeellipse(
self.projected_types.astype(np.float32))
prim.colors = colors * color_scale
features = dict(ambient_light=0,
directional_light=(0, 0, -1),
additive_rendering=True)
scene = draw.Scene([box_prim, prim], features=features)
else:
colors = np.ones((len(self.positions), 4))
colors[:, :3] = plato.cmap.cubeellipse(
self.types.astype(np.float32))
prim = draw.Spheres(positions=self.positions, colors=colors)
scene = draw.Scene([box_prim, prim])
scene.size = size
scene.pixel_scale = pixel_scale
return scene
def colored_spheres(num_per_side=6):
xs = np.arange(num_per_side).astype(np.float32)
rs = np.array(list(itertools.product(*(3*[xs]))))
rs = np.concatenate([rs, rs + .5], axis=0)
colors = np.ones((rs.shape[0], 4))
colors[:, :3] = rs/(num_per_side - 1)
diameters = np.ones((rs.shape[0],))/np.sqrt(2)
rs -= np.mean(rs, axis=0, keepdims=True)
prim = draw.Spheres(positions=rs, colors=colors, diameters=diameters, outline=.02)
features = dict(ambient_light=.25)
features['directional_light'] = .5*np.array([(.5, .25, -.5), (0, -.25, -.25)])
rotation = [0.43797198, -0.4437895 , 0.08068451, 0.7776423]
return draw.Scene(prim, features=features, zoom=4, rotation=rotation)
def test_add_remove_prims(self):
prim1 = draw.Disks()
prim2 = draw.Spheres()
scene = draw.Scene(prim1)
self.assertIn(prim1, scene)
self.assertNotIn(prim2, scene)
scene.add_primitive(prim2)
self.assertIn(prim2, scene)
scene.remove_primitive(prim2)
self.assertNotIn(prim2, scene)
with self.assertRaises(ValueError):
scene.remove_primitive(prim2)
# shouldn't encounter exception with strict mode disabled
scene.remove_primitive(prim2, strict=False)
# prim1 should still be there
self.assertIn(prim1, scene)