Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Feb 2, 2024
1 parent 30bdb09 commit 1c81c37
Showing 1 changed file with 8 additions and 6 deletions.
14 changes: 8 additions & 6 deletions aiida_nanotech_empa/workflows/cp2k/reftraj_md_workchain.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,9 @@ def create_batches(trajectory, batch_size):
"""Create batches of the trajectory."""
ase_structures = trajectory.get_ase()
if not check_cell(ase_structures):
raise ValueError("The cell of the structures in the trajectory is not the same.")
raise ValueError(
"The cell of the structures in the trajectory is not the same."
)
split_structures = split_structures(ase_structures)
batches = []
for ss in split_structures:
Expand Down Expand Up @@ -120,17 +122,18 @@ def first_scf(self):
future = self.submit(builder)
self.report(f"Submitted SCF of the initial geometry: {future.pk}")
self.ToContext(initial_scf=future)

def run_reftraj_batches(self):

for i_batch, trajectory_batch in enumerate(create_batches(self.inputs.trajectory, batch_size=self.inputs.batch_size)):
def run_reftraj_batches(self):
for i_batch, trajectory_batch in enumerate(
create_batches(self.inputs.trajectory, batch_size=self.inputs.batch_size)
):
self.report(f"Running batch of {len(trajectory_batch)} structures")
# create the input for the reftraj workchain
builder = Cp2kBaseWorkChain.get_builder()
builder.cp2k.trajectory = trajectory_batch
builder.cp2k.code = self.inputs.code
builder.cp2k.parameters = orm.Dict(dict=self.ctx.input_dict)

builder.cp2k.parent_calc_folder = self.ctx.initial_scf.outputs.remote_folder
future = self.submit(Cp2kBaseWorkChain, **reftraj_input)
self.report(f"Submitted reftraj batch: {i_batch} with pk: {future.pk}")
Expand All @@ -143,4 +146,3 @@ def merge_batches_output(self):
for i_batch in range(self.ctx.n_batches):
merged_traj.extend(self.ctx[f"reftraj_batch_{i_batch}"].outputs.trajectory)
return merged_traj

0 comments on commit 1c81c37

Please sign in to comment.