Validation and prediction at different time snaps

[RoboticAutonomy]

@rezaakb Hello, I was trying to predict at different time snaps, but I don't know what I should modify at the prediction function and loop also the plotting function. Here is my config file for NS:

pred_dataset:

• mesh_sampler:
  target: pinnstf2.data.MeshSampler
  partial: true
  idx_t: 50
  solution:
  - u
  - v
  - p

• mesh_sampler:
  target: pinnstf2.data.MeshSampler
  partial: true
  idx_t: 100
  solution:
  - u
  - v
  - p

• mesh_sampler:
  target: pinnstf2.data.MeshSampler
  partial: true
  idx_t: 150
  solution:
  - u
  - v
  - p
  def plot_navier_stokes(mesh, preds, train_datasets, val_dataset, file_name):
  """Plot Navier-Stokes continuous inverse PDE."""

  x, t, u = train_datasets[0][:]
  p_star = mesh.solution["p"][:, 100]
  p_pred = preds["p"].reshape(p_star.shape)
  X_star = np.hstack(mesh.spatial_domain)
  lb = X_star.min(0)
  ub = X_star.max(0)
  nn = 200
  x = np.linspace(lb[0], ub[0], nn)
  y = np.linspace(lb[1], ub[1], nn)
  X, Y = np.meshgrid(x, y)
  x_star = X_star[:, 0:1]
  y_star = X_star[:, 1:2]

  PP_star = griddata(X_star, p_pred.flatten(), (X, Y), method="cubic")
  P_exact = griddata(X_star, p_star.flatten(), (X, Y), method="cubic")