Improve export speed for USD scatter

glue_ar/common/scatter_usd.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from typing import List, Optional, Tuple
 
 from glue_vispy_viewers.scatter.layer_state import ScatterLayerState
@@ -14,8 +15,7 @@
 from glue_ar.common.shapes import cone_triangles, cone_points, cylinder_points, cylinder_triangles, \
                                   normalize, rectangular_prism_triangulation, sphere_triangles
 from glue_ar.utils import Viewer3DState, export_label_for_layer, iterable_has_nan, hex_to_components, \
-                          layer_color, xyz_for_layer, Bounds, NoneType
-from glue_ar.usd_utils import material_for_color
+                          layer_color, offset_triangles, xyz_for_layer, Bounds, NoneType
 
 try:
     from glue_jupyter.common.state3d import ViewerState3D
@@ -131,37 +131,49 @@
         normalized = [max(min((cval - layer_state.cmap_vmin) / crange, 1), 0) for cval in cmap_vals]
         colors = [tuple(int(256 * c) for c in cmap(norm)[:3]) for norm in normalized]
 
-    # If we're in fixed-size mode, we can reuse the same prim and translate it
-    if fixed_size:
-        first_point = data[0]
-        points = points_getter(first_point, radius)
-        mesh = builder.add_mesh(points,
-                                triangles,
-                                color=color_components,
-                                opacity=layer_state.alpha,
-                                identifier=identifier)
-
-        for i in range(1, len(data)):
-            point = data[i]
-            translation = tuple(p - fp for p, fp in zip(point, first_point))
-            if fixed_color:
-                material = None
-            else:
-                material = material_for_color(builder.stage, colors[i], layer_state.alpha)
-            builder.add_translated_reference(mesh,
-                                             translation,
-                                             material=material,
-                                             identifier=identifier)
-
+    # If we're in fixed-color mode, we can use one mesh for everything
+    opacity = float(layer_state.alpha)
+    if fixed_color:
+        points = []
+        tris = []
+        triangle_offset = 0
+        for i, point in enumerate(data):
+            size = radius if fixed_size else sizes[i]
+            pts = points_getter(point, size)
+            points.append(pts)
+            pt_triangles = offset_triangles(triangles, triangle_offset)
+            triangle_offset += len(pts)
+            tris.append(pt_triangles)
+
+        mesh_points = [pt for pts in points for pt in pts]
+        mesh_triangles = [tri for sphere in tris for tri in sphere]
+        builder.add_mesh(mesh_points,
+                         mesh_triangles,
+                         color=color_components,
+                         opacity=opacity,
+                         identifier=identifier)
     else:
+        points_by_color = defaultdict(list)
+        triangles_by_color = defaultdict(list)
+        triangle_offsets = defaultdict(int)
         for i, point in enumerate(data):
-            points = points_getter(point, sizes[i])
-            color = color_components
-            if not fixed_color:
-                cval = cmap_vals[i]
-                normalized = max(min((cval - layer_state.cmap_vmin) / crange, 1), 0)
-                color = tuple(int(256 * c) for c in cmap(normalized)[:3])
-            builder.add_mesh(points, triangles, color=color, opacity=layer_state.alpha)
+            color = colors[i]
+            size = radius if fixed_size else sizes[i]
+            pts = points_getter(point, size)
+            pt_triangles = offset_triangles(triangles, triangle_offsets[color])
+            triangle_offsets[color] += len(pts)
+            points_by_color[color].append(pts)
+            triangles_by_color[color].append(pt_triangles)
+
+        for color, points in points_by_color.items():
+            tris = triangles_by_color[color]
+            mesh_points = [pt for pts in points for pt in pts]
+            mesh_triangles = [tri for sphere in tris for tri in sphere]
+            builder.add_mesh(mesh_points,
+                             mesh_triangles,
+                             color=color,
+                             opacity=opacity,
+                             identifier=identifier)
 
     for axis in ("x", "y", "z"):
         if getattr(layer_state, f"{axis}err_visible", False):

glue_ar/common/shapes.py

@@ -4,7 +4,7 @@
 
 from numpy import cross, pi
 
