added docs to twoslab_index.py and small bug fixes to scripts

python/twoslab_index.py

@@ -1,3 +1,37 @@
+"""
+This module provides a class for generating index files for a system with
+two slabs.
+
+The TwoSlabIdxGenerator class takes an input index file and an input gro file
+as arguments and generates an index file for a system with two slabs. It
+provides methods for updating the original index, adding the second slab to
+groups matching patterns, sorting the indices, and saving the index file.
+
+Example usage:
+    generator = TwoSlabIdxGenerator(input_idx, input_gro)
+    generator.generate(output_idx)
+
+Attributes:
+    input_idx (Path): The path to the input index file.
+    input_gro (Path): The path to the input gro file.
+    verbose (bool): Whether to print verbose output.
+
+Methods:
+    __init__(input_idx: Path, input_gro: Path, verbose: bool = True):
+        Initializes the TwoSlabIdxGenerator object.
+    __repr__(): Returns a string representation of the TwoSlabIdxGenerator
+        object.
+    print(): Prints the index groups and the number of atoms in each group.
+    _update_original_index(min_index: int, offset: int): Updates the original
+        index for all groups with an offset from the second slab.
+    _add_second_slab(min_index: int): Adds the second slab to all groups
+        matching patterns.
+    _sort(): Sorts the indices in each group.
+    generate(output_top: Path): Generates the index file for the system with
+        two slabs.
+    save(output_top: Path): Saves the index file to the specified path.
+"""
+
 # standard library
 from argparse import ArgumentParser
 from pathlib import Path
@@ -8,6 +42,34 @@
 
 
 class TwoSlabIdxGenerator:
+    """
+    A class for generating index files for a system with two slabs.
+
+    Args:
+        input_idx (Path): The path to the input index file.
+        input_gro (Path): The path to the input gro file.
+        verbose (bool, optional): Whether to print verbose output. Defaults to True.
+
+    Attributes:
+        ndx (gmx.fileformats.ndx.NDX): The parsed input index file.
+        input_gro (Path): The path to the input gro file.
+        gro (str): The content of the input gro file.
+        n_atoms (int): The total number of atoms in the system.
+        n_crystal_atoms (int): The number of crystal atoms in the first slab.
+        verbose (bool): Whether to print verbose output.
+
+    Methods:
+        __repr__(): Returns a string representation of the TwoSlabIdxGenerator object.
+        print(): Prints the index groups and the number of atoms in each group.
+        _update_original_index(min_index: int, offset: int): Updates the
+            original index for all groups with an offset from the second slab.
+        _add_second_slab(min_index: int): Adds the second slab to all groups
+            matching patterns.
+        _sort(): Sorts the indices in each group.
+        generate(output_top: Path): Generates the index file for the system
+            with two slabs.
+        save(output_top: Path): Saves the index file to the specified path.
+    """
 
     def __init__(self, input_idx: Path, input_gro: Path, verbose: bool = True) -> None:
         # parse input index file
@@ -35,8 +97,35 @@ def __init__(self, input_idx: Path, input_gro: Path, verbose: bool = True) -> No
             print(f"Number of crystal atoms in first slab: {self.n_crystal_atoms}")
             print(f"Index groups: {self.ndx.keys()}")
 
+    def __repr__(self) -> str:
+        return (
+            f"TwoSlabIdxGenerator(input_gro={self.input_gro},"
+            + f" n_atoms={self.n_atoms},"
+            + f" n_crystal_atoms={self.n_crystal_atoms},"
+            + f" verbose={self.verbose})"
+        )
+
+    def print(self) -> None:
+        """
+        Prints the index groups and the number of atoms in each group.
+        """
+        for key, value in self.ndx.items():
+            # left justify key with 25 columns
+            print(f"  {key:<25}: {value} atoms")
+
     def _update_original_index(self, min_index: int, offset: int) -> None:
-        # update original index for all groups with offset from second slab
+        """
+        Updates the original index for all groups with an offset from the
+        second slab.
+
+        Args:
+            min_index (int): The minimum index of the second slab.
+            offset (int): The offset to be added to the indices.
+
+        Raises:
+            ValueError: If the number of atoms in the System group does not
+            match the total number of atoms.
+        """
         for key, value in self.ndx.items():
             updated = [i + offset if i > min_index else i for i in value]
             self.ndx[key] = np.array(updated, dtype=int)
