Airdrop delivery with A* pathfinding

[This article was first published on Laurens Geffert, and kindly contributed to R-bloggers]. (You can report issue about the content on this page here)
Want to share your content on R-bloggers? click here if you have a blog, or here if you don't.

This post is an event report and a quick walk through to a submission that I
developed with a group of participants at an Alibaba / Met Office UK hackathon.
We are using the A* algorithm with a couple of tweaks to route cargo balloons
from London to a number of cities in the UK.

It’s the year 2050. The invention of anti-gravity engines has led to the
creation of unmanned balloons that travel the UK, delivering goods. However,
unpredictable weather conditions mean that these balloons are often delayed,
damaged or even destroyed(!), so we need your help. We’re inviting you to join
our contest AND our hackathon to create an algorithm which allows these balloons
to get to their route safely and effectively.

In January of this year, the Met Office teamed
up with Alibaba Cloud in organising a hackathon
at Huckletree Shoreditch.
Here is a short video that gives a good impression of the event

The Problem

The hackathon was organised as a “Future Challenge”, a fictitious scenario for
the year 2050. Obviously, drone delivery would be considered so 2020 and
completely outdated by then. Instead, people are relying on anti-gravity
balloons to deliver goods to cities across the UK via air drop.
These anti-gravity balloons are reliable and efficient, but have one major
shortcoming: They crash when travelling in areas with high wind speeds.
This is where the Met Office comes in. The task was to navigate balloons from
origin to destination while avoiding storms by using the Met Office forecasts.

I’ll spare you the full run through of rules, terms, and conditions. You can
find them on the competition page. The main facts:

  • forecast data provides projected wind speed for fields of a grid
  • forecasts are for hourly intervals
  • balloons can move up, down, left, right, or stay in place
  • balloons move one field per 2 minutes
  • balloons crash when entering a field when the wind speed is ≥15

So the big question is: How do we safely get the balloons from origin to
destination while avoiding stormy areas? I teamed up with a nice bunch of
people, mostly undergrad or master level university students. It was great to
see their enthusiasm for working on this problem!

The Data

The data that was provided to us was:

  1. the coordinates of cities (origin and destinations)
  2. weather data, separated into:
        * a training set with 7 days of weather forecasts from
    10 models plus observations of the actual conditions that manifested on these days
        * a holdout set with 5 days of weather forecasts

You should still be able to get the data here
in case you’d like to have a go yourself. Note that the weather forecasts came
in at a rather inconvenient file size of 2 x 800 MB, and download speeds were
not that great either.

See the map below for illustration purposes. The map shows gridded forecasts
of wind speed for a one hour time slice, as well as city locations
(origin in yellow, destinations in red).
Weather Forecast

Weather Prediction

We started by looking at the weather predictions. Our initial plan was to use
the first week, for which both forecasts and observations were available, to
train a classifier that would identify the likelihood of high wind speeds in a
given area at a particular time. This, however, turned out to be a bit of a red
herring. The Met Office predictions were so good that averaging them and
using a simple threshold of 15 resulted in close to zero false negatives when
trying to detect cells with storms.

Lesson learned: Do your EDA properly to check which areas are worth
investigating in detail and where you can use a simple ad-hoc solution.

<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="n">np</span>
<span class="kn">import</span> <span class="nn">pandas</span> <span class="k">as</span> <span class="n">pd</span>
<span class="kn">import</span> <span class="nn">matplotlib.pyplot</span> <span class="k">as</span> <span class="n">plt</span>
<span class="kn">import</span> <span class="nn">h5py</span>

