Fill out virtual_sites position functionality (#934)

* Fill out virtual_sites position functionality

* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci

* move expression evaluation for position to gmso expression.py, and add and test for different exceptions

* Identify units from gmso site positions, instead of hard coded as nm

* Remove unncessary variable

---------

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

Author: 3 people

gmso/core/virtual_site.py

@@ -1,11 +1,13 @@
+import string
 from typing import Callable, List, Optional, Union
 
 import unyt as u
 from pydantic import ConfigDict, Field
+from sympy import Matrix, symbols
 
 from gmso.abc.abstract_site import Site
 from gmso.core.virtual_type import VirtualType
-from gmso.exceptions import MissingPotentialError, NotYetImplementedWarning
+from gmso.exceptions import MissingPotentialError
 
 
 class VirtualSite(Site):
@@ -58,8 +60,8 @@ def parent_sites(self) -> List[Site]:
         """Reminder that the order of sites is fixed, such that site index 1 corresponds to ri in the self.virtual_type.virtual_position expression."""
         return self.__dict__.get("parent_sites_", [])
 
-    def position(self) -> str:
-        """Not yet implemented function to get position from virtual_type.virtual_position and parent_sites."""
+    def position(self) -> u.unyt_array:
+        """On the fly position evaluation from virtual_type.virtual_position and parent_sites."""
         if not self.virtual_type:
             raise MissingPotentialError(
                 "No VirtualType associated with this VirtualSite."
@@ -68,10 +70,26 @@ def position(self) -> str:
             raise MissingPotentialError(
                 "No VirtualPositionType associated with this VirtualType."
             )
-        # TODO: validate parent atoms matches virtual_type.virtual_position in terms of independent variables ri, rj, etc.
-        # TODO: Generate position from atoms of parent_atoms and self.virtual_type.virtual_position.expression.
-        raise NotYetImplementedWarning(
-            "Need a functional to call from self.virtual_type.virtual_position, and plug in ri, rj, rk etc."
+
+        independent_namespace = {}
+        for _, symbol in zip(range(len(self.parent_sites)), string.ascii_lowercase[8:]):
+            x, y, z = symbols(f"r{symbol}1 r{symbol}2 r{symbol}3")
+            independent_namespace[f"r{symbol}"] = Matrix([x, y, z])
+
+        independent_parameters = {}
+        for symbol, site in zip(string.ascii_lowercase[8:], self.parent_sites):
+            for i, pos in enumerate(site.position):
+                independent_parameters[f"r{symbol}{i + 1}"] = float(pos.value)
+
+        # get units from parent sites
+        unitsUnyt = self.parent_sites[0].position.units
+
+        # perform expression evaluation
+        return (
+            self.virtual_type.virtual_position.potential_expression.evaluate(
+                independent_namespace, independent_parameters
+            )
+            * unitsUnyt
         )
 
     def __repr__(self):

gmso/tests/test_expression.py

@@ -1,7 +1,9 @@
+import numpy as np
 import pytest
 import sympy
 import unyt as u
 
+from gmso.exceptions import GMSOError
 from gmso.tests.base_test import BaseTest
 from gmso.utils.expression import PotentialExpression, _are_equal_parameters
 
@@ -331,3 +333,68 @@ def test_are_equal_parameters(self):
         u1 = {"a": 2.0 * u.nm, "b": 3.5 * u.nm}
         u2 = {"c": 2.0 * u.nm, "d": 3.5 * u.nm}
         assert _are_equal_parameters(u1, u2) is False
+
+    def test_evaluate_expression(self):
+        import sympy
+
+        from gmso.utils.expression import PotentialExpression
+
+        expression = PotentialExpression(
+            expression="4*epsilon*((sigma/r)**12 - (sigma/r)**6)",
+            parameters={"sigma": 1 * u.nm, "epsilon": -32.507936507936506 * u.kJ},
+            independent_variables="r",
+        )
+        r = sympy.symbols("r")
+        independent_namespace = {"r": r}
+        independent_parameters = {"r": 2}
+        assert expression.evaluate(independent_namespace, independent_parameters) == 2
+
+        # test use of norm in expression on vectors
+        expression = PotentialExpression(
+            expression="norm(a)*r",
+            parameters={"a": [3, 4, 0] * u.nm},  # norm is 5
+            independent_variables="r",
+        )
+        x, y, z = sympy.symbols("x y z")
+        independent_namespace["r"] = sympy.Matrix([x, y, z])
+        independent_parameters = {"x": 0 * u.nm, "y": 1 * u.nm, "z": 2 * u.nm}
+        output = expression.evaluate(independent_namespace, independent_parameters)
+        assert all(output == np.array([0, 1, 2]) * 5)
+
+        with pytest.raises(GMSOError):  # fail with missing r parameters
+            expression = PotentialExpression(
+                expression="4*epsilon*((sigma/r)**12 - (sigma/r)**6)",
+                parameters={"sigma": 1 * u.nm, "epsilon": -32.507936507936506 * u.kJ},
+                independent_variables="r",
+            )
+            expression.evaluate()
+
+        with pytest.raises(GMSOError):  # fail with divide by 0
+            expression = PotentialExpression(
+                expression="a/b",
+                parameters={"a": 1 * u.nm, "b": 0 * u.kJ},
+                independent_variables=[],
+            )
+            expression.evaluate()
+
+        with pytest.raises(GMSOError):  # fail in vector divide by 0
+            expression = PotentialExpression(
+                expression="a/r",
+                parameters={"a": 1 * u.nm},
+                independent_variables="r",
+            )
+            x, y, z = sympy.symbols("x y z")
+            independent_namespace["r"] = sympy.Matrix([x, y, z])
+            independent_parameters = {"x": 1, "y": 0, "z": 0}
+            expression.evaluate(independent_namespace, independent_parameters)
+
+        with pytest.raises(GMSOError):  # fail with bad expression
+            expression = PotentialExpression(
+                expression="normal(a/r)",
+                parameters={"a": 1 * u.nm},
+                independent_variables="r",
+            )
+            x, y, z = sympy.symbols("x y z")
+            independent_namespace["r"] = sympy.Matrix([x, y, z])
+            independent_parameters = {"x": 1, "y": 0, "z": 0}
+            expression.evaluate(independent_namespace, independent_parameters)

gmso/tests/test_virtual_site.py

@@ -1,3 +1,4 @@
+import numpy as np
 import pytest
 import unyt as u
 from sympy import symbols, sympify
@@ -16,7 +17,7 @@
 class TestVirturalSite(BaseTest):
     @pytest.fixture(scope="session")
     def virtual_site(self):
-        site = Atom()
+        site = Atom(position=[1, 1, 1])
         return VirtualSite(parent_sites=[site])
 
     @pytest.fixture(scope="session")
@@ -34,8 +35,30 @@ def test_new_site(self, water_system):
         for site in v_site.parent_sites:
             assert site in water_system.sites
 
-    def test_virtual_position(self):
-        # TODO: Check position as a function of virtual_position_type
+    def test_virtual_position(self, virtual_site):
+        # Check position as a function of virtual_position_type
+        # TODO: check for arrays for b, sin norm
+        # TODO: checkk all gromacs potential forms
+
+        v_pot = VirtualPotentialType(
+            expression="5*a*b",
+            independent_variables={"a"},
+            parameters={"b": 1 * u.kJ},
+        )
+        v_pos = VirtualPositionType(
+            expression="ri*cos(b)",
+            independent_variables=["ri"],
+            parameters={"b": np.pi * u.radian},
+        )
+        assert v_pos
+        v_type = VirtualType(virtual_potential=v_pot, virtual_position=v_pos)
+        virtual_site.virtual_type = v_type  # assign virtual type
+        assert_allclose_units(virtual_site.position(), -1 * ([1, 1, 1] * u.nm))
+
+    def test_tip4p_water(self):
+        pass
+
+    def test_tip5p_water(self):
         pass
 
     def test_virtual_type(self):

gmso/utils/expression.py

@@ -5,14 +5,32 @@
 from functools import lru_cache
 from typing import Dict
 
+import numpy as np
 import sympy
 import unyt as u
+from sympy import Add, Function, Mul, Symbol, lambdify, sympify
+
+from gmso.exceptions import GMSOError
 
 logger = logging.getLogger(__name__)
 
 __all__ = ["PotentialExpression"]
 
 
+class norm(Function):
+    """Sympy functions for use in lambdify"""
+
+    @classmethod
+    def eval(cls, arg):
+        return None
+
+
+# Evaluate vector norm
+def norm_evaluation(matrix_arg):
+    """Evaluation method for norm of a sympy matrix"""
+    return np.linalg.norm(matrix_arg)
+
+
 def _are_equal_parameters(u1, u2):
     """Compare two parameters of unyt quantities/arrays.
 
@@ -407,10 +425,8 @@ def from_non_parametric(
         ----------
         non_parametric: PotentialExpression
             The non-parametric potential expression to create the parametric one from
-
         parameters: dict
             The dictionary of parameters for the newly created parametric expression.
-
         valid: bool, default=False
             Whether to validate expression/independent_variables and with the parameters.
 
@@ -442,3 +458,103 @@ def from_non_parametric(
                 independent_variables=deepcopy(non_parametric.independent_variables),
                 verify_validity=not valid,
             )
+
+    def evaluate(
+        self, independent_namespace: dict = None, independent_parameters: dict = None
+    ):
+        """Evaluate the sympy expression with the given parameters
+
+        Parameters
+        ----------
+        independent_namespace: dict, default None
+            Dictionary with keys that are the string of the symbol to evaluate and values are the sympy Symbol of Function object
+        independent_parameters: dict, default None
+            Keys are strings, and values are the unyt value that corresponds to that symbol in the expression.
+
+        Notes
+        -----
+        Evaluate the LJ expression as follows:
+            ```python
+            from gmso.utils.expression import PotentialExpression
+            import sympy
+            expression = PotentialExpression(
+                expression="4*epsilon*((sigma/r)**12 - (sigma/r)**6)",
+                parameters={"sigma":1*u.nm, "epsilon":-32.507936507936506*u.kJ},
+                independent_variables="r"
+            )
+            r = sympy.symbols(f"r")
+            independent_namespace = {"r": r} # key to symbol
+            independent_parameters = {"r": 2} # key to value
+            expression.evaluate(independent_namespace, independent_parameters)
+            ```
+        The input "norm" in the expression will be evaluated as the Matrix normal of the variable, i.e. the result from np.linalg.norm. See
+        class norm(Function) in the module for the definition and evaluation procedure for this symbol, and possibility to add more shorthand
+        function evaluations in the lambdify namespace.
+
+        Returns
+        -------
+        result: numpy.ndarray or float
+            The expression evaluated as a function with parameters plugged in.
+        """
+        # prep namespace for sympify
+        if independent_namespace is None:
+            namespace = {}
+            args = []
+        else:  # grab symbols from namespace
+            args = [Symbol(key) for key in independent_parameters]
+            namespace = (
+                independent_namespace.copy()
+            )  # make copies to not overwrite dictionary objects
+        namespace.update(
+            {sym: Symbol(sym) for sym in self.parameters}
+        )  # expression parameters
+        namespace["norm"] = norm  # handle Matrix/Vector normalization
+        args.extend([namespace[sym] for sym in self.parameters])  # args to lambdify
+        expr_string = str(self.expression)
+        if independent_parameters is None:
+            parameters = {}
+        else:
+            parameters = independent_parameters.copy()
+
+        # parse expression
+        try:
+            expr = sympify(expr_string, locals=namespace)
+        except (ValueError, TypeError):
+            raise GMSOError(
+                f"Expression {expr_string=} was not viable in sympy for PotentialExpression object:{self}."
+            )
+
+        f = lambdify(args, expr, modules=[{"norm": norm_evaluation}, "numpy"])
+
+        # evaluate
+        parameters.update(
+            {param: val.to_value() for param, val in self.parameters.items()}
+        )
+        # apply parameters here into lamdify'ed object
+        try:
+            result = f(**parameters)
+        except ZeroDivisionError:
+            raise GMSOError(
+                f"PotentialExpression {self} is unabel to be evaluated since the result was a divide by 0 issue."
+            )
+        except NameError as e:
+            raise GMSOError(f"PotentialExpresion {self} raise: " + str(e))
+
+        # error handling
+        if isinstance(result, float):
+            return result  # TODO: Attach units here also processed via sympy expression
+        elif isinstance(result, (Symbol, Add, Mul)) or not np.issubdtype(
+            result.dtype, np.floating
+        ):
+            raise GMSOError(
+                f"{self=} expression was not able to be fully evaluated. Unknown parameters left defined in {expr_string=}. Position evaluation is: {result=}"
+            )
+        elif any(np.isnan(result)):
+            raise GMSOError(
+                f"Failed evaluation of {self=} with {expr_string=} and {parameters=}."
+            )
+        if len(result) > 1:  # return a vector
+            result = np.array(result).T[0]
+        else:
+            result = np.array(result)
+        return result