Merge pull request #5 from pynbody/multiwin

Adding tests

.github/workflows/build-test.yaml

@@ -0,0 +1,40 @@
+name: Build and Test
+
+on: [push, pull_request]
+
+defaults:
+  run:
+    shell: bash
+
+jobs:
+
+  build:
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ["3.10", "3.11"]
+    runs-on: ${{ matrix.os }}
+
+    steps:
+    - name: Install llvmpipe and lavapipe for offscreen canvas
+      if: matrix.os == 'ubuntu-latest'
+      run: |
+        sudo apt-get update -y -qq
+        sudo apt install -y libegl1-mesa libgl1-mesa-dri libxcb-xfixes0-dev mesa-vulkan-drivers
+    - name: Install Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+    - uses: actions/checkout@v2
+    - name: Install
+      run: |
+        pip install .[test]
+    - name: Run all tests
+      working-directory: tests
+      run: python -m pytest
+    - uses: actions/upload-artifact@v3
+      if: always()
+      with:
+        name: Outputs from tests on Python ${{ matrix.python-version }}
+        path: tests/output/

Dockerfile

@@ -0,0 +1,20 @@
+FROM ubuntu:latest
+
+
+RUN apt update && apt install -y python3 python3-pip libegl1-mesa libgl1-mesa-dri libxcb-xfixes0-dev \
+    libglib2.0-0 mesa-vulkan-drivers libgl1-mesa-glx libxkbcommon0 libdbus-1-3
+
+# manually install some dependencies to speed up
+RUN pip3 install numpy pynbody matplotlib pillow wgpu jupyter_rfb tqdm opencv-python PySide6
+
+COPY src /app/src
+COPY tests /app/tests
+COPY pyproject.toml /app/
+COPY README.md /app/
+
+WORKDIR /app
+
+
+RUN pip3 install .[test]
+
+ENTRYPOINT ["/bin/bash"]

README.md

@@ -1,6 +1,8 @@
 topsy
 =====
 
