text: import numpy as np
  from matplotlib import pyplot as plt
  from sklearn.preprocessing import PolynomialFeatures
  from ngboost.evaluation import *
  from ngboost.ngboost import NGBoost
  from ngboost.learners import default_linear_learner, default_tree_learner
  from ngboost.distns import Normal
  from ngboost.scores import MLE, CRPS
  from sklearn.neighbors import KNeighborsRegressor as KNR
  from argparse import ArgumentParser
  
  np.random.seed(1)
  
  default_knr_learner = lambda: KNR()
  
  
  def gen_data(n=50, bound=1, deg=3, beta=1, noise=0.9, intcpt=-1):
      x = np.linspace(-bound, bound, n)[:, np.newaxis]
      h = np.linspace(-bound, bound, n)[:, np.newaxis]
      e = np.random.randn(*x.shape) * (0.1 + 10 * np.abs(x))
      y = 50 * (x ** deg) + h * beta + noise * e + intcpt
      return x, y.squeeze(), np.c_[h, np.ones_like(h)]
  
  
  BLK = 2
  if __name__ == "__main__":
  
      argparser = ArgumentParser()
      argparser.add_argument("--n-estimators", type=int, default=(1 + BLK * 100))
      argparser.add_argument("--lr", type=float, default=0.03)
      argparser.add_argument("--minibatch-frac", type=float, default=0.1)
      argparser.add_argument("--natural", action="store_true")
      args = argparser.parse_args()
  
      x_tr, y_tr, _ = gen_data(n=100)
  
      poly_transform = PolynomialFeatures(1)
      x_tr = poly_transform.fit_transform(x_tr)
  
      ngb = NGBoost(
          Base=default_tree_learner,
          Dist=Normal,
          Score=MLE(),
          n_estimators=args.n_estimators,
          learning_rate=args.lr,
          natural_gradient=args.natural,
          minibatch_frac=args.minibatch_frac,
          verbose=True,
      )
  
      blk = int(args.n_estimators / 100)
      ngb.fit(x_tr, y_tr)
  
      x_te, y_te, _ = gen_data(n=1000, bound=1.3)
      x_te = poly_transform.transform(x_te)
      preds = ngb.pred_dist(x_te)
  
      pctles, obs, _, _ = calibration_regression(preds, y_te)
  
      filenames = []
      all_preds = ngb.staged_pred_dist(x_te)
      for i, preds in enumerate(all_preds):
          if i % blk != 0:
              continue
          plt.figure(figsize=(6, 3))
          plt.scatter(x_tr[:, 1], y_tr, color="black", marker=".", alpha=0.5)
          plt.plot(x_te[:, 1], preds.loc, color="black", linestyle="-", linewidth=1)
          plt.plot(
              x_te[:, 1],
              preds.loc - 1.96 * preds.scale,
              color="black",
              linestyle="--",
              linewidth=0.3,
          )
          plt.plot(
              x_te[:, 1],
              preds.loc + 1.96 * preds.scale,
              color="black",
              linestyle="--",
              linewidth=0.3,
          )
          plt.ylim([-75, 75])
          plt.axis("off")
          plt.title(
              "%s Gradient: %02d%% fit"
              % ("Natural" if args.natural else "Ordinary", i / blk)
          )
          plt.tight_layout()
  
          filenames.append("./figures/anim/toy%d.png" % i)
          print("Saving ./figures/anim/toy%d" % i)
          plt.savefig("./figures/anim/toy%d.png" % i)
          if i % 100 == 0:
              plt.savefig("./figures/anim/toy%d.pdf" % i)
  
      import imageio
  
      images = []
      for filename in filenames:
          images.append(imageio.imread(filename))
      for _ in range(10):
          images.append(imageio.imread(filename))
      imageio.mimsave("./figures/toy.gif", images, fps=5)
  