Support for Quantum ESPRESSO in exatomic

import pandas as pd
import exa, exatomic
from exatomic import qe
from exatomic.base import resource, list_resources   # Helper to find example files available in exatomic

# list_resources shows all available example files; here are some QE related ones
[i for i in list_resources() if "qe" in i]

['qe.vel', 'qe.cel', 'qe.inp', 'qe.for', 'qe.evp', 'qe.eig', 'qe.pos']

inppath = resource("qe.inp")
pospath = resource("qe.pos")
velpath = resource("qe.vel")
forpath = resource("qe.for")
evppath = resource("qe.evp")

exa.Editor(velpath).head()    # We can use the Editor to view the actual file contents

     0:   86300 10.43755843
     1:      0.26582218192461E-03     0.34762151661751E-03    -0.37055692323044E-03
     2:     -0.28357409799207E-03    -0.84385574934985E-03     0.61301203789432E-03
     3:      0.43446778391328E-03    -0.13368379507875E-03     0.22638614726700E-04
     4:     -0.23558263497270E-03     0.11887988312728E-02     0.25475038796956E-03
     5:      0.28858553030712E-03    -0.44466330835073E-03     0.56058154685670E-04
     6:     -0.17782953131705E-03    -0.28382892061277E-04    -0.14986330474008E-03
     7:     -0.33509839767085E-03    -0.58101031888086E-03     0.98300831028985E-03
     8:     -0.45531989084626E-03     0.78072761126234E-04    -0.16013602508294E-03
     9:     -0.72505618884722E-03     0.53608322941604E-03     0.11241564898308E-02

exa.Editor(inppath)[34:38]  # Likewise for the input file...
# Note that the atom order of the pos, vel, etc files is the
# order of the ATOMIC_SPECIES block

['ATOMIC_SPECIES',
 'H  2.01355d0   H.pbe-rrkjus_psl.1.0.0.UPF',
 'O  15.9994d0   O.pbe-n-rrkjus_psl.1.0.0.UPF',
 'ATOMIC_POSITIONS (angstrom)']

# This is one way to access the docs of a specific function
# in the Jupyter notebook or IPython terminal
qe.parse_symbols_from_input?

# Here we get and order the symbol list
xyz = qe.parse_symbols_from_input(inppath)
xyz   # pandas DataFrame

	symbol	x	y	z
0	H	1.032386	6.784453	2.604526
1	H	1.803768	2.861700	5.597986
2	H	1.488886	1.586134	1.667751
3	H	1.304391	7.581787	6.658065
4	H	2.687215	5.624346	7.935414
5	H	3.611814	0.867768	6.808235
6	H	3.455902	3.536857	1.089975
7	H	3.677819	7.840515	2.902969
8	H	5.150307	4.829348	3.958955
9	H	5.359657	1.344457	3.427006
10	H	6.068226	3.004010	8.041704
11	H	6.089307	8.290729	6.112190
12	H	6.270537	6.295901	1.342648
13	H	7.779489	6.219615	6.898890
14	H	7.669570	2.655679	2.535040
15	H	8.631402	2.379311	7.750683
16	O	1.661344	3.288553	0.895464
17	O	2.252441	7.890762	1.787531
18	O	2.193323	6.051175	6.263781
19	O	3.121875	1.568882	5.210278
20	O	6.867725	7.678591	7.634948
21	O	6.482595	5.721000	3.118545
22	O	7.361162	3.199850	6.768000
23	O	6.554401	1.278874	2.030744

# Turns out the symbols are in the right order but if we needed to reorder them...
def mapper(sym):
    return {'H': 0, 'O': 1}[sym]

xyz['order'] = xyz['symbol'].map(mapper)
symbols = xyz.sort_values('order')['symbol'].tolist()    # symbols = xyz['symbol'].tolist()
symbols

