Python trajectory.PickleTrajectory类代码示例

 def accept_move(self):
     e = self.atoms.get_potential_energy()
     self.lists['energy'].append(e)
     self.lists['counts'].append(1)
     self.lists['coordinations'].append(self.sumcoordinations(self.surfmove.coordinations))
     self.lists['types'].append(self.sumcoordinations(self.surfmove.types)) #See sumcoordinations
     self.lists['vac_coordinations'].append(self.sumcoordinations(self.surfmove.vacant_coordinations))
     self.lists['vac_types'].append(self.sumcoordinations(self.surfmove.vacant_types))
     
     
     if self.mc_step>=self.total_steps/2:
         if self.id_relax == None:
     	    self.id_relax = self.n_accept
         unrelax_energy, relax_energy = self.relax()
         self.lists['unrelaxed_energies'].append(unrelax_energy)
         self.lists['relaxed_energies'].append(relax_energy)
         e += relax_energy - unrelax_energy
     
         if e < self.ground_state_energy:
             #now store .amc in tmpfn[1] and 0000* in tmpfn[0]
             tmpfn = self.filename.split(".")
             traj = PickleTrajectory(tmpfn[0]+".traj", 'w', self.atoms, master=True, backup=False)
             traj.write()
             traj.close()
             #write(filename=tmpfn[0]+".traj",images=self.atoms)
         self.ground_state_energy = e
     else:
         self.lists['unrelaxed_energies'].append(np.nan)
         self.lists['relaxed_energies'].append(np.nan)
     self.mc_step += 1
     self.n_accept += 1

 def __init__(self, filename, mode='r', atoms=None, master=None,
              backup=True):
     if master is not None:
         if atoms.get_comm().rank != 0:
             raise NotImplementedError("It is required that the cpu with rank 0 is the master")
     _PickleTrajectory.__init__(self, filename, mode, atoms, master,
                                backup=backup)

 def eos(self, atoms, name):
     opts = self.opts
     
     traj = PickleTrajectory(self.get_filename(name, 'traj'), 'w', atoms)
     eps = 0.01
     strains = np.linspace(1 - eps, 1 + eps, 5)
     v1 = atoms.get_volume()
     volumes = strains**3 * v1
     energies = []
     cell1 = atoms.cell
     for s in strains:
         atoms.set_cell(cell1 * s, scale_atoms=True)
         energies.append(atoms.get_potential_energy())
         traj.write(atoms)
     traj.close()
     eos = EquationOfState(volumes, energies, opts.eos_type)
     v0, e0, B = eos.fit()
     atoms.set_cell(cell1 * (v0 / v1)**(1 / 3), scale_atoms=True)
     data = {'volumes': volumes,
             'energies': energies,
             'fitted_energy': e0,
             'fitted_volume': v0,
             'bulk_modulus': B,
             'eos_type': opts.eos_type}
     return data

def vasp2xyz():
    try:
        from lxml import etree
    except ImportError:
        print "You need to install python-lxml."
    print "start parse"
    xml = etree.parse("vasprun.xml")
    calculations = xml.xpath('//calculation')
    print 'len(calculations)=', len(calculations)
    from ase.io.trajectory import PickleTrajectory
    atoms = io.read('POSCAR')
    traj = PickleTrajectory('a.traj', 'w', atoms)

    print "start find forces and positions"
    allforce = []
    allpos = []

    def toarr(v):
        x = v.text.split()
        return map(float, x)
    for i, u in enumerate(calculations):
        print "step : ", i
        forces = map(toarr, u.xpath('./varray[@name="forces"]/v'))
        positions = map(toarr, u.xpath(
            './structure/varray[@name="positions"]/v'))
        atoms.set_scaled_positions(positions)
        allforce.append(forces)
        allpos.append(positions)
        traj.write()
    np.save('allforce.npy', allforce)
    np.save('allpos.npy', allpos)
    passthru("ase-gui a.traj -o a.xyz ")

def breed(particles,objects,gen):
	sortedIndices = numpy.argsort(numpy.asarray(objects))
	particlestobreed = []
	explambda = numpy.log(10) / (NPsPerGen - 1)
	for i in range(len(particles)):
		trialvar = numpy.random.random()
		if (trialvar <= numpy.exp(-1 * explambda * i)):
			particlestobreed.append(particles[sortedIndices[i]])
	NewGen = []
	for i in range(NPsPerGen):
		first = numpy.random.randint(len(particlestobreed))
		second = numpy.random.randint(len(particlestobreed))
		gene1 = numpy.random.randint(2)
		gene2 = numpy.random.randint(2)
		if (gene1 == 0):
			type1 = particlestobreed[first].gene[0]
		else:	
			type1 = particlestobreed[second].gene[0]
		if (gene2 == 0):
			type2 = particlestobreed[first].gene[1]
		else:
			type2 = particlestobreed[second].gene[1]
		if (particlestobreed[first].gene[2] < particlestobreed[second].gene[2]):
			minval = particlestobreed[first].gene[2]
			maxval = particlestobreed[second].gene[2]
		else:
			minval = particlestobreed[second].gene[2]
			maxval = particlestobreed[first].gene[2]
		numbertype1 = numpy.random.randint(minval, maxval + 1)
		NewGen.append(createparticle(type1,type2,numbertype1))
		traj = PickleTrajectory('generation'+str(gen)+ '_' + str(i) + '.traj','w')
		NewGen[i].gene = (type1,type2,numbertype1)
        	traj.write(ParticleList[i])
		traj.close()
	return NewGen

    def fit_volume(self, name, atoms, data=None):
        N, x = self.fit
        cell0 = atoms.get_cell()
        v = atoms.get_volume()
        if x > 0:
            strains = np.linspace(1 - x, 1 + x, N)
        else:
            strains = np.linspace(1 + x / v, 1 - x / v,  N)**(1./3)
        energies = []
        traj = PickleTrajectory(self.get_filename(name, 'fit.traj'), 'w')
        for s in strains:
            atoms.set_cell(cell0 * s, scale_atoms=True)
            energies.append(atoms.get_potential_energy())
            traj.write(atoms)

        traj.close()

        if data is not None:
            data['strains'] = strains
            data['energies'] = energies
        else:
            assert N % 2 == 1
            data = {'energy': energies[N // 2],
                    'strains': strains,
                    'energies': energies}

        return data

    def _test_timestepping(self, t):
        #XXX DEBUG START
        if debug and os.path.isfile('%s_%d.gpw' % (self.tdname, t)):
            return
        #XXX DEBUG END

        timestep = self.timesteps[t]
        self.assertAlmostEqual(self.duration % timestep, 0.0, 12)
        niter = int(self.duration / timestep)
        ndiv = 1 #XXX

        traj = PickleTrajectory('%s_%d.traj' % (self.tdname, t),
                                'w', self.tdcalc.get_atoms())

        t0 = time.time()
        f = paropen('%s_%d.log' % (self.tdname, t), 'w')
        print >>f, 'propagator: %s, duration: %6.1f as, timestep: %5.2f as, ' \
            'niter: %d' % (self.propagator, self.duration, timestep, niter)

        for i in range(1, niter+1):
            # XXX bare bones propagation without all the nonsense
            self.tdcalc.propagator.propagate(self.tdcalc.time,
                                             timestep * attosec_to_autime)
            self.tdcalc.time += timestep * attosec_to_autime
            self.tdcalc.niter += 1

            if i % ndiv == 0:
                rate = 60 * ndiv / (time.time()-t0)
                ekin = self.tdcalc.atoms.get_kinetic_energy()
                epot = self.tdcalc.get_td_energy() * Hartree
                F_av = np.zeros((len(self.tdcalc.atoms), 3))
                print >>f, 'i=%06d, time=%6.1f as, rate=%6.2f min^-1, ' \
                    'ekin=%13.9f eV, epot=%13.9f eV, etot=%13.9f eV' \
                    % (i, timestep * i, rate, ekin, epot, ekin + epot)
                t0 = time.time()

                # Hack to prevent calls to GPAW::get_potential_energy when saving
                spa = self.tdcalc.get_atoms()
                spc = SinglePointCalculator(epot, F_av, None, None, spa)
                spa.set_calculator(spc)
                traj.write(spa)
        f.close()
        traj.close()
        self.tdcalc.write('%s_%d.gpw' % (self.tdname, t), mode='all')

        # Save density and wavefunctions to binary
        gd, finegd = self.tdcalc.wfs.gd, self.tdcalc.density.finegd
        if world.rank == 0:
            big_nt_g = finegd.collect(self.tdcalc.density.nt_g)
            np.save('%s_%d_nt.npy' % (self.tdname, t), big_nt_g)
            del big_nt_g

            big_psit_nG = gd.collect(self.tdcalc.wfs.kpt_u[0].psit_nG)
            np.save('%s_%d_psit.npy' % (self.tdname, t), big_psit_nG)
            del big_psit_nG
        else:
            finegd.collect(self.tdcalc.density.nt_g)
            gd.collect(self.tdcalc.wfs.kpt_u[0].psit_nG)
        world.barrier()

    def do_calculations(self, formulas):
        """Perform calculation on molecules, write results to .gpw files."""
        atoms = {}
        for formula in formulas:
            for symbol in string2symbols(formula.split('_')[0]):
                atoms[symbol] = None
        formulas = formulas + atoms.keys()

        for formula in formulas:
            if path.isfile(formula + '.gpw'):
                continue

            barrier()
            open(formula + '.gpw', 'w')
            s = molecule(formula)
            s.center(vacuum=self.vacuum)
            cell = s.get_cell()
            h = self.h
            s.set_cell((cell / (4 * h)).round() * 4 * h)
            s.center()
            calc = GPAW(h=h,
                        xc=self.xc,
                        eigensolver=self.eigensolver,
                        setups=self.setups,
                        basis=self.basis,
                        fixmom=True,
                        txt=formula + '.txt')

            if len(s) == 1:
                calc.set(hund=True)

            s.set_calculator(calc)

            if formula == 'BeH':
                calc.initialize(s)
                calc.nuclei[0].f_si = [(1, 0, 0.5, 0, 0),
                                       (0.5, 0, 0, 0, 0)]

            if formula in ['NO', 'ClO', 'CH']:
                s.positions[:, 1] += h * 1.5

            try:
                energy = s.get_potential_energy()
            except (RuntimeError, ConvergenceError):
                if rank == 0:
                    print >> sys.stderr, 'Error in', formula
                    traceback.print_exc(file=sys.stderr)
            else:
                print >> self.txt, formula, repr(energy)
                self.txt.flush()
                calc.write(formula)

            if formula in diatomic and self.calculate_dimer_bond_lengths:
                traj = PickleTrajectory(formula + '.traj', 'w')
                d = diatomic[formula][1]
                for x in range(-2, 3):
                    s.set_distance(0, 1, d * (1.0 + x * 0.02))
                    traj.write(s)

def mainBasinLoop(symbol1, symbol2, elementN1, elementN2, numberOfType1, numberOfType2, radius, percentType1):
  """symbol1 and symbol2 should be STRINGS.
  elementN1 and elementN2 should be ATOMIC NUMBERS of those elements.
  numberOfType1 and numberOfType2 are the actual # of those atoms in the molecule
  radius can be a decimal
  percent should be expressed as a number between 0 and 1.

  After that, it will create five new sets of 400 molecules,
  each set being the result of that first set having one of our small perturbations
  performed on it, which are also then minimized, perturbed, etc., until 50,000
  total original orientations have been minimized."""

  bigKickResults, round1PE = [], []
  global finalList

  creationString = symbol1 + str(numberOfType1) + symbol2 + str(numberOfType2)
  creationString2 = creationString + '.traj'

  baseAtom = nearlySphericalAtom(str(creationString),radius,numberOfType1+numberOfType2)
  baseAtom.set_pbc((1,1,1))
  baseAtom.set_cell((100,100,100))
  calc = QSC()
  baseAtom.set_calculator(calc)


  baseAtom = makeBimetallic(baseAtom,numberOfType1+numberOfType2,elementN1,elementN2,percentType1)
  baseAtom = preventExplosions(baseAtom)
  baseAtom = preventExplosions(baseAtom)

  for x in range(listLength):
    bigKickResults.append(baseAtom.copy())

  for x in range(listLength/2):
    bigKickResults[x] = shake(bigKickResults[x])

  for x in range(listLength/2, listLength):
    bigKickResults[x] = switchAtoms(bigKickResults[x])

  for x in range(len(bigKickResults)):
    bigKickResults[x] = optimizeMolecule(bigKickResults[x],FIREMoves,creationString2)
    round1PE.append(bigKickResults[x].get_potential_energy())

  minimumPE = min(round1PE)
  minimumIndex = round1PE.index(minimumPE)
  bestAtom = bigKickResults[minimumIndex]

  minimaList = PickleTrajectory(str(creationString2),mode='a')
  minimaList.write(bestAtom)
  minimaList.close()

  finalList.append(bestAtom.copy())

  smallKicks(bigKickResults,0,creationString2)

  veryBestAtom = returnFinalBestAtom(str(creationString2), str(creationString))

  return veryBestAtom

 def write_optimized(self, atoms, energy):
     magmoms = None  # XXX we could do better ...
     forces = np.zeros((len(atoms), 3))
     calc = SinglePointCalculator(energy, forces, np.zeros((3, 3)),
                                  magmoms, atoms)
     atoms.set_calculator(calc)
     filename = '%s%s-optimized.traj' % (self.name, self.tag)
     traj = PickleTrajectory(filename, 'w', backup=False)
     traj.write(atoms)
     traj.close()

 def write_mode(self, n, kT=units.kB * 300, nimages=30):
     """Write mode to trajectory file."""
     mode = self.get_mode(n) * sqrt(kT / self.hnu[n])
     p = self.atoms.positions.copy()
     n %= 3 * len(self.indices)
     traj = PickleTrajectory('%s.%d.traj' % (self.name, n), 'w')
     calc = self.atoms.get_calculator()
     self.atoms.set_calculator()
     for x in np.linspace(0, 2 * pi, nimages, endpoint=False):
         self.atoms.set_positions(p + sin(x) * mode)
         traj.write(self.atoms)
     self.atoms.set_positions(p)
     self.atoms.set_calculator(calc)
     traj.close()

    def __init__(self, atoms,
                 temperature=100 * kB,
                 optimizer=FIRE,
                 fmax=0.1,
                 dr=0.1,
                 logfile='-', 
                 trajectory='lowest.traj',
                 optimizer_logfile='-',
                 local_minima_trajectory='local_minima.traj',
                 adjust_cm=True):
        Dynamics.__init__(self, atoms, logfile, trajectory)
        self.kT = temperature
        self.optimizer = optimizer
        self.fmax = fmax
        self.dr = dr
        if adjust_cm:
            self.cm = atoms.get_center_of_mass()
        else:
            self.cm = None

        self.optimizer_logfile = optimizer_logfile
        self.lm_trajectory = local_minima_trajectory
        if isinstance(local_minima_trajectory, str):
            self.lm_trajectory = PickleTrajectory(local_minima_trajectory,
                                                  'w', atoms)

        self.initialize()

    def fit_bond_length(self, name, atoms, data):
        N, x = self.fit
        assert N % 2 == 1
        d0 = atoms.get_distance(0, 1)
        distances = np.linspace(d0 * (1 - x), d0 * (1 + x), N)
        energies = []
        traj = PickleTrajectory(self.get_filename(name, 'fit.traj'), 'w')
        for d in distances:
            atoms.set_distance(0, 1, d)
            energies.append(atoms.get_potential_energy())
            self.check_occupation_numbers(atoms)
            traj.write(atoms)

        traj.close()

        data['distances'] = distances
        data['energies'] = energies

def optimizeMolecule(molecule,NMoves,creationString):

  bestEnergy = 0.
  totalMinimaFound = 0
  sinceLastFind = 0

  calc = QSC()
  molecule.set_calculator(calc)

  for i in range(NMoves):
        molecule = newMove(molecule)
        molecule = preventExplosions(molecule)
        molecule.center()
        sinceLastFind += 1
        # do a last optimization of the structure
        dyn = FIRE(molecule)
        dyn.run(steps = 2000)
        newEnergy = molecule.get_potential_energy()

        if (newEnergy < bestEnergy):
                optimizedMolecule = molecule
                bestEnergy = newEnergy
                line = str(totalMinimaFound) + "  " + str(molecule.get_potential_energy()) + "  " + str(i) +"\n"
                print line
                f = open('EnergyList.txt','a')
                f.write(line)
                f.close()
                minimaList = PickleTrajectory(str(creationString),mode='a')
                minimaList.write(optimizedMolecule)
                totalMinimaFound += 1
                minimaList.close()
                sinceLastFind = 0
        elif (sinceLastFind < 200):
          break


  minimaList.close()

  minimaList = PickleTrajectory(str(creationString),mode='r')

  atomslist = [atom for atom in minimaList]
  ase.io.write('movie.xyz',atomslist,format='xyz') # write a movie file of our dynamics

  minimaList.close()

  return optimizedMolecule

    def MD(self):
        """Molecular Dynamic"""
        from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
        from ase import units
        from ase.md import MDLogger
        from ase.io.trajectory import PickleTrajectory
        from ase.md.langevin import Langevin
        from ase.md.verlet import VelocityVerlet

        dyndrivers = {
            'Langevin': Langevin,
            'None': VelocityVerlet,
        }

        useAsap = False

        mol = self.mol
        temperature = self.definedParams['temperature']
        init_temperature = self.definedParams['init_temperature']
        time_step = self.definedParams['time_step']
        nstep = self.definedParams['nstep']
        nprint = self.definedParams['nprint']
        thermostat = self.definedParams['thermostat']
        prop_file = os.path.join(self.definedParams['workdir'],
                                 self.definedParams['output_prefix']+'.out')
        traj_file = os.path.join(self.definedParams['workdir'],
                                 self.definedParams['output_prefix']+'.traj')

        MaxwellBoltzmannDistribution(mol,init_temperature*units.kB)

        if thermostat == 'None':
            dyn = VelocityVerlet(mol, time_step*units.fs)
        elif thermostat == 'Langevin':
            dyn = Langevin(mol, time_step*units.fs, temperature*units.kB, 0.01 )
        else:
            raise ImplementationError(method,'Thermostat is not implemented in the MD function')

        #Function to print the potential, kinetic and total energy
        traj = PickleTrajectory(traj_file,"a",mol)
        dyn.attach(MDLogger(dyn,mol,prop_file),interval=nprint)
        dyn.attach(traj.write, interval = nprint)

        dyn.run(nstep)
        traj.close()