<span class="k">def</span> <span class="nf">convert_forecast</span><span class="p">(</span><span class="n">data_cube</span><span class="p">):</span>
    <span class="c"># take mean of forecasts</span>
    <span class="n">arr_world</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="nb">min</span><span class="p">(</span><span class="n">data_cube</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">2</span><span class="p">)</span>
    <span class="c"># binarize to storm (1) or safe (0)</span>
    <span class="n">arr_world</span> <span class="o">=</span> <span class="n">arr_world</span> <span class="o">>=</span> <span class="mi">15</span>
    <span class="c"># from boolean to binary</span>
    <span class="n">arr_world</span> <span class="o">=</span> <span class="n">arr_world</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="nb">int</span><span class="p">)</span>
    <span class="c"># swap axes so x=0, y=1, z=2, day=3</span>
    <span class="n">arr_world</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">swapaxes</span><span class="p">(</span><span class="n">arr_world</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="mi">2</span><span class="p">)</span>
    <span class="n">arr_world</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">swapaxes</span><span class="p">(</span><span class="n">arr_world</span><span class="p">,</span><span class="mi">1</span><span class="p">,</span><span class="mi">3</span><span class="p">)</span>
    <span class="n">arr_world</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">swapaxes</span><span class="p">(</span><span class="n">arr_world</span><span class="p">,</span><span class="mi">2</span><span class="p">,</span><span class="mi">3</span><span class="p">)</span>

    <span class="k">return</span><span class="p">(</span><span class="n">arr_world</span><span class="p">)</span>

<span class="n">data_cube</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">load</span><span class="p">(</span><span class="s">'../data/5D_test.npy'</span><span class="p">)</span>
<span class="c"># convert forecast to world array</span>
<span class="n">forecast</span> <span class="o">=</span> <span class="n">convert_forecast</span><span class="p">(</span><span class="n">data_cube</span><span class="p">)</span>

Balloon Navigation

We wanted balloons to take the shortest path from origin to destination without
passing into storms. That means storms can be viewed as obstacles in our path
search problem, because we would never, ever, want to pass through them, even
if that means a massive detour for the balloon. We therefore chose to use an
A* path search algorithm. This algorithm finds the shortest path around
obstacles in a reasonable amount of time and is quite straight forward to
implement.

The basic approach is to start from origin and generate a frontier possible next
moves from a list of valid neighbouring fields. Fields with obstacles are
excluded, as well as fields that have been visited before. For each field that
is part of this frontier we log which field we came from and calculate a
heuristic cost of our movement to far. As soon as the destination field becomes
part of the frontier, we can recursively follow the trail laid out in our log,
back to the origin, to find the shortest path.

In case you would like more details or want to compare A* to other pathfinding
algorithms, this page
has the best summary of pathfinding algorithms I have seen, including a great
section on A* with interactive simulations.

In our case, there is one additional complication: our search is two-dimensional
(geographical space with latitude and longitude), but our “obstacles” change
over time. My way around that was to make the search three-dimensional, with
each movement time step (2 minutes) as a slice in the third dimension.

<span class="c"># repeat time slices x30</span>
<span class="n">forecast_stack</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="n">forecast</span><span class="p">,</span> <span class="n">repeats</span><span class="o">=</span><span class="mi">30</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">2</span><span class="p">)</span>

I then forced the search algorithm
to always take a step forward in time by restricting the valid neighbours to
[..., ..., z+1]. I tried to illustrate this schematically in the diagramm below:
3D A* schematic

The code for my A* implementation in python using heapq can be found below.

Note that I allow for neighbours that have the same x and y
coordinates. This essentially allows balloons to “hover in place” to wait out
unfavourable weather conditions in the area ahead, should that be the most
promising course of action.

Another thing to mention is that this approach massively inflates our frontier.
Usually, an advantage of the A* algorithm is that fields that have been visited
before do not need to be considered again. In my approach, field [0,0,0] is
different from field [0,0,1] (the same latitude and longitude, but a different
time step). As a result, computation becomes a lot more resource intense, but it
is still feasible to run this on your local machine and in the fast-paced
setting of a hackathon I think that prioritising developer time over computing
time was the right call.

<span class="kn">import</span> <span class="nn">heapq</span>