['H',
 'H',
 'H',
 'H',
 'H',
 'H',
 'H',
 'H',
 'H',
 'H',
 'H',
 'H',
 'H',
 'H',
 'H',
 'H',
 'O',
 'O',
 'O',
 'O',
 'O',
 'O',
 'O',
 'O']

pos = qe.parse_xyz(pospath, symbols)
vel = qe.parse_xyz(velpath, symbols)
vel.columns = ("vx", "vy", "vz", "frame", "symbol")
force = qe.parse_xyz(forpath, symbols)
force.columns = ("fx", "fy", "fz", "frame", "symbol")

pos.head()

	x	y	z	frame	symbol
0	1.386029	10.148803	2.527119	86300	H
1	6.298432	0.408980	6.614086	86300	H
2	2.705972	4.812221	1.851418	86300	H
3	0.925476	7.572369	3.711328	86300	H
4	-0.712621	7.029148	6.158907	86300	H

vel.head()

	vx	vy	vz	frame	symbol
0	0.000266	0.000348	-0.000371	86300	H
1	-0.000284	-0.000844	0.000613	86300	H
2	0.000434	-0.000134	0.000023	86300	H
3	-0.000236	0.001189	0.000255	86300	H
4	0.000289	-0.000445	0.000056	86300	H

force.head()

	fx	fy	fz	frame	symbol
0	-0.014144	-0.001770	-0.010127	86300	H
1	-0.003695	0.015956	-0.013014	86300	H
2	-0.003526	-0.001231	0.003031	86300	H
3	0.017270	-0.005110	-0.001918	86300	H
4	0.002262	0.006420	0.002992	86300	H

atom = pd.concat((pos, vel[["vx", "vy", "vz"]], force[["fx", "fy", "fz"]]), axis=1)

# Some extra steps to set the column names
evp = qe.parse_evp(evppath, symbols)
columns = evp.iloc[0, :].tolist()[1:] + ["None"]    # Add a dummy column
evp = evp.drop("#", axis=0)                         # Delete first row
evp.columns = columns
evp['atom_count'] = len(symbols)
evp = evp.dropna(how='all', axis=1)                 # Remove dummy column
evp.head()

	time(ps)	ekinc	T_cell(K)	Tion(K)	etot	enthal	econs	econt	Volume	Pressure(GPa)	atom_count
86300	1.043756E+01	1.165034E-03	0.000000E+00	1.991763E+02	-139.02493778	-139.02493778	-139.00223062	-139.04934698	7.272825E+02	-0.00000	24
86310	1.043877E+01	1.177623E-03	0.000000E+00	1.951384E+02	-139.02452520	-139.02452520	-139.00227838	-139.00110075	7.272825E+02	-0.00000	24
86320	1.043998E+01	2.023734E-03	0.000000E+00	1.817591E+02	-139.02385374	-139.02385374	-139.00313223	-139.00110849	7.272825E+02	-0.00000	24
86330	1.044119E+01	1.223379E-03	0.000000E+00	1.886373E+02	-139.02383875	-139.02383875	-139.00233309	-139.00110971	7.272825E+02	-0.00000	24
86340	1.044240E+01	2.015737E-03	0.000000E+00	1.967750E+02	-139.02555726	-139.02555726	-139.00312386	-139.00110813	7.272825E+02	-0.00000	24

# Make a universe
uni = exatomic.Universe(atom=atom, frame=evp)

len(uni)

10001

uni.atom.dtypes

x          float64
y          float64
z          float64
frame     category
symbol    category
vx         float64
vy         float64
vz         float64
fx         float64
fy         float64
fz         float64
dtype: object

%time uni.compute_atom_two()    # Compute interatomic distances

CPU times: user 6.69 s, sys: 298 ms, total: 6.99 s
Wall time: 7.04 s

uni.atom.shape

(240024, 11)

uni.atom_two.shape

(1874130, 4)

uni.molecule.shape     # Automatically computes molecules from  atom-atom distances

(80030, 3)

exatomic.UniverseWidget(uni)    # Visualize the universe!