git-svn-id: svn://svn.icms.temple.edu/lammps-ro/trunk@16053 f3b2605a-c512-4ea7-a41b...
[lammps.git] / tools / i-pi / ipi / engine / forces.py
blob58987f1e4dcb8fd5ee1e4ed56d2f75f7f341bd6b
1 """Contains the classes that connect the driver to the python code.
3 Copyright (C) 2013, Joshua More and Michele Ceriotti
5 This program is free software: you can redistribute it and/or modify
6 it under the terms of the GNU General Public License as published by
7 the Free Software Foundation, either version 3 of the License, or
8 (at your option) any later version.
10 This program is distributed in the hope that it will be useful,
11 but WITHOUT ANY WARRANTY; without even the implied warranty of
12 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13 GNU General Public License for more details.
15 You should have received a copy of the GNU General Public License
16 along with this program. If not, see <http.//www.gnu.org/licenses/>.
19 Communicates with the driver code, obtaining the force, virial and potential.
20 Deals with creating the jobs that will be sent to the driver, and
21 returning the results to the python code.
23 Classes:
24 ForceField: Base forcefield class with the generic methods and attributes.
25 FFSocket: Deals with a single replica of the system
26 ForceBeads: Deals with the parallelization of the force calculation over
27 different beads.
28 Forces: Deals with the parallelizatoin of the force calculation over
29 different forcefields.
30 """
32 __all__ = ['ForceField', 'ForceBeads', 'Forces', 'FFSocket']
34 import numpy as np
35 import time
36 from ipi.utils.softexit import softexit
37 from ipi.utils.messages import verbosity, warning
38 from ipi.utils.depend import *
39 from ipi.utils.nmtransform import nm_rescale
40 from ipi.interfaces.sockets import InterfaceSocket
41 from ipi.engine.beads import Beads
43 class ForceField(dobject):
44 """Base forcefield class.
46 Gives the standard methods and quantities needed in all the forcefield
47 classes.
49 Attributes:
50 atoms: An Atoms object containing all the atom positions.
51 cell: A Cell object containing the system box.
53 Depend objects:
54 ufvx: A list of the form [pot, f, vir]. These quantities are calculated
55 all at one time by the driver, so are collected together. Each separate
56 object is then taken from the list. Depends on the atom positions and
57 the system box.
58 extra: A string containing some formatted output returned by the client. Depends on ufvx.
59 pot: A float giving the potential energy of the system. Depends on ufvx.
60 f: An array containing all the components of the force. Depends on ufvx.
61 fx: A slice of f containing only the x components of the forces.
62 fy: A slice of f containing only the y components of the forces.
63 fz: A slice of f containing only the z components of the forces.
64 vir: An array containing the components of the virial tensor in upper
65 triangular form, not divided by the volume. Depends on ufvx.
66 """
68 def __init__(self):
69 """Initialises ForceField."""
71 # ufvx is a list [ u, f, vir, extra ] which stores the results of the force
72 #calculation
73 dset(self,"ufvx", depend_value(name="ufvx", func=self.get_all))
75 def copy(self):
76 """Creates a deep copy without the bound objects.
78 Used in ForceBeads to create a ForceField for each replica of the system.
80 Returns:
81 A ForceField object without atoms or cell attributes.
82 """
84 return type(self)(self.nbeads, self.weight)
86 def bind(self, atoms, cell):
87 """Binds atoms and cell to the forcefield.
89 This takes an atoms object and a cell object and makes them members of
90 the forcefield. It also then creates the objects that will hold the data
91 that the driver returns and the dependency network.
93 Args:
94 atoms: The Atoms object from which the atom positions are taken.
95 cell: The Cell object from which the system box is taken.
96 """
98 # stores a reference to the atoms and cell we are computing forces for
99 self.atoms = atoms
100 self.cell = cell
102 # ufv depends on the atomic positions and on the cell
103 dget(self,"ufvx").add_dependency(dget(self.atoms,"q"))
104 dget(self,"ufvx").add_dependency(dget(self.cell,"h"))
106 # potential and virial are to be extracted very simply from ufv
107 dset(self,"pot",
108 depend_value(name="pot", func=self.get_pot,
109 dependencies=[dget(self,"ufvx")]))
111 dset(self,"vir",
112 depend_array(name="vir", value=np.zeros((3,3),float),func=self.get_vir,
113 dependencies=[dget(self,"ufvx")]))
115 # NB: the force requires a bit more work, to define shortcuts to xyz
116 # slices without calculating the force at this point.
117 fbase = np.zeros(atoms.natoms*3, float)
118 dset(self,"f",
119 depend_array(name="f", value=fbase, func=self.get_f,
120 dependencies=[dget(self,"ufvx")]))
122 dset(self,"extra",
123 depend_value(name="extra", func=self.get_extra,
124 dependencies=[dget(self,"ufvx")]))
126 dset(self,"fx", depend_array(name="fx", value=fbase[0:3*atoms.natoms:3]))
127 dset(self,"fy", depend_array(name="fy", value=fbase[1:3*atoms.natoms:3]))
128 dset(self,"fz", depend_array(name="fz", value=fbase[2:3*atoms.natoms:3]))
129 depcopy(self,"f", self,"fx")
130 depcopy(self,"f", self,"fy")
131 depcopy(self,"f", self,"fz")
133 def queue(self):
134 """Dummy queueing method."""
136 pass
138 def stop(self):
139 """Dummy queueing method."""
141 pass
143 def run(self):
144 """Dummy queueing method."""
146 pass
148 def get_all(self):
149 """Dummy driver routine.
151 Returns:
152 A list of the form [potential, force, virial] where the potential
153 and all components of the force and virial have been set to zero.
156 return [0.0, np.zeros(3*self.atoms.natoms), np.zeros((3,3),float), ""]
158 def get_pot(self):
159 """Calls get_all routine of forcefield to update potential.
161 Returns:
162 Potential energy.
165 return self.ufvx[0]
167 def get_f(self):
168 """Calls get_all routine of forcefield to update force.
170 Returns:
171 An array containing all the components of the force.
174 return depstrip(self.ufvx[1])
176 def get_vir(self):
177 """Calls get_all routine of forcefield to update virial.
179 Returns:
180 An array containing the virial in upper triangular form, not divided
181 by the volume.
184 vir = depstrip(self.ufvx[2])
185 vir[1,0] = 0.0
186 vir[2,0:2] = 0.0
187 return vir
189 def get_extra(self):
190 """Calls get_all routine of forcefield to update potential.
192 Returns:
193 A string containing all formatted additional output that the
194 client might have produced.
197 return self.ufvx[3]
200 class FFSocket(ForceField):
201 """Interface between the PIMD code and the socket for a single replica.
203 Deals with an individual replica of the system, obtaining the potential
204 force and virial appropriate to this system. Deals with the distribution of
205 jobs to the interface.
207 Attributes:
208 parameters: A dictionary of the parameters used by the driver. Of the
209 form {'name': value}.
210 socket: The interface object which contains the socket through which
211 communication between the forcefield and the driver is done.
212 request: During the force calculation step this holds a dictionary
213 containing the relevant data for determining the progress of the step.
214 Of the form {'atoms': atoms, 'cell': cell, 'pars': parameters,
215 'status': status, 'result': result, 'id': bead id,
216 'start': starting time}.
219 def __init__(self, pars=None, interface=None):
220 """Initialises FFSocket.
222 Args:
223 pars: Optional dictionary, giving the parameters needed by the driver.
224 interface: Optional Interface object, which contains the socket.
227 # a socket to the communication library is created or linked
228 super(FFSocket,self).__init__()
229 if interface is None:
230 self.socket = InterfaceSocket()
231 else:
232 self.socket = interface
234 if pars is None:
235 self.pars = {}
236 else:
237 self.pars = pars
238 self.request = None
240 def bind(self, atoms, cell):
241 """Pass on the binding request from ForceBeads.
243 Also makes sure to set the socket's softexit.
245 Args:
246 atoms: Atoms object from which the bead positions are taken.
247 cell: Cell object from which the system box is taken.
250 super(FFSocket,self).bind(atoms, cell)
252 def copy(self):
253 """Creates a deep copy without the bound objects.
255 Used in ForceBeads to create a FFSocket for each replica of the system.
257 Returns:
258 A FFSocket object without atoms or cell attributes.
261 # does not copy the bound objects
262 # (i.e., the returned forcefield must be bound before use)
263 return type(self)(self.pars, self.socket)
265 def get_all(self):
266 """Driver routine.
268 When one of the force, potential or virial are called, this sends the
269 atoms and cell to the driver through the interface, requesting that the
270 driver does the calculation. This then waits until the driver is finished,
271 and then returns the ufvx list.
273 Returns:
274 A list of the form [potential, force, virial, extra].
277 # this is converting the distribution library requests into [ u, f, v ] lists
278 if self.request is None:
279 self.request = self.socket.queue(self.atoms, self.cell, pars=self.pars, reqid=-1)
280 while self.request["status"] != "Done":
281 if self.request["status"] == "Exit":
282 break
283 time.sleep(self.socket.latency)
284 if self.request["status"] == "Exit":
285 softexit.trigger(" @Force: Requested returned a Exit status")
287 # data has been collected, so the request can be released and a slot
288 #freed up for new calculations
289 self.socket.release(self.request)
290 result = self.request["result"]
291 self.request = None
293 return result
295 def queue(self, reqid=-1):
296 """Sends the job to the interface queue directly.
298 Allows the ForceBeads object to ask for the ufvx list of each replica
299 directly without going through the get_all function. This allows
300 all the jobs to be sent at once, allowing them to be parallelized.
302 Args:
303 reqid: An optional integer that indentifies requests of the same type,
304 e.g. the bead index.
307 if self.request is None and dget(self,"ufvx").tainted():
308 self.request = self.socket.queue(self.atoms, self.cell, pars=self.pars, reqid=reqid)
310 def run(self):
311 """Makes the socket start looking for driver codes.
313 Tells the interface code to start the thread that looks for
314 connection from the driver codes in a loop. Until this point no
315 jobs can be queued.
318 if not self.socket.started():
319 self.socket.start_thread()
321 def stop(self):
322 """Makes the socket stop looking for driver codes.
324 Tells the interface code to stop the thread that looks for
325 connection from the driver codes in a loop. After this point no
326 jobs can be queued.
329 if self.socket.started():
330 self.socket.end_thread()
333 class ForceBeads(dobject):
334 """Class that gathers the forces for each replica together.
336 Deals with splitting the bead representation into
337 separate replicas, and collecting the data from each replica.
339 Attributes:
340 natoms: An integer giving the number of atoms.
341 nbeads: An integer giving the number of beads.
342 f_model: A model used to create the forcefield objects for each replica
343 of the system.
344 _forces: A list of the forcefield objects for all the replicas.
345 weight: A float that will be used to weight the contribution of this
346 forcefield to the total force.
348 Depend objects:
349 f: An array containing the components of the force. Depends on each
350 replica's ufvx list.
351 pots: A list containing the potential energy for each system replica.
352 Depends on each replica's ufvx list.
353 virs: A list containing the virial tensor for each system replica.
354 Depends on each replica's ufvx list.
355 pot: The sum of the potential energy of the replicas.
356 vir: The sum of the virial tensor of the replicas.
357 extras: Strings containing some formatted output returned by the client.
358 Depends on each replica's ufvx list.
361 def __init__(self, model, nbeads=0, weight=1.0):
362 """Initializes ForceBeads
364 Args:
365 model: A model to be used to create the forcefield objects for all
366 the replicas of the system.
367 nbeads: The number of replicas.
368 weight: A relative weight to be given to the values obtained with this
369 forcefield. When the contribution of all the forcefields is
370 combined to give a total force, the contribution of this forcefield
371 will be weighted by this factor.
374 self.f_model = model
375 self.nbeads = nbeads
376 self.weight = weight
378 def copy(self):
379 """Creates a deep copy without the bound objects.
381 Used so that we can create multiple Forces objects from the same
382 Forcebeads model, without binding a particular ForceBeads object twice.
384 Returns:
385 A ForceBeads object without beads or cell attributes.
388 # does not copy the bound objects (i.e., the returned forcefield must be bound before use)
389 return type(self)(self.f_model, self.nbeads, self.weight)
392 def bind(self, beads, cell):
393 """Binds beads, cell and force to the forcefield.
395 Takes the beads, cell objects and makes them members of the forcefield.
396 Also takes the force object and copies it once for each replica of the
397 system, then binds each replica to one of the copies so that the force
398 calculation can be parallelized. Creates the objects that will
399 hold the data that the driver returns and the dependency network.
401 Args:
402 beads: Beads object from which the bead positions are taken.
403 cell: Cell object from which the system box is taken.
406 # stores a copy of the number of atoms and of beads
407 #!TODO! make them read-only properties
408 self.natoms = beads.natoms
409 if (self.nbeads != beads.nbeads):
410 raise ValueError("Binding together a Beads and a ForceBeads objects with different numbers of beads")
412 # creates an array of force objects, which are bound to the beads
413 #and the cell
414 self._forces = [];
415 for b in range(self.nbeads):
416 new_force = self.f_model.copy()
417 new_force.bind(beads[b], cell)
418 self._forces.append(new_force)
420 # f is a big array which assembles the forces on individual beads
421 dset(self,"f",
422 depend_array(name="f",value=np.zeros((self.nbeads,3*self.natoms)),
423 func=self.f_gather,
424 dependencies=[dget(self._forces[b],"f") for b in range(self.nbeads)]))
426 # collection of pots and virs from individual beads
427 dset(self,"pots",
428 depend_array(name="pots", value=np.zeros(self.nbeads,float),
429 func=self.pot_gather,
430 dependencies=[dget(self._forces[b],"pot") for b in range(self.nbeads)]))
431 dset(self,"virs",
432 depend_array(name="virs", value=np.zeros((self.nbeads,3,3),float),
433 func=self.vir_gather,
434 dependencies=[dget(self._forces[b],"vir") for b in range(self.nbeads)]))
435 dset(self,"extras",
436 depend_value(name="extras", value=np.zeros(self.nbeads,float),
437 func=self.extra_gather,
438 dependencies=[dget(self._forces[b],"extra") for b in range(self.nbeads)]))
440 # total potential and total virial
441 dset(self,"pot",
442 depend_value(name="pot", func=(lambda: self.pots.sum()),
443 dependencies=[dget(self,"pots")]))
444 dset(self,"vir",
445 depend_array(name="vir", func=self.get_vir, value=np.zeros((3,3)),
446 dependencies=[dget(self,"virs")]))
448 def run(self):
449 """Makes the socket start looking for driver codes.
451 Tells the interface code to start the thread that looks for
452 connection from the driver codes in a loop. Until this point no
453 jobs can be queued.
456 for b in range(self.nbeads):
457 self._forces[b].run()
459 def stop(self):
460 """Makes the socket stop looking for driver codes.
462 Tells the interface code to stop the thread that looks for
463 connection from the driver codes in a loop. After this point no
464 jobs can be queued.
467 for b in range(self.nbeads):
468 self._forces[b].stop()
470 def queue(self):
471 """Submits all the required force calculations to the interface."""
473 # this should be called in functions which access u,v,f for ALL the beads,
474 # before accessing them. it is basically pre-queueing so that the
475 # distributed-computing magic can work
476 for b in range(self.nbeads):
477 self._forces[b].queue(reqid=b)
479 def pot_gather(self):
480 """Obtains the potential energy for each replica.
482 Returns:
483 A list of the potential energy of each replica of the system.
486 self.queue()
487 return np.array([b.pot for b in self._forces], float)
489 def extra_gather(self):
490 """Obtains the potential energy for each replica.
492 Returns:
493 A list of the potential energy of each replica of the system.
496 self.queue()
497 return [b.extra for b in self._forces]
499 def vir_gather(self):
500 """Obtains the virial for each replica.
502 Returns:
503 A list of the virial of each replica of the system.
506 self.queue()
507 return np.array([b.vir for b in self._forces], float)
509 def f_gather(self):
510 """Obtains the force vector for each replica.
512 Returns:
513 An array with all the components of the force. Row i gives the force
514 array for replica i of the system.
517 newf = np.zeros((self.nbeads,3*self.natoms),float)
519 self.queue()
520 for b in range(self.nbeads):
521 newf[b] = depstrip(self._forces[b].f)
523 return newf
525 #serial
526 # for b in range(self.nbeads): newf[b]=self._forces[b].f
527 # threaded
528 # bthreads=[]
529 # print "starting threads"
530 # for b in range(self.nbeads):
531 # thread=threading.Thread(target=self._getbead, args=(b,newf,))
532 # thread.start()
533 # bthreads.append(thread)
535 # print "waiting threads"
536 # for b in range(self.nbeads): bthreads[b].join()
537 # print "threads joined in"
539 def get_vir(self):
540 """Sums the virial of each replica.
542 Not the actual system virial, as it has not been divided by either the
543 number of beads or the cell volume.
545 Returns:
546 Virial sum.
549 vir = np.zeros((3,3))
550 for v in depstrip(self.virs):
551 vir += v
552 return vir
554 def __len__(self):
555 """Length function.
557 This is called whenever the standard function len(forcebeads) is used.
559 Returns:
560 The number of beads.
563 return self.nbeads
565 def __getitem__(self,index):
566 """Overwrites standard getting function.
568 This is called whenever the standard function forcebeads[index] is used.
569 Returns the force on bead index.
571 Args:
572 index: The index of the replica of the system to be accessed.
574 Returns:
575 The forces acting on the replica of the system given by the index.
578 return self._forces[index]
581 class Forces(dobject):
582 """Class that gathers all the forces together.
584 Collects many forcefield instances and parallelizes getting the forces
585 in a PIMD environment.
587 Attributes:
588 natoms: An integer giving the number of atoms.
589 nbeads: An integer giving the number of beads.
590 nforces: An integer giving the number of ForceBeads objects.
591 mforces: A list of all the forcefield objects.
592 mbeads: A list of all the beads objects. Some of these may be contracted
593 ring polymers, with a smaller number of beads than of the simulation.
594 mweights: A list of the weights of all the forcefields.
595 mrpc: A list of the objects containing the functions required to
596 contract the ring polymers of the different forcefields.
598 Depend objects:
599 f: An array containing the components of the force. Depends on each
600 replica's ufvx list.
601 pots: A list containing the potential energy for each system replica.
602 Depends on each replica's ufvx list.
603 virs: A list containing the virial tensor for each system replica.
604 Depends on each replica's ufvx list.
605 extras: A list containing the "extra" strings for each replica.
606 pot: The sum of the potential energy of the replicas.
607 vir: The sum of the virial tensor of the replicas.
610 def bind(self, beads, cell, flist):
612 self.natoms = beads.natoms
613 self.nbeads = beads.nbeads
614 self.nforces = len(flist)
616 # flist should be a list of tuples containing ( "name", forcebeads)
617 self.mforces = []
618 self.mbeads = []
619 self.mweights = []
620 self.mrpc = []
622 # a "function factory" to generate functions to automatically update
623 #contracted paths
624 def make_rpc(rpc, beads):
625 return lambda: rpc.b1tob2(depstrip(beads.q))
627 # creates new force objects, possibly acting on contracted path
628 #representations
629 for (ftype, fbeads) in flist:
631 # creates an automatically-updated contracted beads object
632 newb = fbeads.nbeads
633 newforce = fbeads.copy()
634 newweight = fbeads.weight
636 # if the number of beads for this force component is unspecified,
637 #assume full force evaluation
638 if newb == 0:
639 newb = beads.nbeads
640 newforce.nbeads = newb
642 newbeads = Beads(beads.natoms, newb)
643 newrpc = nm_rescale(beads.nbeads, newb)
645 dget(newbeads,"q")._func = make_rpc(newrpc, beads)
646 for b in newbeads:
647 # must update also indirect access to the beads coordinates
648 dget(b,"q")._func = dget(newbeads,"q")._func
650 # makes newbeads.q depend from beads.q
651 dget(beads,"q").add_dependant(dget(newbeads,"q"))
653 #now we create a new forcebeads which is bound to newbeads!
654 newforce.bind(newbeads, cell)
656 #adds information we will later need to the appropriate lists.
657 self.mweights.append(newweight)
658 self.mbeads.append(newbeads)
659 self.mforces.append(newforce)
660 self.mrpc.append(newrpc)
662 #now must expose an interface that gives overall forces
663 dset(self,"f",
664 depend_array(name="f",value=np.zeros((self.nbeads,3*self.natoms)),
665 func=self.f_combine,
666 dependencies=[dget(ff, "f") for ff in self.mforces] ) )
668 # collection of pots and virs from individual ff objects
669 dset(self,"pots",
670 depend_array(name="pots", value=np.zeros(self.nbeads,float),
671 func=self.pot_combine,
672 dependencies=[dget(ff, "pots") for ff in self.mforces]) )
674 # must take care of the virials!
675 dset(self,"virs",
676 depend_array(name="virs", value=np.zeros((self.nbeads,3,3),float),
677 func=self.vir_combine,
678 dependencies=[dget(ff, "virs") for ff in self.mforces]) )
680 dset(self,"extras",
681 depend_value(name="extras", value=np.zeros(self.nbeads,float),
682 func=self.extra_combine,
683 dependencies=[dget(ff, "extras") for ff in self.mforces]))
685 # total potential and total virial
686 dset(self,"pot",
687 depend_value(name="pot", func=(lambda: self.pots.sum()),
688 dependencies=[dget(self,"pots")]))
689 dset(self,"vir",
690 depend_array(name="vir", func=self.get_vir, value=np.zeros((3,3)),
691 dependencies=[dget(self,"virs")]))
693 def run(self):
694 """Makes the socket start looking for driver codes.
696 Tells the interface code to start the thread that looks for
697 connection from the driver codes in a loop. Until this point no
698 jobs can be queued.
701 for ff in self.mforces:
702 ff.run()
704 def stop(self):
705 """Makes the socket stop looking for driver codes.
707 Tells the interface code to stop the thread that looks for
708 connection from the driver codes in a loop. After this point no
709 jobs can be queued.
712 for ff in self.mforces:
713 ff.stop()
715 def queue(self):
716 """Submits all the required force calculations to the forcefields."""
718 for ff in self.mforces:
719 ff.queue()
721 def get_vir(self):
722 """Sums the virial of each forcefield.
724 Not the actual system virial.
726 Returns:
727 Virial sum.
730 vir = np.zeros((3,3))
731 for v in depstrip(self.virs):
732 vir += v
733 return vir
735 def f_combine(self):
736 """Obtains the total force vector."""
738 self.queue()
739 rf = np.zeros((self.nbeads,3*self.natoms),float)
740 for k in range(self.nforces):
741 # "expand" to the total number of beads the forces from the
742 #contracted one
743 rf += self.mweights[k]*self.mrpc[k].b2tob1(depstrip(self.mforces[k].f))
744 return rf
746 def pot_combine(self):
747 """Obtains the potential energy for each forcefield."""
749 self.queue()
750 rp = np.zeros(self.nbeads,float)
751 for k in range(self.nforces):
752 # "expand" to the total number of beads the potentials from the
753 #contracted one
754 rp += self.mweights[k]*self.mrpc[k].b2tob1(self.mforces[k].pots)
755 return rp
757 def extra_combine(self):
758 """Obtains the potential energy for each forcefield."""
760 self.queue()
761 rp = [ "" for b in range(self.nbeads) ]
762 for k in range(self.nforces):
763 # "expand" to the total number of beads the potentials from the
764 #contracted one
765 for b in range(self.nbeads):
766 rp[b] += self.mforces[k].extras[b]
767 return rp
769 def vir_combine(self):
770 """Obtains the virial tensor for each forcefield."""
772 self.queue()
773 rp = np.zeros((self.nbeads,3,3),float)
774 for k in range(self.nforces):
775 virs = depstrip(self.mforces[k].virs)
776 # "expand" to the total number of beads the virials from the
777 #contracted one, element by element
778 for i in range(3):
779 for j in range(3):
780 rp[:,i,j] += self.mweights[k]*self.mrpc[k].b2tob1(virs[:,i,j])
781 return rp