<span class="k">def</span> <span class="nf">heuristic_function</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">):</span>
    <span class="k">return</span> <span class="p">(</span><span class="n">b</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">a</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">**</span> <span class="mi">2</span> <span class="o">+</span> <span class="p">(</span><span class="n">b</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="n">a</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span> <span class="o">**</span> <span class="mi">2</span>

<span class="k">def</span> <span class="nf">astar_3D</span><span class="p">(</span><span class="n">space</span><span class="p">,</span> <span class="n">origin_xy</span><span class="p">,</span> <span class="n">destination_xy</span><span class="p">):</span>
    <span class="c"># make origin 3D with timeslice 0</span>
    <span class="n">origin</span> <span class="o">=</span> <span class="n">origin_xy</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">origin_xy</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="mi">0</span>
    <span class="c"># logs the path</span>
    <span class="n">came_from</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="c"># holds the legal next moves in order of priority</span>
    <span class="n">frontier</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="c"># define legal moves:</span>
    <span class="c"># up, down, left, right, stay in place.</span>
    <span class="c"># no diagonals and always move forward one time step (z)</span>
    <span class="n">neighbours</span> <span class="o">=</span> <span class="p">[(</span><span class="mi">0</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="mi">1</span><span class="p">),(</span><span class="mi">0</span><span class="p">,</span><span class="mi">1</span><span class="p">,</span><span class="mi">1</span><span class="p">),(</span><span class="mi">0</span><span class="p">,</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span><span class="mi">1</span><span class="p">),(</span><span class="mi">1</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="mi">1</span><span class="p">),(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span><span class="mi">0</span><span class="p">,</span><span class="mi">1</span><span class="p">)]</span>

    <span class="n">cost_so_far</span> <span class="o">=</span> <span class="p">{</span><span class="n">origin</span><span class="p">:</span> <span class="mi">0</span><span class="p">}</span>
    <span class="n">priority</span> <span class="o">=</span> <span class="p">{</span><span class="n">origin</span><span class="p">:</span> <span class="n">heuristic_function</span><span class="p">(</span><span class="n">origin_xy</span><span class="p">,</span> <span class="n">destination_xy</span><span class="p">)}</span>
    <span class="n">heapq</span><span class="o">.</span><span class="n">heappush</span><span class="p">(</span><span class="n">frontier</span><span class="p">,</span> <span class="p">(</span><span class="n">priority</span><span class="p">[</span><span class="n">origin</span><span class="p">],</span> <span class="n">origin</span><span class="p">))</span>

    <span class="c"># While there is still options to explore</span>
    <span class="k">while</span> <span class="n">frontier</span><span class="p">:</span>

        <span class="n">current</span> <span class="o">=</span> <span class="n">heapq</span><span class="o">.</span><span class="n">heappop</span><span class="p">(</span><span class="n">frontier</span><span class="p">)[</span><span class="mi">1</span><span class="p">]</span>

        <span class="c"># if current position is destination,</span>
        <span class="c"># break the loop and find the path that lead here</span>
        <span class="k">if</span> <span class="p">(</span><span class="n">current</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">current</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span> <span class="o">==</span> <span class="n">destination_xy</span><span class="p">:</span>
            <span class="n">data</span> <span class="o">=</span> <span class="p">[]</span>
            <span class="k">while</span> <span class="n">current</span> <span class="ow">in</span> <span class="n">came_from</span><span class="p">:</span>
                <span class="n">data</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">current</span><span class="p">)</span>
                <span class="n">current</span> <span class="o">=</span> <span class="n">came_from</span><span class="p">[</span><span class="n">current</span><span class="p">]</span>
            <span class="k">return</span> <span class="n">data</span>

        <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">neighbours</span><span class="p">:</span>
            <span class="n">move</span> <span class="o">=</span> <span class="n">current</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="n">i</span><span class="p">,</span> <span class="n">current</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="n">j</span><span class="p">,</span> <span class="n">current</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">+</span> <span class="n">k</span>

            <span class="c"># check that move is legal</span>
            <span class="k">if</span> <span class="p">((</span><span class="mi">0</span> <span class="o"><=</span> <span class="n">move</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o"><</span> <span class="n">space</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">&</span>
                <span class="p">(</span><span class="mi">0</span> <span class="o"><=</span> <span class="n">move</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o"><</span> <span class="n">space</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span> <span class="o">&</span>
                <span class="p">(</span><span class="mi">0</span> <span class="o"><=</span> <span class="n">move</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o"><</span> <span class="n">space</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">])):</span>

                <span class="k">if</span> <span class="n">space</span><span class="p">[</span><span class="n">move</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">move</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">move</span><span class="p">[</span><span class="mi">2</span><span class="p">]]</span> <span class="o">!=</span> <span class="mi">1</span><span class="p">:</span>

                    <span class="n">new_cost</span> <span class="o">=</span> <span class="mi">1</span>
                    <span class="n">new_total</span> <span class="o">=</span> <span class="n">cost_so_far</span><span class="p">[</span><span class="n">current</span><span class="p">]</span> <span class="o">+</span> <span class="n">new_cost</span>

                    <span class="k">if</span> <span class="n">move</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">cost_so_far</span><span class="p">:</span>
                        <span class="n">cost_so_far</span><span class="p">[</span><span class="n">move</span><span class="p">]</span> <span class="o">=</span> <span class="n">new_total</span>
                        <span class="c"># calculate total cost</span>
                        <span class="n">priority</span><span class="p">[</span><span class="n">move</span><span class="p">]</span> <span class="o">=</span> <span class="n">new_total</span> <span class="o">+</span> <span class="n">heuristic_function</span><span class="p">(</span><span class="n">move</span><span class="p">,</span> <span class="n">destination_xy</span><span class="p">)</span>
                        <span class="c"># update frontier</span>
                        <span class="n">heapq</span><span class="o">.</span><span class="n">heappush</span><span class="p">(</span><span class="n">frontier</span><span class="p">,</span> <span class="p">(</span><span class="n">priority</span><span class="p">[</span><span class="n">move</span><span class="p">],</span> <span class="n">move</span><span class="p">))</span>
                        <span class="c"># log this move</span>
                        <span class="n">came_from</span><span class="p">[</span><span class="n">move</span><span class="p">]</span> <span class="o">=</span> <span class="n">current</span>

    <span class="k">return</span> <span class="s">'no solution found :('</span>

<span class="c"># get city data</span>
<span class="n">cities</span> <span class="o">=</span> <span class="n">pd</span><span class="o">.</span><span class="n">read_csv</span><span class="p">(</span><span class="s">'../data/CityData.csv'</span><span class="p">,</span> <span class="n">header</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
<span class="c"># run algorithm</span>
<span class="n">x</span> <span class="o">=</span> <span class="n">astar_3D</span><span class="p">(</span><span class="n">space</span><span class="o">=</span><span class="n">arr_world_big</span><span class="p">[:,:,:,</span><span class="mi">0</span><span class="p">],</span>
             <span class="n">origin_xy</span><span class="o">=</span><span class="n">origin</span><span class="p">,</span>
             <span class="n">destination_xy</span><span class="o">=</span><span class="n">destination</span><span class="p">)</span>

Visualising the Results

We have a predicted optimal route now. That’s great, but it would be even
better to visualise these results in a way that allows us to develop some
intuition about how our solution is doing and where we could improve it further.
I thought that an animation of the time slices with the paths generated would
be ideal for this. So I used matplotlib.pyplot to create an image of each
time slice and then combined them into an animated gif. Output and code below:

Conditions and routes for day 11

You can see that, for this day, the solution for most cities is relatively
straight-forward because of low wind speeds in the majority of the area.
However, the A* pathfinding algorithm can be seen nicely at work in the later
timeslices and the centre-right of the map, where the purple balloon pauses
for a timeslice to wait out unfavourable conditions ahead and then winds around
patches of high wind speed towards its target.

<span class="k">def</span> <span class="nf">plot_solution</span><span class="p">(</span><span class="n">world</span><span class="p">,</span> <span class="n">cities</span><span class="p">,</span> <span class="n">solution</span><span class="p">,</span> <span class="n">day</span><span class="p">):</span>
    <span class="n">timesteps</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">540</span><span class="p">,</span> <span class="mi">30</span><span class="p">))</span>
    <span class="n">solution</span> <span class="o">=</span> <span class="n">solution</span><span class="o">.</span><span class="n">loc</span><span class="p">[</span><span class="n">solution</span><span class="o">.</span><span class="n">day</span> <span class="o">==</span> <span class="n">day</span><span class="p">,:]</span>
    <span class="c"># colour map for cities</span>
    <span class="n">cmap</span> <span class="o">=</span> <span class="n">plt</span><span class="o">.</span><span class="n">cm</span><span class="o">.</span><span class="n">cool</span>
    <span class="n">norm</span> <span class="o">=</span> <span class="n">matplotlib</span><span class="o">.</span><span class="n">colors</span><span class="o">.</span><span class="n">Normalize</span><span class="p">(</span><span class="n">vmin</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">vmax</span><span class="o">=</span><span class="mi">10</span><span class="p">)</span>
    <span class="c"># colour map for weather</span>
    <span class="n">cm</span> <span class="o">=</span> <span class="n">matplotlib</span><span class="o">.</span><span class="n">colors</span><span class="o">.</span><span class="n">LinearSegmentedColormap</span><span class="o">.</span><span class="n">from_list</span><span class="p">(</span><span class="s">'grid'</span><span class="p">,</span> <span class="p">[(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="p">(</span><span class="mf">0.5</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">)],</span> <span class="n">N</span><span class="o">=</span><span class="mi">2</span><span class="p">)</span>

    <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="n">timesteps</span><span class="p">:</span>
        <span class="n">timeslice</span> <span class="o">=</span> <span class="n">world</span><span class="p">[:,:,</span><span class="n">t</span><span class="p">]</span>
        <span class="n">moves_sofar</span> <span class="o">=</span> <span class="n">solution</span><span class="o">.</span><span class="n">loc</span><span class="p">[</span><span class="n">solution</span><span class="o">.</span><span class="n">z</span> <span class="o"><=</span> <span class="n">t</span><span class="p">,:]</span>
        <span class="n">moves_new</span> <span class="o">=</span> <span class="n">solution</span><span class="o">.</span><span class="n">loc</span><span class="p">[(</span><span class="n">t</span> <span class="o"><=</span> <span class="n">solution</span><span class="o">.</span><span class="n">z</span><span class="p">)</span> <span class="o">&</span> <span class="p">(</span><span class="n">solution</span><span class="o">.</span><span class="n">z</span> <span class="o"><=</span> <span class="n">t</span> <span class="o">+</span> <span class="mi">30</span><span class="p">),:]</span>

        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">solution_subset</span><span class="p">)</span> <span class="o">></span> <span class="mi">0</span><span class="p">:</span>    
            <span class="n">plt</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mi">5</span><span class="p">,</span><span class="mi">5</span><span class="p">))</span>
            <span class="n">plt</span><span class="o">.</span><span class="n">imshow</span><span class="p">(</span><span class="n">timeslice</span><span class="p">[:,:]</span><span class="o">.</span><span class="n">T</span><span class="p">,</span> <span class="n">aspect</span><span class="o">=</span><span class="s">'equal'</span><span class="p">,</span> <span class="n">cmap</span> <span class="o">=</span> <span class="n">cm</span><span class="p">)</span>

            <span class="c"># plot old moves</span>
            <span class="k">for</span> <span class="n">city</span> <span class="ow">in</span> <span class="n">moves_sofar</span><span class="o">.</span><span class="n">city</span><span class="o">.</span><span class="n">unique</span><span class="p">():</span>
                <span class="n">moves_sofar_city</span> <span class="o">=</span> <span class="n">moves_sofar</span><span class="o">.</span><span class="n">loc</span><span class="p">[</span><span class="n">moves_sofar</span><span class="o">.</span><span class="n">city</span> <span class="o">==</span> <span class="n">city</span><span class="p">,:]</span>
                <span class="n">x</span> <span class="o">=</span> <span class="n">moves_sofar_city</span><span class="o">.</span><span class="n">x</span>
                <span class="n">y</span> <span class="o">=</span> <span class="n">moves_sofar_city</span><span class="o">.</span><span class="n">y</span>
                <span class="n">z</span> <span class="o">=</span> <span class="n">moves_sofar_city</span><span class="o">.</span><span class="n">z</span>
                <span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="nb">list</span><span class="p">(</span><span class="n">y</span><span class="p">),</span> <span class="n">linestyle</span><span class="o">=</span><span class="s">'-'</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s">'black'</span><span class="p">)</span>

            <span class="c"># plot new moves</span>
            <span class="k">for</span> <span class="n">city</span> <span class="ow">in</span> <span class="n">moves_new</span><span class="o">.</span><span class="n">city</span><span class="o">.</span><span class="n">unique</span><span class="p">():</span>
                <span class="n">moves_new_city</span> <span class="o">=</span> <span class="n">moves_new</span><span class="o">.</span><span class="n">loc</span><span class="p">[</span><span class="n">moves_new</span><span class="o">.</span><span class="n">city</span> <span class="o">==</span> <span class="n">city</span><span class="p">,:]</span>
                <span class="n">x</span> <span class="o">=</span> <span class="n">moves_new_city</span><span class="o">.</span><span class="n">x</span>
                <span class="n">y</span> <span class="o">=</span> <span class="n">moves_new_city</span><span class="o">.</span><span class="n">y</span>
                <span class="n">z</span> <span class="o">=</span> <span class="n">moves_new_city</span><span class="o">.</span><span class="n">z</span>
                <span class="n">plt</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="nb">list</span><span class="p">(</span><span class="n">y</span><span class="p">),</span> <span class="n">linestyle</span><span class="o">=</span><span class="s">'-'</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="n">cmap</span><span class="p">(</span><span class="n">norm</span><span class="p">(</span><span class="n">city</span><span class="p">)))</span>

            <span class="c"># plot cities</span>
            <span class="k">for</span> <span class="n">city</span><span class="p">,</span><span class="n">x</span><span class="p">,</span><span class="n">y</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">cities</span><span class="o">.</span><span class="n">cid</span><span class="p">,</span> <span class="n">cities</span><span class="o">.</span><span class="n">xid</span><span class="p">,</span> <span class="n">cities</span><span class="o">.</span><span class="n">yid</span><span class="p">):</span>
                <span class="k">if</span> <span class="n">city</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">plt</span><span class="o">.</span><span class="n">scatter</span><span class="p">([</span><span class="n">x</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="n">y</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="n">c</span><span class="o">=</span><span class="s">'black'</span><span class="p">)</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="c"># balloon still en-route?</span>
                    <span class="k">if</span> <span class="n">city</span> <span class="ow">in</span> <span class="n">moves_new</span><span class="o">.</span><span class="n">city</span><span class="o">.</span><span class="n">unique</span><span class="p">():</span>
                        <span class="n">plt</span><span class="o">.</span><span class="n">scatter</span><span class="p">([</span><span class="n">x</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="n">y</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="n">c</span><span class="o">=</span><span class="n">cmap</span><span class="p">(</span><span class="n">norm</span><span class="p">(</span><span class="n">city</span><span class="p">)))</span>
                    <span class="k">else</span><span class="p">:</span>
                        <span class="n">plt</span><span class="o">.</span><span class="n">scatter</span><span class="p">([</span><span class="n">x</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="n">y</span><span class="o">-</span><span class="mi">1</span><span class="p">],</span> <span class="n">c</span><span class="o">=</span><span class="s">'black'</span><span class="p">)</span>

            <span class="c"># save and display</span>
            <span class="n">plt</span><span class="o">.</span><span class="n">savefig</span><span class="p">(</span><span class="s">'img_day'</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">day</span><span class="p">)</span> <span class="o">+</span> <span class="s">'_timestep_'</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">t</span><span class="p">)</span> <span class="o">+</span> <span class="s">'.png'</span><span class="p">)</span>
            <span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>

To leave a comment for the author, please follow the link and comment on their blog: Laurens Geffert.

R-bloggers.com offers daily e-mail updates about R news and tutorials about learning R and many other topics. Click here if you're looking to post or find an R/data-science job.
Want to share your content on R-bloggers? click here if you have a blog, or here if you don't.

Never miss an update!
Subscribe to R-bloggers to receive
e-mails with the latest R posts.
(You will not see this message again.)

Click here to close (This popup will not appear again)