Enable partial loading with swift snapshots

Author: apontzen

src/topsy/__init__.py

@@ -30,12 +30,19 @@ def parse_args(args=None):
     argparser.add_argument("--particle", "-p", help="Specify the particle type to visualise",
                             default="dm", type=str)
     argparser.add_argument("--center", "-c", help="Specify the centering method: 'halo-<N>', 'all', 'zoom' or 'none'",
-                           default="halo-1", type=str)
+                           default="none", type=str)
     argparser.add_argument("--quantity", "-q", help="Specify a quantity to render instead of density",
                            default=None, type=str)
     argparser.add_argument("--tile", "-t", help="Wrap and tile the simulation box using its periodicity",
                            default=False, action="store_true")
 
+    argparser.add_argument("--load-sphere", nargs=4, help="Load a sphere of particles with the given "
+                                                          "radius and centre in simulation units, "
+                                                          "e.g. --load-sphere 0.2 0.3 0.4 0.5 to load a sphere of "
+                                                          "particles centre (0.3, 0.4, 0.5), radius 0.2. " 
+                                                          "Supported only for swift simulations",
+                            default=None, type=float)
+
     if args is None:
         args = sys.argv[1:]
     arg_batches = []
@@ -77,8 +84,13 @@ def main():
             logger.info(f"Using test data with {n_part} particles")
             loader_args = (n_part,)
         else:
+            import pynbody
             loader_class = loader.PynbodyDataLoader
-            loader_args = (args.filename, args.center, args.particle)
+            if args.load_sphere is not None:
+                loader_args = (args.filename, args.center, args.particle,
+                               pynbody.filt.Sphere(args.load_sphere[0], args.load_sphere[1:]))
+            else:
+                loader_args = (args.filename, args.center, args.particle)
 
 
         vis = visualizer.Visualizer(data_loader_class=loader_class,

src/topsy/loader.py

@@ -2,7 +2,9 @@
 import wgpu
 import logging
 import pynbody
-import pickle 
+import pickle
+
+from typing import Optional
 
 from . import config
 # ABC support:
@@ -126,10 +128,15 @@ def get_filename(self):
 
 class PynbodyDataLoader(PynbodyDataInMemory):
     """Literal data loader for pynbody (starts from just a filename)"""
-    def __init__(self, device: wgpu.GPUDevice, filename: str, center: str, particle: str):
+    def __init__(self, device: wgpu.GPUDevice, filename: str, center: str, particle: str,
+                 take_region: Optional[pynbody.filt.Filter] = None):
 
         logger.info(f"Data filename = {filename}, center = {center}, particle = {particle}")
-        snapshot = pynbody.load(filename)
+        if take_region is None:
+            snapshot = pynbody.load(filename)
+        else:
+            snapshot = pynbody.load(filename, take_region=take_region)
+
         snapshot.physical_units()
         self.filename = filename
 

src/topsy/sph.py

@@ -255,7 +255,11 @@ def encode_render_pass(self, command_encoder):
 
     def _get_kernel_at_resolution(self, n_samples):
         if self._kernel is None:
-            self._kernel = pynbody.sph.Kernel2D()
+            try:
+                self._kernel = pynbody.sph.Kernel2D()
+            except AttributeError:
+                # pynbody v2:
+                self._kernel = pynbody.sph.kernels.Kernel2D()
 
         # sph kernel is sampled at the centre of the pixels, and the full grid ranges from -2 to 2.
         # thus the left hand most pixel is at -2+2/n_samples, and the right hand most pixel is at 2-2/n_samples.