make mass optional

marrink-lab · Nov 23, 2023 · 737b45c · 737b45c
1 parent 5207801
commit 737b45c
Showing 1 changed file with 13 additions and 10 deletions.
diff --git a/polyply/src/ffoutput.py b/polyply/src/ffoutput.py
@@ -159,16 +159,19 @@ def write_atoms_block(self, nodes):
 
         for idx, (node, attrs) in enumerate(nodes, start=1):
             write_attrs = {attr: str(attrs[attr]) for attr in self.normal_order_block_atoms if attr in attrs}
-            self.stream.write('{idx:>{max_length[idx]}} '
-                              '{atype:<{max_length[atype]}} '
-                              '{resid:>{max_length[resid]}} '
-                              '{resname:<{max_length[resname]}} '
-                              '{atomname:<{max_length[atomname]}} '
-                              '{charge_group:>{max_length[charge_group]}} '
-                              '{charge:>{max_length[charge]}} '
-                              '{mass:>{max_length[mass]}}\n'.format(idx=idx,
-                                                                    max_length=max_length,
-                                                                    **write_attrs))
+            template = ('{idx:>{max_length[idx]}} '
+                        '{atype:<{max_length[atype]}} '
+                        '{resid:>{max_length[resid]}} '
+                        '{resname:<{max_length[resname]}} '
+                        '{atomname:<{max_length[atomname]}} '
+                        '{charge_group:>{max_length[charge_group]}} '
+                        '{charge:>{max_length[charge]}} ')
+            if 'mass' in write_attrs:
+                template += '{mass:>{max_length[mass]}}\n'
+            else:
+                template += '\n'
+
+            self.stream.write(template.format(idx=idx, max_length=max_length, **write_attrs))
 
     def write_atoms_link(self, nodes, nometa=False):
         """