index | modules | gitlab

# Source code for ase.optimize.gpmin.gpmin

    import warnings

    import numpy as np
    from scipy.optimize import minimize

    from ase.io.jsonio import write_json
    from ase.optimize.gpmin.gp import GaussianProcess
    from ase.optimize.gpmin.kernel import SquaredExponential
    from ase.optimize.gpmin.prior import ConstantPrior
    from ase.optimize.optimize import Optimizer
    from ase.parallel import world

    class GPMin(Optimizer, GaussianProcess):
        def __init__(self, atoms, restart=None, logfile='-', trajectory=None,
                     prior=None, kernel=None, noise=None, weight=None, scale=None,
                     batch_size=None, bounds=None, update_prior_strategy="maximum",
                     update_hyperparams=False, **kwargs):
            """Optimize atomic positions using GPMin algorithm, which uses both
            potential energies and forces information to build a PES via Gaussian
            Process (GP) regression and then minimizes it.

            Default behaviour:
            --------------------
            The default values of the scale, noise, weight, batch_size and bounds
            parameters depend on the value of update_hyperparams. In order to get
            the default value of any of them, they should be set up to None.
            Default values are:

            update_hyperparams = True
                scale : 0.3
                noise : 0.004
                weight: 2.
                bounds: 0.1
                batch_size: 1

            update_hyperparams = False
                scale : 0.4
                noise : 0.005
                weight: 1.
                bounds: irrelevant
                batch_size: irrelevant

            Parameters
            ----------

            atoms: :class:`~ase.Atoms`
                The Atoms object to relax.

            restart: str
                JSON file used to store the training set. If set, file with
                such a name will be searched and the data in the file incorporated
                to the new training set, if the file exists.

            logfile: file object or str
                If *logfile* is a string, a file with that name will be opened.
                Use '-' for stdout

            trajectory: str
                File used to store trajectory of atomic movement.

            prior: Prior object or None
                Prior for the GP regression of the PES surface
                See ase.optimize.gpmin.prior
                If *prior* is None, then it is set as the
                ConstantPrior with the constant being updated
                using the update_prior_strategy specified as a parameter

            kernel: Kernel object or None
                Kernel for the GP regression of the PES surface
                See ase.optimize.gpmin.kernel
                If *kernel* is None the SquaredExponential kernel is used.
                Note: It needs to be a kernel with derivatives!!!!!

            noise: float
                Regularization parameter for the Gaussian Process Regression.

            weight: float
                Prefactor of the Squared Exponential kernel.
                If *update_hyperparams* is False, changing this parameter
                has no effect on the dynamics of the algorithm.

            update_prior_strategy: str
                Strategy to update the constant from the ConstantPrior
                when more data is collected. It does only work when
                Prior = None

                options:
                    'maximum': update the prior to the maximum sampled energy
                    'init' : fix the prior to the initial energy
                    'average': use the average of sampled energies as prior

            scale: float
                scale of the Squared Exponential Kernel

            update_hyperparams: bool
                Update the scale of the Squared exponential kernel
                every batch_size-th iteration by maximizing the
                marginal likelihood.

            batch_size: int
                Number of new points in the sample before updating
                the hyperparameters.
                Only relevant if the optimizer is executed in update_hyperparams
                mode: (update_hyperparams = True)

            bounds: float, 0<bounds<1
                Set bounds to the optimization of the hyperparameters.
                Let t be a hyperparameter. Then it is optimized under the
                constraint (1-bound)*t_0 <= t <= (1+bound)*t_0
                where t_0 is the value of the hyperparameter in the previous
                step.
                If bounds is False, no constraints are set in the optimization of
                the hyperparameters.

            kwargs : dict, optional
                Extra arguments passed to
                :class:`~ase.optimize.optimize.Optimizer`.

            .. warning:: The memory of the optimizer scales as O(n²N²) where
                         N is the number of atoms and n the number of steps.
                         If the number of atoms is sufficiently high, this
                         may cause a memory issue.
                         This class prints a warning if the user tries to
                         run GPMin with more than 100 atoms in the unit cell.
            """
            # Warn the user if the number of atoms is very large
            if len(atoms) > 100:
                warning = ('Possible Memory Issue. There are more than '
                           '100 atoms in the unit cell. The memory '
                           'of the process will increase with the number '
                           'of steps, potentially causing a memory issue. '
                           'Consider using a different optimizer.')

                warnings.warn(warning)

            # Give it default hyperparameters
            if update_hyperparams:       # Updated GPMin
                if scale is None:
                    scale = 0.3
                if noise is None:
                    noise = 0.004
                if weight is None:
                    weight = 2.
                if bounds is None:
                    self.eps = 0.1
                elif bounds is False:
                    self.eps = None
                else:
                    self.eps = bounds
                if batch_size is None:
                    self.nbatch = 1
                else:
                    self.nbatch = batch_size
            else:                        # GPMin without updates
                if scale is None:
                    scale = 0.4
                if noise is None:
                    noise = 0.001
                if weight is None:
                    weight = 1.
                if bounds is not None:
                    warning = ('The parameter bounds is of no use '
                               'if update_hyperparams is False. '
                               'The value provided by the user '
                               'is being ignored.')
                    warnings.warn(warning, UserWarning)
                if batch_size is not None:
                    warning = ('The parameter batch_size is of no use '
                               'if update_hyperparams is False. '
                               'The value provided by the user '
                               'is being ignored.')
                    warnings.warn(warning, UserWarning)

                # Set the variables to something anyways
                self.eps = False
                self.nbatch = None

            self.strategy = update_prior_strategy
            self.update_hp = update_hyperparams
            self.function_calls = 1
            self.force_calls = 0
            self.x_list = []      # Training set features
            self.y_list = []      # Training set targets

            Optimizer.__init__(self, atoms, restart=restart, logfile=logfile,
                               trajectory=trajectory, **kwargs)
            if prior is None:
                self.update_prior = True
                prior = ConstantPrior(constant=None)
            else:
                self.update_prior = False

            if kernel is None:
                kernel = SquaredExponential()
            GaussianProcess.__init__(self, prior, kernel)
            self.set_hyperparams(np.array([weight, scale, noise]))

        def acquisition(self, r):
            e = self.predict(r)
            return e[0], e[1:]

        def update(self, r, e, f):
            """Update the PES

            Update the training set, the prior and the hyperparameters.
            Finally, train the model
            """
            # update the training set
            self.x_list.append(r)
            f = f.reshape(-1)
            y = np.append(np.array(e).reshape(-1), -f)
            self.y_list.append(y)

            # Set/update the constant for the prior
            if self.update_prior:
                if self.strategy == 'average':
                    av_e = np.mean(np.array(self.y_list)[:, 0])
                    self.prior.set_constant(av_e)
                elif self.strategy == 'maximum':
                    max_e = np.max(np.array(self.y_list)[:, 0])
                    self.prior.set_constant(max_e)
                elif self.strategy == 'init':
                    self.prior.set_constant(e)
                    self.update_prior = False

            # update hyperparams
            if (self.update_hp and self.function_calls % self.nbatch == 0 and
               self.function_calls != 0):
                self.fit_to_batch()

            # build the model
            self.train(np.array(self.x_list), np.array(self.y_list))

        def relax_model(self, r0):
            result = minimize(self.acquisition, r0, method='L-BFGS-B', jac=True)
            if result.success:
                return result.x
            else:
                self.dump()
                raise RuntimeError("The minimization of the acquisition function "
                                   "has not converged")

        def fit_to_batch(self):
            """Fit hyperparameters keeping the ratio noise/weight fixed"""
            ratio = self.noise / self.kernel.weight
            self.fit_hyperparameters(np.array(self.x_list),
                                     np.array(self.y_list), eps=self.eps)
            self.noise = ratio * self.kernel.weight

        def step(self, f=None):
            optimizable = self.optimizable
            if f is None:
                f = optimizable.get_forces()

            r0 = optimizable.get_positions().reshape(-1)
            e0 = optimizable.get_potential_energy()
            self.update(r0, e0, f)

            r1 = self.relax_model(r0)
            optimizable.set_positions(r1.reshape(-1, 3))
            e1 = optimizable.get_potential_energy()
            f1 = optimizable.get_forces()
            self.function_calls += 1
            self.force_calls += 1
            count = 0
            while e1 >= e0:
                self.update(r1, e1, f1)
                r1 = self.relax_model(r0)

                optimizable.set_positions(r1.reshape(-1, 3))
                e1 = optimizable.get_potential_energy()
                f1 = optimizable.get_forces()
                self.function_calls += 1
                self.force_calls += 1
                if self.converged(f1):
                    break

                count += 1
                if count == 30:
                    raise RuntimeError("A descent model could not be built")
            self.dump()

        def dump(self):
            """Save the training set"""
            if world.rank == 0 and self.restart is not None:
                with open(self.restart, 'w') as fd:
                    write_json(fd, (self.x_list, self.y_list))

        def read(self):
            self.x_list, self.y_list = self.load()