<jupyter_start><jupyter_text>Copyright 2021 The Balloon Learning Environment Authors.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.0Unless required by applicable law or agreed to in writing, softwaredistributed 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 andlimitations under the License.  BLE View generated wind fieldsThis notebook allows you to sample from our trained WindField VAE and generate an animation of the resulting wind fields.Simply run the cells in order!<jupyter_code># @title Install necessary libraries
!pip install --upgrade pip
!pip install -U balloon_learning_environment
# @title Necessary imports
import datetime as dt
import numpy as np
import jax
import matplotlib.pyplot as plt
from matplotlib import animation
from matplotlib import rc
rc('animation', html='jshtml')

from balloon_learning_environment.env import generative_wind_field
# @title Generate a wind field
wind_field = generative_wind_field.GenerativeWindField()
rng = jax.random.PRNGKey(0)
wind_field.reset_forecast(rng, dt.datetime.now())
# @title Generate animation object (with a test still)
lats = np.linspace(-150, 150, 21)
lngs = np.linspace(-150, 150, 21)
pressure_levels = np.linspace(5000, 14000, 10)
x, y, z = np.meshgrid(lngs, lats, pressure_levels)
xx, yy = np.meshgrid(lats, lngs)

time_step = 0
direction = 1
episode_num = 0
elev = 10
azim = -60

fig = plt.figure(figsize=(20, 15))

def frame(w):
  global time_step
  global direction
  global episode_num
  global azim

  u = wind_field.field[:, :, :, time_step, 0]
  v = wind_field.field[:, :, :, time_step, 1]

  ax = fig.add_subplot(projection='3d', elev=elev, azim=azim)
  ax.quiver(x, y, z, u, v, np.zeros_like(u), length=1.5,
            arrow_length_ratio=0.01, normalize=False)
  ax.set_xlabel('Longitude displacement (km)')
  ax.set_ylabel('Latitude displacement (km)')
  ax.set_zlabel('Pressure (Pa)')
  ax.invert_zaxis()

  time_step += direction
  azim += 1
  if time_step == 9:
    time_step = 7
    direction = -1
    episode_num += 1
  elif time_step == -1:
    time_step = 1
    direction = 1
    episode_num += 1
    rng = jax.random.PRNGKey(episode_num)
    wind_field.reset_forecast(rng, dt.datetime.now())

  return ax
  
anim = animation.FuncAnimation(fig, frame, frames=40, blit=False, repeat=True)<jupyter_output><empty_output><jupyter_text>Run one of the two cells below. To avoid large notebooks we are not embedding the generated videos, but you can find them in our [generated directory](https://github.com/google/balloon-learning-environment/tree/master/balloon_learning_environment/generated).<jupyter_code># @title Display the animation in this notebook.
anim
# @title Save the animation (you can access it from the left panel).
filename = 'wind_field.mp4'  # @param {type: 'string'}
anim.save(filename)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.0956188039654621, 'fraction_numerical': 0.021426287176207228, 'mean_word_length': 4.890772128060264, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '575754', 'license': 'Apache-2.0', 'n_tokens_mistral': 1053, 'n_tokens_neox': 989, 'n_words': 329, 'path': 'balloon_learning_environment/colab/BLE_Generative_Wind_Field.ipynb', 'repo_name': 'johannah/balloon-learning-environment', 'size': '3127'}