candidate_generation.py 5.62 KB
Newer Older
Malthe Kjær Bisbo's avatar
update  
Malthe Kjær Bisbo committed
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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
import numpy as np
from abc import ABC, abstractmethod
from ase.data import covalent_radii
from ase.geometry import get_distances

from ase.visualize import view

class OffspringOperation(ABC):
    """Baseclass for mutation and crossover operations.
    """
    def __init__(self, blmin=0.7, blmax=1.3):
        self.blmin = blmin
        self.blmax = blmax
        self.description = 'Unspecified'

    @abstractmethod
    def get_new_candidate(self):
        pass

    def check_valid_bondlengths(self, a, indices=None):
        bl = self.get_distances_as_fraction_of_covalent(a,indices)

        # Filter away self interactions.
        bl = bl[bl > 1e-6].reshape(bl.shape[0],bl.shape[1]-1)

        # Check if atoms are too close
        tc = np.any(bl < self.blmin)

        # Check if there are isolated atoms
        if self.blmax is not None:
            isolated = np.any(np.all(bl > self.blmax, axis=1))
        else:
            isolated = False

        is_valid = not tc and not isolated
        return is_valid

    def get_covalent_distance_from_atom_numbers(self, a, indices=None):
        r_cov = np.array([covalent_radii[n] for n in a.get_atomic_numbers()])
        if indices is None:
            r_cov_sub = r_cov
        else:
            r_cov_sub = r_cov[indices]
        cd_mat = r_cov_sub.reshape(-1,1) + r_cov.reshape(1,-1)
        return cd_mat
    
    def get_distances_as_fraction_of_covalent(self, a, indices=None, covalent_distances=None):
        if indices is None:
            indices = np.arange(len(a))

        if covalent_distances is None:
            cd = self.get_covalent_distance_from_atom_numbers(a, indices=indices)
        else:
            cd = covalent_distances[indices,:]
        _, d = get_distances(a[indices].positions,
                             a.positions,
                             cell=a.get_cell(),
                             pbc=a.get_pbc())
        bl = d/cd
        
        return bl

    def finalize(self, a):
        a.info['key_value_pairs'] = {'origin': self.description}
        valid_bondlengths = self.check_valid_bondlengths(a)
        assert valid_bondlengths, 'bondlengths are not valid'
        return a

def pos_add_sphere(rattle_strength):
    # Rattle within a sphere
    r = rattle_strength * np.random.rand()**(1/3)
    theta = np.random.uniform(low=0, high=2*np.pi)
    phi = np.random.uniform(low=0, high=np.pi)
    pos_add = r * np.array([np.cos(theta)*np.sin(phi),
                            np.sin(theta)*np.sin(phi),
                            np.cos(phi)])
    return pos_add

def pos_add_sphere_shell(rmin, rmax):
    # Rattle within a sphere
    r = np.random.uniform(rmin**3, rmax**3)**(1/3)
    theta = np.random.uniform(low=0, high=2*np.pi)
    phi = np.random.uniform(low=0, high=np.pi)
    pos_add = r * np.array([np.cos(theta)*np.sin(phi),
                            np.sin(theta)*np.sin(phi),
                            np.cos(phi)])
    return pos_add


class CandidateGenerator():
    """Class to generate new candidates by applying candidate generation
    operations drawn randomly from a list of possible operations
    
    operations : "list" or "list of lists" of mutations/crossovers.

    probabilities : probability for each of the mutations/crossovers
    in operations. Must have the same dimensions as operations.
    """
    def __init__(self, probabilities, operations):
        cond1 = isinstance(operations[0], list)
        cond2 = isinstance(probabilities[0], list)
        if not cond1 and not cond2:
            operations = [operations]
            probabilities = [probabilities]
        element_count_operations = [len(op_list) for op_list in operations]
        element_count_probabilities = [len(prob_list)
                                       for prob_list in probabilities]
        assert element_count_operations == element_count_probabilities
        self.operations = operations
        self.rho = [np.cumsum(prob_list) for prob_list in probabilities]

    def __get_index__(self, rho):
        """Draw from the cumulative probalility distribution, rho,
        to return the index of which operation to use"""
        v = np.random.random() * rho[-1]
        for i in range(len(rho)):
            if rho[i] > v:
                return i
        
    def get_new_candidate(self, parents):
        """Generate new candidate by applying a randomly drawn
        operation on the structures. This is done successively for
        each list of operations, if multiple are present.
        """
        for op_list, rho_list in zip(self.operations, self.rho):
            to_use = self.__get_index__(rho_list)
            anew = op_list[to_use].get_new_candidate(parents)
            parents[0] = anew
        return anew

if __name__ == '__main__':
    from ase.io import read
    from ase.visualize import view

    from candidate_operations.basic_mutations import RattleMutation, RattleMutation2, PermutationMutation
    
    np.random.seed(1)
    
    #a = read('/home/mkb/DFT/gpLEA/Si3x3/ref/gm_unrelaxed_done.traj', index='0')
    #a = read('si3x3.traj', index='0')
    #a = read('c6h6.traj', index='0')
    a = read('sn2o3.traj', index='0')
    
    rattle = RattleMutation(n_top=16, Nrattle=3)
    rattle2 = RattleMutation2(n_top=16, Nrattle=0.1)
    permut = PermutationMutation(n_top=16, Npermute=2)

    candidategenerator = CandidateGenerator([0., 1., 0.], [rattle, rattle2, permut])
    #candidategenerator = CandidateGenerator([[1],[1]], [[rattle2], [permut]])

    traj = []
    for i in range(100):
        a0 = a.copy()
        anew = candidategenerator.get_new_candidate([a0,a0])
        traj += [a0, anew]

    view(traj)
    """
    a_mut = rattle.get_new_candidate([a])
    view([a,a_mut])
    """