text: import sys
  sys.path.insert(0, "../")
  
  from PINN_Survey.problems.burgers.v1 import Burgers, Burgers_Sphere_Mesh
  from PINN_Survey.problems.helmholtz.v1 import Helmholtz, Helmholtz_Sphere_Mesh
  from PINN_Survey.problems.burgers.data.load import load_burgers_bounds
  from PINN_Survey.problems.helmholtz.data.load import load_helmholtz_bounds
  from PINN_Survey.viz.loss_viz import viz_2d_layer_norm, viz_base, viz_mesh, save_as_heightmap
  from PINN_Base.util import bounds_from_data, random_choices, rel_error, percent_noise
  import numpy as np
  import tensorflow as tf
  
  
  @viz_base
  class Burgers_Viz(Burgers):
      pass
  
  
  @viz_base
  class Helmholtz_Viz(Helmholtz):
      pass
  
  
  @viz_mesh
  class Burgers_Sphere_Mesh_Viz(Burgers_Sphere_Mesh):
      pass
  
  
  @viz_mesh
  class Helmholtz_Sphere_Mesh_Viz(Helmholtz_Sphere_Mesh):
      pass
  
  
  MAX_THREADS = 32
  config = tf.ConfigProto(
      intra_op_parallelism_threads=MAX_THREADS
  )
  
  
  def quick_setup(loader, size=1000):
  
      X_true, U_true, _, _, _ = loader()
  
      [X, U] = random_choices(X_true, U_true, size=size)
  
      lower_bound, upper_bound = bounds_from_data(X_true)
  
      return X_true, U_true, X, U, lower_bound, upper_bound
  
  
  def burgers_regularization_base(data_noise, df_multiplier):
  
      X_true, U_true, X, U, lower_bound, upper_bound = quick_setup(
          load_burgers_bounds)
  
      if data_noise != 0:
          U = percent_noise(U, data_noise)
  
      layers = [2, 20, 20, 20, 20, 20, 20, 20, 20, 1]
      nu = 0.01 / np.pi
      model = Burgers(lower_bound, upper_bound, layers,
                      nu, use_differential_points=False, session_config=config, df_multiplier=df_multiplier)
      model.train_BFGS(X, U, print_loss=True)
      U_hat = model.predict(X_true)
      print(rel_error(U_true, U_hat))
  
      w0 = model.get_all_weights()
  
      model_viz = Burgers_Viz(lower_bound, upper_bound,
                              layers, nu, use_differential_points=False, session_config=config, df_multiplier=df_multiplier)
  
      _, _, [_, _, loss] = viz_2d_layer_norm(model_viz, X, U, None, w0, 512)
  
      save_as_heightmap(
          f"burgers-regularization-base-N({data_noise})-L({df_multiplier}).png", np.log(loss))
  
  
  def burgers_regularization_mesh(data_noise):
  
      X_true, U_true, X, U, lower_bound, upper_bound = quick_setup(
          load_burgers_bounds)
  
      if data_noise != 0:
          U = percent_noise(U, data_noise)
  
      layers_mesh = [2, 20, 20, 20, 20, 20, 20]
      layers_approx = [2, 20, 20, 1]
      nu = 0.01 / np.pi
      model = Burgers_Sphere_Mesh(lower_bound, upper_bound, layers_approx, layers_mesh,
                                  nu, use_differential_points=False, session_config=config)
      model.train_BFGS(X, U, print_loss=True)
      U_hat = model.predict(X_true)
      print(rel_error(U_true, U_hat))
  
      w0 = model.get_all_weights()
  
      model_viz = Burgers_Sphere_Mesh_Viz(lower_bound, upper_bound, layers_approx, layers_mesh,
                                          nu, use_differential_points=False, session_config=config)
  
      _, _, [_, _, loss] = viz_2d_layer_norm(model_viz, X, U, None, w0, 512)
  
      save_as_heightmap(
          f"burgers-regularization-mesh-{data_noise}.png", np.log(loss))
  
  
  def helmholtz_regularization_base(data_noise, df_multiplier):
  
      X_true, U_true, X, U, lower_bound, upper_bound = quick_setup(
          load_helmholtz_bounds)
  
      if data_noise != 0:
          U = percent_noise(U, data_noise)
  
      layers = [2, 20, 20, 20, 20, 20, 20, 20, 20, 1]
      model = Helmholtz(lower_bound, upper_bound, layers,
                        1, 4, use_differential_points=False, session_config=config, df_multiplier=df_multiplier)
      model.train_BFGS(X, U, print_loss=True)
      U_hat = model.predict(X_true)
      print(rel_error(U_true, U_hat))
  
      w0 = model.get_all_weights()
  
      model_viz = Helmholtz_Viz(lower_bound, upper_bound,
                                layers, 1, 4, use_differential_points=False, session_config=config, df_multiplier=df_multiplier)
  
      _, _, [_, _, loss] = viz_2d_layer_norm(model_viz, X, U, None, w0, 512)
  
      save_as_heightmap(
          f"helmholtz-regularization-base-N({data_noise})-L({df_multiplier}).png", np.log(loss))
  
  
  def helmholtz_regularization_mesh(data_noise):
  
      X_true, U_true, X, U, lower_bound, upper_bound = quick_setup(
          load_helmholtz_bounds)
  
      if data_noise != 0:
          U = percent_noise(U, data_noise)
  
      layers_mesh = [2, 20, 20, 20, 20, 20, 20]
      layers_approx = [2, 20, 20, 1]
      model = Helmholtz_Sphere_Mesh(lower_bound, upper_bound, layers_approx, layers_mesh,
                                    1, 4, use_differential_points=False, session_config=config)
      model.train_BFGS(X, U, print_loss=True)
      U_hat = model.predict(X_true)
      print(rel_error(U_true, U_hat))
  
      w0 = model.get_all_weights()
  
      model_viz = Helmholtz_Sphere_Mesh_Viz(lower_bound, upper_bound, layers_approx, layers_mesh,
                                            1, 4, use_differential_points=False, session_config=config)
  
      _, _, [_, _, loss] = viz_2d_layer_norm(model_viz, X, U, None, w0, 512)
  
      save_as_heightmap(
          f"helmholtz-regularization-mesh-{data_noise}.png", np.log(loss))
  
  
  if __name__ == "__main__":
      if int(sys.argv[1]) == 0:
          burgers_regularization_base(.1, 0)
          burgers_regularization_base(.1, 1e-5)
          burgers_regularization_base(.1, 1e-4)
          burgers_regularization_base(.1, 1e-3)
          burgers_regularization_base(.1, 1e-2)
          burgers_regularization_base(.1, 1e-1)
          burgers_regularization_base(.1, 1)
          burgers_regularization_base(.1, 10)
          burgers_regularization_base(.1, 100)
      elif int(sys.argv[1]) == 1:
          helmholtz_regularization_base(.1, 0)
          helmholtz_regularization_base(.1, 1e-5)
          helmholtz_regularization_base(.1, 1e-4)
          helmholtz_regularization_base(.1, 1e-3)
          helmholtz_regularization_base(.1, 1e-2)
          helmholtz_regularization_base(.1, 1e-1)
          helmholtz_regularization_base(.1, 1)
          helmholtz_regularization_base(.1, 10)
          helmholtz_regularization_base(.1, 100)
  