animate.R

# Parse command line arguments
library(argparser)
arg_parser <- arg_parser("Animate compiled neutral atom instructions")
arg_parser <- add_argument(arg_parser, "input",
  help = "CSV file generated by the compile.py script.",
  type = "character"
)
arg_parser <- add_argument(arg_parser, "output",
  help = "Output movie file [.mp4]",
  default = "animation.mp4",
  type="character"
)
arg_parser <- add_argument(arg_parser, "--layout",
  help = "Layout CSV file",
  type = "character"
)
arg_parser <- add_argument(arg_parser, "--ratio",
  help = "Milliseconds per step",
  type = "numeric",
  default = 0
)
args <- parse_args(arg_parser)
print(paste("Input file:", args$input))
print(paste("Output file:", args$output))
data <- read.csv(file=args$input, sep = ';')
layout <- read.csv(file=args$layout, sep = ';')
ratio <- args$ratio
# load other packages
library(ggplot2)
library(gganimate)
library(ggnewscale)
library(hrbrthemes)
library(sqldf)
# use a minimalistic theme, adjust the margin to the minimum
thm <- theme_ipsum(base_family = "CMU Sans Serif") +
  theme(plot.margin = margin(0, 8, 0, 4, unit = "pt"))
zone_colors <- c("#CCD6DD", "#7F98AA", "#345A78")
gate_colors <- c("#A2AD00", "#64A0C8", "#CCDCE8", "#80A7C6", "#E37222", "#cc0000")

# compute appropriate size of block radius and atom size
xlim = range(layout$x) + c(-1,1)
ylim = rev(range(layout$y) + c(-1,1))
blockSize =  720 / (ylim[1]-ylim[2]) /16 * 9 * max(data$marginSize)
atomSize = 720 / (ylim[1]-ylim[2]) /40
# add column if marginSize is not present and set to 1
if (!"marginSize" %in% colnames(data)) {
  data$marginSize <- 1
}
# add column if atomSize is not present
if (!"atomSize" %in% colnames(layout)) {
  layout$atomSize <- 3
}


anim <- ggplot() +
  geom_point(
    data = layout,
    mapping = aes(
      x = x,
      y = y,
      color = factor(color),
      size = atomSize,
      fill = factor(color)),
    shape = 21,
    stroke = 0.5) +
  scale_fill_manual(
    values = zone_colors,
    labels = c('Entangling', 'Storage', 'Readout'),
    name = "Zone") +
  guides(fill = guide_legend(override.aes = list(shape = 21, stroke = 0, size=3))) +
  scale_color_manual(
    values = zone_colors,
    guide = 'none') +
  new_scale_fill() +
  new_scale_color() +
  scale_size(
    guide = 'none',
    range = c(2,3) * atomSize
  ) +
  new_scale('size') +
  geom_point(
    data = sqldf('select time, x, y, id || "-" || marginId as gid, marginSize from data where margin = 1'),
    mapping = aes(
      x = x,
      y = y,
      size = marginSize,
      group = gid),
    color = '#E37222',
    alpha = .2) +
  scale_size(
    guide = 'none',
    range = c(blockSize , blockSize)) +
  new_scale("size") +
  geom_hline(
    data = sqldf('select time, y, min(id) || "-" || axesId as gid from data where axes = 1 group by time, y'),
    mapping = aes(
      yintercept = y,
      group = gid),
    linetype = 'dashed',
    color = '#333333',
    linewidth = .2) +
  geom_vline(
    data = sqldf('select time, x, min(id) || "-" || axesId as gid from data where axes = 1 group by time, x'),
    mapping = aes(
      xintercept = x,
      group = gid),
    linetype = 'dashed',
    color = '#333333',
    linewidth = .2) +
  geom_point(
    data = data,
    mapping = aes(
      x = x,
      y = y,
      color = factor(color),
      fill = factor(fill),
      size = size,
      group = id),
    shape = 21,
    stroke = 0.5) +
  geom_text(
    data = data,
    mapping = aes(
      x = x,
      y = y,
      label = id,
      group = id
    ),
    alpha = .7,
    size = 2
  ) +
  scale_fill_manual(
    values = gate_colors,
    breaks = c(0, 1, 2, 3, 4),
    labels = c('SLM', 'AOD', 'RY', 'RZ', 'CZ'),
    name = "Gates",
    guide = guide_legend(override.aes = list(shape = 21, stroke = 0, size = 2))) +
  scale_color_manual(
    values = gate_colors,
    breaks = c(0, 1, 2, 3, 4, 5),
    guide = 'none') +
  scale_size(
    guide = 'none',
    range = c(2, 3) * atomSize) +
  labs(caption = "Time: {format(round(frame_time, 1), nsmall = 1)} µs") +
  coord_fixed() +
  scale_y_reverse() +
  thm +
  transition_components(time, enter_length = 1, exit_length = 1) +
  enter_fade() +
  exit_fade()
animate(anim,
    duration = ceiling(ratio / 1000 * (max(data$time) - min(data$time))),
    fps = 30,
    height = 720,
    width = 1280,
    res = 100,
    start_pause = 20,
    renderer = av_renderer(args$output)
  )
print("Successfully created neutral atom animation.")
print(paste("Input file:", args$input))
print(paste("Output file:", args$output))