<commit_before># This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import logging
from collections import namedtuple
from collections.abc import Generator
from itertools import chain

import numpy as np
from numpy import dot, eye
from numpy.linalg import norm

from . import Math
from .coords import InternalCoords

__version__ = '0.3.0'
__all__ = ['Berny']

log = logging.getLogger(__name__)

defaults = {
    'gradientmax': 0.45e-3,
    'gradientrms': 0.15e-3,
    'stepmax': 1.8e-3,
    'steprms': 1.2e-3,
    'trust': 0.3,
    'dihedral': True,
    'superweakdih': False,
}
"""
``gradientmax``, ``gradientrms``, ``stepmax``, ``steprms``
    Convergence criteria in atomic units ("step" refers to the step in
    internal coordinates, assuming radian units for angles).

``trust``
    Initial trust radius in atomic units. It is the maximum RMS of the
    quadratic step (see below).

``dihedral``
    Form dihedral angles.

``superweakdih``
    Form dihedral angles containing two or more noncovalent bonds.
"""


class BernyAdapter(logging.LoggerAdapter):
    def process(self, msg, kwargs):
        return '{} {}'.format(self.extra['step'], msg), kwargs


class Berny(Generator):
    """Generator that receives energy and gradients and yields the next geometry.

    Args:
        geom (:class:`~berny.Geometry`): geometry to start with
        debug (bool): if :data:`True`, the generator yields debug info on receiving
            the energy and gradients, otherwise it yields :data:`None`
        restart (dict): start from a state saved from previous run
            using ``debug=True``
        maxsteps (int): abort after maximum number of steps
        params: parameters that override the :data:`~berny.berny.defaults`

    The Berny object is to be used as follows::

        optimizer = Berny(geom)
        for geom in optimizer:
            # calculate energy and gradients (as N-by-3 matrix)
            debug = optimizer.send((energy, gradients))
    """

    class State(object):
        pass

    Point = namedtuple('Point', 'q E g')

    def __init__(self, geom, debug=False, restart=None, maxsteps=100, **params):
        self._debug = debug
        self._maxsteps = maxsteps
        self._converged = False
        self._n = 0
        self._log_extra = {'step': self._n}
        self._log = BernyAdapter(log, self._log_extra)
        s = self._state = Berny.State()
        if restart:
            vars(s).update(restart)
            return
        s.geom = geom
        s.params = dict(chain(defaults.items(), params.items()))
        s.trust = s.params['trust']
        s.coords = InternalCoords(
            s.geom, dihedral=s.params['dihedral'], superweakdih=s.params['superweakdih']
        )
        s.H = s.coords.hessian_guess(s.geom)
        s.weights = s.coords.weights(s.geom)
        s.future = Berny.Point(s.coords.eval_geom(s.geom), None, None)
        s.first = True
        for line in str(s.coords).split('\n'):
            self._log.info(line)

    def __next__(self):
        assert self._n <= self._maxsteps
        if self._n == self._maxsteps or self._converged:
            raise StopIteration
        self._n += 1
        return self._state.geom

    @property
    def trust(self):
        """Current trust radius."""
        return self._state.trust

    def send(self, energy_and_gradients):  # noqa: D102
        log = self._log.info
        self._log_extra['step'] = self._n
        s = self._state
        energy, gradients = energy_and_gradients
        gradients = np.array(gradients)
        log('Energy: {:.12}'.format(energy))
        B = s.coords.B_matrix(s.geom)
        B_inv = B.T.dot(Math.pinv(np.dot(B, B.T), log=log))
        current = Berny.Point(s.future.q, energy, dot(B_inv.T, gradients.reshape(-1)))
        if not s.first:<commit_msg>add Berny.converged property<commit_after># This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import logging
from collections import namedtuple
from collections.abc import Generator
from itertools import chain

import numpy as np
from numpy import dot, eye
from numpy.linalg import norm

from . import Math
from .coords import InternalCoords

__version__ = '0.3.1'
__all__ = ['Berny']

log = logging.getLogger(__name__)

defaults = {
    'gradientmax': 0.45e-3,
    'gradientrms': 0.15e-3,
    'stepmax': 1.8e-3,
    'steprms': 1.2e-3,
    'trust': 0.3,
    'dihedral': True,
    'superweakdih': False,
}
"""
``gradientmax``, ``gradientrms``, ``stepmax``, ``steprms``
    Convergence criteria in atomic units ("step" refers to the step in
    internal coordinates, assuming radian units for angles).

``trust``
    Initial trust radius in atomic units. It is the maximum RMS of the
    quadratic step (see below).

``dihedral``
    Form dihedral angles.

``superweakdih``
    Form dihedral angles containing two or more noncovalent bonds.
"""


class BernyAdapter(logging.LoggerAdapter):
    def process(self, msg, kwargs):
        return '{} {}'.format(self.extra['step'], msg), kwargs


class Berny(Generator):
    """Generator that receives energy and gradients and yields the next geometry.

    Args:
        geom (:class:`~berny.Geometry`): geometry to start with
        debug (bool): if :data:`True`, the generator yields debug info on receiving
            the energy and gradients, otherwise it yields :data:`None`
        restart (dict): start from a state saved from previous run
            using ``debug=True``
        maxsteps (int): abort after maximum number of steps
        params: parameters that override the :data:`~berny.berny.defaults`

    The Berny object is to be used as follows::

        optimizer = Berny(geom)
        for geom in optimizer:
            # calculate energy and gradients (as N-by-3 matrix)
            debug = optimizer.send((energy, gradients))
    """

    class State(object):
        pass

    Point = namedtuple('Point', 'q E g')

    def __init__(self, geom, debug=False, restart=None, maxsteps=100, **params):
        self._debug = debug
        self._maxsteps = maxsteps
        self._converged = False
        self._n = 0
        self._log_extra = {'step': self._n}
        self._log = BernyAdapter(log, self._log_extra)
        s = self._state = Berny.State()
        if restart:
            vars(s).update(restart)
            return
        s.geom = geom
        s.params = dict(chain(defaults.items(), params.items()))
        s.trust = s.params['trust']
        s.coords = InternalCoords(
            s.geom, dihedral=s.params['dihedral'], superweakdih=s.params['superweakdih']
        )
        s.H = s.coords.hessian_guess(s.geom)
        s.weights = s.coords.weights(s.geom)
        s.future = Berny.Point(s.coords.eval_geom(s.geom), None, None)
        s.first = True
        for line in str(s.coords).split('\n'):
            self._log.info(line)

    def __next__(self):
        assert self._n <= self._maxsteps
        if self._n == self._maxsteps or self._converged:
            raise StopIteration
        self._n += 1
        return self._state.geom

    @property
    def trust(self):
        """Current trust radius."""
        return self._state.trust

    @property
    def converged(self):
        """Whether the optimized has converged."""
        return self._converged

    def send(self, energy_and_gradients):  # noqa: D102
        log = self._log.info
        self._log_extra['step'] = self._n
        s = self._state
        energy, gradients = energy_and_gradients
        gradients = np.array(gradients)
        log('Energy: {:.12}'.format(energy))
        B = s.coords.B_matrix(s.geom)
        B_inv = B.T.dot(Math.pinv(np.dot(B, B.T), log=log))
        current = Berny.Point(s.future.q, energy, dot(B_inv.T, gradients.reshape(-1)))
        if not s.first:
{'fraction_non_alphanumeric': 0.11851759731376695, 'fraction_numerical': 0.008705384902375327, 'mean_word_length': 2.59499329459097, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5479972', 'n_tokens_mistral': 2555, 'n_tokens_neox': 2457, 'n_words': 737}