1+ import numpy as np
2+ import tqdm
3+ from typing import Callable , Optional
4+
5+ from sisl import BrillouinZone , DensityMatrix , get_distribution , unit_convert
6+ from ._compute_dm import add_cnc_diag_spin , add_cnc_nc
7+
8+ def compute_dm (bz : BrillouinZone , occ_distribution : Optional [Callable ] = None ,
9+ occtol : float = 1e-9 , fermi_dirac_T : float = 300. , eta : bool = True ):
10+ """Computes the DM from the eigenstates of a Hamiltonian along a BZ.
11+
12+ Parameters
13+ ----------
14+ bz: BrillouinZone
15+ The brillouin zone object containing the Hamiltonian of the system
16+ and the k-points to be sampled.
17+ occ_distribution: function, optional
18+ The distribution that will determine the occupations of states. It will
19+ receive an array of energies (in eV, referenced to fermi level) and it should
20+ return an array of floats.
21+ If not provided, a fermi_dirac distribution will be considered, being the
22+ fermi_dirac_T parameter the electronic temperature.
23+ occtol: float, optional
24+ Threshold below which the contribution of a state is not even added to the
25+ DM.
26+ fermi_dirac_T: float, optional
27+ If an occupation distribution is not provided, a fermi-dirac distribution centered
28+ at the chemical potential is assumed. This argument controls the electronic temperature (in K).
29+ eta: bool, optional
30+ Whether a progress bar should be displayed or not.
31+ """
32+ # Get the hamiltonian
33+ H = bz .parent
34+
35+ # Geometry
36+ geom = H .geometry
37+
38+ # Sparsity pattern information
39+ row_orbs , col_orbs = H .nonzero ()
40+ col_orbs_uc = H .osc2uc (col_orbs )
41+ col_isc = col_orbs // H .no
42+ sc_offsets = H .sc_off .dot (H .cell )
43+
44+ # Initialize the density matrix using the sparsity pattern of the Hamiltonian.
45+ last_dim = H .dim
46+ S = None
47+ if not H .orthogonal :
48+ last_dim -= 1
49+ S = H .tocsr (dim = last_dim )
50+ DM = DensityMatrix .fromsp (geom , [H .tocsr (dim = idim ) for idim in range (last_dim )], S = S )
51+ # Keep a reference to its data array so that we can have
52+ # direct access to it (instead of through orbital indexing).
53+ vals = DM ._csr .data
54+ # And set all values to 0
55+ if DM .orthogonal :
56+ vals [:, :] = 0
57+ else :
58+ # Don't touch the overlap values
59+ vals [:, :- 1 ] = 0
60+
61+ # For spin polarized calculations, we need to iterate over the two spin components.
62+ # If spin is unpolarized, we will multiply the contributions by 2.
63+ if DM .spin .is_polarized :
64+ spin_iterator = (0 , 1 )
65+ spin_factor = 1
66+ else :
67+ spin_iterator = (0 ,)
68+ spin_factor = 2
69+
70+ # Set the distribution that will compute occupations (or more generally, weights)
71+ # for eigenstates. If not provided, use a fermi-dirac
72+ if occ_distribution is None :
73+ kT = unit_convert ("K" , "eV" ) * fermi_dirac_T
74+ occ_distribution = get_distribution ("fermi_dirac" , smearing = kT , x0 = 0 )
75+
76+ # Loop over spins
77+ for ispin in spin_iterator :
78+ # Create the eigenstates generator
79+ eigenstates = bz .apply .eigenstate (spin = ispin )
80+
81+ # Zip it with the weights so that we can scale the contribution of each k point.
82+ k_it = zip (bz .weight , eigenstates )
83+ # Provide progress bar if requested
84+ if eta :
85+ k_it = tqdm .tqdm (k_it , total = len (bz .weight ))
86+
87+ # Now, loop through all k points
88+ for k_weight , k_eigs in k_it :
89+ # Get the k point for which this state has been calculated (in fractional coordinates)
90+ k = k_eigs .info ['k' ]
91+ # Convert the k points to 1/Ang
92+ k = k .dot (geom .rcell )
93+
94+ # Ensure R gauge so that we can use supercell phases. Much faster and less memory requirements
95+ # than using the r gauge, because we just have to compute the phase one time for each sc index.
96+ k_eigs .change_gauge ("R" )
97+
98+ # Calculate all phases, this will be a (nnz, ) shaped array.
99+ sc_phases = sc_offsets .dot (k )
100+ phases = sc_phases [col_isc ]
101+ phases = np .exp (- 1j * phases )
102+
103+ # Now find out the occupations for each wavefunction
104+ occs = k_eigs .occupation (occ_distribution )
105+
106+ state = k_eigs .state
107+
108+ if DM .spin .is_diagonal :
109+ # Calculate the matrix elements contributions for this k point.
110+ DM_kpoint = np .zeros (row_orbs .shape [0 ], dtype = k_eigs .state .dtype )
111+ add_cnc_diag_spin (state , row_orbs , col_orbs_uc , occs , DM_kpoint , occtol = occtol )
112+
113+ # Apply phases
114+ DM_kpoint = DM_kpoint * phases
115+
116+ # Take only the real part, weighting the contribution
117+ vals [:, ispin ] += k_weight * DM_kpoint .real * spin_factor
118+
119+ else :
120+ # Non colinear eigenstates contain an array of coefficients
121+ # of shape (n_wfs, no * 2), where n_wfs is also no * 2.
122+ # However, we only have "no" basis orbitals. The extra factor of 2 accounts for a hidden dimension
123+ # corresponding to spin "up"/"down". We reshape the array to uncover this extra dimension.
124+ state = state .reshape (- 1 , state .shape [1 ] // 2 , 2 )
125+
126+ # Calculate the matrix elements contributions for this k point. For each matrix element
127+ # we allocate a 2x2 spin box.
128+ DM_kpoint = np .zeros ((row_orbs .shape [0 ], 2 , 2 ), dtype = np .complex128 )
129+ add_cnc_nc (state , row_orbs , col_orbs_uc , occs , DM_kpoint , occtol = occtol )
130+
131+ # Apply phases
132+ DM_kpoint *= phases .reshape (- 1 , 1 , 1 )
133+
134+ # Now, each matrix element is a 2x2 spin box of complex numbers. That is, 4 complex numbers
135+ # i.e. 8 real numbers. What we do is to store these 8 real numbers separately in the DM.
136+ # However, in the non-colinear case (no spin orbit), since H is spin box hermitian we can force
137+ # the DM to also be spin-box hermitian. This means that DM[:, 0, 1] and DM[:, 1, 0] are complex
138+ # conjugates and we can store only 4 numbers while keeping the same information.
139+ # Here is how the spin-box can be reconstructed from the stored values:
140+ # D[j, i] =
141+ # NON-COLINEAR
142+ # [[ D[j, i, 0], D[j, i, 2] -i D[j, i, 3] ],
143+ # [ D[j, i, 2] + i D[j, i, 3], D[j, i, 1] ]]
144+ # SPIN-ORBIT
145+ # [[ D[j, i, 0], D[j, i, 6] + i D[j, i, 7]],
146+ # [ D[j, i, 2] -i D[j, i, 3], D[j, i, 1] ]]
147+
148+ # Force DM spin-box to be hermitian in the non-colinear case.
149+ if DM .spin .is_noncolinear :
150+ DM_kpoint [:, 1 , 0 ] = 0.5 * (DM_kpoint [:, 1 , 0 ] + DM_kpoint [:, 0 , 1 ].conj ())
151+
152+ # Add each contribution to its location
153+ vals [:, 0 ] += DM_kpoint [:, 0 , 0 ].real * k_weight
154+ vals [:, 1 ] += DM_kpoint [:, 1 , 1 ].real * k_weight
155+ vals [:, 2 ] += DM_kpoint [:, 1 , 0 ].real * k_weight
156+ vals [:, 3 ] -= DM_kpoint [:, 1 , 0 ].imag * k_weight
157+
158+ if DM .spin .is_spinorbit :
159+ vals [:, 4 ] -= DM_kpoint [:, 0 , 0 ].imag * k_weight
160+ vals [:, 5 ] -= DM_kpoint [:, 1 , 1 ].imag * k_weight
161+ vals [:, 6 ] += DM_kpoint [:, 0 , 1 ].real * k_weight
162+ vals [:, 7 ] -= DM_kpoint [:, 0 , 1 ].imag * k_weight
163+
164+ return DM
0 commit comments