Fix Lattice system class for coarse-grained molecules (#129)

* add test that makes lattice with bead spring polymer

* add warning in align z axis method

* add notes to Lattice class

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix precommit issue

* fix typo

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: chrisjonesBSU

flowermd/base/molecule.py

@@ -3,6 +3,7 @@
 import itertools
 import os.path
 import random
+import warnings
 from typing import List
 
 import mbuild as mb
@@ -211,13 +212,20 @@ def _align_backbones_z_axis(self, heavy_atoms_only=False):
         backbone_direction = np.array([0, 0, 1])
         for mol in self.molecules:
             if heavy_atoms_only:
-                positions = np.array(
-                    [
-                        p.xyz[0]
-                        for p in mol.particles()
-                        if p.element.symbol != "H"
-                    ]
-                )
+                try:
+                    positions = np.array(
+                        [
+                            p.xyz[0]
+                            for p in mol.particles()
+                            if p.element.symbol != "H"
+                        ]
+                    )
+                except AttributeError:
+                    positions = mol.xyz
+                    warnings.warn(
+                        "No element information found."
+                        "Using all particle positions to fit backbone axis."
+                    )
             else:
                 positions = mol.xyz
             backbone = get_backbone_vector(positions)

flowermd/base/system.py

@@ -712,6 +712,13 @@ class Lattice(System):
     basis_vector : array-like, default [0.5, 0.5, 0]
         The vector between points in the unit cell.
 
+    Notes
+    -----
+    The system is built in a way that the long axis of the
+    molecules is aligned with the z direction, and the
+    lattice is made by repeating and translating in the
+    x and y directions.
+
     """
 
     def __init__(

flowermd/tests/base/test_system.py

@@ -10,7 +10,12 @@
 
 from flowermd import Lattice, Pack
 from flowermd.internal.exceptions import ForceFieldError, ReferenceUnitError
-from flowermd.library import OPLS_AA, OPLS_AA_DIMETHYLETHER, OPLS_AA_PPS
+from flowermd.library import (
+    OPLS_AA,
+    OPLS_AA_DIMETHYLETHER,
+    OPLS_AA_PPS,
+    LJChain,
+)
 from flowermd.tests import BaseTest
 
 
@@ -811,6 +816,10 @@ def test_lattice_molecule(self, benzene_molecule):
         assert system.n_particles == system.hoomd_snapshot.particles.N
         assert system.reference_values.keys() == {"energy", "length", "mass"}
 
+    def test_lattice_bead_spring(self):
+        chains = LJChain(lengths=10, num_mols=32)
+        Lattice(molecules=chains, x=1, y=1, n=4)
+
     def test_scale_charges(self, pps):
         pps_mol = pps(num_mols=5, lengths=5)
         no_scale = Pack(

flowermd/tests/library/test_polymers.py

@@ -47,12 +47,14 @@ def test_lj_chain(self):
         cg_chain = LJChain(
             lengths=3,
             num_mols=1,
-            bead_sequence=["A"],
-            bead_mass={"A": 100},
-            bond_lengths={"A-A": 1.5},
+            bead_sequence=["_A"],
+            bead_mass={"_A": 100},
+            bond_lengths={"_A-_A": 1.5},
         )
         assert cg_chain.n_particles == 3
         assert cg_chain.molecules[0].mass == 300
+        with pytest.warns():
+            cg_chain._align_backbones_z_axis(heavy_atoms_only=True)
 
     def test_lj_chain_sequence(self):
         cg_chain = LJChain(