diff --git a/src/topsy/__init__.py b/src/topsy/__init__.py index 52ca033..7ae4f90 100644 --- a/src/topsy/__init__.py +++ b/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-', '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, diff --git a/src/topsy/loader.py b/src/topsy/loader.py index b474f0a..8893b49 100644 --- a/src/topsy/loader.py +++ b/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 diff --git a/src/topsy/sph.py b/src/topsy/sph.py index ad0a522..b0393b2 100644 --- a/src/topsy/sph.py +++ b/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.