bfgslinesearch_constrained.py 4.54 KB
Newer Older
1
2
3
4
5
6
7
8
import numpy as np
from ase.optimize.bfgslinesearch import BFGSLineSearch
from ase.constraints import FixAtoms
from ase.data import covalent_radii
from ase.ga.utilities import get_mic_distance

from utils import check_valid_bondlengths

9
10
import traceback

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

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:
                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]):
                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:
                valid_bondlengths = False
        return valid_bondlengths
Malthe Kjær Bisbo's avatar
Malthe Kjær Bisbo committed
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

def relax(structure, calc, Fmax=0.05, steps_max=200, dmax_cov=None):
    a = structure.copy()
    # Set calculator 
    a.set_calculator(calc)
    pos_init = a.get_positions()

    # Catch if linesearch fails
    try:
        dyn = BFGSLineSearch_constrained(a,
                                         logfile=None,
                                         pos_init=pos_init,
                                         dmax_cov=dmax_cov)
        dyn.run(fmax = Fmax, steps = steps_max)
    except Exception as err:
        print('Error in surrogate-relaxation:', err, flush=True)
        traceback.print_exc()
        traceback.print_exc(file=sys.stderr)
    return a