PAMD/ptmd.py at public_version · christiansenh/PAMD · GitHub

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
import numpy as np

from mpi4py import MPI

import sys


class Simulation:
    '''Main Class for the simulation with population annealing.'''
    def __init__(self, temperatures, seed_offset=0, comm=MPI.COMM_WORLD, gpu_sim_object=None):
        self.comm = comm
        self.gpu_sim_object = gpu_sim_object

        self.temperatures = np.array(temperatures) #index to temp
        self.equilibration_steps = 10
        self.simulation_steps = 10
        self.steps = 100

        self.simulations = None
        self.boltzmann = None

        self.rank = self.comm.Get_rank()
        self.size = self.comm.Get_size()
        self.simulations_per_rank = int(len(self.temperatures)/self.size)

        np.random.seed(self.rank+seed_offset)
        assert len(self.temperatures) % self.size == 0, "The number of temperatures must be dividable by the number of threads"

    def finalize(self):
        self.comm.Barrier()
        for sim in self.simulations:
            sim.finalize()

    def before_sim(self):
        '''
        Actions that should be performed on the
        simulation before starting the pamd-scheme.
        '''
        for sim in self.simulations:
            sim.before_sim()

    def after_sim(self):
        '''
        Actions that should be performed on the
        simulation after the pamd-scheme.
        '''
        for sim in self.simulations:
            sim.after_sim()

    def before_step(self):
        '''
        Actions that should be performed on the
        simulation before running the intermediate
        MD equilibration.
        '''
        for sim in self.simulations:
            sim.before_step()

    def after_step(self):
        '''
        Actions that should be performed on the
        simulation before running the intermediate
        MD equilibration.
        '''
        for sim in self.simulations:
            sim.after_step()

    def run(self, simulations):
        '''Run PAMD simulation, pass array of initial configurations'''
        self._setup_run(simulations)
        self.equilibrate()

        self.before_sim()

        for i in range(self.steps):
            self.before_step()
            for simulation in self.simulations:
                simulation.run_steps(self.simulation_steps, self.gpu_sim_object)
            self.comm.Barrier()

            pot_energies = np.empty(self.simulations_per_rank)
            for idx,sim in enumerate(self.simulations):
                pot_energies[idx]=sim.get_potential_energy()

            #get potential energies to rank 0
            if self.rank == 0:
                all_pot_energies = np.empty(len(self.temperatures))
            else:
                all_pot_energies = None

            self.comm.Gather(sendbuf=pot_energies, recvbuf=(all_pot_energies, self.simulations_per_rank), root=0)
            if self.rank==0:
                update_informations = []
                for i in range(len(self.temperatures)):
                    idx1,idx2 = np.random.choice(self.temperatures.shape[0], 2, replace=False)
                    temp1,temp2 = self.temperatures[idx1], self.temperatures[idx2]
                    d_beta = 1.0/temp2-1.0/temp1
                    eng1,eng2 = all_pot_energies[idx1], all_pot_energies[idx2]
                    d_eng = eng2-eng1
                    swap_prob = np.exp(self.boltzmann*d_beta*d_eng)
                    if np.random.uniform()<swap_prob:
                        update_informations.append([idx1,self.temperatures[idx1],self.temperatures[idx2]])
                        update_informations.append([idx2,self.temperatures[idx2],self.temperatures[idx1]])
                        self.temperatures[idx1],self.temperatures[idx2] = self.temperatures[idx2],self.temperatures[idx1]
                        all_pot_energies[idx1], all_pot_energies[idx2] = all_pot_energies[idx2], all_pot_energies[idx1]
            else:
                update_informations = None

            update_informations = self.comm.bcast(update_informations, root=0)
            for update_information in update_informations:
                if self.rank == int(update_information[0]/self.simulations_per_rank):
                    rescale_factor = np.sqrt(update_information[2]/update_information[1])
                    simulations[update_information[0]%self.simulations_per_rank].rescale_velocities(rescale_factor)
                    simulations[update_information[0]%self.simulations_per_rank].set_temperature(update_information[2], self.gpu_sim_object)
            self.after_step()

        self.after_sim()
        self.finalize()

    def __del__(self):
        MPI.Finalize()

    def equilibrate(self):
        for idx,simulation in enumerate(self.simulations):
            self.boltzmann = simulation.equilibrate(float(self.temperatures[int(idx+self.rank/self.size*len(self.temperatures))]), self.equilibration_steps, self.gpu_sim_object)

    def _setup_run(self, simulations):
        self.simulations = simulations