ethanfetaya / NRI

Neural relational inference for interacting systems - pytorch
MIT License
732 stars 157 forks source link

How to plot the beautiful trajectory like your figure1? #8

Closed Reeshark closed 5 years ago

Reeshark commented 5 years ago

Thanks for your amazing code. How to plot the beautiful trajectory like your figure1?

ethanfetaya commented 5 years ago

Thanks! The figure was done by Thomas, so I am not 100% sure, but I think he used omnigraffle.

tkipf commented 5 years ago

Thanks! It was just a very simple matplotlib script in python:

from matplotlib import cm
from numpy import linspace

start = 0.0
stop = 1.0
num_colors = 10
cm_subsection = linspace(start, stop, num_colors) 

colors = [ cm.Set1(x) for x in cm_subsection ]

for i in range(loc.shape[-1]):
    for t in range(loc.shape[0]):
        # Plot fading tail for past locations.
        plt.plot(loc[t, 0, i], loc[t, 1, i], 'o', markersize=3, 
                 color=colors[i], alpha=1-(float(t)/loc.shape[0]))
    # Plot final location.
    plt.plot(loc[0, 0, i], loc[0, 1, i], 'o', color=colors[I])

where loc is the num_timesteps x num_dims x num_objects matrix of particle trajectories. The particle trajectories are from the spring simulation.