Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding additional QE.out parsing #546

Merged
214 changes: 213 additions & 1 deletion mala/targets/target.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from scipy.integrate import simps

from mala.common.parameters import Parameters, ParametersTargets
from mala.common.parallelizer import printout, parallel_warn
from mala.common.parallelizer import printout, parallel_warn, get_rank
from mala.targets.calculation_helpers import fermi_function
from mala.common.physical_data import PhysicalData
from mala.descriptors.atomic_density import AtomicDensity
Expand Down Expand Up @@ -119,6 +119,13 @@ def __init__(self, params):
self.band_energy_dft_calculation = None
self.total_energy_dft_calculation = None
self.entropy_contribution_dft_calculation = None
self.total_energy_contributions_dft_calculation = {
"one_electron_contribution": None,
"hartree_contribution": None,
"xc_contribution": None,
"ewald_contribution": None,
}
self.atomic_forces_dft = None
self.atoms = None
self.electrons_per_atom = None
self.qe_input_data = {
Expand Down Expand Up @@ -319,6 +326,13 @@ def read_additional_calculation_data(self, data, data_type=None):
self.band_energy_dft_calculation = None
self.total_energy_dft_calculation = None
self.entropy_contribution_dft_calculation = None
self.total_energy_contributions_dft_calculation = {
"one_electron_contribution": None,
"hartree_contribution": None,
"xc_contribution": None,
"ewald_contribution": None,
}
self.atomic_forces_dft = None
self.grid_dimensions = [0, 0, 0]
self.atoms = None

Expand All @@ -328,12 +342,26 @@ def read_additional_calculation_data(self, data, data_type=None):
self.fermi_energy_dft = (
self.atoms.get_calculator().get_fermi_level()
)
# The forces may not have been computed. If they are indeed
# not computed, ASE will by default throw an PropertyNotImplementedError
# error
try:
self.atomic_forces_dft = self.atoms.get_forces()
except ase.calculators.calculator.PropertyNotImplementedError:
pass

# Parse the file for energy values.
total_energy = None
past_calculation_part = False
bands_included = True

# individual energy contributions.
entropy_contribution = None
one_electron_contribution = None
hartree_contribution = None
xc_contribution = None
ewald_contribution = None

with open(data) as out:
pseudolinefound = False
lastpseudo = None
Expand Down Expand Up @@ -390,6 +418,32 @@ def read_additional_calculation_data(self, data, data_type=None):
entropy_contribution = float(
(line.split("=")[1]).split("Ry")[0]
)
if (
"one-electron contribution" in line
and past_calculation_part
):
if one_electron_contribution is None:
one_electron_contribution = float(
(line.split("=")[1]).split("Ry")[0]
)
if (
"hartree contribution" in line
and past_calculation_part
):
if hartree_contribution is None:
hartree_contribution = float(
(line.split("=")[1]).split("Ry")[0]
)
if "xc contribution" in line and past_calculation_part:
if xc_contribution is None:
xc_contribution = float(
(line.split("=")[1]).split("Ry")[0]
)
if "ewald contribution" in line and past_calculation_part:
if ewald_contribution is None:
ewald_contribution = float(
(line.split("=")[1]).split("Ry")[0]
)
if "set verbosity='high' to print them." in line:
bands_included = False

Expand All @@ -413,6 +467,25 @@ def read_additional_calculation_data(self, data, data_type=None):
self.entropy_contribution_dft_calculation = (
entropy_contribution * Rydberg
)
if one_electron_contribution is not None:
self.total_energy_contributions_dft_calculation[
"one_electron_contribution"
] = (one_electron_contribution * Rydberg)

if hartree_contribution is not None:
self.total_energy_contributions_dft_calculation[
"hartree_contribution"
] = (hartree_contribution * Rydberg)

if xc_contribution is not None:
self.total_energy_contributions_dft_calculation[
"xc_contribution"
] = (xc_contribution * Rydberg)

if ewald_contribution is not None:
self.total_energy_contributions_dft_calculation[
"ewald_contribution"
] = (ewald_contribution * Rydberg)

# Calculate band energy, if the necessary data is included in
# the output file.
Expand Down Expand Up @@ -446,6 +519,13 @@ def read_additional_calculation_data(self, data, data_type=None):
self.band_energy_dft_calculation = None
self.total_energy_dft_calculation = None
self.entropy_contribution_dft_calculation = None
self.total_energy_contributions_dft_calculation = {
"one_electron_contribution": None,
"hartree_contribution": None,
"xc_contribution": None,
"ewald_contribution": None,
}
self.atomic_forces_dft = None
self.grid_dimensions = [0, 0, 0]
self.atoms: ase.Atoms = data[0]

Expand Down Expand Up @@ -490,6 +570,13 @@ def read_additional_calculation_data(self, data, data_type=None):
self.voxel = None
self.band_energy_dft_calculation = None
self.total_energy_dft_calculation = None
self.total_energy_contributions_dft_calculation = {
"one_electron_contribution": None,
"hartree_contribution": None,
"xc_contribution": None,
"ewald_contribution": None,
}
self.atomic_forces_dft = None
self.entropy_contribution_dft_calculation = None
self.grid_dimensions = [0, 0, 0]
self.atoms = None
Expand All @@ -505,6 +592,24 @@ def read_additional_calculation_data(self, data, data_type=None):
self.qe_input_data["degauss"] = json_dict["degauss"]
self.qe_pseudopotentials = json_dict["pseudopotentials"]

energy_contribution_ids = [
"one_electron_contribution",
"hartree_contribution",
"xc_contribution",
"ewald_contribution",
]
for key in energy_contribution_ids:
if key in json_dict:
self.total_energy_contributions_dft_calculation[key] = (
json_dict[key]
)

# Not always read from DFT files.
if "atomic_forces_dft" in json_dict:
self.atomic_forces_dft = np.array(
json_dict["atomic_forces_dft"]
)

else:
raise Exception("Unsupported auxiliary file type.")

Expand Down Expand Up @@ -539,6 +644,18 @@ def write_additional_calculation_data(self, filepath, return_string=False):
"degauss": self.qe_input_data["degauss"],
"pseudopotentials": self.qe_pseudopotentials,
"entropy_contribution_dft_calculation": self.entropy_contribution_dft_calculation,
"one_electron_contribution": self.total_energy_contributions_dft_calculation[
"one_electron_contribution"
],
"hartree_contribution": self.total_energy_contributions_dft_calculation[
"hartree_contribution"
],
"xc_contribution": self.total_energy_contributions_dft_calculation[
"xc_contribution"
],
"ewald_contribution": self.total_energy_contributions_dft_calculation[
"ewald_contribution"
],
}
if self.voxel is not None:
additional_calculation_data["voxel"] = self.voxel.todict()
Expand All @@ -560,6 +677,11 @@ def write_additional_calculation_data(self, filepath, return_string=False):
additional_calculation_data["atoms"]["pbc"] = (
additional_calculation_data["atoms"]["pbc"].tolist()
)
if self.atomic_forces_dft is not None:
additional_calculation_data["atomic_forces_dft"] = (
self.atomic_forces_dft.tolist()
)

