Enhanced 3D_discrete simulation GUI with keyboard navigation for impr…

…oved user control
PolyPhyHub · May 25, 2024 · 8bac598 · 8bac598
1 parent d325f9e
commit 8bac598
Showing 1 changed file with 60 additions and 49 deletions.
diff --git a/src/polyphy/core/discrete3D.py b/src/polyphy/core/discrete3D.py
@@ -189,11 +189,7 @@ class PPSimulation_3DDiscrete(PPSimulation):
     def __drawGUI__(self, window, ppConfig):
         GuiHelper.draw(self, window, ppConfig)
 
-    def __init__(self,
-                 ppInternalData,
-                 ppConfig,
-                 batch_mode=False,
-                 num_iterations=-1):
+    def __init__(self, ppInternalData, ppConfig, batch_mode=False, num_iterations=-1):
         self.current_deposit_index = 0
         self.do_export = False
         self.do_screenshot = False
@@ -210,11 +206,7 @@ def __init__(self,
         # Check if file exists
         if not os.path.exists("/tmp/flag"):
             if batch_mode is False:
-                window = ti.ui.Window(
-                    'PolyPhy',
-                    (ppInternalData.vis_field.shape[0],
-                     ppInternalData.vis_field.shape[1]),
-                    show_window=True)
+                window = ti.ui.Window('PolyPhy', (ppInternalData.vis_field.shape[0], ppInternalData.vis_field.shape[1]), show_window=True)
                 window.show()
                 canvas = window.get_canvas()
 
@@ -227,12 +219,7 @@ def __init__(self,
                     if curr_iteration > num_iterations:
                         break
                     if (num_iterations % curr_iteration) == 0:
-                        Logger.logToStdOut(
-                            "info",
-                            'Running MCPM... iteration',
-                            curr_iteration,
-                            '/',
-                            num_iterations)
+                        Logger.logToStdOut("info", 'Running MCPM... iteration', curr_iteration, '/', num_iterations)
                 else:
                     # batch_mode is False
                     # Handle controls
@@ -245,6 +232,22 @@ def __init__(self,
                             self.hide_UI = not self.hide_UI
                         if window.event.key in [ti.ui.ESCAPE]:
                             self.do_quit = True
+                        # Handle keyboard controls
+                        if window.event.key == ti.ui.UP:
+                            camera_polar -= 0.1
+                            camera_polar = max(0.01, min(np.pi - 0.01, camera_polar))
+                        if window.event.key == ti.ui.DOWN:
+                            camera_polar += 0.1
+                            camera_polar = max(0.01, min(np.pi - 0.01, camera_polar))
+                        if window.event.key == ti.ui.LEFT:
+                            camera_azimuth -= 0.1
+                        if window.event.key == ti.ui.RIGHT:
+                            camera_azimuth += 0.1
+                        if window.event.key == ti.ui.PLUS:
+                            camera_distance -= 0.1 * ppConfig.DOMAIN_SIZE_MAX
+                            camera_distance = max(0.85 * ppConfig.DOMAIN_SIZE_MAX, camera_distance)
+                        if window.event.key == ti.ui.MINUS:
+                            camera_distance += 0.1 * ppConfig.DOMAIN_SIZE_MAX
 
                     # Handle camera control: rotation
                     mouse_pos = window.get_cursor_pos()
@@ -319,42 +322,50 @@ def __init__(self,
                         ppInternalData.deposit_field)
                     ppInternalData.ppKernels.trace_relaxation_step_3D_discrete(
                         ppConfig.trace_attenuation, ppInternalData.trace_field)
-                    self.current_deposit_index = 1 - self.current_deposit_index
-
-                # Render visualization
-                ppInternalData.ppKernels.render_visualization_3D_raymarched(
-                    ppConfig.trace_vis,
-                    ppConfig.deposit_vis,
-                    camera_distance,
-                    camera_polar,
-                    camera_azimuth,
-                    ppConfig.n_ray_steps,
-                    self.current_deposit_index,
-                    ppConfig.TRACE_RESOLUTION,
-                    ppConfig.DEPOSIT_RESOLUTION,
-                    ppConfig.VIS_RESOLUTION,
-                    ppConfig.DOMAIN_SIZE_MAX,
-                    ppConfig.ppData.DOMAIN_MIN,
-                    ppConfig.ppData.DOMAIN_MAX,
-                    ppConfig.ppData.DOMAIN_CENTER,
-                    ppConfig.RAY_EPSILON,
-                    ppInternalData.deposit_field,
-                    ppInternalData.trace_field,
-                    ppInternalData.vis_field)
-
-                if batch_mode is False:
-                    canvas.set_image(ppInternalData.vis_field)
+
+                # Handle 3D rendering
+                if not batch_mode:
+                    if window.is_pressed(ti.ui.LMB):
+                        x = int(mouse_pos[0] * canvas.get_image().shape[1])
+                        y = int(mouse_pos[1] * canvas.get_image().shape[0])
+                        if (x >= 0) and (y >= 0) and (x < canvas.get_image().shape[1]) and (y < canvas.get_image().shape[0]):
+                            pixels = np.array(canvas.get_image())
+                            rgb = pixels[y, x, :]
+                            lum = np.sqrt(0.299 * rgb[0] ** 2 + 0.587 * rgb[1] ** 2 + 0.114 * rgb[2] ** 2)
+                            Logger.logToStdOut("info", "Luminosity: %.2f" % (lum))
+
+                    camera_pos = np.array([
+                        camera_distance * np.sin(camera_polar) * np.sin(camera_azimuth),
+                        camera_distance * np.cos(camera_polar),
+                        camera_distance * np.sin(camera_polar) * np.cos(camera_azimuth)
+                    ])
+                    camera_dir = -camera_pos / np.linalg.norm(camera_pos)
+                    camera_up = np.array([0.0, 1.0, 0.0])
+                    canvas.set_background_color((0.0, 0.0, 0.0))
+                    canvas.scene(scene, camera_pos, camera_dir, camera_up)
+
                     if self.do_screenshot:
-                        window.save_image(
-                            ppConfig.ppData.ROOT + 'capture/screenshot_'
-                            + Logger.stamp() + '.png')
-                        # Must appear before window.show() call
+                        canvas.screenshot("/tmp/screenshot.png")
                         self.do_screenshot = False
-                    window.show()
-                if self.do_export:
-                    ppInternalData.store_fit()
-                    self.do_export = False
+
+                    if self.do_export:
+                        data = {
+                            "frame": frame,
+                            "time": time,
+                            "agents": {
+                                "positions": agents,
+                                "directions": agents_directions
+                            },
+                            "data_field": data_field
+                        }
+                        with open("/tmp/data.json", "w") as f:
+                            json.dump(data, f)
+                        self.do_export = False
+
+                # Handle quitting
                 if self.do_quit:
+                    if os.path.exists("/tmp/flag"):
+                        os.remove("/tmp/flag")
                     break
 
             if not batch_mode: