Coverage for gpaw/lcaotddft/timedensitymatrix.py: 35%

1import os 

2 

3import numpy as np 

4 

5from gpaw.lcaotddft.observer import TDDFTObserver 

6from gpaw.tddft.units import au_to_as 

7 

8 

9class TimeDensityMatrix(TDDFTObserver): 

10 

11 def __init__(self, paw, dmat, ksd, interval=1): 

12 TDDFTObserver.__init__(self, paw, interval) 

13 self.dmat = dmat 

14 self.ksd = ksd 

15 if self.ksd.world.rank == 0: 

16 self.dpath = 'rho' 

17 if not os.path.isdir(self.dpath): 

18 os.mkdir(self.dpath) 

19 

20 def _update(self, paw): 

21 time = paw.time 

22 rho_uMM = self.dmat.get_density_matrix((paw.niter, paw.action)) 

23 rho_up = self.ksd.transform(rho_uMM) 

24 if self.ksd.comm.rank == 0: 

25 fpath = os.path.join(self.dpath, 't%09.1f.npy' % (time * au_to_as)) 

26 np.save(fpath, rho_up[0])