+[![Build Status](https://github.com/pynbody/topsy/actions/workflows/build-test.yaml/badge.svg)](https://github.com/pynbody/topsy/actions)
+
 This package visualises simulations, and is an add-on to the [pynbody](https://github.com/pynbody/pynbody) analysis package.
 Its name nods to the [TIPSY](https://github.com/N-BodyShop/tipsy) project.
 It is built using [wgpu](https://wgpu.rs), which is a future-facing GPU standard (with thanks to the [python wgpu bindings](https://wgpu-py.readthedocs.io/en/stable/guide.html)).

src/topsy/__init__.py

@@ -1,8 +1,8 @@
-from __future__ import annotations
-
 """topsy - An astrophysics simulation visualization package based on webgpu, using pynbody for reading data"""
 
-__version__ = "0.3.0"
+from __future__ import annotations
+
+__version__ = "0.3.1"
 
 import argparse
 import logging
@@ -101,8 +101,9 @@ def topsy(snapshot: pynbody.snapshot.SimSnap, quantity: str | None = None, **kwa
     vis.quantity_name = quantity
     return vis
 
-def _test(nparticle=config.TEST_DATA_NUM_PARTICLES_DEFAULT):
+def _test(nparticle=config.TEST_DATA_NUM_PARTICLES_DEFAULT, **kwargs):
     from . import visualizer, loader
     vis = visualizer.Visualizer(data_loader_class=loader.TestDataLoader,
-                                data_loader_args=(nparticle,))
+                                data_loader_args=(nparticle,),
+                                **kwargs)
     return vis

src/topsy/canvas/offscreen.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import numpy as np
+from wgpu.gui.offscreen import WgpuManualOffscreenCanvas, call_later
+
+from . import VisualizerCanvasBase
+
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from ..visualizer import Visualizer
+
+
+class VisualizerCanvas(VisualizerCanvasBase, WgpuManualOffscreenCanvas):
+
+    @classmethod
+    def call_later(cls, delay, fn, *args):
+        call_later(delay, fn, *args)
+

src/topsy/loader.py

@@ -204,13 +204,16 @@ def _evaluate_density(self, pos):
     def _generate_samples(self):
         # simple gaussian mixture model
         pos = np.empty((self._n_particles, 3), dtype=np.float32)
-        offset = 0
-        for i in range(len(self._gmm_weights)):
-            cpt_len = int(self._n_particles*self._gmm_weights[i])
-            pos[offset:offset+cpt_len] = \
-                np.random.normal(size=(cpt_len, 3), scale=1.0).astype(np.float32) * self._gmm_std[np.newaxis,i,:] + self._gmm_means[i]
-            offset += cpt_len
-        assert offset == self._n_particles
+        if self._n_particles==1:
+            pos[0] = self._gmm_means[0]
+        else:
+            offset = 0
+            for i in range(len(self._gmm_weights)):
+                cpt_len = int(self._n_particles*self._gmm_weights[i])
+                pos[offset:offset+cpt_len] = \
+                    np.random.normal(size=(cpt_len, 3), scale=1.0).astype(np.float32) * self._gmm_std[np.newaxis,i,:] + self._gmm_means[i]
+                offset += cpt_len
+            assert offset == self._n_particles
         return np.random.permutation(pos)
 
     def get_positions(self):
@@ -222,7 +225,7 @@ def get_smooth(self):
         return sm
 
     def get_mass(self):
-        return np.random.uniform(0.01, 1.0, size=(self._n_particles)).astype(np.float32)*1e-8
+        return np.repeat(np.float32(1e-8), self._n_particles)
 
     def get_named_quantity(self, name):
         if name=="test-quantity":

src/topsy/sph.py

@@ -20,6 +20,7 @@ def __init__(self, visualizer: Visualizer, render_texture: wgpu.GPUTexture,
         self._render_texture = render_texture
         self._device = visualizer.device
         self._wrapping = wrapping
+        self._kernel = None
 
         self._setup_shader_module()
         self._setup_transform_buffer()
@@ -252,59 +253,53 @@ def encode_render_pass(self, command_encoder):
         sph_render_pass.end()
 
 
-    def _setup_kernel_texture(self, n_samples=64, n_mip_levels = 4):
-        if hasattr(SPH, "_kernel_texture"):
-            # we only do this once, even if multiple SPH objects (i.e. multi-resolution) is in play
-            return
+    def _get_kernel_at_resolution(self, n_samples):
+        if self._kernel is None:
+            self._kernel = pynbody.sph.Kernel2D()
 
-        pynbody_sph_kernel = pynbody.sph.Kernel2D()
-        x, y = np.meshgrid(np.linspace(-2, 2, n_samples), np.linspace(-2, 2, n_samples))
+        # 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.
+        pixel_centres = np.linspace(-2+2./n_samples, 2-2./n_samples, n_samples)
+        x, y = np.meshgrid(pixel_centres, pixel_centres)
         distance = np.sqrt(x ** 2 + y ** 2)
-        kernel_im = np.array([pynbody_sph_kernel.get_value(d) for d in distance.flatten()]).reshape(n_samples, n_samples)
+
+        # TODO: the below could easily be optimized
+        kernel_im = np.array([self._kernel.get_value(d) for d in distance.flatten()]).reshape(n_samples, n_samples)
 
         # make kernel explicitly mass conserving; naive pixelization makes it not automatically do this.
         # It should be normalized such that the integral over the kernel is 1/h^2. We have h=1 here, and the
         # full width is 4h, so the width of a pixel is dx=4/n_samples. So we need to multiply by dx^2=(n_samples/4)^2.
         # This results in a correction of a few percent, typically; not huge but not negligible either.
         #
         # (Obviously h!=1 generally, so the h^2 normalization occurs within the shader later.)
-        kernel_im *= (n_samples/4)**2 / kernel_im.sum()
+        kernel_im *= (n_samples / 4) ** 2 / kernel_im.sum()
+
+        return kernel_im
+
+    def _setup_kernel_texture(self, n_samples=64, n_mip_levels = 4):
+        if hasattr(SPH, "_kernel_texture"):
+            # we only do this once, even if multiple SPH objects (i.e. multi-resolution) is in play
+            return
 
         SPH._kernel_texture = self._device.create_texture(
             label="kernel_texture",
             size=(n_samples, n_samples, 1),
             dimension=wgpu.TextureDimension.d2,
             format=wgpu.TextureFormat.r32float,
-            mip_level_count=2,
+            mip_level_count=n_mip_levels,
             sample_count=1,
             usage=wgpu.TextureUsage.COPY_DST | wgpu.TextureUsage.TEXTURE_BINDING,
         )
 
-        self._device.queue.write_texture(
-            {
-                "texture": self._kernel_texture,
-                "mip_level": 0,
-                "origin": (0, 0, 0),
-            },
-            kernel_im.astype(np.float32).tobytes(),
-            {
-                "offset": 0,
-                "bytes_per_row": 4*n_samples,
-            },
-            (n_samples, n_samples, 1)
-        )
-
-        kt_mip = kernel_im
 
-        for i in range(1, n_mip_levels):
-            kt_mip = kt_mip[::2, ::2]
+        for i in range(0, n_mip_levels):
             self._device.queue.write_texture(
                 {
                     "texture": self._kernel_texture,
-                    "mip_level": 1,
+                    "mip_level": i,
                     "origin": (0, 0, 0),
                 },
-                kt_mip.astype(np.float32).tobytes(),
+                self._get_kernel_at_resolution(n_samples//2**i).astype(np.float32).tobytes(),
                 {
                     "offset": 0,
                     "bytes_per_row": 4 * n_samples // 2**i,

src/topsy/visualizer.py

@@ -30,7 +30,7 @@ class VisualizerBase:
 
     def __init__(self, data_loader_class = loader.TestDataLoader, data_loader_args = (),
                  *, render_resolution = config.DEFAULT_RESOLUTION, periodic_tiling = False,
-                 colormap_name = config.DEFAULT_COLORMAP):
+                 colormap_name = config.DEFAULT_COLORMAP, canvas_class = canvas.VisualizerCanvas):
         self._colormap_name = colormap_name
         self._render_resolution = render_resolution
         self.crosshairs_visible = False
@@ -41,7 +41,7 @@ def __init__(self, data_loader_class = loader.TestDataLoader, data_loader_args =
         self.show_colorbar = True
         self.show_scalebar = True
 
-        self.canvas = canvas.VisualizerCanvas(visualizer=self, title="topsy")
+        self.canvas = canvas_class(visualizer=self, title="topsy")
 
         self._setup_wgpu()
 
@@ -360,6 +360,26 @@ def get_sph_image(self) -> np.ndarray:
             im = np_im[:,:,0]
         return im
 
+    def get_presentation_image(self) -> np.ndarray:
+        texture_view = self.context.get_current_texture()
+        size = texture_view.size
+        bytes_per_pixel = 4 # NB this might be wrong in principle!
+        data = self.device.queue.read_texture(
+            {
+                "texture": texture_view.texture,
+                "mip_level": 0,
+                "origin": (0, 0, 0),
+            },
+            {
+                "offset": 0,
+                "bytes_per_row": bytes_per_pixel * size[0],
+                "rows_per_image": size[1],
+            },
+            size,
+        )
+
+        return np.frombuffer(data, np.uint8).reshape(size[1], size[0], 4)
+
 
     def save(self, filename='output.pdf'):
         image = self.get_sph_image()

tests/run_tests_in_docker.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+# It may be useful to run the tests in docker to track down issues using a standardised environment
+# This script will build the docker image, run the tests and copy the output to the local machine
+
+rm -rf docker_test_output
+docker build .. -t topsy
+docker run --name running-tests topsy -c 'pytest'
+docker cp running-tests:/app/tests/output ./docker_test_output
+docker rm running-tests
+

tests/test_arg_parse.py

@@ -7,9 +7,6 @@ def test_simple_arg_parse():
     args = args[0]
     assert args.filename == "test://1000"
     assert args.quantity == "test-quantity"
-    assert args.center is None
-    assert args.particle is None
-    assert args.tile is False
     assert args.resolution == topsy.config.DEFAULT_RESOLUTION
     assert args.colormap == topsy.config.DEFAULT_COLORMAP