Spaces:
Sleeping
Sleeping
| from typing import Set | |
| import numpy as np | |
| from warp.fem import DiscreteField | |
| def plot_grid_surface(field, axes=None): | |
| import matplotlib.pyplot as plt | |
| from matplotlib import cm | |
| if axes is None: | |
| fig, axes = plt.subplots(subplot_kw={"projection": "3d"}) | |
| node_positions = field.space.node_grid() | |
| # Make data. | |
| X = node_positions[0] | |
| Y = node_positions[1] | |
| Z = field.dof_values.numpy().reshape(X.shape) | |
| # Plot the surface. | |
| return axes.plot_surface(X, Y, Z, cmap=cm.coolwarm, linewidth=0, antialiased=False) | |
| def plot_tri_surface(field, axes=None): | |
| import matplotlib.pyplot as plt | |
| from matplotlib import cm | |
| from matplotlib.tri.triangulation import Triangulation | |
| if axes is None: | |
| fig, axes = plt.subplots(subplot_kw={"projection": "3d"}) | |
| node_positions = field.space.node_positions().numpy() | |
| triangulation = Triangulation( | |
| x=node_positions[:, 0], y=node_positions[:, 1], triangles=field.space.node_triangulation() | |
| ) | |
| Z = field.dof_values.numpy() | |
| # Plot the surface. | |
| return axes.plot_trisurf(triangulation, Z, cmap=cm.coolwarm, linewidth=0, antialiased=False) | |
| def plot_scatter_surface(field, axes=None): | |
| import matplotlib.pyplot as plt | |
| from matplotlib import cm | |
| if axes is None: | |
| fig, axes = plt.subplots(subplot_kw={"projection": "3d"}) | |
| X, Y = field.space.node_positions().numpy().T | |
| # Make data. | |
| Z = field.dof_values.numpy().reshape(X.shape) | |
| # Plot the surface. | |
| return axes.scatter(X, Y, Z, c=Z, cmap=cm.coolwarm) | |
| def plot_surface(field, axes=None): | |
| if hasattr(field.space, "node_grid"): | |
| return plot_grid_surface(field, axes) | |
| elif hasattr(field.space, "node_triangulation"): | |
| return plot_tri_surface(field, axes) | |
| else: | |
| return plot_scatter_surface(field, axes) | |
| def plot_grid_color(field, axes=None): | |
| import matplotlib.pyplot as plt | |
| from matplotlib import cm | |
| if axes is None: | |
| fig, axes = plt.subplots() | |
| node_positions = field.space.node_grid() | |
| # Make data. | |
| X = node_positions[0] | |
| Y = node_positions[1] | |
| Z = field.dof_values.numpy().reshape(X.shape) | |
| # Plot the surface. | |
| return axes.pcolormesh(X, Y, Z, cmap=cm.coolwarm) | |
| def plot_velocities(field, axes=None): | |
| import matplotlib.pyplot as plt | |
| if axes is None: | |
| fig, axes = plt.subplots() | |
| node_positions = field.space.node_positions().numpy() | |
| # Make data. | |
| X = node_positions[:, 0] | |
| Y = node_positions[:, 1] | |
| vel = field.dof_values.numpy() | |
| u = np.ascontiguousarray(vel[:, 0]) | |
| v = np.ascontiguousarray(vel[:, 1]) | |
| u = u.reshape(X.shape) | |
| v = v.reshape(X.shape) | |
| return axes.quiver(X, Y, u, v) | |
| def plot_grid_streamlines(field, axes=None): | |
| import matplotlib.pyplot as plt | |
| if axes is None: | |
| fig, axes = plt.subplots() | |
| node_positions = field.space.node_grid() | |
| # Make data. | |
| X = node_positions[0][:, 0] | |
| Y = node_positions[1][0, :] | |
| vel = field.dof_values.numpy() | |
| u = np.ascontiguousarray(vel[:, 0]) | |
| v = np.ascontiguousarray(vel[:, 1]) | |
| u = np.transpose(u.reshape(node_positions[0].shape)) | |
| v = np.transpose(v.reshape(node_positions[0].shape)) | |
| splot = axes.streamplot(X, Y, u, v, density=2) | |
| splot.axes = axes | |
| return splot | |
| def plot_3d_scatter(field, axes=None): | |
| import matplotlib.pyplot as plt | |
| from matplotlib import cm | |
| if axes is None: | |
| fig, axes = plt.subplots(subplot_kw={"projection": "3d"}) | |
| X, Y, Z = field.space.node_positions().numpy().T | |
| # Make data. | |
| f = field.dof_values.numpy().reshape(X.shape) | |
| # Plot the surface. | |
| return axes.scatter(X, Y, Z, c=f, cmap=cm.coolwarm) | |
| def plot_3d_velocities(field, axes=None): | |
| import matplotlib.pyplot as plt | |
| if axes is None: | |
| fig, axes = plt.subplots(subplot_kw={"projection": "3d"}) | |
| X, Y, Z = field.space.node_positions().numpy().T | |
| vel = field.dof_values.numpy() | |
| u = np.ascontiguousarray(vel[:, 0]) | |
| v = np.ascontiguousarray(vel[:, 1]) | |
| w = np.ascontiguousarray(vel[:, 2]) | |
| u = u.reshape(X.shape) | |
| v = v.reshape(X.shape) | |
| w = w.reshape(X.shape) | |
| return axes.quiver(X, Y, Z, u, v, w, length=1.0 / X.shape[0], normalize=False) | |
| class Plot: | |
| def __init__(self, stage=None, default_point_radius=0.01): | |
| self.default_point_radius = default_point_radius | |
| self._surfaces = {} | |
| self._surface_vectors = {} | |
| self._volumes = {} | |
| if stage is not None: | |
| from warp.render import UsdRenderer | |
| self._usd_renderer = UsdRenderer(stage) | |
| self._plt = None | |
| else: | |
| self._usd_renderer = None | |
| def begin_frame(self, time): | |
| if self._usd_renderer is not None: | |
| self._usd_renderer.begin_frame(time=time) | |
| def end_frame(self): | |
| if self._usd_renderer is not None: | |
| self._usd_renderer.end_frame() | |
| def add_surface(self, name: str, field: DiscreteField): | |
| if self._usd_renderer is not None: | |
| points_2d = field.space.node_positions().numpy() | |
| values = field.dof_values.numpy() | |
| points_3d = np.hstack((points_2d, values.reshape(-1, 1))) | |
| if hasattr(field.space, "node_triangulation"): | |
| indices = field.space.node_triangulation() | |
| self._usd_renderer.render_mesh(name, points=points_3d, indices=indices) | |
| else: | |
| self._usd_renderer.render_points(name, points=points_3d, radius=self.default_point_radius) | |
| if name not in self._surfaces: | |
| field_clone = field.space.make_field(space_partition=field.space_partition) | |
| self._surfaces[name] = (field_clone, []) | |
| self._surfaces[name][1].append(field.dof_values.numpy()) | |
| def add_surface_vector(self, name: str, field: DiscreteField): | |
| if self._usd_renderer is not None: | |
| points_2d = field.space.node_positions().numpy() | |
| values = field.dof_values.numpy() | |
| points_3d = np.hstack((points_2d + values, np.zeros_like(points_2d[:, 0]).reshape(-1, 1))) | |
| if hasattr(field.space, "node_triangulation"): | |
| indices = field.space.node_triangulation() | |
| self._usd_renderer.render_mesh(name, points=points_3d, indices=indices) | |
| else: | |
| self._usd_renderer.render_points(name, points=points_3d, radius=self.default_point_radius) | |
| if name not in self._surface_vectors: | |
| field_clone = field.space.make_field(space_partition=field.space_partition) | |
| self._surface_vectors[name] = (field_clone, []) | |
| self._surface_vectors[name][1].append(field.dof_values.numpy()) | |
| def add_volume(self, name: str, field: DiscreteField): | |
| if self._usd_renderer is not None: | |
| points_3d = field.space.node_positions().numpy() | |
| values = field.dof_values.numpy() | |
| self._usd_renderer.render_points(name, points_3d, radius=values) | |
| if name not in self._volumes: | |
| field_clone = field.space.make_field(space_partition=field.space_partition) | |
| self._volumes[name] = (field_clone, []) | |
| self._volumes[name][1].append(field.dof_values.numpy()) | |
| def plot(self, streamlines: Set[str] = []): | |
| return self._plot_matplotlib(streamlines) | |
| def _plot_matplotlib(self, streamlines: Set[str] = []): | |
| import matplotlib.pyplot as plt | |
| import matplotlib.animation as animation | |
| def make_animation(ax, field, values, plot_func, num_frames: int): | |
| def animate(i): | |
| ax.clear() | |
| field.dof_values = values[i] | |
| return plot_func(field, axes=ax) | |
| return animation.FuncAnimation( | |
| ax.figure, | |
| animate, | |
| interval=30, | |
| blit=False, | |
| frames=len(values), | |
| ) | |
| for name, (field, values) in self._surfaces.items(): | |
| field.dof_values = values[0] | |
| ax = plot_surface(field).axes | |
| if len(values) > 1: | |
| anim = make_animation(ax, field, values, plot_func=plot_surface, num_frames=len(values)) | |
| for name, (field, values) in self._surface_vectors.items(): | |
| field.dof_values = values[0] | |
| if name in streamlines and hasattr(field.space, "node_grid"): | |
| ax = plot_grid_streamlines(field).axes | |
| else: | |
| ax = plot_velocities(field).axes | |
| if len(values) > 1: | |
| anim = make_animation(ax, field, values, plot_func=plot_velocities, num_frames=len(values)) | |
| for name, (field, values) in self._volumes.items(): | |
| field.dof_values = values[0] | |
| ax = plot_3d_scatter(field).axes | |
| plt.show() | |