Fixed pytest and Centre of Mass for issue MDAnalysis#3603

Author: Echo Fia

Diff for: testsuite/MDAnalysisTests/analysis/test_lineardensity.py

@@ -1,3 +1,25 @@
+# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
+# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
+#
+# MDAnalysis --- https://www.mdanalysis.org
+# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
+# (see the file AUTHORS for the full list of names)
+#
+# Released under the Lesser GNU Public Licence, v2.1 or any higher version
+#
+# Please cite your use of MDAnalysis in published work:
+#
+# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
+# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
+# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
+# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
+# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
+# doi: 10.25080/majora-629e541a-00e
+#
+# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
+# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
+# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
+#
 import MDAnalysis as mda
 import numpy as np
 import pytest
@@ -8,9 +30,11 @@
 from numpy.testing import assert_allclose
 from MDAnalysis.core._get_readers import get_reader_for
 from MDAnalysisTests.util import no_deprecated_call
+
 from MDAnalysis.units import constants
 
 
+
 def test_invalid_grouping():
     """Invalid groupings raise AttributeError"""
     universe = mda.Universe(waterPSF, waterDCD)
@@ -21,26 +45,6 @@ def test_invalid_grouping():
         ld = LinearDensity(selection, grouping="centroid", binsize=5)
         ld.run()
 
-### Initialising Variables ###
-
-expected_masses_atoms = None
-expected_charges_atoms = None
-expected_xmass_atoms = None
-expected_xcharge_atoms = None
-
-expected_masses_residues = None
-expected_charges_residues = None
-expected_xmass_residues = None
-expected_xcharge_residues = None
-
-expected_masses_segments = None
-expected_charges_segments = None
-expected_xmass_segments = None
-expected_xcharge_segments = None
-
-'''
-FRAGMENTS!!!! Can't find the documentation for Fragment
-'''
 
 ### Creating the Test Universes ###
 
