""" The window the globe is shown in: camera, lighting, the click readout, and saving or exporting what was built. Only this module talks to a render window, so the geometry in `geometry.py` and the atmosphere in `shell.py` stay testable without one. """ import contextlib import os import numpy as np from ...figure import VIDEO_SUFFIXES from ..movie import open_encoder try: import vtkmodules.all as vtk except ImportError: # pragma: no cover try: import vtk except ImportError: vtk = None try: from vedo import Light, Line, Plotter, Text2D, Volume, merge, settings vedo_available = True except ImportError: vedo_available = False Volume = () # isinstance() against an empty tuple is False # Mars' axial tilt, used to place the sun from a solar longitude. OBLIQUITY_DEG = 25.19 # Depth-peeling passes. VTK's default of 4 resolves four translucent surfaces # along a ray and gives up on the rest, which for a stack of concentric shells - # two crossings each - means it stops after the second one and the blending # below that is whatever order the actors happened to be in. DEPTH_PEELS = 16 # Which way the cutting plane faces before anyone drags it. VTK keeps the side # the normal points to, and the default camera sits on +X. # # Halfway between facing the camera and lying edge-on to it, because the two # extremes each lose half the picture: square-on, the cross-section fills the # view and the shells behind it are never seen; edge-on, the cut face is a line # and only the shells are left. At 45 degrees the quarter that was removed opens # onto the face, and both are legible at once. CUT_NORMAL = (-0.7071, 0.7071, 0.0) # How long one --spin turn lasts, in seconds of finished video. The frame # count follows from this and --fps, so asking for more frames per second # makes the turn smoother rather than faster. SPIN_SECONDS = 8.0 # How far the camera sits from the centre, in radii of whatever the scene turns # out to be - see GlobeScene.camera. CAMERA_DISTANCE = 4.5 # How far the things drawn on top of the surface are lifted off it, in metres of # elevation - exaggerated with the relief, like every other height here. Ordered # so nothing z-fights with anything else. # # Small, because they no longer have to cover a modelling error. These used to # be ten times larger, to clear the gap between the smooth surface the lines # were placed on and the flat triangles they had to cross; `mesh_elevation_at` # puts them on the triangles themselves, so what is left to buy is depth-buffer # clearance and an order between the three. 200 m is 2 km once exaggerated, # invisible against a globe 3390 km across and ample against a depth buffer. CONTOUR_LIFT = 200.0 GRATICULE_LIFT = 300.0 TERMINATOR_LIFT = 400.0 # How far apart the vertices of a line laid on the surface may be, in degrees. # The render grid is one degree, so half of it keeps every segment inside a # single triangle's neighbourhood; see `geometry.densify`. LINE_STEP_DEG = 0.5 def subsolar_point(sun=None, ls=None): """ (longitude, latitude) of the sun on the planet, or None when neither was asked for. `sun` is an explicit "lon,lat" pair. `ls` is a solar longitude, which fixes the sub-solar *latitude* through the obliquity but says nothing about the longitude - that depends on the time of day - so the longitude defaults to 0 and the caller can override it with `sun`. """ if sun is not None: return float(sun[0]), float(sun[1]) if ls is None: return None lat = np.rad2deg(np.arcsin(np.sin(np.deg2rad(OBLIQUITY_DEG)) * np.sin(np.deg2rad(ls)))) return 0.0, float(lat) def sun_direction(lon, lat): """ Unit vector from the planet centre towards the sun. """ phi, theta = np.deg2rad(90.0 - lat), np.deg2rad(lon) return np.array([np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta), np.cos(phi)]) def terminator_points(globe, lon, lat, samples=1441): """ The day/night great circle: the surface points at right angles to the sun. Built in cartesian space from two vectors spanning the plane normal to the sun direction, then pushed back through the relief so the line follows the terrain like the graticule does. A quarter of a degree of arc between samples rather than a whole one, so that consecutive points land in neighbouring cells of the relief and the line follows the ground between them. Subdividing afterwards, the way a contour is subdivided, would not do: an interpolation in latitude and longitude leaves the great circle, and this line is one. """ s = sun_direction(lon, lat) # Any axis not parallel to the sun gives a first perpendicular seed = np.array([0.0, 0.0, 1.0]) if abs(s[2]) < 0.9 else np.array([1.0, 0.0, 0.0]) u = np.cross(s, seed) u /= np.linalg.norm(u) v = np.cross(s, u) t = np.linspace(0.0, 2.0 * np.pi, samples) ring = np.outer(np.cos(t), u) + np.outer(np.sin(t), v) lats, lons, _ = globe.to_geographic(ring) return globe.on_surface(lats, lons, TERMINATOR_LIFT) class GlobeScene: """ Collects the actors of one globe view and puts them on screen or on disk. """ def __init__(self, globe, title="3D globe view", bg="bb", size=(1200, 900)): self.globe = globe self.title = title self.bg = bg self.size = size self.actors = [] self.surface = None self.shell = None self._readout = None self._on_click = None self._plotter = None self._animation = None self._caption = None self._caption_lines = [] self._frame_line = None self._depth_peeling = False self._cut_plane = None self._cut_widget = None self._slicer = None self._slice_actor = None self._spin = None def add(self, *actors): """ Add actors, ignoring the Nones that merge() returns for empty inputs. """ self.actors.extend(a for a in actors if a is not None) return self def add_lines(self, segments, color='k', width=1.0): """ Merge many line segments into one actor. Hundreds of separate Line objects are what made the globe sluggish to rotate. """ segments = [s for s in segments if s is not None] if not segments: return self return self.add(merge(*segments).c(color).lw(width)) def add_sun(self, lon, lat, show_terminator=True, lit=True): """ Light the globe from the sub-solar point and mark the day/night line. The terminator is geometry rather than illumination, so it is drawn even when `lit` is false and the scene carries no light of its own. """ if lit: direction = sun_direction(lon, lat) self.add(Light(pos=(direction * self.globe.radius * 20).tolist(), focal_point=(0, 0, 0), intensity=1.0)) if show_terminator: self.add(Line(terminator_points(self.globe, lon, lat), c='orange', lw=2)) return self def add_cutter(self, slicer=None, normal=CUT_NORMAL): """ Slice the whole scene with a plane through the planet's centre, so the view looks into it instead of at its outermost surface. A stack of concentric shells cannot be seen into from outside: the outermost one covers every other, and no amount of transparency fixes that because VTK sorts translucent actors by centroid and every shell here shares the globe's. Cutting is the way out. Everything is cut, not only the shell. Removing the near half of the atmosphere and leaving the planet whole only swaps one lid for another - the globe fills the opening and the stack is as hidden as it was. `slicer` is what makes the opening worth having. Clipping on its own exposes nothing, because a stack of shells is a stack of *surfaces*: there is no interior behind them, only the far shells seen from inside. The cross-section is the face of the cut - altitude up it, coloured by the field - and it is the part that actually shows how the field is stacked. The plane is handed to each mapper, which clips on the GPU - no geometry is rebuilt, and it works on a Volume as readily as on a Mesh. `show` then puts a handle on it if there is an interactor to hang one from; without one the cut is simply static, which is what a saved still wants anyway. """ if vtk is None: return self plane = vtk.vtkPlane() plane.SetOrigin(0.0, 0.0, 0.0) plane.SetNormal(*normal) cut_any = False for actor in self.actors: # Text2D and Light carry no mapper, and a scalar bar clipped in # world space would simply vanish mapper = getattr(actor, 'mapper', None) if mapper is None or not hasattr(mapper, 'AddClippingPlane'): continue mapper.AddClippingPlane(plane) cut_any = True if not cut_any: return self self._cut_plane = plane self._slicer = slicer self._refresh_slice() return self def _refresh_slice(self): """ Rebuild the cross-section for wherever the plane is now. A slice costs about ten milliseconds against the seconds the volume took to sample, so this can run on every drag. The old actor has to leave the renderer as well as the list, or each drag leaves another one behind. """ if self._slicer is None or self._cut_plane is None: return # Sliced just clear of the clip, so the face is not fighting the very # plane that cut it for the same depth normal = np.asarray(self._cut_plane.GetNormal(), dtype=float) origin = np.asarray(self._cut_plane.GetOrigin(), dtype=float) fresh = self._slicer(origin - normal * self.globe.radius * 1e-3, normal) if self._slice_actor is not None: if self._slice_actor in self.actors: self.actors.remove(self._slice_actor) if self._plotter is not None: self._plotter.remove(self._slice_actor) self._slice_actor = fresh if fresh is not None: self.add(fresh) if self._plotter is not None: self._plotter.add(fresh) def _add_cut_handle(self): """ Put a draggable handle on the cutting plane. The representation is stripped back to the one gesture that matters here: the plane itself is not drawn (it would hide exactly what the cut reveals), the outline cannot be dragged off the planet, and it cannot be scaled. What is left is the normal, which sweeps the cut around. VTK drops a widget the moment nothing references it, so it is kept on the scene rather than in a local. """ if self._cut_plane is None or self._plotter is None: return interactor = getattr(self._plotter, 'interactor', None) if interactor is None: return radius = self.globe.radius * 1.2 rep = vtk.vtkImplicitPlaneRepresentation() rep.SetPlaceFactor(1.0) rep.PlaceWidget((-radius, radius, -radius, radius, -radius, radius)) rep.SetOrigin(*self._cut_plane.GetOrigin()) rep.SetNormal(*self._cut_plane.GetNormal()) rep.DrawPlaneOff() rep.OutlineTranslationOff() rep.ScaleEnabledOff() rep.GetPlaneProperty().SetOpacity(0.15) def on_move(widget, _event): # Copy the handle's plane onto the one the mappers already hold, # rather than swapping in a new object they are not watching widget.GetRepresentation().GetPlane(self._cut_plane) self._refresh_slice() widget = vtk.vtkImplicitPlaneWidget2() widget.SetRepresentation(rep) widget.SetInteractor(interactor) widget.AddObserver('InteractionEvent', on_move) widget.On() self._cut_widget = widget def want_depth_peeling(self, wanted=True): """ Ask for order-independent transparency while this scene is on screen. vedo reads `settings.use_depth_peeling` when a Plotter is built, and it is process-global: setting it where the translucent actor is created leaked into every later view, including flat ones. The scene turns it on around its own window and puts the previous value back. """ self._depth_peeling = self._depth_peeling or bool(wanted) return self def caption(self, lines): """ A standing description of the view, in the top-left corner. The globe carries several independent choices at once - what shapes it, what colours it, where the shell sits and what threshold drew it - and none of them are visible in the picture. Without this the only record is the terminal the run was launched from. """ self._caption_lines = [str(line) for line in lines if line] if self._caption is None and self._caption_lines: self._caption = Text2D('', pos='top-left', c='white', bg='black', alpha=0.6, s=0.75) self.add(self._caption) self._refresh_caption() return self def set_frame_label(self, text): """ The line of the caption that changes as an animation plays. """ self._frame_line = text self._refresh_caption() return self def _refresh_caption(self): if self._caption is None: return lines = list(self._caption_lines) if self._frame_line: lines.append(self._frame_line) self._caption.text('\n'.join(lines)) def attach_readout(self, mesh, varname, units=None): """ Report longitude, latitude, elevation and value wherever the user clicks, which is what `attach_format_coord` already gives the flat maps. """ label = f"{varname}" + (f" [{units}]" if units else "") self._readout = Text2D("click the globe for a value", pos='bottom-left', c='white', bg='black', alpha=0.6, s=0.8) self.add(self._readout) def on_click(event): if event.actor is not mesh or event.picked3d is None: return point = np.asarray(event.picked3d, dtype=float) lat, lon, height = self.globe.to_geographic(point) text = f"lon={lon:.2f} lat={lat:.2f} elev={height / 1000:.2f} km" value = self._value_at(mesh, point) if value is not None: text += f"\n{label} = {value:.4g}" self._readout.text(text) self._render() self._on_click = on_click return self @staticmethod def _value_at(mesh, point): """ The scalar carried by the vertex nearest a picked point, if there is one. """ scalars = mesh.pointdata['Scalars'] if scalars is None: return None values = np.asarray(scalars).ravel() index = int(mesh.closest_point(point, return_point_id=True)) if not 0 <= index < values.size or not np.isfinite(values[index]): return None return float(values[index]) def _render(self): if self._plotter is not None: self._plotter.render() def camera(self, distance=None, position=(1, 0, 0), viewup=(0, 0, 1)): """ Camera arguments for show(). These have to go through show(): a Plotter defaults to resetcam=True and would otherwise refit the camera to the bounds and discard the distance. 4.5 radii keeps the whole disc inside VTK's default 30 degree view angle, with room for the exaggerated relief. At 3 radii the globe subtends 41 degrees and is visibly cropped. Radii of *the scene*, though, not of the planet. An atmosphere drawn at thirty times its own depth reaches 1.43 R, while 4.5 R frames 1.21 R, so a fixed distance cut the outer half of every stack of layers and the whole rim of every cloud. `distance` is in units of the planet radius and overrides the fit. """ if distance is None: distance = CAMERA_DISTANCE * max(1.0, self._scene_radius() / self.globe.radius) pos = np.asarray(position, dtype=float) pos = pos / np.linalg.norm(pos) * self.globe.radius * distance return dict(pos=tuple(pos), focal_point=(0, 0, 0), viewup=tuple(viewup)) def _scene_radius(self): """ How far the furthest thing on screen reaches from the planet's centre. Measured per axis rather than as the corner of the bounding box: a cloud is ray-cast from a cube that circumscribes the atmosphere, and taking its corner would push the camera out by another root three to frame air that is not there. The cross-section of a cutaway is left out of it. It is a face cut through that same box, so it reaches into the corners the box has and the atmosphere does not - 3.4 radii against the shell's 1.4 - and framing it would shrink the planet to a third of its size to make room for a triangle of transparent air. """ radius = self.globe.radius for actor in self.actors: bounds = getattr(actor, 'bounds', None) if not callable(bounds) or actor is self._slice_actor: continue values = np.asarray(bounds(), dtype=float) if values.size < 6 or not np.all(np.isfinite(values)): continue radius = max(radius, float(np.abs(values).max())) return radius def animate(self, count, on_frame, label='frame', fps=12, describe=None): """ Drive `count` frames through `on_frame(index)`. Only the scalars change from frame to frame, never the geometry, so a frame costs one re-colouring rather than a rebuild. Interactively this becomes a slider; with a video path it becomes a movie. `describe(index)` returns the caption line for a frame - the coordinate value rather than the index, where the file gives one. """ if count > 1: self._animation = dict(count=count, on_frame=on_frame, label=label, fps=fps, describe=describe) return self def spin(self, turns=1.0, fps=12, seconds=SPIN_SECONDS): """ Turn the camera around the planet. Worth having because a fixed camera cannot tell a feature that moves with the planet from one that moves against it, and because a globe is a sphere: half of it is always facing away. To a file this becomes an animation, so a still globe can be written straight out as a movie; on top of an existing one it rides along and the data steps while the view turns. On screen it is a timer instead, which leaves the mouse free - the globe keeps turning while it is dragged. """ frames = max(2, int(round(fps * float(seconds)))) self._spin = dict(step=360.0 * float(turns) / frames, frames=frames, fps=fps, turns=float(turns)) if self._animation is None: self.animate(frames, lambda index: None, label='turn', fps=fps, describe=lambda index: f"turned {index * self._spin['step']:.0f} degrees") spec = self._animation inner = spec['on_frame'] def on_frame(index): inner(index) # Frame 0 is drawn where the camera already is, so the first step # comes after it and a whole turn ends where it began if index and self._plotter is not None: self._plotter.azimuth(self._spin['step']) spec['on_frame'] = on_frame return self def _add_spin_timer(self): """ Keep the globe turning on screen without holding the interactor. A timer callback rather than a loop, so dragging, clicking the readout and the cut handle all keep working while it turns. """ if self._spin is None or self._plotter is None: return def on_tick(_event): self._plotter.azimuth(self._spin['step']) self._plotter.render() self._plotter.add_callback('timer', on_tick) self._plotter.timer_callback('create', dt=int(1000 / max(1, self._spin['fps']))) def _show_frame(self, index): """ Draw one frame and keep the caption in step with it. """ spec = self._animation spec['on_frame'](index) describe = spec.get('describe') self.set_frame_label(describe(index) if describe else f"{spec['label']} {index + 1}/{spec['count']}") def show(self, output_path=None, camera=None): """ Open the window, or render offscreen and write `output_path`. """ camera = camera or self.camera() if self._animation: # Frame 0 is already drawn; this only puts its label in the caption self._show_frame(0) if output_path: if self._animation and output_path.lower().endswith(VIDEO_SUFFIXES): return self._save_video(output_path, camera) return self._save(output_path, camera) with self._peeling(): self._plotter = Plotter(title=self.title, bg=self.bg, axes=0, size=self.size) try: if self._readout is not None: self._plotter.add_callback('mouse click', self._on_click) if self._animation: self._add_frame_slider() self._add_cut_handle() self._add_spin_timer() self._plotter.show(*self.actors, camera=camera) finally: # Without this the VTK render window outlives the call: the # interactor returns but the window is never torn down, so in # the interactive loop it lingers and the next globe opens # behind it. self._close(self._plotter) self._plotter = None return 0 @contextlib.contextmanager def _peeling(self): """ Apply, then restore, the global depth-peeling settings. The peel count travels with the flag: both are process-global, and turning peeling on while leaving it at VTK's four passes only resolves the first two shells of a stack. """ previous = settings.use_depth_peeling previous_peels = settings.max_number_of_peels settings.use_depth_peeling = self._depth_peeling or previous if self._depth_peeling: settings.max_number_of_peels = max(previous_peels, DEPTH_PEELS) try: yield finally: settings.use_depth_peeling = previous settings.max_number_of_peels = previous_peels @staticmethod def _close(plotter): """ Tear a plotter down without letting VTK's own teardown mask a real error. """ if plotter is None: return try: plotter.close() except Exception: pass def _add_frame_slider(self): """ A slider along the bottom that scrubs the animated dimension. """ spec = self._animation last = spec['count'] - 1 def on_slide(widget, _event): self._show_frame(int(round(widget.value))) self._plotter.add_slider(on_slide, 0, last, value=0, pos=((0.15, 0.06), (0.85, 0.06)), title=f"{spec['label']} (0-{last})", show_value=True, delayed=True) def _save_video(self, output_path, camera): """ Render every frame offscreen into a movie. Frames go straight down a pipe into ffmpeg as raw pixels. They used to go through vedo's Video, which writes every frame to disk as a PNG and reads it back to encode - three passes over each frame, two of them compression, for output that is about to be compressed again. Piping removes all of it, and the temporary directory with it. """ spec = self._animation plotter = None try: with self._peeling(): plotter = Plotter(offscreen=True, bg=self.bg, axes=0, size=self.size) plotter.show(*self.actors, camera=camera) # A frame callback may need to reach the plotter - the spin # turns the camera between frames - and this is the only handle # on it that is not local to this method self._plotter = plotter first = self._grab(plotter, 0) encoder = open_encoder(output_path, spec['fps'], first.shape) if encoder is None: return 1 with encoder: encoder.add(first) for index in range(1, spec['count']): encoder.add(self._grab(plotter, index)) except Exception as err: print(f"Could not render the 3D globe offscreen: {err}") print("Offscreen rendering needs a GL context; try 'xvfb-run -a '.") return 1 finally: self._close(plotter) self._plotter = None if not os.path.isfile(output_path) or os.path.getsize(output_path) == 0: print(f"Rendered {spec['count']} frames but '{output_path}' was not written.") return 1 print(f"Saved {spec['count']} frames to {output_path}") return 0 def _grab(self, plotter, index): """ Draw one frame and hand back its pixels. """ self._show_frame(index) plotter.render() return plotter.screenshot(asarray=True) def _save(self, output_path, camera): """ Offscreen render. VTK needs a GL context even with no window, so a headless machine without EGL or OSMesa fails here; say so plainly instead of letting a VTK traceback out. """ from ...figure import ensure_directory if not ensure_directory(output_path): return 1 plotter = None try: with self._peeling(): plotter = Plotter(offscreen=True, bg=self.bg, axes=0, size=self.size) plotter.show(*self.actors, camera=camera) plotter.screenshot(output_path) except Exception as err: print(f"Could not render the 3D globe offscreen: {err}") print("Offscreen rendering needs a GL context; try 'xvfb-run -a '.") return 1 finally: self._close(plotter) print(f"Saved to {output_path}") return 0 def export(self, path): """ Write the meshes to a 3D file (.vtk, .vtp, .ply, .obj, .stl). A shell is written beside the surface as '_shell', the same convention the polar views use for their two hemispheres. A cloud is a volume rather than a mesh, so it goes to '.vti' whatever the surface asked for: none of the mesh formats can hold voxels. """ written = [] if self.surface is not None: self.surface.write(path) written.append(path) if self.shell is not None: stem, ext = os.path.splitext(path) if isinstance(self.shell, Volume): ext = '.vti' shell_path = f"{stem}_shell{ext or '.vtk'}" self.shell.write(shell_path) written.append(shell_path) for item in written: print(f"Exported {item}") return written