@@ -50,12 +139,22 @@ def _update_original_index(self, min_index: int, offset: int) -> None:
         total_atoms = self.n_atoms
         if total_atoms != len(self.ndx["System"]):
             raise ValueError(
-                f"Number of atoms mismatch in System group: {total_atoms} != {len(self.ndx['System'])}"
+                "Number of atoms mismatch in System group:"
+                + f" {total_atoms} != {len(self.ndx['System'])}"
             )
 
     def _add_second_slab(self, min_index: int) -> None:
-        # add second slab to all groups matching patterns
-        patterns = ["System", "non-Water", "Frozen", "Mobile", "Crystal"]
+        """
+        Adds the second slab to all groups matching patterns.
+
+        Args:
+            min_index (int): The minimum index of the second slab.
+
+        Raises:
+            ValueError: If the number of atoms in the System group does not
+            match the total number of atoms.
+        """
+        patterns = ["System", "non-Water", "non-Protein", "Frozen", "Mobile", "Crystal"]
         for key, value in self.ndx.items():
             if any(p in key for p in patterns):
                 group_new = np.array(
@@ -71,25 +170,40 @@ def _add_second_slab(self, min_index: int) -> None:
         total_atoms = self.n_atoms + self.n_crystal_atoms
         if total_atoms != len(self.ndx["System"]):
             raise ValueError(
-                f"Number of atoms mismatch in System group: {total_atoms} != {len(self.ndx['System'])}"
+                "Number of atoms mismatch in System group:"
+                + f" {total_atoms} != {len(self.ndx['System'])}"
             )
 
     def _sort(self) -> None:
+        """
+        Sorts the indices in each group.
+        """
         for key, value in self.ndx.items():
             self.ndx[key] = np.sort(value)
 
     def generate(self, output_top: Path) -> None:
+        """
+        Generates the index file for the system with two slabs.
+
+        Args:
+            output_top (Path): The path to save the generated index file.
+
+        Raises:
+            ValueError: If the number of atoms in the System group does not
+            match the total number of atoms.
+        """
         self._update_original_index(self.n_crystal_atoms, self.n_crystal_atoms)
         self._add_second_slab(self.n_crystal_atoms)
         self._sort()
-
-        assert self.n_atoms + self.n_crystal_atoms == max(
-            self.ndx["System"]
-        ), f"Number of atoms mismatch in System group: {self.n_atoms} != {max(self.ndx['System'])}"
-
         self.save(output_top)
 
     def save(self, output_top: Path) -> None:
+        """
+        Saves the index file to the specified path.
+
+        Args:
+            output_top (Path): The path to save the index file.
+        """
         self.ndx.write(output_top)
 
 
@@ -112,7 +226,13 @@ def main():
     args = parser.parse_args()
 
     generator = TwoSlabIdxGenerator(args.input_idx, args.input_gro, args.verbose)
+    if args.verbose:
+        print("Initial index groups:")
+        generator.print()
     generator.generate(args.output_idx)
+    if args.verbose:
+        print("Final index groups:")
+        generator.print()
 
 
 if __name__ == "__main__":

scripts/method/equilibration.sh

@@ -102,8 +102,6 @@ else
             sed -i 's/^ewald-geometry .*/ewald-geometry            = 3dc/g' "${sim_name}.mdp" || exit 1
             sed -i 's/^pbc .*/pbc                       = xy/g' "${sim_name}.mdp" || exit 1
             sed -i 's/^nwall .*/nwall                     = 2/g' "${sim_name}.mdp" || exit 1
-            # FIXME: wall-atomtype is leading to instability, may need to be removed
-            # sed -i 's/^wall-atomtype             = WR WL.*/wall-atomtype             = WL WL/g' "${sim_name}.mdp" || exit 1
             echo "wall-r-linpot = 0.1" >>"${sim_name}.mdp" || exit 1
         fi
 
@@ -206,9 +204,6 @@ else
             sed -i 's/^ewald-geometry .*/ewald-geometry            = 3dc/g' "${sim_name}.mdp" || exit 1
             sed -i 's/^pbc .*/pbc                       = xy/g' "${sim_name}.mdp" || exit 1
             sed -i 's/^nwall .*/nwall                     = 2/g' "${sim_name}.mdp" || exit 1
-            # FIXME: wall-atomtype is leading to instability, may need to be removed
-            # sed -i 's/^wall-atomtype             = WR WL.*/wall-atomtype             = WL WL/g' "${sim_name}.mdp" || exit 1
-            # echo "wall-r-linpot = 0.1" >>"${sim_name}.mdp" || exit 1
         fi
 
         # make tpr file

scripts/method/initialization.sh

@@ -111,8 +111,6 @@ echo "INFO: Copying input files to working directory"
         sed -i 's/^ewald-geometry .*/ewald-geometry            = 3dc/g' "mdin.mdp"
         sed -i 's/^pbc .*/pbc                       = xy/g' "mdin.mdp"
         sed -i 's/^nwall .*/nwall                     = 2/g' "mdin.mdp"
-        # FIXME: this is not working
-        # sed -i 's/^wall-atomtype             = WR WL.*/wall-atomtype             = WL WL/g' "mdin.mdp"
     fi
 } >>"${log_file}" 2>&1
 

