def get_dos(
model,
posinp,
device="cpu",
supercell=(6, 6, 6),
qpoints=[30, 30, 30],
npts=1000,
width=0.004,
):
if isinstance(posinp, str):
atoms = posinp_to_ase_atoms(Posinp.from_file(posinp))
elif isinstance(posinp, Posinp):
atoms = posinp_to_ase_atoms(posinp)
else:
raise ValueError("The posinp variable is not recognized.")
if isinstance(model, str):
model = load_model(model, map_location=device)
elif isinstance(model, torch.nn.Module):
pass
else:
raise ValueError("The model variable is not recognized.")
# Bugfix to make older models work with PyTorch 1.6
# Hopefully temporary
for mod in model.modules():
if not hasattr(mod, "_non_persistent_buffers_set"):
mod._non_persistent_buffers_set = set()
assert len(supercell) == 3, "Supercell should be a length 3 object."
assert len(qpoints) == 3, "Qpoints should be a length 3 object."
supercell = tuple(supercell)
cutoff = float(model.state_dict()
["representation.interactions.0.cutoff_network.cutoff"])
calculator = SpkCalculator(
model,
device=device,
energy="energy",
forces="forces",
environment_provider=AseEnvironmentProvider(cutoff),
)
ph = Phonons(atoms, calculator, supercell=supercell, delta=0.02)
ph.run()
ph.read(acoustic=True)
dos = ph.get_dos(kpts=qpoints).sample_grid(npts=npts, width=width)
ph.clean()
return Dos(dos.energy * 8065.6, dos.weights[0])
def main(args):
function = get_function(args.function)
distances = np.linspace(args.range[0], args.range[1], args.ndata)
with connect(args.dbname) as db:
for d in distances:
pos_dict = {
"units":
"angstroem",
"boundary_conditions":
"free",
"positions": [
{
args.element: [0, 0, 0]
},
{
args.element: [d, 0, 0]
},
],
}
posinp = Posinp.from_dict(pos_dict)
atoms = posinp_to_ase_atoms(posinp)
energy = function.value(d)
forces = np.array([
[-1.0 * function.first_derivative(-d), 0, 0],
[-1.0 * function.first_derivative(d), 0, 0],
])
db.write(atoms, data={"energy": energy, "forces": forces})
def determine_unique_configurations(configurations):
cutoff = float(np.max(configurations[0].cell.array) / 2 + 1)
unique_reps, unique_config, reps, count_configs = [], [], [], []
schnet = SchNet(n_atom_basis=32,
n_filters=32,
n_interactions=1,
cutoff=cutoff,
cutoff_network=CosineCutoff)
env = AseEnvironmentProvider(cutoff=cutoff)
data = [posinp_to_ase_atoms(pos) for pos in configurations]
data = SchnetPackData(data=data,
environment_provider=env,
collect_triples=False)
data_loader = AtomsLoader(data, batch_size=1)
for batch in data_loader:
reps.append(torch.squeeze(schnet(batch)))
for i, rep in enumerate(reps):
for j, uni in enumerate(unique_reps):
if compare_reps(rep, uni):
count_configs[j] += 1
break
else:
unique_reps.append(rep)
unique_config.append(configurations[i])
count_configs.append(1)
return unique_config, count_configs
def main(args):
with connect(args.dbname) as db:
db.metadata = DEFAULT_METADATA
if args.run_mode == "bigdft":
files = [f for f in os.listdir() if f.endswith(".xyz")]
for f in files:
atoms = posinp_to_ase_atoms(Posinp.from_file(f))
with open(f, "r") as posinp_file:
energy = float(
posinp_file.readline().split()[2]) * 27.21138602
forces = None
for line in posinp_file:
if "forces" not in line:
continue
else:
forces = []
for _ in range(len(atoms)):
forces.append([
float(force) for force in
posinp_file.readline().split()[-3:]
])
forces = np.array(
forces) * 27.21138602 / 0.529177249
break
db.write(atoms, data={"energy": energy, "forces": forces})
elif args.run_mode == "abinit":
files = [f for f in os.listdir() if f.endswith(".out")]
for f in files:
atoms = read(f, format="abinit-out")
about = AbinitOutputFile(f)
energy = float(about.final_vars_global["etotal"]) * 27.21138602
forces = (np.array([
float(g) for g in about.final_vars_global["fcart"].split()
]).reshape(-1, 3) * 27.21138602 / 0.529177249)
db.write(atoms, data={"energy": energy, "forces": forces})
def run(
self,
property,
posinp=None,
batch_size=128,
):
r"""
Central method to use when making a calculation with
the calculator.
Parameters
----------
property : str
Property to be predicted by the calculator
posinp : Posinp
Atomic configuration to pass to the model
batch_size : int
Batch sizes. Default is 128.
Returns
-------
predictions : :class:`numpy.ndarray`
Corresponding prediction by the model.
"""
init_property, out_name, derivative, wrt = get_derivative_names(
property, self.available_properties)
if abs(derivative) >= 1:
self.model.output_modules[0].create_graph = True
if len(posinp) > 1 and derivative:
batch_size = 1
data = [posinp_to_ase_atoms(pos) for pos in posinp]
pbc = True if any(pos.pbc.any() for pos in data) else False
environment_provider = (AseEnvironmentProvider(
cutoff=self.cutoff) if pbc else SimpleEnvironmentProvider())
data = SchnetPackData(
data=data,
environment_provider=environment_provider,
collect_triples=self.model_type == "wacsf",
)
data_loader = AtomsLoader(data, batch_size=batch_size)
pred = []
if derivative == 0:
if self.model.output_modules[0].derivative is not None:
for batch in data_loader:
batch = {k: v.to(self.device) for k, v in batch.items()}
pred.append(self.model(batch))
else:
with torch.no_grad():
for batch in data_loader:
batch = {
k: v.to(self.device)
for k, v in batch.items()
}
pred.append(self.model(batch))
if abs(derivative) == 1:
for batch in data_loader:
batch = {k: v.to(self.device) for k, v in batch.items()}
batch[wrt[0]].requires_grad_()
results = self.model(batch)
deriv1 = torch.unsqueeze(
torch_derivative(results[init_property], batch[wrt[0]]), 0)
if derivative < 0:
deriv1 = -1.0 * deriv1
pred.append({out_name: deriv1})
if abs(derivative) == 2:
for batch in data_loader:
batch = {k: v.to(self.device) for k, v in batch.items()}
for inp in set(wrt):
batch[inp].requires_grad_()
results = self.model(batch)
deriv2 = torch.unsqueeze(
torch_derivative(
torch_derivative(
results[init_property],
batch[wrt[0]],
create_graph=True,
),
batch[wrt[0]],
),
0,
)
if derivative < 0:
deriv2 = -1.0 * deriv2
pred.append({out_name: deriv2})
predictions = {}
if self.md:
for p in ["energy", "forces"]:
predictions[p] = np.concatenate(
[batch[p].cpu().detach().numpy() for batch in pred])
else:
if derivative:
predictions[property] = np.concatenate(
[batch[out_name].cpu().detach().numpy() for batch in pred])
else:
predictions[property] = np.concatenate([
batch[init_property].cpu().detach().numpy()
for batch in pred
])
return predictions
def run(
self,
property,
posinp=None,
batch_size=1,
):
r"""
Central method to use when making a calculation with
the calculator.
Parameters
----------
property : str
Property to be predicted by the calculator
posinp : Posinp
Atomic configuration to pass to the model
Returns
-------
predictions : :class:`numpy.ndarray`
Corresponding prediction by the model.
"""
# Initial setup
assert (
len(posinp) == 1
), "Use the PatchSPCalculator for one configuration at a time."
atoms = posinp_to_ase_atoms(posinp[0])
if property == "hessian" and any(self.subgrid == 2):
raise warnings.warn(
"""
The hessian matrix can have some bad values with a grid of
size 2 because the same atom can be copied multiple times
in the buffers of the same subcell. Use a larger grid.
"""
)
init_property, out_name, derivative, wrt = get_derivative_names(
property, self.available_properties
)
if abs(derivative) >= 1:
self.model.output_modules[0].create_graph = True
pbc = True if atoms.pbc.any() else False
environment_provider = (
AseEnvironmentProvider(cutoff=self.cutoff)
if pbc
else SimpleEnvironmentProvider()
)
# Split the configuration according to the subgrid
at_to_patches = AtomsToPatches(
cutoff=self.cutoff, n_interaction=self.n_interaction, grid=self.subgrid
)
(
subcells,
subcells_main_idx,
original_cell_idx,
complete_subcell_copy_idx,
) = at_to_patches.split_atoms(atoms)
# Pass each subcell independantly
results = []
for subcell in subcells:
data = SchnetPackData(
data=[subcell],
environment_provider=environment_provider,
collect_triples=self.model_type == "wacsf",
)
data_loader = AtomsLoader(data, batch_size=1)
if derivative == 0:
if self.model.output_modules[0].derivative is not None:
for batch in data_loader:
batch = {k: v.to(self.device) for k, v in batch.items()}
results.append(self.model(batch))
else:
with torch.no_grad():
for batch in data_loader:
batch = {k: v.to(self.device) for k, v in batch.items()}
results.append(self.model(batch))
if abs(derivative) == 1:
for batch in data_loader:
batch = {k: v.to(self.device) for k, v in batch.items()}
batch[wrt[0]].requires_grad_()
forward_results = self.model(batch)
deriv1 = torch_derivative(
forward_results[init_property], batch[wrt[0]]
)
if derivative < 0:
deriv1 = -1.0 * deriv1
results.append({out_name: deriv1})
if abs(derivative) == 2:
raise NotImplementedError()
predictions = {}
if property == "energy":
predictions["energy"] = np.sum(
[
patch["individual_energy"][subcells_main_idx[i]]
.detach()
.cpu()
.numpy()
for i, patch in enumerate(results)
]
)
elif property == "forces":
forces = np.zeros((len(atoms), 3))
for i in range(len(results)):
forces[original_cell_idx[i]] = (
results[i]["forces"]
.detach()
.squeeze()
.cpu()
.numpy()[subcells_main_idx[i]]
)
predictions["forces"] = forces
elif property == "hessian":
hessian = np.zeros((3 * len(atoms), 3 * len(atoms)))
for i in range(len(results)):
(
hessian_original_cell_idx_0,
hessian_original_cell_idx_1,
) = prepare_hessian_indices(
original_cell_idx[i], complete_subcell_copy_idx[i]
)
(
hessian_subcells_main_idx_0,
hessian_subcells_main_idx_1,
) = prepare_hessian_indices(
subcells_main_idx[i],
np.arange(0, len(complete_subcell_copy_idx[i])),
)
hessian[hessian_original_cell_idx_0, hessian_original_cell_idx_1] = (
results[i]["hessian"]
.detach()
.squeeze()
.cpu()
.numpy()[hessian_subcells_main_idx_0, hessian_subcells_main_idx_1]
)
predictions["hessian"] = hessian
else:
raise NotImplementedError()
return predictions
def test_periodic(self):
pos1 = Posinp.from_file(os.path.join(pos_folder, "periodic.xyz"))
atoms = posinp_to_ase_atoms(pos1)
pos2 = Posinp.from_ase(atoms)
assert pos1 == pos2
def test_surface(self):
pos1 = Posinp.from_file(os.path.join(pos_folder, "surface2.xyz"))
atoms = posinp_to_ase_atoms(pos1)
pos2 = Posinp.from_ase(atoms)
assert pos1 == pos2