if return_string is False:
with open(filepath, "w", encoding="utf-8") as f:
json.dump(
Expand Down Expand Up @@ -1306,6 +1428,8 @@ def _process_additional_metadata(self, additional_metadata):
)

def _set_openpmd_attribtues(self, iteration, mesh):
import openpmd_api as io

super(Target, self)._set_openpmd_attribtues(iteration, mesh)

# If no atoms have been read, neither have any of the other
Expand All @@ -1321,6 +1445,7 @@ def _set_openpmd_attribtues(self, iteration, mesh):
and key is not None
and key != "pseudopotentials"
and additional_calculation_data[key] is not None
and key != "atomic_forces_dft"
):
iteration.set_attribute(key, additional_calculation_data[key])
if key == "pseudopotentials":
Expand All @@ -1334,6 +1459,42 @@ def _set_openpmd_attribtues(self, iteration, mesh):
],
)

# If the data contains atomic forces from a DFT calculation, we need
# to process it in much the same fashion as the atoms.
if "atomic_forces_dft" in additional_calculation_data:
atomic_forces = additional_calculation_data["atomic_forces_dft"]
if atomic_forces is not None:
# This data is equivalent across the ranks, so just write it once
atomic_forces_dft_openpmd = iteration.particles[
"atomic_forces_dft"
]
forces = io.Dataset(
# Need bugfix https://github.com/openPMD/openPMD-api/pull/1357
(
np.array(atomic_forces[0]).dtype
if io.__version__ >= "0.15.0"
else io.Datatype.DOUBLE
),
np.array(atomic_forces[0]).shape,
)
for atom in range(0, np.shape(atomic_forces)[0]):
atomic_forces_dft_openpmd["force_compopnents"][
str(atom)
].reset_dataset(forces)

