code: def main(args):
      funsor.set_backend("torch")
  
      # download and pre-process EEG data if not in test mode
      if not args.test:
          download_data()
          N_val, N_test = 149, 200
          data = np.loadtxt("eeg.dat", delimiter=",", skiprows=19)
          print(f"[raw data shape] {data.shape}")
          data = data[::20, :]
          print(f"[data shape after thinning] {data.shape}")
          eye_state = [int(d) for d in data[:, -1].tolist()]
          data = torch.tensor(data[:, :-1]).float()
      # in test mode (for continuous integration on github) so create fake data
      else:
          data = torch.randn(10, 3)
          N_val, N_test = 2, 2
  
      T, obs_dim = data.shape
      N_train = T - N_test - N_val
  
      np.random.seed(0)
      rand_perm = np.random.permutation(N_val + N_test)
      val_indices = rand_perm[0:N_val]
      test_indices = rand_perm[N_val:]
  
      data_mean = data[0:N_train, :].mean(0)
      data -= data_mean
      data_std = data[0:N_train, :].std(0)
      data /= data_std
  
      print(
          f"Length of time series T: {T}   Observation dimension: {obs_dim}"
      )
      print(f"N_train: {N_train}  N_val: {N_val}  N_test: {N_test}")
  
      torch.manual_seed(args.seed)
  
      # set up model
      slds = SLDS(
          num_components=args.num_components,
          hidden_dim=args.hidden_dim,
          obs_dim=obs_dim,
          fine_observation_noise=args.fon,
          fine_transition_noise=args.ftn,
          fine_observation_matrix=args.fom,
          fine_transition_matrix=args.ftm,
          moment_matching_lag=args.moment_matching_lag,
      )
  
      # set up optimizer
      adam = torch.optim.Adam(
          slds.parameters(),
          lr=args.learning_rate,
          betas=(args.beta1, 0.999),
          amsgrad=True,
      )
      scheduler = torch.optim.lr_scheduler.ExponentialLR(
          adam, gamma=args.gamma
      )
      ts = [time.time()]
  
      report_frequency = 1
  
      # training loop
      for step in range(args.num_steps):
          nll = -slds.log_prob(data[0:N_train, :]) / N_train
          nll.backward()
  
          if step == 5:
              scheduler.base_lrs[0] *= 0.20
  
          adam.step()
          scheduler.step()
          adam.zero_grad()
  
          if step % report_frequency == 0 or step == args.num_steps - 1:
              step_dt = ts[-1] - ts[-2] if step > 0 else 0.0
              pred_mse, pred_LLs = slds.filter_and_predict(
                  data[0 : N_train + N_val + N_test, :]
              )
              val_mse = pred_mse[val_indices].mean().item()
              test_mse = pred_mse[test_indices].mean().item()
              val_ll = pred_LLs[val_indices].mean().item()
              test_ll = pred_LLs[test_indices].mean().item()
  
              stats = "[step %03d] train_nll: %.5f val_mse: %.5f val_ll: %.5f test_mse: %.5f test_ll: %.5f\t(dt: %.2f)"
              print(
                  stats
                  % (
                      step,
                      nll.item(),
                      val_mse,
                      val_ll,
                      test_mse,
                      test_ll,
                      step_dt,
                  )
              )
  
          ts.append(time.time())
  
      # plot predictions and smoothed means
      if args.plot:
          assert not args.test
          (
              predicted_mse,
              LLs,
              pred_means,
              pred_vars,
              smooth_means,
              smooth_probs,
          ) = slds.filter_and_predict(data, smoothing=True)
  
          pred_means = pred_means.data.numpy()
          pred_stds = pred_vars.sqrt().data.numpy()
          smooth_means = smooth_means.data.numpy()
          smooth_probs = smooth_probs.data.numpy()
  
          import matplotlib
  
          matplotlib.use("Agg")  # noqa: E402
          import matplotlib.pyplot as plt
  
          f, axes = plt.subplots(4, 1, figsize=(12, 8), sharex=True)
          T = data.size(0)
          N_valtest = N_val + N_test
          to_seconds = 117.0 / T
  
          for k, ax in enumerate(axes[:-1]):
              which = [0, 4, 10][k]
              ax.plot(
                  to_seconds * np.arange(T),
                  data[:, which],
                  "ko",
                  markersize=2,
              )
              ax.plot(
                  to_seconds * np.arange(N_train),
                  smooth_means[:N_train, which],
                  ls="solid",
                  color="r",
              )
  
              ax.plot(
                  to_seconds * (N_train + np.arange(N_valtest)),
                  pred_means[-N_valtest:, which],
                  ls="solid",
                  color="b",
              )
              ax.fill_between(
                  to_seconds * (N_train + np.arange(N_valtest)),
                  pred_means[-N_valtest:, which]
                  - 1.645 * pred_stds[-N_valtest:, which],
                  pred_means[-N_valtest:, which]
                  + 1.645 * pred_stds[-N_valtest:, which],
                  color="lightblue",
              )
              ax.set_ylabel(f"$y_{which + 1}$", fontsize=20)
              ax.tick_params(axis="both", which="major", labelsize=14)
  
          axes[-1].plot(
              to_seconds * np.arange(T), eye_state, "k", ls="solid"
          )
          axes[-1].plot(
              to_seconds * np.arange(T), smooth_probs, "r", ls="solid"
          )
          axes[-1].set_xlabel("Time (s)", fontsize=20)
          axes[-1].set_ylabel("Eye state", fontsize=20)
          axes[-1].tick_params(axis="both", which="major", labelsize=14)
  
          plt.tight_layout(pad=0.7)
          plt.savefig("eeg.pdf")
  
imports: import argparse
  import time
  from collections import OrderedDict
  from os.path import exists
  from urllib.request import urlopen
  import numpy
  import pyro
  import torch
  import torch.nn
  import funsor
  import funsor.ops
  import funsor.torch.distributions
  from funsor.pyro.convert import funsor_to_cat_and_mvn, funsor_to_mvn, matrix_and_mvn_to_funsor, mvn_to_funsor
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types