High-dimensional Sobol’ Function

We consider a map \(f: \mathbb{R}^{5}\to\mathbb{R}^{4}\) expressed as

\[f(\boldsymbol{z}) = \boldsymbol{A}\,\boldsymbol{g}(e^{\boldsymbol{z}}),\]

where \(g_i(\boldsymbol{r}) = (2\cdot |2\,a_{i} - 1| + r_i) / (1 + r_i)\) with \(r_i > 0\) for \(i=1,\dots,5\) is the Sobol function [Sobol03] and \(\boldsymbol{A}\) is a \(4\times5\) matrix. We also set

\[\begin{split}\boldsymbol{a} = (0.084, 0.229, 0.913, 0.152, 0.826)^T \mbox{ and }\boldsymbol{A} = \frac{1}{\sqrt{2}} \begin{pmatrix} 1 & 1 & 0 & 0 & 0\\ 0 & 1 & 1 & 0 & 0\\ 0 & 0 & 1 & 1 & 0\\ 0 & 0 & 0 & 1 & 1\\ \end{pmatrix}.\end{split}\]

The true parameter vector is set at \(\boldsymbol{z}^{*} = (2.75,-1.5, 0.25,-2.5,1.75)^T\). While the Sobol function is bijective and analytic, \(f\) is over-parameterized and non identifiabile.

This is also confirmed by the fact that the curve segment \(\gamma(t) = g^{-1}(g(\boldsymbol z^*) + \boldsymbol v\,t)\in Z\) gives the same model solution as \(\boldsymbol{x}^{*} = f(\boldsymbol{z}^{*}) = f(\gamma(t)) \approx (1.4910,1.6650,1.8715,1.7011)^T\) for \(t \in (-0.0153, 0.0686]\), where \(\boldsymbol v = (1,-1,1,-1,1)^T\). This is consistent with the one-dimensional null-space of the matrix \(\boldsymbol A\). We also generate synthetic observations from the Gaussian distribution

\[\boldsymbol{x} = \boldsymbol{x}^{*} + 0.01\cdot |\boldsymbol{x}^{*}| \odot \boldsymbol{x}_{0},\,\,\text{and}\,\,\boldsymbol{x}_{0} \sim \mathcal{N}(0,\boldsymbol I_5).\]

Results are shown in Fig. 2.

Fig. 2 Results from the high-dimensional model. Loss profile and representative plots for the posterior samples and posterior predictive distribution are shown.