individual_force = atomic_forces_dft_openpmd[
"force_compopnents"
][str(atom)]
if get_rank() == 0:
individual_force.store_chunk(
np.array(atomic_forces)[atom]
)

# Positions are stored in Angstrom.
atomic_forces_dft_openpmd["force_compopnents"][
str(atom)
].unit_SI = 1.0e-10

def _process_openpmd_attributes(self, series, iteration, mesh):
super(Target, self)._process_openpmd_attributes(
series, iteration, mesh
Expand Down Expand Up @@ -1374,6 +1535,20 @@ def _process_openpmd_attributes(self, series, iteration, mesh):
"periodic_boundary_conditions_z"
)

# Forces may not necessarily have been read (and therefore written)

try:
atomic_forces_dft = iteration.particles["atomic_forces_dft"]
nr_atoms = len(atomic_forces_dft["force_compopnents"])
self.atomic_forces_dft = np.zeros((nr_atoms, 3))
for i in range(0, nr_atoms):
atomic_forces_dft["force_compopnents"][str(i)].load_chunk(
self.atomic_forces_dft[i, :]
)
series.flush()
except IndexError:
pass

# Process all the regular meta info.
self.fermi_energy_dft = self._get_attribute_if_attribute_exists(
iteration, "fermi_energy_dft", default_value=self.fermi_energy_dft
Expand Down Expand Up @@ -1447,6 +1622,43 @@ def _process_openpmd_attributes(self, series, iteration, mesh):
iteration.get_attribute(attribute)
)

self.total_energy_contributions_dft_calculation[
"one_electron_contribution"
] = self._get_attribute_if_attribute_exists(
iteration,
"one_electron_contribution",
default_value=self.total_energy_contributions_dft_calculation[
"one_electron_contribution"
],
)
self.total_energy_contributions_dft_calculation[
"hartree_contribution"
] = self._get_attribute_if_attribute_exists(
iteration,
"hartree_contribution",
default_value=self.total_energy_contributions_dft_calculation[
"hartree_contribution"
],
)
self.total_energy_contributions_dft_calculation["xc_contribution"] = (
self._get_attribute_if_attribute_exists(
iteration,
"xc_contribution",
default_value=self.total_energy_contributions_dft_calculation[
"xc_contribution"
],
)
)
self.total_energy_contributions_dft_calculation[
"ewald_contribution"
] = self._get_attribute_if_attribute_exists(
iteration,
"ewald_contribution",
default_value=self.total_energy_contributions_dft_calculation[
"ewald_contribution"
],
)

def _set_geometry_info(self, mesh):
# Geometry: Save the cell parameters and angles of the grid.
if self.atoms is not None:
Expand Down