diff --git a/doc/sphinx/installation.rst b/doc/sphinx/installation.rst index 198e90d2f8..adb639a6c5 100644 --- a/doc/sphinx/installation.rst +++ b/doc/sphinx/installation.rst @@ -132,7 +132,7 @@ To install the ZnDraw visualizer: .. code-block:: bash - python3 -m pip install --user -c requirements.txt 'zndraw==0.4.5' + python3 -m pip install --user -c requirements.txt 'zndraw==0.4.6' .. _Nvidia GPU acceleration: diff --git a/doc/sphinx/visualization.rst b/doc/sphinx/visualization.rst index b421f49aca..652013b560 100644 --- a/doc/sphinx/visualization.rst +++ b/doc/sphinx/visualization.rst @@ -303,7 +303,7 @@ ZnDraw visualizer |es| supports the ZnDraw visualizer :cite:`elijosius24a` in Jupyter Notebooks. With ZnDraw [1]_, you can visualize your simulation live in a notebook or -web browser. The visualizer is based on ``THREE.js``. +web browser. .. _ZnDraw General usage: @@ -316,6 +316,9 @@ colors like ``red``, ``black`` etc., but one can also use hex colors like ``#ff0 Then write your integration loop in a separate function, and call the update function of the visualizer to capture the current state of the system and visualize it. Note that the visualizer needs to be started by pressing space. +Note that due to the need of a server running in the background, the visualizer is currently not able to run in multiple notebooks at once. +Make sure only one kernel is running the visualizer at a time. + Example code:: import espressomd @@ -351,6 +354,10 @@ as a criterium. The ``normalize`` boolean which normalizes the color to the larg ``normalize`` is false and describes the range to what the colorrange is applied to. ``scale_vector_thickness`` is a boolean and changes the thickness scaling of the vectors and ``opacity`` is a float value that sets the opacity of the vectors. +The ``ZnDraw``-object used for visualization is wrapped inside the visualizer object and is exposed by the attribute ``vis.zndraw``. + +When the ``jupyter`` visualization-window comes up as a blank page, make sure to allow third party cookies in your browser settings. + An example code snippet containing the :class:`~espressomd.zn.LBField` object:: import espressomd.zn diff --git a/src/python/espressomd/zn.py b/src/python/espressomd/zn.py index 7cb887cd5c..9520664230 100644 --- a/src/python/espressomd/zn.py +++ b/src/python/espressomd/zn.py @@ -29,6 +29,7 @@ import time import urllib.parse import typing as t +import scipy.spatial.transform # Standard colors @@ -417,10 +418,13 @@ class Visualizer(): """ + SERVER_PORT = None + SOCKET_PORT = None + def __init__(self, system: espressomd.system.System = None, port: int = 1234, - token: str = secrets.token_hex(4), + token: str = None, folded: bool = True, colors: dict = None, radii: dict = None, @@ -440,13 +444,18 @@ def __init__(self, self.url = "http://127.0.0.1" self.frame_count = 0 - self.port = port - self.token = token + if token is None: + self.token = secrets.token_hex(4) + else: + self.token = token # A server is started in a subprocess, and we have to wait for it - print("Starting ZnDraw server, this may take a few seconds") - self._start_server() - time.sleep(10) + if self.SERVER_PORT is None: + print("Starting ZnDraw server, this may take a few seconds") + self.port = port + self._start_server() + time.sleep(10) + self._start_zndraw() time.sleep(1) @@ -470,9 +479,8 @@ def __init__(self, if jupyter: self._show_jupyter() else: - # Problems with server being terminated at the end of the script - raise NotImplementedError("Only Jupyter is supported for now") - # webbrowser.open_new_tab(self.address) + raise NotImplementedError( + "Only Jupyter notebook is supported at the moment") def _start_server(self): """ @@ -480,6 +488,9 @@ def _start_server(self): """ self.socket_port = zndraw.utils.get_port(default=6374) + Visualizer.SERVER_PORT = self.port + Visualizer.SOCKET_PORT = self.socket_port + self.server = subprocess.Popen(["zndraw", "--no-browser", f"--port={self.port}", f"--storage-port={self.socket_port}"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL @@ -500,12 +511,12 @@ def _start_zndraw(self): while True: try: self.r = znsocket.Client( - address=f"{self.url}:{self.socket_port}") + address=f"{self.url}:{self.SOCKET_PORT}") break except BaseException: time.sleep(0.5) - url = f"{self.url}:{self.port}" + url = f"{self.url}:{self.SERVER_PORT}" self.zndraw = zndraw.zndraw.ZnDrawLocal( r=self.r, url=url, token=self.token, timeout=config) parsed_url = urllib.parse.urlparse( @@ -532,16 +543,28 @@ def update(self): ) # Catch when the server is initializing an empty frame - if self.frame_count == 0 and len(self.zndraw) > 0: - self.zndraw.__setitem__(0, data) - else: + # len(self.zndraw) is a expensive socket call, so we try to avoid it + if self.frame_count != 0 or len(self.zndraw) == 0: self.zndraw.append(data) + else: + self.zndraw.__setitem__(0, data) - if self.params["vector_field"] is not None and self.frame_count == 0: + if self.frame_count == 0: self.zndraw.socket.sleep(1) - for key, value in self.arrow_config.items(): - setattr(self.zndraw.config.arrows, key, value) + x = self.system.box_l[0] / 2 + y = self.system.box_l[1] / 2 + z = self.system.box_l[2] / 2 + + z_dist = max([1.5 * y, 1.5 * x, 1.5 * z]) + + self.zndraw.camera = {'position': [ + x, y, z_dist], 'target': [x, y, z]} + self.zndraw.config.scene.frame_update = False + + if self.params["vector_field"] is not None: + for key, value in self.arrow_config.items(): + setattr(self.zndraw.config.arrows, key, value) self.frame_count += 1 @@ -580,18 +603,40 @@ def draw_constraints(self, shapes: list): dist = shape.dist normal = np.array(shape.normal) - rotation_angles = zndraw.utils.direction_to_euler(normal) - - position = (dist * normal).tolist() - # Not optimal, but ensures its always larger than the box. - wall_width = wall_height = 2 * max(self.system.box_l) - - objects.append(zndraw.draw.Plane( - position=position, - rotation=rotation_angles, - width=wall_width, - height=wall_height, - material=mat)) + position = dist * normal + helper = WallIntersection( + plane_normal=normal, plane_point=position, box_l=self.system.box_l) + corners = helper.get_intersections() + + base_position = np.copy(corners[0]) + corners -= base_position + + # Rotate plane to align with z-axis, Custom2DShape only works + # in the xy-plane + unit_z = np.array([0, 0, 1]) + r, _ = scipy.spatial.transform.Rotation.align_vectors( + [unit_z], [normal]) + rotated_corners = r.apply(corners) + + # Sort corners in a clockwise order, except the first corner + angles = np.arctan2( + rotated_corners[1:, 1], rotated_corners[1:, 0]) + sorted_indices = np.argsort(angles) + sorted_corners = rotated_corners[1:][sorted_indices] + sorted_corners = np.vstack( + [rotated_corners[0], sorted_corners])[:, :2] + + r, _ = scipy.spatial.transform.Rotation.align_vectors( + [normal], [unit_z]) + euler_angles = r.as_euler("xyz") + + # invert the z-axis, unsure why this is needed, maybe + # different coordinate systems + euler_angles[2] *= -1. + + objects.append(zndraw.draw.Custom2DShape( + position=base_position, rotation=euler_angles, + points=sorted_corners, material=mat)) elif shape_type == "Sphere": center = shape.center @@ -621,4 +666,64 @@ def draw_constraints(self, shapes: list): raise NotImplementedError( f"Shape of type {shape_type} isn't available in ZnDraw") - self.zndraw.geometries = objects + self.zndraw.geometries = objects + + +class WallIntersection: + """ + Simple helper to calculate all Box edges that intersect with a plane. + """ + + def __init__(self, plane_point, plane_normal, box_l): + self.plane_point = plane_point + self.plane_normal = plane_normal / np.linalg.norm(plane_normal) + self.box_l = box_l + + # Create 8 vertices of the bounding box + self.vertices = np.array([ + [0, 0, 0], + [0, 0, box_l[2]], + [0, box_l[1], 0], + [0, box_l[1], box_l[2]], + [box_l[0], 0, 0], + [box_l[0], 0, box_l[2]], + [box_l[0], box_l[1], 0], + [box_l[0], box_l[1], box_l[2]] + ]) + + self.edges = [ + (self.vertices[0], self.vertices[1]), + (self.vertices[0], self.vertices[2]), + (self.vertices[0], self.vertices[4]), + (self.vertices[1], self.vertices[3]), + (self.vertices[1], self.vertices[5]), + (self.vertices[2], self.vertices[3]), + (self.vertices[2], self.vertices[6]), + (self.vertices[3], self.vertices[7]), + (self.vertices[4], self.vertices[5]), + (self.vertices[4], self.vertices[6]), + (self.vertices[5], self.vertices[7]), + (self.vertices[6], self.vertices[7]) + ] + + def plane_intersection_with_line(self, line_point1, line_point2): + # Calculate the intersection point of a line and a plane + line_dir = line_point2 - line_point1 + denom = np.dot(self.plane_normal, line_dir) + + if np.abs(denom) > 1e-6: # Avoid division by zero + t = np.dot(self.plane_normal, + (self.plane_point - line_point1)) / denom + if 0 <= t <= 1: + return line_point1 + t * line_dir + return None + + def get_intersections(self): + intersections = [] + + for edge in self.edges: + intersection = self.plane_intersection_with_line(edge[0], edge[1]) + if intersection is not None: + intersections.append(intersection) + + return np.array(intersections) diff --git a/testsuite/scripts/importlib_wrapper.py b/testsuite/scripts/importlib_wrapper.py index f3dd8e5b44..5bc879338c 100644 --- a/testsuite/scripts/importlib_wrapper.py +++ b/testsuite/scripts/importlib_wrapper.py @@ -394,7 +394,7 @@ class GetEspressomdVisualizerImports(ast.NodeVisitor): """ def __init__(self): - self.visualizers = {"visualization"} + self.visualizers = {"visualization", "zn"} self.namespace_visualizers = { "espressomd." + x for x in self.visualizers} self.visu_items = {}