@@ -52,8 +56,8 @@ def make_Universe(coords, charges, masses, n_atoms, n_frames, atomsPerRes, resPe
     n_segments = n_residues // resPerSegs # Arbitrarily 4 residues per segment
 
     # Indexing atoms into residues & residues into segments
-    atom_resindex = np.array([[i] * (n_atoms // n_residues) for i in range(n_residues)]).flatten()
-    residue_segindex=np.array([[i] * (n_residues // n_segments) for i in range(n_segments)]).flatten()
+    atom_resindex = np.repeat(np.arange(n_residues), atomsPerRes)
+    residue_segindex = np.repeat(np.arange(n_segments), resPerSegs)
 
     # Creating the universe
     u = mda.Universe.empty(n_atoms=n_atoms,
@@ -75,6 +79,14 @@ def make_Universe(coords, charges, masses, n_atoms, n_frames, atomsPerRes, resPe
 
     return u
 
+
+### Defining the Systems ###
+
+# The masses are all set to 1
+# The charges are dependent on each system
+# The positions are randomly taken from a uniform distribution between 0 and 1
+# NOTE: The positions at each time frame are independent on the previous time step
+
 def neutral_Particles(n_atoms, n_frames, atomsPerRes, resPerSegs):
     charges = np.zeros(n_atoms)
     masses = np.ones(n_atoms)
@@ -99,7 +111,6 @@ def charged_Dimers(n_dimers, n_frames, dimersPerRes, resPerSegs, dimerLength = 0
     charges = np.random.random(n_atoms) * 2 - np.ones(n_atoms) # Between -1 and 1
     masses = np.ones(n_atoms)
 
-    # Setting each position to be random for each timestep (independent of previous timestep)
     shape = (n_frames, n_dimers, 3)
     coords = np.random.random(shape) * 0.9 + np.ones(shape) * 0.05
     # Puts in the same coordinate twice per dimer
@@ -116,6 +127,7 @@ def charged_Dimers(n_dimers, n_frames, dimersPerRes, resPerSegs, dimerLength = 0
 
     return make_Universe(coords, charges, masses, n_atoms, n_frames, dimersPerRes * 2, resPerSegs)
 
+
 ### Calculating the Expected Values ###
 
 def calc_Prop(u, prop = 'masses'): # Property can be 'masses' or 'charges'
@@ -125,15 +137,14 @@ def calc_Prop(u, prop = 'masses'): # Property can be 'masses' or 'charges'
 
     return expected_atoms, expected_residues, expected_segments
 
-def find_Centres(groups, prop): 
+def find_Centres(groups): # Always based on Centre of Mass (NOT Centre of Charge)
     centres = []
     for group in groups:
-        # NOTE: Absolute is taken for charges
-        total_Prop = sum(abs(eval(f'group.atoms.{prop}')))
-        if total_Prop != 0:
-            centres.append(np.sum(group.atoms.positions.transpose() * abs(eval(f'group.atoms.{prop}')), axis = 1) / total_Prop)
-        elif total_Prop == 0:
-            centres.append(np.sum(group.atoms.positions.transpose() * abs(eval(f'group.atoms.{prop}')), axis = 1) / len(group.atoms))
+        total_Mass = sum(group.atoms.masses)
+        if total_Mass != 0:
+            centres.append(np.sum(group.atoms.positions.transpose() * abs(group.atoms.masses), axis = 1) / total_Mass)
+        elif total_Mass == 0:
+            centres.append(np.sum(group.atoms.positions.transpose() * abs(group.atoms.masses), axis = 1) / len(group.atoms)) ############## NOT NEDEDEDDED 
 
     return np.array(centres)
 
@@ -153,7 +164,7 @@ def calc_Densities(u, prop = 'masses', spliceLen = 0.25): # Property can be 'mas
     _,residue_Totals,segment_Totals = calc_Prop(u, prop) # Total of Charge OR Mass
     ### Residues
     expected_residues = np.zeros((3, int(u.dimensions[0] // spliceLen))).astype(float)
-    residue_Centres = find_Centres(u.residues, prop = prop)
+    residue_Centres = find_Centres(u.residues)
 
 
     for i in range(len(residue_Centres)):
@@ -162,15 +173,14 @@ def calc_Densities(u, prop = 'masses', spliceLen = 0.25): # Property can be 'mas
 
     ### Segments
     expected_segments = np.zeros((3, int(u.dimensions[0] // spliceLen))).astype(float)
-    segment_Centres = find_Centres(u.segments, prop = prop)
+    segment_Centres = find_Centres(u.segments)
 
 
     for i in range(len(segment_Centres)):
         for j in range(3):
             expected_segments[j][int(segment_Centres[i][j] // spliceLen)] += segment_Totals[i]
 
 
-
     # Scaling based on splice volumes & converting units
     for i in range(3):
         expected_atoms[i,:] /= spliceLen * u.dimensions[(i + 1) % 3] * u.dimensions[(i + 2) % 3]
@@ -181,84 +191,72 @@ def calc_Densities(u, prop = 'masses', spliceLen = 0.25): # Property can be 'mas
     expected_segments /= constants['N_Avogadro'] * 1e-24 # To be consistent with lineardensity.py
 
     return expected_atoms, expected_residues, expected_segments
-    
 
+### Creating the Expected Values ###
+
+spliceLen = 0.25
+
+universes = []
+groupings = ['atoms', 'residues', 'segments']
+# Will be [[Atoms, Residues, Segments (of Universe 1)], [... of Universe 2], ...]
+expected_masses = [] 
+expected_charges = []
+expected_xmass = []
+expected_xcharge = []
 
-####
+for system in test_Systems:
 
+    universe = eval(f'{system}(100, 1, 5, 4)')
+    universes.append(universe)
 
+##    expected_masses_atoms, expected_masses_residues, expected_masses_segments = calc_Prop(universe, 'masses')
+##    expected_charges_atoms, expected_charges_residues, expected_charges_segments = calc_Prop(universe, 'charges')
+##    expected_xmass_atoms, expected_xmass_residues, expected_xmass_segments = calc_Densities(universe, 'masses', spliceLen)
+##    expected_xcharge_atoms, expected_xcharge_residues, expected_xcharge_segments = calc_Densities(universe, 'charges', spliceLen)
+
+    expected_masses.append(calc_Prop(universe, 'masses'))
+    expected_charges.append(calc_Prop(universe, 'charges'))
+    expected_xmass.append(calc_Densities(universe, 'masses', spliceLen))
+    expected_xcharge.append(calc_Densities(universe, 'charges', spliceLen))
+
+### Parametrizing for Pytest ###
+
+# NOTE: There is an extra [0] after the expected_xmass and expected_xcharge as the original data has all 3 co-ords, but only comparing to x here
+params = [(universes[i], groupings[j], expected_masses[i][j], expected_charges[i][j], expected_xmass[i][j][0], expected_xcharge[i][j][0]) for j in range(len(groupings)) for i in range(len(universes))]
 @pytest.mark.parametrize(
-    "grouping, expected_masses, expected_charges, expected_xmass, expected_xcharge",
-    [
-        (
-            "atoms",
-            expected_masses_atoms,
-            expected_charges_atoms,
-            expected_xmass_atoms,
-            expected_xcharge_atoms,
-        ),
-        (
-            "residues",
-            expected_masses_residues,
-            expected_charges_residues,
-            expected_xmass_residues,
-            expected_xcharge_residues,
-        ),
-        (
-            "segments",
-            expected_masses_segments,
-            expected_charges_segments,
-            expected_xmass_segments,
-            expected_xcharge_segments,
-        ),
-##        (
-##            "fragments",
-##            expected_masses_fragments,
-##            expected_charges_fragments,
-##            expected_xmass_fragments,
-##            expected_xcharge_fragments,
-##        ),
-    ],
+    "universe, grouping, expected_masses, expected_charges, expected_xmass, expected_xcharge",
+    params,
 )
 
 def test_lineardensity(
-##    universe,
+    universe,
     grouping,
-##    expected_masses,
-##    expected_charges,
-##    expected_xmass,
-##    expected_xcharge,
+    expected_masses,
+    expected_charges,
+    expected_xmass,
+    expected_xcharge,
 ):
+    sel_string = "all"
+    selection = universe.select_atoms(sel_string)
+    ld = LinearDensity(selection, grouping, binsize=0.25).run()
+    assert_allclose(ld.masses, expected_masses)
+    assert_allclose(ld.charges, expected_charges)
+    # rtol changed here due to floating point imprecision
+    assert_allclose(ld.results.x.mass_density, expected_xmass, rtol=1e-06)
+    assert_allclose(ld.results.x.charge_density, expected_xcharge)
 
-    spliceLen = 0.25
-    for system in test_Systems:
-        universe = eval(f'{system}(100, 1, 1, 1)')
-
-        expected_masses_atoms, expected_masses_residues, expected_masses_segments = calc_Prop(universe, 'masses')
-        expected_charges_atoms, expected_charges_residues, expected_charges_segments = calc_Prop(universe, 'charges')
-        expected_xmass_atoms, expected_xmass_residues, expected_xmass_segments = calc_Densities(universe, 'masses', spliceLen)
-        expected_xcharge_atoms, expected_xcharge_residues, expected_xcharge_segments = calc_Densities(universe, 'charges', spliceLen)
-
-        if grouping == 'atoms':
-            expected_masses, expected_charges, expected_xmass, expected_xcharge = expected_masses_atoms, expected_charges_atoms, expected_xmass_atoms, expected_xcharge_atoms
-
-        elif grouping == 'residues':
-            expected_masses, expected_charges, expected_xmass, expected_xcharge = expected_masses_residues, expected_charges_residues, expected_xmass_residues, expected_xcharge_residues
-
-        elif grouping == 'segments':
-            expected_masses, expected_charges, expected_xmass, expected_xcharge = expected_masses_segments, expected_charges_segments, expected_xmass_segments, expected_xcharge_segments
-    
-  
-        sel_string = "all"
-        selection = universe.select_atoms(sel_string)
-        ld = LinearDensity(selection, grouping, binsize=spliceLen).run()
-        assert_allclose(ld.masses, expected_masses)
-        assert_allclose(ld.charges, expected_charges)
-        # rtol changed here due to floating point imprecision
-        assert_allclose(ld.results.x.mass_density, expected_xmass[0], rtol=1e-06)
-        assert_allclose(ld.results.x.charge_density, expected_xcharge[0])
 
-##test_lineardensity('atoms')
+for i in range(len(params)):
+    print(i)
+    universe = params[i][0]
+    sel_string = "all"
+    selection = universe.select_atoms(sel_string)
+    ld = LinearDensity(selection, params[i][1], binsize=0.25).run()
+    assert_allclose(ld.masses, params[i][2])
+    assert_allclose(ld.charges, params[i][3])
+    # rtol changed here due to floating point imprecision
+    assert_allclose(ld.results.x.mass_density, params[i][4], rtol=1e-06)
+    assert_allclose(ld.results.x.charge_density, params[i][5])
 
 
 @pytest.fixture(scope="module")
@@ -362,3 +360,4 @@ def test_parallel_analysis(testing_Universe):
     assert_allclose(
         ld1.results.x.mass_density, ld_whole.results.x.mass_density
     )
+