code: def main():
      # Dataset.
      ts_, ts_ext_, ts_vis_, ts, ts_ext, ts_vis, ys, ys_ = make_data()
  
      # Plotting parameters.
      vis_batch_size = 1024
      ylims = (-1.75, 1.75)
      alphas = [
          0.05,
          0.10,
          0.15,
          0.20,
          0.25,
          0.30,
          0.35,
          0.40,
          0.45,
          0.50,
          0.55,
      ]
      percentiles = [
          0.999,
          0.99,
          0.9,
          0.8,
          0.7,
          0.6,
          0.5,
          0.4,
          0.3,
          0.2,
          0.1,
      ]
      vis_idx = np.random.permutation(vis_batch_size)
      # From https://colorbrewer2.org/.
      if args.color == "blue":
          sample_colors = ("#8c96c6", "#8c6bb1", "#810f7c")
          fill_color = "#9ebcda"
          mean_color = "#4d004b"
          num_samples = len(sample_colors)
      else:
          sample_colors = ("#fc4e2a", "#e31a1c", "#bd0026")
          fill_color = "#fd8d3c"
          mean_color = "#800026"
          num_samples = len(sample_colors)
  
      eps = torch.randn(vis_batch_size, 1).to(
          device
      )  # Fix seed for the random draws used in the plots.
      bm = torchsde.BrownianInterval(
          t0=ts_vis[0],
          t1=ts_vis[-1],
          size=(vis_batch_size, 1),
          device=device,
          levy_area_approximation="space-time",
      )  # We need space-time Levy area to use the SRK solver
  
      # Model.
      model = LatentSDE().to(device)
      optimizer = optim.Adam(model.parameters(), lr=1e-2)
      scheduler = optim.lr_scheduler.ExponentialLR(
          optimizer, gamma=0.999
      )
      kl_scheduler = LinearScheduler(iters=args.kl_anneal_iters)
  
      logpy_metric = EMAMetric()
      kl_metric = EMAMetric()
      loss_metric = EMAMetric()
  
      if args.show_prior:
          with torch.no_grad():
              zs = model.sample_p(
                  ts=ts_vis, batch_size=vis_batch_size, eps=eps, bm=bm
              ).squeeze()
              ts_vis_, zs_ = ts_vis.cpu().numpy(), zs.cpu().numpy()
              zs_ = np.sort(zs_, axis=1)
  
              img_dir = os.path.join(args.train_dir, "prior.png")
              plt.subplot(frameon=False)
              for alpha, percentile in zip(alphas, percentiles):
                  idx = int((1 - percentile) / 2.0 * vis_batch_size)
                  zs_bot_ = zs_[:, idx]
                  zs_top_ = zs_[:, -idx]
                  plt.fill_between(
                      ts_vis_,
                      zs_bot_,
                      zs_top_,
                      alpha=alpha,
                      color=fill_color,
                  )
  
              # `zorder` determines who's on top; the larger the more at the top.
              plt.scatter(
                  ts_, ys_, marker="x", zorder=3, color="k", s=35
              )  # Data.
              plt.ylim(ylims)
              plt.xlabel("$t$")
              plt.ylabel("$Y_t$")
              plt.tight_layout()
              plt.savefig(img_dir, dpi=args.dpi)
              plt.close()
              logging.info(f"Saved prior figure at: {img_dir}")
  
      for global_step in tqdm.tqdm(range(args.train_iters)):
          # Plot and save.
          if global_step % args.pause_iters == 0:
              img_path = os.path.join(
                  args.train_dir, f"global_step_{global_step}.png"
              )
  
              with torch.no_grad():
                  zs = model.sample_q(
                      ts=ts_vis,
                      batch_size=vis_batch_size,
                      eps=eps,
                      bm=bm,
                  ).squeeze()
                  samples = zs[:, vis_idx]
                  ts_vis_, zs_, samples_ = (
                      ts_vis.cpu().numpy(),
                      zs.cpu().numpy(),
                      samples.cpu().numpy(),
                  )
                  zs_ = np.sort(zs_, axis=1)
                  plt.subplot(frameon=False)
  
                  if args.show_percentiles:
                      for alpha, percentile in zip(alphas, percentiles):
                          idx = int(
                              (1 - percentile) / 2.0 * vis_batch_size
                          )
                          zs_bot_, zs_top_ = zs_[:, idx], zs_[:, -idx]
                          plt.fill_between(
                              ts_vis_,
                              zs_bot_,
                              zs_top_,
                              alpha=alpha,
                              color=fill_color,
                          )
  
                  if args.show_mean:
                      plt.plot(
                          ts_vis_, zs_.mean(axis=1), color=mean_color
                      )
  
                  if args.show_samples:
                      for j in range(num_samples):
                          plt.plot(
                              ts_vis_,
                              samples_[:, j],
                              color=sample_colors[j],
                              linewidth=1.0,
                          )
  
                  if args.show_arrows:
                      num, dt = 12, 0.12
                      t, y = torch.meshgrid(
                          [
                              torch.linspace(0.2, 1.8, num).to(device),
                              torch.linspace(-1.5, 1.5, num).to(device),
                          ]
                      )
                      t, y = t.reshape(-1, 1), y.reshape(-1, 1)
                      fty = model.f(t=t, y=y).reshape(num, num)
                      dt = torch.zeros(num, num).fill_(dt).to(device)
                      dy = fty * dt
                      dt_, dy_, t_, y_ = (
                          dt.cpu().numpy(),
                          dy.cpu().numpy(),
                          t.cpu().numpy(),
                          y.cpu().numpy(),
                      )
                      plt.quiver(
                          t_,
                          y_,
                          dt_,
                          dy_,
                          alpha=0.3,
                          edgecolors="k",
                          width=0.0035,
                          scale=50,
                      )
  
                  if args.hide_ticks:
                      plt.xticks([], [])
                      plt.yticks([], [])
  
                  plt.scatter(
                      ts_, ys_, marker="x", zorder=3, color="k", s=35
                  )  # Data.
                  plt.ylim(ylims)
                  plt.xlabel("$t$")
                  plt.ylabel("$Y_t$")
                  plt.tight_layout()
                  plt.savefig(img_path, dpi=args.dpi)
                  plt.close()
                  logging.info(f"Saved figure at: {img_path}")
  
                  if args.save_ckpt:
                      torch.save(
                          {
                              "model": model.state_dict(),
                              "optimizer": optimizer.state_dict(),
                              "scheduler": scheduler.state_dict(),
                              "kl_scheduler": kl_scheduler,
                          },
                          os.path.join(
                              ckpt_dir,
                              f"global_step_{global_step}.ckpt",
                          ),
                      )
  
          # Train.
          optimizer.zero_grad()
          zs, kl = model(ts=ts_ext, batch_size=args.batch_size)
          zs = zs.squeeze()
          zs = zs[
              1:-1
          ]  # Drop first and last which are only used to penalize out-of-data region and spread uncertainty.
  
          likelihood_constructor = {
              "laplace": distributions.Laplace,
              "normal": distributions.Normal,
          }[args.likelihood]
          likelihood = likelihood_constructor(loc=zs, scale=args.scale)
          logpy = likelihood.log_prob(ys).sum(dim=0).mean(dim=0)
  
          loss = -logpy + kl * kl_scheduler.val
          loss.backward()
  
          optimizer.step()
          scheduler.step()
          kl_scheduler.step()
  
          logpy_metric.step(logpy)
          kl_metric.step(kl)
          loss_metric.step(loss)
  
          logging.info(
              f"global_step: {global_step}, "
              f"logpy: {logpy_metric.val:.3f}, "
              f"kl: {kl_metric.val:.3f}, "
              f"loss: {loss_metric.val:.3f}"
          )
  
imports: import argparse
  import logging
  import math
  import os
  import random
  from collections import namedtuple
  from typing import Optional, Union
  import matplotlib.pyplot
  import numpy
  import torch
  import tqdm
  from torch import distributions, nn, optim
  import torchsde
label_desc: Use image node_id 7 for calling a global function with example usage: main() without return types