80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
|
import sys
|
||
|
|
import datashader as ds
|
||
|
|
import pandas as pd
|
||
|
|
import datashader.transfer_functions as tf
|
||
|
|
from datashader.utils import export_image
|
||
|
|
from datashader.transfer_functions import set_background
|
||
|
|
import subprocess as SP
|
||
|
|
import os
|
||
|
|
import itertools as ITT
|
||
|
|
from glob import glob
|
||
|
|
|
||
|
|
print("Loading stars...")
|
||
|
|
stars = pd.read_csv("stars.csv", usecols=["id", "x", "z", "mult"], index_col=0)
|
||
|
|
stars.loc[stars.mult == 1.0, "mult"] = float("nan")
|
||
|
|
|
||
|
|
steps = int(sys.argv[1])
|
||
|
|
size = 1080
|
||
|
|
mode = "eq_hist"
|
||
|
|
|
||
|
|
cvs = ds.Canvas(plot_width=size, plot_height=size)
|
||
|
|
|
||
|
|
print("Plotting density")
|
||
|
|
density_agg = cvs.points(stars, "x", "z")
|
||
|
|
density = tf.shade(density_agg, cmap=["black", "white"], how=mode)
|
||
|
|
|
||
|
|
print("Plotting neutrons")
|
||
|
|
neutrons_agg = cvs.points(stars, "x", "z", agg=ds.count("mult"))
|
||
|
|
neutrons = tf.shade(neutrons_agg, cmap=["darkblue", "lightblue"], how=mode)
|
||
|
|
|
||
|
|
base = tf.stack(density, neutrons)
|
||
|
|
|
||
|
|
# ffplay = SP.Popen([
|
||
|
|
# "ffplay","-f","image2pipe","-"
|
||
|
|
# ],stdin=SP.PIPE,bufsize=0)
|
||
|
|
|
||
|
|
|
||
|
|
for rh_fn in ITT.chain.from_iterable(map(glob, sys.argv[2:])):
|
||
|
|
basename = os.path.splitext(os.path.split(rh_fn)[-1])[0]
|
||
|
|
filename = "img/{}_{}_{}.mkv".format(basename, size, mode)
|
||
|
|
ffmpeg = SP.Popen(
|
||
|
|
[
|
||
|
|
"ffmpeg",
|
||
|
|
"-y",
|
||
|
|
"-f",
|
||
|
|
"image2pipe",
|
||
|
|
"-i",
|
||
|
|
"-",
|
||
|
|
"-crf",
|
||
|
|
"17",
|
||
|
|
"-r",
|
||
|
|
"25",
|
||
|
|
"-pix_fmt",
|
||
|
|
"yuv420p",
|
||
|
|
filename,
|
||
|
|
],
|
||
|
|
stdin=SP.PIPE,
|
||
|
|
bufsize=0,
|
||
|
|
)
|
||
|
|
|
||
|
|
print("Loading", rh_fn)
|
||
|
|
route_hist = pd.read_csv(
|
||
|
|
rh_fn, names=["id", "d"], index_col=0, dtype={"d": int}, low_memory=False,
|
||
|
|
)
|
||
|
|
exp_span = [route_hist.d.min(), route_hist.d.max()]
|
||
|
|
stars["d"] = float("nan")
|
||
|
|
rng = range(route_hist.d.min(), route_hist.d.max() + 1, steps)
|
||
|
|
if steps == 0:
|
||
|
|
rng = [route_hist.d.max() + 1]
|
||
|
|
for n in rng:
|
||
|
|
stars['d'] = route_hist[route_hist.d < n] # slow
|
||
|
|
explored_agg = cvs.points(stars, "x", "z", agg=ds.mean("d")) # slow
|
||
|
|
explored = tf.shade(
|
||
|
|
explored_agg, cmap=["darkred", "lightpink"], how="linear", span=exp_span
|
||
|
|
)
|
||
|
|
img = set_background(tf.stack(base, explored), "black").to_pil()
|
||
|
|
img.save(ffmpeg.stdin, "png")
|
||
|
|
img.save("img/current.png")
|
||
|
|
ffmpeg.stdin.close()
|
||
|
|
ffmpeg.wait()
|