mirror of
https://github.com/google-deepmind/deepmind-research.git
synced 2026-02-06 12:02:08 +08:00
68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
# pylint: disable=g-bad-file-header
|
|
# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ============================================================================
|
|
"""Plots a CFD trajectory rollout."""
|
|
|
|
import pickle
|
|
|
|
from absl import app
|
|
from absl import flags
|
|
from matplotlib import animation
|
|
from matplotlib import tri as mtri
|
|
import matplotlib.pyplot as plt
|
|
|
|
FLAGS = flags.FLAGS
|
|
flags.DEFINE_string('rollout_path', None, 'Path to rollout pickle file')
|
|
|
|
|
|
def main(unused_argv):
|
|
with open(FLAGS.rollout_path, 'rb') as fp:
|
|
rollout_data = pickle.load(fp)
|
|
|
|
fig, ax = plt.subplots(1, 1, figsize=(12, 8))
|
|
skip = 10
|
|
num_steps = rollout_data[0]['gt_velocity'].shape[0]
|
|
num_frames = len(rollout_data) * num_steps // skip
|
|
|
|
# compute bounds
|
|
bounds = []
|
|
for trajectory in rollout_data:
|
|
bb_min = trajectory['gt_velocity'].min(axis=(0, 1))
|
|
bb_max = trajectory['gt_velocity'].max(axis=(0, 1))
|
|
bounds.append((bb_min, bb_max))
|
|
|
|
def animate(num):
|
|
step = (num*skip) % num_steps
|
|
traj = (num*skip) // num_steps
|
|
ax.cla()
|
|
ax.set_aspect('equal')
|
|
ax.set_axis_off()
|
|
vmin, vmax = bounds[traj]
|
|
pos = rollout_data[traj]['mesh_pos'][step]
|
|
faces = rollout_data[traj]['faces'][step]
|
|
velocity = rollout_data[traj]['pred_velocity'][step]
|
|
triang = mtri.Triangulation(pos[:, 0], pos[:, 1], faces)
|
|
ax.tripcolor(triang, velocity[:, 0], vmin=vmin[0], vmax=vmax[0])
|
|
ax.triplot(triang, 'ko-', ms=0.5, lw=0.3)
|
|
ax.set_title('Trajectory %d Step %d' % (traj, step))
|
|
return fig,
|
|
|
|
_ = animation.FuncAnimation(fig, animate, frames=num_frames, interval=100)
|
|
plt.show(block=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
app.run(main)
|