few more steps

aiida_nanotech_empa/workflows/cp2k/protocols/md_reftraj_protocol.yml

@@ -2,20 +2,31 @@ standard:
     GLOBAL:
         EXTENDED_FFT_LENGTHS: ''
         PRINT_LEVEL: MEDIUM
-        RUN_TYPE: GEO_OPT
-        WALLTIME: 600
+        RUN_TYPE: MD
+        WALLTIME: '600'
         ELPA_KERNEL: AVX2_BLOCK2
     MOTION:
         MD:
+            STEPS: '100'
             ENSEMBLE: 'REFTRAJ'
-                REFTRAJ:
-                    EVAL_FORCES: '.TRUE.'
-                TRAJ_FILE_NAME: 'aiida-traj.xyz'
+            REFTRAJ:
+                FIRST_SNAPSHOT: '1'
+                LAST_SNAPSHOT: '100'
+                STRIDE: '1'
+                EVAL_FORCES: '.TRUE.'
+                TRAJ_FILE_NAME: 'reftraj.xyz'
+                VARIABLE_VOLUME: .TRUE.
+                CELL_FILE_NAME: 'reftraj.cell'
         PRINT:
             RESTART_HISTORY:
                 _: 'OFF'
             FORCES:
+                EACH:
+                    MD: '1'
                 FORMAT: 'XYZ'
+            CELL:
+                EACH:
+                    MD: '1'
     FORCE_EVAL:
         METHOD: Quickstep
         SUBSYS:
@@ -24,50 +35,49 @@ standard:
                 SYMMETRY: ORTHORHOMBIC
         DFT:
             UKS: .FALSE.
-            MULTIPLICITY: 0
-            CHARGE: 0
+            MULTIPLICITY: '0'
+            CHARGE: '0'
             BASIS_SET_FILE_NAME: BASIS_MOLOPT
             POTENTIAL_FILE_NAME: POTENTIAL
             RESTART_FILE_NAME: ./parent_calc/aiida-RESTART.wfn
             MGRID:
-                CUTOFF: 600
-                NGRIDS: 5
+                CUTOFF: '600'
+                NGRIDS: '5'
             POISSON:
                 PERIODIC: XYZ
                 POISSON_SOLVER: PERIODIC
             QS:
-                EPS_DEFAULT: 1.0E-14
+                EPS_DEFAULT: '1.0E-14'
                 EXTRAPOLATION: ASPC
-                EXTRAPOLATION_ORDER: 3
+                EXTRAPOLATION_ORDER: '3'
                 METHOD: GPW
             SCF:
                 SCF_GUESS: RESTART
-                EPS_SCF: 1.0E-7
-                MAX_SCF: 40
+                EPS_SCF: '1.0E-7'
+                MAX_SCF: '40'
                 OT:
                     MINIMIZER: CG
                     PRECONDITIONER: FULL_SINGLE_INVERSE
                 OUTER_SCF:
-                    EPS_SCF: 1.0E-7
-                    MAX_SCF: 50
+                    EPS_SCF: '1.0E-7'
+                    MAX_SCF: '50'
                 PRINT:
                     RESTART:
                         ADD_LAST: NUMERIC
                         EACH:
-                            GEO_OPT: 1
-                            QS_SCF: 0
+                            MD: '1'
                         FILENAME: RESTART
                     RESTART_HISTORY:
                         _: 'OFF'
-                        BACKUP_COPIES: 0
+                        BACKUP_COPIES: '0'
             XC:
                 VDW_POTENTIAL:
                     DISPERSION_FUNCTIONAL: PAIR_POTENTIAL
                     PAIR_POTENTIAL:
                         CALCULATE_C9_TERM: .TRUE.
                         PARAMETER_FILE_NAME: dftd3.dat
                         REFERENCE_FUNCTIONAL: PBE
-                        R_CUTOFF: 15
+                        R_CUTOFF: '15'
                         TYPE: DFTD3
                 XC_FUNCTIONAL:
                         _: PBE
@@ -76,7 +86,7 @@ low_accuracy:
         EXTENDED_FFT_LENGTHS: ''
         PRINT_LEVEL: MEDIUM
         RUN_TYPE: GEO_OPT
-        WALLTIME: 600
+        WALLTIME: '600'
         ELPA_KERNEL: AVX2_BLOCK2
     MOTION:
         MD:

aiida_nanotech_empa/workflows/cp2k/reftraj_md_workchain.py

@@ -9,41 +9,14 @@
 TrajectoryData = plugins.DataFactory("array.trajectory")
 
 
-# function to check if all ase structures have the same cell
-def check_cell(ase_structures, cell_tolerance=1e-4):
-    """Check if all the structures in the trajectory have the same cell."""
-    cell = ase_structures[0].get_cell()
-    for s in ase_structures[1:]:
-        if not np.allclose(cell, s.get_cell(), atol=cell_tolerance):
-            return False
-    return True
-
-
-# fuction that takes a list of ase structures and returns different leists where in each list teh structures have the same cell and the same formula
-def split_structures(ase_structures):
-    """Split the structures in the trajectory in different lists based on the cell and the formula."""
-    formulas = [s.get_chemical_formula() for s in ase_structures]
-    unique_formulas = list(set(formulas))
-    split_structures = [
-        [s for s, f in zip(ase_structures, formulas) if f == uf]
-        for uf in unique_formulas
-    ]
-    return split_structures
-
-
 @calcfunction
 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."
-        )
-    split_structures = split_structures(ase_structures)
+    """Create a list of touples with integers as batches indeces. Start from 1 for CP2K input."""
+    n_structures = trajectory.get_array("positions").shape[0]
     batches = []
-    for ss in split_structures:
-        for i in range(0, len(ss), batch_size):
-            batches.append(ss[i : i + batch_size])
+    for i_batch in range(1, n_structures, batch_size):
+        batches.append((i_batch, min(i_batch + batch_size - 1, n_structures)))
+
     return batches
 
 
@@ -92,6 +65,13 @@ def define(cls, spec):
     def setup(self):
         """Initialize the workchain process."""
         self.report("Inspecting input and setting up things")
+
+        # create an ase Atoms object from the first structure of the trajectory
+        cell = self.inputs.trajectory.get_array("cells")[0]
+        positions = self.inputs.trajectory.get_array("positions")[0]
+        symbols = self.inputs.trajectory.symbols
+        ase_structure = Atoms(symbols, positions=positions, cell=cell)
+
         return engine.ExitCode(0)
 
     def first_scf(self):