Coverage for rhodent/writers/density.py: 98%

1from __future__ import annotations 

2 

3from time import strftime 

4import numpy as np 

5from numpy.typing import NDArray 

6from ase.io.cube import write_cube 

7 

8from ..density_matrices.frequency import FrequencyDensityMatrices 

9from ..density_matrices.base import WorkMetadata 

10from ..density_matrices.time import ConvolutionDensityMatrices 

11from ..utils import Result, get_gaussian_pulse_values 

12from ..voronoi import VoronoiWeights, EmptyVoronoiWeights 

13from .writer import Writer, ResultsCollector 

14 

15 

16class DensityWriter(Writer): 

17 

18 """ Calculate density contributions 

19 

20 Parameters 

21 ---------- 

22 collector 

23 ResultsCollector object 

24 """ 

25 

26 def __init__(self, 

27 collector: ResultsCollector): 

28 super().__init__(collector) 

29 if isinstance(self.density_matrices, ConvolutionDensityMatrices): 

30 self._ulm_tag = 'Time Density' 

31 assert len(self.density_matrices.pulses) == 1, 'Only one pulse allowed' 

32 else: 

33 assert isinstance(self.density_matrices, FrequencyDensityMatrices) 

34 self._ulm_tag = 'Frequency Density' 

35 

36 @property 

37 def voronoi(self) -> VoronoiWeights: 

38 return EmptyVoronoiWeights() 

39 

40 @property 

41 def common_arrays(self) -> dict[str, NDArray[np.float64] | NDArray[np.int64] | int | float]: 

42 from ..calculators import DensityCalculator 

43 

44 common = super().common_arrays 

45 common.pop('eig_i') 

46 common.pop('eig_a') 

47 common.pop('eig_n') 

48 common.pop('imin') 

49 common.pop('imax') 

50 common.pop('amin') 

51 common.pop('amax') 

52 

53 assert isinstance(self.calc, DensityCalculator) 

54 common['N_c'] = self.calc.N_c 

55 common['cell_cv'] = self.calc.cell_cv 

56 

57 atoms = self.density_matrices.ksd.atoms 

58 common['atom_numbers_i'] = atoms.get_atomic_numbers() 

59 common['atom_positions_iv'] = atoms.get_positions() 

60 

61 if isinstance(self.density_matrices, ConvolutionDensityMatrices): 

62 common['time_t'] = self.density_matrices.times * 1e-3 

63 else: 

64 assert isinstance(self.density_matrices, FrequencyDensityMatrices) 

65 common['freq_w'] = self.density_matrices.frequencies 

66 

67 if isinstance(self.density_matrices, ConvolutionDensityMatrices): 

68 # If pulse is Gaussian pulse, then get dictionary with 'pulsefreq' and 'pulsefwhm' 

69 common.update(**get_gaussian_pulse_values(self.density_matrices.pulses[0])) 

70 

71 return common 

72 

73 def fill_ulm(self, 

74 writer, 

75 work: WorkMetadata, 

76 result: Result): 

77 writer.fill(result['rho_g']) 

78 

79 def write_empty_arrays_ulm(self, writer): 

80 if isinstance(self.density_matrices, ConvolutionDensityMatrices): 

81 Nt = len(self.density_matrices.times) 

82 writer.add_array('rho_tg', (Nt, ) + self.calc.gdshape, dtype=float) 

83 else: 

84 assert isinstance(self.density_matrices, FrequencyDensityMatrices) 

85 Nw = len(self.density_matrices.frequencies) 

86 writer.add_array('rho_wg', (Nw, ) + self.calc.gdshape, dtype=float) 

87 

88 

89def write_density(out_fname: str, 

90 atoms, 

91 data, 

92 comment: str | None = None): 

93 """ Calculate density contribution and save in Gaussian cube file format. 

94 

95 Parameters 

96 ---------- 

97 out_fname 

98 File name of the resulting cube file. 

99 comment 

100 Comment line in the cube file. 

101 """ 

102 

103 if comment is None: 

104 comment = f'Cube file from rhodent, written on {strftime("%c")}' 

105 else: 

106 comment = comment.strip() 

107 

108 assert data is not None 

109 with open(out_fname, 'w') as fp: 

110 write_cube(fp, atoms, data=data, comment=comment)