Source code for pyiron_atomistics.interactive.scipy_minimizer

# coding: utf-8
# Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
# Distributed under the terms of "New BSD License", see the LICENSE file.

import warnings
from typing import Iterable, List

import numpy as np
import scipy
from pyiron_base import HasStorage
from scipy.optimize import minimize

from pyiron_atomistics.atomistics.job.interactive import GenericInteractiveOutput
from pyiron_atomistics.atomistics.job.interactivewrapper import InteractiveWrapper

__author__ = "Osamu Waseda"
__copyright__ = (
    "Copyright 2021, Max-Planck-Institut für Eisenforschung GmbH - "
    "Computational Materials Design (CM) Department"
)
__version__ = "1.0"
__maintainer__ = "Osamu Waseda"
__email__ = "waseda@mpie.de"
__status__ = "development"
__date__ = "Sep 1, 2018"

GPa_to_eV_by_A3 = (
    1e21 / scipy.constants.physical_constants["joule-electron volt relationship"][0]
)




class ScipyMinimizer(InteractiveWrapper):
    """
    Structure optimization class based on Scipy minimizers.

    Example I:

    # Position optimization
    >>> pr = Project('position')
    >>> job = pr.create_job('SomeAtomisticJob', 'atomistic')
    >>> job.structure = pr.create_structure('Al', 'fcc', 4.)
    >>> # it works also in the static mode, but interactive is recommended
    >>> job.server.run_mode.interactive = True
    >>> minim = pr.create_job('ScipyMinimizer', 'scipy')
    >>> minim.ref_job = job
    >>> minim.run()

    Example II:

    # Volume optimization:
    >>> pr = Project('volume')
    >>> job = pr.create_job('SomeAtomisticJob', 'atomistic')
    >>> job.structure = pr.create_structure('Al', 'fcc', 4.)
    >>> # it works also in the static mode, but interactive is recommended
    >>> job.server.run_mode.interactive = True
    >>> minim = pr.create_job('ScipyMinimizer', 'scipy')
    >>> minim.ref_job = job
    >>> minim.calc_minimize(pressure=0, volume_only=True)
    >>> minim.run()

    By setting `volume_only`, positions are not updated, so that only the
    pressures are minimized.

    It is possible to optimize both the volume and the positions, but since
    changing the cell also changes the definition of coordinates, it is a
    mathematically ill-defined problem and therefore it might end up in a
    physically undefined state. For this reason, it is strongly recommended to
    launch several jobs using the Murnaghan class (cf. `Murnaghan`).

    In order to perform volume optimization, the child job must have
    3x3-pressure output.
    """



    def __init__(self, project, job_name):
        super(ScipyMinimizer, self).__init__(project, job_name)

        self._ref_job = None
        self.input = ScipyMinimizerInput()
        self.output = ScipyMinimizerOutput(job=self)
        self.interactive_cache = {}
        self._delete_existing_job = True




    def set_input_to_read_only(self):
        """
        This function enforces read-only mode for the input classes, but it
        has to be implement in the individual classes.
        """
        super(ScipyMinimizer, self).set_input_to_read_only()
        self.input.read_only = True




    def write_input(self):
        pass


    def _initialize_structure(self):
        self._original_cell = self.ref_job.structure.cell.copy()
        self._current_strain = np.zeros(6)



    def run_static(self):
        self.ref_job_initialize()
        self._logger.debug("cg status: " + str(self.status))
        self._initialize_structure()
        if self.ref_job.server.run_mode.interactive:
            self._delete_existing_job = False
        self.ref_job.run(delete_existing_job=self._delete_existing_job)
        self.status.running = True
        if self.input.pressure is not None:
            x0 = np.zeros(sum(self.input.pressure != None))
            if not self.input.volume_only:
                x0 = np.append(
                    x0, self.ref_job.structure.get_scaled_positions().flatten()
                )
        else:
            x0 = self.ref_job.structure.positions.flatten()
        self.output._result = minimize(
            method=self.input.minimizer,
            fun=self._get_value,
            x0=x0,
            jac=self._get_gradient if self.input.use_pressure else None,
            tol=self.input.ionic_energy_tolerance,
            options={"maxiter": self.input.ionic_steps, "return_all": True},
        )
        self.status.collect = True
        self.collect_output()
        if self.ref_job.server.run_mode.interactive:
            self.ref_job.interactive_close()
        if self["output/convergence"] > 0:
            self.status.finished = True
        else:
            self.status.not_converged = True


    @staticmethod
    def _tensor_to_voigt(s, strain=False):
        ss = 0.5 * (s + s.T)
        ss = ss.flatten()[[0, 4, 8, 5, 2, 1]]
        if strain:
            ss[3:] *= 2
        return ss

    @staticmethod
    def _voigt_to_tensor(s, strain=False):
        ss = np.array(s).copy()
        if not strain:
            ss[:3] /= 2
        ss = np.array([[ss[0], ss[5], ss[4]], [0, ss[1], ss[3]], [0, 0, ss[2]]])
        ss += ss.T
        return ss

    def _update(self, x):
        rerun = False
        if self.input.pressure is not None:
            if not np.allclose(x[: len(self.input.pressure)], self._current_strain):
                if len(self.input.pressure) == 1:
                    self._current_strain[:3] = x[0]
                else:
                    self._current_strain[self.input.pressure != None] = x[
                        : len(self.input.pressure)
                    ]
                cell = np.matmul(
                    self._voigt_to_tensor(self._current_strain, strain=True)
                    + np.eye(3),
                    self._original_cell,
                )
                self.ref_job.structure.set_cell(cell, scale_atoms=True)
                rerun = True
            if not self.input.volume_only and not np.allclose(
                x[len(self.input.pressure) :],
                self.ref_job.structure.get_scaled_positions().flatten(),
            ):
                self.ref_job.structure.set_scaled_positions(
                    x[len(self.input.pressure) :].reshape(-1, 3)
                )
                rerun = True
        elif not np.allclose(x, self.ref_job.structure.positions.flatten()):
            self.ref_job.structure.positions = x.reshape(-1, 3)
            rerun = True
        if rerun:
            self.ref_job.run(delete_existing_job=self._delete_existing_job)

    def check_convergence(self):
        if self.input.ionic_energy_tolerance > 0:
            if len(self.ref_job.output.energy_pot) < 2:
                return False
            elif (
                np.absolute(np.diff(self.ref_job.output.energy_pot)[-1])
                > self.input.ionic_energy_tolerance
            ):
                return False
            if self.input.ionic_force_tolerance == 0:
                return True
        max_force = np.linalg.norm(self.ref_job.output.forces[-1], axis=-1).max()
        if self.input.pressure is None:
            if max_force > self.input.ionic_force_tolerance:
                return False
        elif self.input.volume_only:
            if self.input.use_pressure and (
                np.absolute(self._get_pressure() - self.input.pressure).max()
                > self.input.pressure_tolerance
            ):
                return False
        else:
            if max_force > self.input.ionic_force_tolerance:
                return False
            if self.input.use_pressure and (
                np.absolute(self._get_pressure() - self.input.pressure).max()
                > self.input.pressure_tolerance
            ):
                return False
        return True

    def _get_pressure(self):
        if len(self.input.pressure) == 1:
            return [np.mean(np.diagonal(self.ref_job.output.pressures[-1]))]
        else:
            return self._tensor_to_voigt(self.ref_job.output.pressures[-1])[
                self.input.pressure != None
            ]

    def _get_gradient(self, x):
        self._update(x)
        prefactor = 1.0e-1
        if self.check_convergence():
            prefactor = 0
        if self.input.pressure is not None:
            pressure = -(self._get_pressure() - self.input.pressure)
            if self.input.volume_only:
                return pressure * prefactor
            else:
                return (
                    np.append(
                        pressure,
                        -np.einsum(
                            "ij,ni->nj",
                            np.linalg.inv(self.ref_job.structure.cell),
                            self.ref_job.output.forces[-1],
                        ).flatten(),
                    ).flatten()
                    * prefactor
                )
        else:
            return -self.ref_job.output.forces[-1].flatten() * prefactor

    def _get_value(self, x):
        self._update(x)
        return self.ref_job.output.energy_pot[-1]



    def collect_output(self):
        self.output.to_hdf(self._hdf5)




    def to_hdf(self, hdf=None, group_name=None):
        super(ScipyMinimizer, self).to_hdf(hdf=hdf, group_name=group_name)
        self.output.to_hdf(self.project_hdf5)




    def calc_minimize(
        self,
        max_iter=100,
        pressure=None,
        algorithm="CG",
        ionic_energy_tolerance=0,
        ionic_force_tolerance=1.0e-2,
        pressure_tolerance=1.0e-3,
        volume_only=False,
    ):
        """
        Args:
            algorithm (str): scipy algorithm (currently only 'CG' and 'BFGS' run reliably)
            pressure (float/list/numpy.ndarray): target pressures
            max_iter (int): maximum number of iterations
            ionic_energy_tolerance (float): convergence goal in terms of
                                  energy (optional)
            ionic_force_tolerance (float): convergence goal in terms of
                                  forces (optional)
            volume_only (bool): Only pressure minimization
        """
        if pressure is None and volume_only:
            raise ValueError("pressure must be specified if volume_only")
        if pressure is not None and not volume_only:
            warnings.warn(
                "Simultaneous optimization of pressures and positions is a"
                + " mathematically ill posed problem - there is no guarantee"
                + " that it converges to the desired structure"
            )
        if pressure is not None:
            pressure = np.array([pressure]).flatten()
            if len(pressure) == 9:
                pressure = self._tensor_to_voigt(pressure.reshape(3, 3))
            if len(pressure) == 3:
                pressure = np.append(pressure, 3 * [None])
        self.input.minimizer = algorithm
        self.input.ionic_steps = max_iter
        self.input.pressure = pressure
        self.input.volume_only = volume_only
        self.input.ionic_force_tolerance = ionic_force_tolerance
        self.input.ionic_energy_tolerance = ionic_energy_tolerance
        self.input.pressure_tolerance = pressure_tolerance






