Skip to content

Commit 7e69933

Browse files
floatingCattyfloatingCatty
committed
add support of abacus md matrix parse for version <3.9
1 parent 4a04e74 commit 7e69933

1 file changed

Lines changed: 79 additions & 28 deletions

File tree

‎dftio/io/abacus/abacus_parser.py‎

Lines changed: 79 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -248,30 +248,54 @@ def get_blocks(self, idx, hamiltonian=True, overlap=False, density_matrix=False)
248248
if hamiltonian:
249249
hamiltonian_dict = []
250250
for i in range(sys.get_nframes()):
251-
hamil, tmp = self.parse_matrix(
252-
matrix_path=os.path.join(self.raw_datas[idx], "OUT.ABACUS", "matrix/"+str(i)+"_data-HR-sparse_SPIN0.csr"),
253-
nsites=nsites,
254-
site_norbits=site_norbits,
255-
orbital_types_dict=orbital_types_dict,
256-
element=element,
257-
factor=13.605698, # Ryd2eV
258-
spinful=spinful
259-
)
251+
if os.path.exists(os.path.join(self.raw_datas[idx], "OUT.ABACUS", "matrix/")):
252+
hamil, tmp = self.parse_matrix(
253+
matrix_path=os.path.join(self.raw_datas[idx], "OUT.ABACUS", "matrix/"+str(i)+"_data-HR-sparse_SPIN0.csr"),
254+
nsites=nsites,
255+
site_norbits=site_norbits,
256+
orbital_types_dict=orbital_types_dict,
257+
element=element,
258+
factor=13.605698, # Ryd2eV
259+
spinful=spinful
260+
)
261+
else:
262+
hamil, tmp = self.parse_matrix(
263+
matrix_path=os.path.join(self.raw_datas[idx], "OUT.ABACUS/data-HR-sparse_SPIN0.csr"),
264+
nsites=nsites,
265+
site_norbits=site_norbits,
266+
orbital_types_dict=orbital_types_dict,
267+
element=element,
268+
factor=13.605698, # Ryd2eV
269+
spinful=spinful,
270+
step=i
271+
)
260272
assert tmp == int(np.sum(site_norbits)) * (1 + spinful)
261273
hamiltonian_dict.append(hamil)
262274

263275
if overlap:
264276
overlap_dict = []
265277
for i in range(sys.get_nframes()):
266-
ovp, tmp = self.parse_matrix(
267-
matrix_path=os.path.join(self.raw_datas[idx], "OUT.ABACUS", "matrix/"+str(i)+"_data-SR-sparse_SPIN0.csr"),
268-
nsites=nsites,
269-
site_norbits=site_norbits,
270-
orbital_types_dict=orbital_types_dict,
271-
element=element,
272-
factor=1,
273-
spinful=spinful
274-
)
278+
if os.path.exists(os.path.join(self.raw_datas[idx], "OUT.ABACUS", "matrix/")):
279+
ovp, tmp = self.parse_matrix(
280+
matrix_path=os.path.join(self.raw_datas[idx], "OUT.ABACUS", "matrix/"+str(i)+"_data-SR-sparse_SPIN0.csr"),
281+
nsites=nsites,
282+
site_norbits=site_norbits,
283+
orbital_types_dict=orbital_types_dict,
284+
element=element,
285+
factor=1,
286+
spinful=spinful
287+
)
288+
else:
289+
ovp, tmp = self.parse_matrix(
290+
matrix_path=os.path.join(self.raw_datas[idx], "OUT.ABACUS/data-SR-sparse_SPIN0.csr"),
291+
nsites=nsites,
292+
site_norbits=site_norbits,
293+
orbital_types_dict=orbital_types_dict,
294+
element=element,
295+
factor=1,
296+
spinful=spinful,
297+
step=i
298+
)
275299
assert tmp == int(np.sum(site_norbits)) * (1 + spinful)
276300

277301
if spinful:
@@ -285,31 +309,58 @@ def get_blocks(self, idx, hamiltonian=True, overlap=False, density_matrix=False)
285309
if density_matrix:
286310
density_matrix_dict = []
287311
for i in range(sys.get_nframes()):
288-
dm, tmp = self.parse_matrix(
289-
matrix_path=os.path.join(self.raw_datas[idx], "OUT.ABACUS", "matrix/"+str(i)+"_data-DMR-sparse_SPIN0.csr"),
290-
nsites=nsites,
291-
site_norbits=site_norbits,
292-
orbital_types_dict=orbital_types_dict,
293-
element=element,
294-
factor=1,
295-
spinful=spinful
296-
)
312+
if os.path.exists(os.path.join(self.raw_datas[idx], "OUT.ABACUS", "matrix/")):
313+
dm, tmp = self.parse_matrix(
314+
matrix_path=os.path.join(self.raw_datas[idx], "OUT.ABACUS", "matrix/"+str(i)+"_data-DMR-sparse_SPIN0.csr"),
315+
nsites=nsites,
316+
site_norbits=site_norbits,
317+
orbital_types_dict=orbital_types_dict,
318+
element=element,
319+
factor=1,
320+
spinful=spinful
321+
)
322+
else:
323+
dm, tmp = self.parse_matrix(
324+
matrix_path=os.path.join(self.raw_datas[idx], "OUT.ABACUS/data-DMR-sparse_SPIN0.csr"),
325+
nsites=nsites,
326+
site_norbits=site_norbits,
327+
orbital_types_dict=orbital_types_dict,
328+
element=element,
329+
factor=1,
330+
spinful=spinful,
331+
step=i
332+
)
333+
297334
assert tmp == int(np.sum(site_norbits)) * (1 + spinful)
298335
density_matrix_dict.append(dm)
299336
else:
300337
raise NotImplementedError("mode {} is not supported.".format(mode))
301338

302339
return hamiltonian_dict, overlap_dict, density_matrix_dict
303340

304-
def parse_matrix(self, matrix_path, nsites, site_norbits, orbital_types_dict, element, factor, spinful=False):
341+
def parse_matrix(self, matrix_path, nsites, site_norbits, orbital_types_dict, element, factor, spinful=False, step=0):
305342
site_norbits_cumsum = np.cumsum(site_norbits)
306343
norbits = int(np.sum(site_norbits))
307344
matrix_dict = dict()
308345
with open(matrix_path, 'r') as f:
309346
line = f.readline() # read "Matrix Dimension of ..."
310347
if not "Matrix Dimension of" in line:
348+
"""In this case, the starting of the file is STEP 0"""
349+
# find the correct step
350+
step_found = False
351+
while line and not step_found:
352+
if "STEP" in line:
353+
stp = int(line.split()[-1])
354+
if stp != step:
355+
line = f.readline()
356+
else:
357+
step_found = True
358+
else:
359+
line = f.readline()
311360
line = f.readline() # ABACUS >= 3.0
312361
assert "Matrix Dimension of" in line
362+
else:
363+
assert step == 0
313364
f.readline() # read "Matrix number of ..."
314365
norbits = int(line.split()[-1])
315366
for line in f:

0 commit comments

Comments
 (0)