Enable partial loading with swift snapshots

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.