* import imagewriter only if needed

* rename variables in `BuildGrid_trimesh`
* speed up grid creation using numpy and vectorization
* remove duplicate imports
* add description to progress bar
* create new function for vectorized point checking

cellpack/autopack/Analysis.py

@@ -34,7 +34,6 @@
 from cellpack.autopack.upy import colors as col
 from cellpack.autopack.upy.colors import map_colors
 from cellpack.autopack.utils import check_paired_key, get_paired_key
-from cellpack.autopack.writers.ImageWriter import ImageWriter
 import concurrent.futures
 import multiprocessing
 
@@ -2354,6 +2353,8 @@ def pack_one_seed(
                 plt.close()  # closes the current figure
 
         if image_export_options is not None:
+            from cellpack.autopack.writers.ImageWriter import ImageWriter
+
             image_writer = ImageWriter(
                 env=self.env,
                 name=seed_basename,

cellpack/autopack/Compartment.py

@@ -1286,7 +1286,15 @@ def BuildGrid_binvox(self, env, grdPos, vSurfaceArea, srfPts, idarray, ray=1):
 
         return self.insidePoints, self.surfacePoints
 
-    def BuildGrid_trimesh(self, env, grdPos, vSurfaceArea, srfPts, idarray, mesh_store):
+    def BuildGrid_trimesh(
+        self,
+        env,
+        master_grid_positions,
+        vSurfaceArea,
+        off_grid_surface_points,
+        compartment_ids,
+        mesh_store,
+    ):
         """Build the compartment grid ie surface and inside points"""
         insidePoints = []
         number = self.number
@@ -1299,42 +1307,46 @@ def BuildGrid_trimesh(self, env, grdPos, vSurfaceArea, srfPts, idarray, mesh_sto
         ).hollow()
         self.log.info("VOXELIZED MESH")
         # the main loop
-        tree = spatial.cKDTree(grdPos, leafsize=10)
-        points_in_encap_sphere = tree.query_ball_point(
-            self.center,
-            self.encapsulating_radius + env.grid.gridSpacing * 2,
-            return_sorted=True,
+        tree = spatial.cKDTree(master_grid_positions, leafsize=10)
+        points_in_encap_sphere = numpy.array(
+            tree.query_ball_point(
+                self.center,
+                self.encapsulating_radius + env.grid.gridSpacing * 2,
+                return_sorted=True,
+            )
         )
         self.log.info(f"GOT POINTS IN SPHERE {len(points_in_encap_sphere)}")
-        for ptInd in points_in_encap_sphere:
-            coord = [
-                grdPos.item((ptInd, 0)),
-                grdPos.item((ptInd, 1)),
-                grdPos.item((ptInd, 2)),
-            ]
-            if idarray[ptInd] > 0:
-                continue
-            if trimesh_grid_surface.is_filled(coord):
-                idarray.itemset(ptInd, number)
-            elif mesh_store.contains_point(self.gname, coord):
-                insidePoints.append(ptInd)
-                idarray.itemset(ptInd, -number)
+
+        point_compartment_ids = compartment_ids[points_in_encap_sphere]
+        point_ids_to_assign = points_in_encap_sphere[point_compartment_ids == 0]
+        point_positions = master_grid_positions[point_ids_to_assign]
+
+        # check surface points
+        points_on_surface = trimesh_grid_surface.is_filled(point_positions)
+        compartment_ids[point_ids_to_assign[points_on_surface]] = number
+
+        # check inside points
+        points_in_mesh = mesh_store.contains_points_mesh(self.gname, point_positions)
+        points_in_mesh = points_in_mesh & ~points_on_surface
+        compartment_ids[point_ids_to_assign[points_in_mesh]] = -number
+        insidePoints = point_ids_to_assign[points_in_mesh]
+
         self.log.info("ASSIGNED INSIDE OUTSIDE")
 
         nbGridPoints = len(env.grid.masterGridPositions)
 
         (
             surface_points_in_bounding_box,
             surfPtsBBNorms,
-        ) = self.filter_surface_pts_to_fill_box(srfPts, env)
+        ) = self.filter_surface_pts_to_fill_box(off_grid_surface_points, env)
 
-        srfPts = surface_points_in_bounding_box
+        off_grid_surface_points = surface_points_in_bounding_box
 
         ex = False  # True if nbGridPoints == len(idarray) else False
 
         surfacePoints, surfacePointsNormals = self.extendGridArrays(
             nbGridPoints,
-            srfPts,
+            off_grid_surface_points,
             surfPtsBBNorms,
             env,
             surfacePointsNormals=self.surfacePointsNormals,
@@ -2541,8 +2553,6 @@ def getSurfaceInnerPoints_kevin(
         It is simply there as a safeguard.
         """
         # Start the timer.
-        from time import time
-
         startTime = time()
 
         gridSpacing = spacing

cellpack/autopack/Environment.py

@@ -2200,6 +2200,7 @@ def pack_grid(
         stime = time()
         if self.show_progress_bar:
             pbar = tqdm(total=totalNumMols, mininterval=0, miniters=1)
+            pbar.set_description(f"Packing {self.name}")
         while nbFreePoints:
             self.log.info(
                 ".........At start of while loop, with vRangeStart = %d", vRangeStart

cellpack/autopack/MeshStore.py

@@ -264,6 +264,12 @@ def contains_points(self, geomname, points):
             return intersector.contains_points(points)
         return [False]
 
+    def contains_points_mesh(self, geomname, points):
+        mesh = self.get_object(geomname)
+        if mesh is not None:
+            return mesh.contains(points)
+        return numpy.full(len(points), False)
+
     def get_smallest_radius(self, geomname, center):
         mesh = self.get_object(geomname)
         if mesh is not None: