bfgslinesearch_constrained.py 4.78 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
from __future__ import print_function

# ******NOTICE***************
# optimize.py module by Travis E. Oliphant
#
# You may copy and use this module as you see fit with no
# guarantee implied provided you keep this notice in all copies.
# *****END NOTICE************

import time
import numpy as np
from numpy import eye, absolute, sqrt, isinf
from ase.utils.linesearch import LineSearch
from ase.optimize.bfgslinesearch import BFGSLineSearch
from ase.constraints import FixAtoms
from ase.utils import basestring
from ase.data import covalent_radii
from ase.ga.utilities import get_mic_distance

from utils import check_valid_bondlengths

# These have been copied from Numeric's MLab.py
# I don't think they made the transition to scipy_core

# Modified from scipy_optimize
abs = absolute
pymin = min
pymax = max
__version__ = '0.1'


class BFGSLineSearch_constrained(BFGSLineSearch):
    def __init__(self, atoms, pos_init=None, restart=None, logfile='-', maxstep=.2,
                 trajectory=None, c1=0.23, c2=0.46, alpha=10.0, stpmax=50.0,
                 master=None, force_consistent=None, zlim=None,
                 blmin=None, blmax=None, dmax_cov=3.0, rk=None):
        """
        add maximum displacement of single atoms to BFGSLineSearch:

        dmax_cov: maximum distance the atom is alowed to move from it's initial position.
        in units of it's covalent distance.
        """

        self.rk = rk  # for testing
        
        self.zlim = zlim
        self.blmin = blmin
        self.blmax = blmax
        
        self.dmax_cov = dmax_cov
        self.d_cov = 2*np.array([covalent_radii[num] for num in atoms.get_atomic_numbers()])
        if pos_init is not None:
            self.pos_init = pos_init
        else:
            self.pos_init = np.copy(atoms.positions)

        self.cell = atoms.get_cell()
        self.pbc = atoms.get_pbc()

        BFGSLineSearch.__init__(self, atoms, restart=restart, logfile=logfile, maxstep=maxstep,
                                trajectory=trajectory, c1=c1, c2=c2, alpha=alpha, stpmax=stpmax,
                                master=master, force_consistent=force_consistent)

    def converged(self, forces=None):
        """Did the optimization converge?"""
        if forces is None:
            forces = self.atoms.get_forces()
        if hasattr(self.atoms, 'get_curvature'):
            return ((forces**2).sum(axis=1).max() < self.fmax**2 and
                    self.atoms.get_curvature() < 0.0)

        # Check constraints
        terminate_due_to_constraints = self.check_constraints()
        if terminate_due_to_constraints:
            return True
        
        return (forces**2).sum(axis=1).max() < self.fmax**2

    def check_constraints(self):
        # Check if stop due to large displacement
        valid_displace = self.check_displacement()
        # Check if stop due to zlim
        valid_z = self.check_zcoord()
        # Check if stop due to invalid bondlengths
        valid_bondlengths = self.check_bondlengths()
        
        if not valid_displace or not valid_z or not valid_bondlengths:
            #self.r0 = self.atoms_prior.get_positions()
            return True
    
    def check_displacement(self):
        valid_displace = True
        if self.dmax_cov is not None:
            d = np.array([get_mic_distance(p1,p2,self.cell,self.pbc) for p1,p2 in zip(self.pos_init,self.atoms.get_positions())])
            max_covDisplace = np.max(d/self.d_cov)
            if max_covDisplace > self.dmax_cov:
                print('RELAX STOPPED on rank {} - max relax dist (>{} covDist) reached.'.format(self.rk, self.dmax_cov), flush=True)
                valid_displace = False
        return valid_displace

    def check_zcoord(self):
        valid_z = True
        if self.zlim is not None:
            # get indices of non-fixed atoms
            indices = np.arange(self.atoms.get_number_of_atoms())
            for constraint in self.atoms.constraints:
                if isinstance(constraint, FixAtoms):
                    indices_fixed = constraint.get_indices()
                    indices = np.delete(np.arange(self.atoms.get_number_of_atoms()), indices_fixed)
            pos_z = self.atoms.positions[indices,2]
            if np.any(pos_z < self.zlim[0]) or np.any(pos_z > self.zlim[1]):
                print('RELAXATION STOPPED on rank {} due to - zlim crossed.'.format(self.rk))
                valid_z = False
        return valid_z

    def check_bondlengths(self):
        valid_bondlengths = True
        if self.blmin is not None or self.blmax is not None:
            valid_bondlengths = check_valid_bondlengths(self.atoms, self.blmin, self.blmax)
            if not valid_bondlengths:
                #print('RELAXATION STOPPED on rank {} due to - invalid bondlengths.'.format(self.rk))
                valid_bondlengths = False
        return valid_bondlengths