scripts/utilities/second_slab.sh

@@ -111,19 +111,17 @@ fi
 # ##############################################################################
 # Remake system topology #######################################################
 # ##############################################################################
-
 # update number of atoms in topology file
 if [[ -f "${archive_dir}/topol.top" ]]; then
     echo "DEBUG: Topology file already updated"
-
 else
     echo "DEBUG: Updating topology file"
     {
         # get all lines [ molecules ] and following from topol.top
         grep -A 1000 "molecules" topol.top >temp.txt
         n_mol_slab="$(grep "Ion_chain_X" temp.txt | awk '{print $2}')"
         n_mol_2slab="$(echo "2 * ${n_mol_slab}" | bc)"
-        sed -i "s/${n_mol_slab}/${n_mol_2slab}/" temp.txt
+        sed -i "s/Ion_chain_X.*/Ion_chain_X         ${n_mol_2slab}/" temp.txt
         # replace [ molecules ] and following in topol.top
         sed -i '/molecules/,$d' topol.top
         cat temp.txt >>topol.top
@@ -138,34 +136,21 @@ fi
 # ##############################################################################
 # Remake system index ##########################################################
 # ##############################################################################
-
-# Current groups:
-
 if [[ -f "${archive_dir}/index.ndx" ]]; then
     echo "DEBUG: Index file already updated"
 else
-    # call python script to update index file
-    echo "DEBUG: Updating index file"
-    cp -np "index.ndx" "${archive_dir}/index_1slab.ndx" || exit 1
-    "${GMX_BIN}" make_ndx \
-        -f "${sim_name}.gro" \
-        -n "index.ndx" \
-        -o "index.ndx" <<EOF
-q
-EOF
-
-    python3 "${project_path}/python/twoslab_index.py" \
-        -i "index.ndx" \
-        -g "${sim_name}.gro" \
-        -o "index.ndx" \
-        -v || exit 1
-
-    cp "index.ndx" "${archive_dir}/"
-    cp -np "index.ndx" "${archive_dir}/index_1slab.ndx" || exit 1
-    "${GMX_BIN}" make_ndx \
-        -f "${sim_name}.gro" \
-        -n "index.ndx" \
-        -o "index.ndx" <<EOF
-q
-EOF
+    {
+        python3 "${project_path}/python/twoslab_index.py" \
+            -i "index.ndx" \
+            -g "${sim_name}.gro" \
+            -o "index.ndx" \
+            -v || exit 1
+        cp "index.ndx" "${archive_dir}/"
+    } >>"${log_file}" 2>&1
 fi
+
+# ##############################################################################
+# Clean up #####################################################################
+# ##############################################################################
+# delete all backup files
+find . -type f -name '#*#' -delete || true