class ScipyMinimizerInput(HasStorage):


    def __init__(self):
        super().__init__()
        self.storage.minimizer = "CG"
        self.storage.ionic_steps = 100
        self.storage.ionic_energy_tolerance = 0
        self.storage.ionic_force_tolerance = 1.0e-2
        self.storage.pressure_tolerance = 1.0e-3
        self.storage.pressure = None
        self.storage.use_pressure = True
        self.storage.volume_only = False


    def _get_hdf_group_name(self) -> str:
        return "parameters"

    @property
    def minimizer(self) -> str:
        """str: name of minimizer to use"""
        return self.storage.minimizer

    @minimizer.setter
    def minimizer(self, value: str):
        self.storage.minimizer = value

    @property
    def ionic_steps(self) -> int:
        """int: maximum number of minimization steps"""
        return self.storage.ionic_steps

    @ionic_steps.setter
    def ionic_steps(self, value: int):
        self.storage.ionic_steps = value

    @property
    def ionic_force_tolerance(self) -> float:
        """float: convergence goal in terms of forces"""
        return self.storage.ionic_force_tolerance

    @ionic_force_tolerance.setter
    def ionic_force_tolerance(self, value: float):
        self.storage.ionic_force_tolerance = value

    @property
    def ionic_energy_tolerance(self) -> float:
        """float: convergence goal in terms of energye"""
        return self.storage.ionic_energy_tolerance

    @ionic_energy_tolerance.setter
    def ionic_energy_tolerance(self, value: float):
        self.storage.ionic_energy_tolerance = value

    @property
    def pressure_tolerance(self) -> float:
        """float: convergence goal in terms of energye"""
        return self.storage.pressure_tolerance

    @pressure_tolerance.setter
    def pressure_tolerance(self, value: float):
        self.storage.pressure_tolerance = value

    @property
    def pressure(self):
        """float: target pressure"""
        if isinstance(self.storage.pressure, Iterable):
            return np.asarray(self.storage.pressure)
        else:
            return self.storage.pressure

    @pressure.setter
    def pressure(self, value: Iterable[float]):
        value = list(value)
        self.storage.pressure = value

    @property
    def use_pressure(self) -> bool:
        """bool: rely on pressures computed by reference job or not"""
        return self.storage.use_pressure

    @use_pressure.setter
    def use_pressure(self, value: bool):
        self.storage.use_pressure = value

    @property
    def volume_only(self) -> bool:
        """bool: only pressure minimization"""
        return self.storage.volume_only

    @volume_only.setter
    def volume_only(self, value: bool):
        self.storage.volume_only = value





class ScipyMinimizerOutput(GenericInteractiveOutput):


    def __init__(self, job):
        super(ScipyMinimizerOutput, self).__init__(job=job)
        self._result = None


    def to_hdf(self, hdf, group_name="output"):
        if self._result is None:
            return
        with hdf.open(group_name) as hdf_output:
            hdf_output["convergence"] = self._result["success"]
            if "hess_inv" in self._result.keys():
                hdf_output["hessian"] = self._result["hess_inv"]