An implementation of this model can be found below.

    def highdim_example(self):

        if "it" in os.environ:
          max_it  = int(os.environ["it"])
          max_pre = 1000
        else:
          max_it  = 25001
          max_pre = 100000

        print('')
        print('--- TEST 2: HIGH DIMENSIONAL SOBOL FUNCTION - NOFAS')
        print('')

        from linfa.models.highdim_models import Highdim
        
        exp = experiment()
        exp.name = "highdim"
        exp.flow_type         = 'realnvp'     # str: Type of flow (default 'realnvp')
        exp.n_blocks          = 15            # int: Number of layers (default 5)
        exp.hidden_size       = 100           # int: Hidden layer size for MADE in each layer (default 100)
        exp.n_hidden          = 1             # int: Number of hidden layers in each MADE (default 1)
        exp.activation_fn     = 'relu'        # str: Actication function used (default 'relu')
        exp.input_order       = 'sequential'  # str: Input order for create_mask (default 'sequential')
        exp.batch_norm_order  = True          # bool: Order to decide if batch_norm is used (default True)
        exp.save_interval     = 5000          # int: How often to sample from normalizing flow

        exp.input_size    = 5       # int: Dimensionality of input (default 2)
        exp.batch_size    = 250     # int: Number of samples generated (default 100)
        exp.true_data_num = 12      # double: number of true model evaluated (default 2)
        # exp.n_iter      = 25001
        exp.n_iter        = max_it  # int: Number of iterations (default 25001)
        exp.lr            = 0.003   # float: Learning rate (default 0.003)
        exp.lr_decay      = 0.9999  # float: Learning rate decay (default 0.9999)
        exp.log_interval  = 10      # int: How often to show loss stat (default 10)

        exp.run_nofas          = True
        exp.surr_pre_it       = max_pre #:int: Number of pre-training iterations for surrogate model
        exp.surr_upd_it       = 6000   #:int: Number of iterations for the surrogate model update

        exp.annealing          = False
        exp.calibrate_interval = 200   # int: How often to update surrogate model (default 1000)
        exp.budget             = 1023  # int: Total number of true model evaluation
        exp.surr_folder        = "./"
        exp.use_new_surr       = True

        exp.output_dir   = './results/' + exp.name
        exp.log_file     = 'log.txt'
        exp.seed         = random.randint(0, 10 ** 9)  # int: Random seed used
        exp.n_sample     = 5000                        # int: Batch size to generate final results/plots
        exp.no_cuda      = True                        # Running on CPU by default but tested on CUDA

        exp.optimizer    = 'RMSprop'
        exp.lr_scheduler = 'ExponentialLR'

        exp.device = torch.device('cuda:0' if torch.cuda.is_available() and not exp.no_cuda else 'cpu')

        # Define transformation
        # One list for each variable
        trsf_info = [['identity',0,0,0,0],
                     ['identity',0,0,0,0],
                     ['identity',0,0,0,0],
                     ['identity',0,0,0,0],
                     ['identity',0,0,0,0]]
        trsf = Transformation(trsf_info)
        exp.transform = trsf

        # Define the model
        model = Highdim(exp.device)
        exp.model = model

        # Read data
        res_path = os.path.abspath(os.path.join(os.path.dirname(linfa.__file__), os.pardir))
        model.data = np.loadtxt(res_path + '/resource/data/data_highdim.txt')

        # Define the surrogate
        exp.surrogate = Surrogate(exp.name, lambda x: model.solve_t(trsf.forward(x)), model.input_num, model.output_num,
                                  model_folder=exp.surr_folder, limits=torch.Tensor([[-3.0, 3.0], [-3.0, 3.0], [-3.0, 3.0], [-3.0, 3.0], [-3.0, 3.0]]), 
                                  memory_len=20, device=exp.device)
        surr_filename = exp.surr_folder + exp.name
        if exp.use_new_surr or (not os.path.isfile(surr_filename + ".sur")) or (not os.path.isfile(surr_filename + ".npz")):
            print("Warning: Surrogate model files: {0}.npz and {0}.npz could not be found. ".format(surr_filename))
            exp.surrogate.gen_grid(gridnum=3)
            exp.surrogate.pre_train(exp.surr_pre_it, 0.03, 0.9999, 500, store=True)
        # Load the surrogate
        exp.surrogate.surrogate_load()

        # Define the log density        
        def log_density(x, model, surrogate, transform):

            # Compute transformation log Jacobian
            adjust = transform.compute_log_jacob_func(x)

            # Eval model or surrogate
            if surrogate:
              modelOut = surrogate.forward(x)
            else:
              modelOut = model.solve_t(transform(x))
            stds = model.defOut * model.stdRatio
            Data = torch.tensor(model.data).to(exp.device)

            # Compute LL
            ll1 = -0.5 * np.prod(model.data.shape) * np.log(2.0 * np.pi)  # a number
            ll2 = (-0.5 * model.data.shape[1] * torch.log(torch.prod(stds))).item()  # a number
            ll3 = 0.0
            for i in range(4):
              ll3 += - 0.5 * torch.sum(((modelOut[:, i].unsqueeze(1) - Data[i, :].unsqueeze(0)) / stds[0, i]) ** 2, dim=1)
            negLL = -(ll1 + ll2 + ll3) 
            res = -negLL.reshape(x.size(0), 1) + adjust
            
            # Return LL
            return res
        
        # Assign log-density model
        exp.model_logdensity = lambda x: log_density(x, model, exp.surrogate,trsf)

        # Run VI
        exp.run()