-from glue_ar.gltf_utils import offset_triangles
+from glue_ar.utils import offset_triangles
 
 __all__ = [
     "rectangular_prism_points",

glue_ar/common/tests/test_scatter_usd.py

@@ -1,7 +1,7 @@
 from sys import platform
 from tempfile import NamedTemporaryFile
 
-from pxr import Sdf, Usd
+from pxr import Usd
 import pytest
 
 from glue_ar.common.export import export_viewer
@@ -44,42 +44,33 @@ def test_basic_export(self, app_type: str, viewer_type: str):
         phi_resolution: int = getattr(options, "phi_resolution", 3)
         sphere_pts_count = sphere_points_count(theta_resolution=theta_resolution, phi_resolution=phi_resolution)
         sphere_tris_count = sphere_triangles_count(theta_resolution=theta_resolution, phi_resolution=phi_resolution)
-        expected_vert_cts = [3] * sphere_tris_count
+        expected_vert_cts = [3] * sphere_tris_count * self.n
 
         color_precision = 5
         color_components = [round(c / 255, color_precision) for c in hex_to_components(layer.state.color)]
 
-        # There should be 2n + 4 total prims:
-        # The xform and the mesh for each point -> 2 * n
+        # There should be 6 total prims:
+        # The xform and the mesh (single color means one mesh)
         # The top-level world prim
         # The light
         # The material
         # The PBR shader
-        assert iterator_count(stage.TraverseAll()) == 2 * self.n + 4
+        assert iterator_count(stage.TraverseAll()) == 6
 
-        original_point_mesh = None
-        for i in range(self.n):
-            point_mesh = stage.GetPrimAtPath(f"/world/xform_{identifier}_{i}/mesh_{identifier}_{i}")
-            assert point_mesh is not None
-            points = list(point_mesh.GetAttribute("points").Get())
-            assert len(points) == sphere_pts_count
-            vertex_counts = list(point_mesh.GetAttribute("faceVertexCounts").Get())
-            assert vertex_counts == expected_vert_cts
-            vertex_indices = list(point_mesh.GetAttribute("faceVertexIndices").Get())
-            assert len(vertex_indices) == sphere_tris_count * 3
+        point_mesh = stage.GetPrimAtPath(f"/world/xform_{identifier}_0/mesh_{identifier}_0")
+        assert point_mesh is not None
+        points = list(point_mesh.GetAttribute("points").Get())
+        assert len(points) == sphere_pts_count * self.n
+        vertex_counts = list(point_mesh.GetAttribute("faceVertexCounts").Get())
+        assert vertex_counts == expected_vert_cts
+        vertex_indices = list(point_mesh.GetAttribute("faceVertexIndices").Get())
+        assert len(vertex_indices) == sphere_tris_count * 3 * self.n
 
-            material = material_for_mesh(point_mesh)
-            pbr_shader = stage.GetPrimAtPath(f"{material.GetPath()}/PBRShader")
-            assert pbr_shader.GetAttribute("inputs:metallic").Get() == 0.0
-            assert pbr_shader.GetAttribute("inputs:roughness").Get() == 1.0
-            assert round(pbr_shader.GetAttribute("inputs:opacity").Get(), color_precision) == \
-                   round(layer.state.alpha, color_precision)
-            assert [round(c, color_precision) for c in pbr_shader.GetAttribute("inputs:diffuseColor").Get()] == \
-                   color_components
-
-            if i == 0:
-                original_point_mesh = point_mesh
-            else:
-                stack_top = point_mesh.GetPrimStack()[0]
-                assert stack_top.referenceList.prependedItems[0] == \
-                       Sdf.Reference(primPath=original_point_mesh.GetPath())
+        material = material_for_mesh(point_mesh)
+        pbr_shader = stage.GetPrimAtPath(f"{material.GetPath()}/PBRShader")
+        assert pbr_shader.GetAttribute("inputs:metallic").Get() == 0.0
+        assert pbr_shader.GetAttribute("inputs:roughness").Get() == 1.0
+        assert round(pbr_shader.GetAttribute("inputs:opacity").Get(), color_precision) == \
+               round(layer.state.alpha, color_precision)
+        assert [round(c, color_precision) for c in pbr_shader.GetAttribute("inputs:diffuseColor").Get()] == \
+               color_components

glue_ar/gltf_utils.py

@@ -8,7 +8,6 @@
     "slope_intercept_between",
     "clip_linear_transformations",
     "bring_into_clip",
-    "offset_triangles",
     "create_material_for_color",
     "add_points_to_bytearray",
     "add_triangles_to_bytearray",
@@ -44,10 +43,6 @@ def bring_into_clip(points, transforms):
     return [tuple(transform[0] * c + transform[1] for transform, c in zip(transforms, pt)) for pt in points]
 
 
-def offset_triangles(triangle_indices, offset):
-    return [tuple(idx + offset for idx in triangle) for triangle in triangle_indices]
-
-
 def create_material_for_color(
     color: List[int],
     opacity: float

glue_ar/utils.py

@@ -25,6 +25,20 @@
 except ImportError:
     NoneType = type(None)
 
+
+__all__ = [
+    "NoneType", "PACKAGE_DIR", "AR_ICON", "RESOURCES_DIR", "data_count",
+    "export_label_for_layer", "layers_to_export", "isomin_for_layer",
+    "isomax_for_layer", "xyz_bounds", "bounds_3d_from_layers",
+    "slope_intercept_between", "layer_color", "bring_into_clip",
+    "mask_for_bounds", "xyz_for_layer", "hex_to_components",
+    "unique_id", "alpha_composite", "data_for_layer", "frb_for_layer",
+    "ndarray_has_nan", "iterable_has_nan", "iterator_count",
+    "is_volume_viewer", "get_resolution", "clamp", "clamped_opacity",
+    "binned_opacity", "offset_triangles",
+]
+
+
 PACKAGE_DIR = dirname(abspath(__file__))
 AR_ICON = abspath(join(dirname(__file__), "ar.png"))
 RESOURCES_DIR = join(PACKAGE_DIR, "resources")
@@ -298,3 +312,7 @@ def clamped_opacity(opacity: float) -> float:
 
 def binned_opacity(raw_opacity: float, resolution: float) -> float:
     return clamped_opacity(round(raw_opacity / resolution) * resolution)
+
+
+def offset_triangles(triangle_indices, offset):
+    return [tuple(idx + offset for idx in triangle) for triangle in triangle_indices]