Merge branch 'lanl:development' into development

.github/workflows/deploy.yml

@@ -31,4 +31,4 @@ jobs:
         python setup.py sdist bdist_wheel
     - name: Publish distribution 📦 to PyPI
       if: startsWith(github.event.ref, 'refs/tags') || github.event_name == 'release'
-      uses: pypa/gh-action-pypi-publish@release/v1
+      uses: pypa/gh-action-pypi-publish@release/v1

hippynn/graphs/nodes/physics.py

@@ -7,7 +7,7 @@
 from .base.node_functions import NodeNotFound
 from .indexers import AtomIndexer, PaddingIndexer, acquire_encoding_padding
 from .pairs import OpenPairIndexer
-from .tags import Encoder, PairIndexer, Charges
+from .tags import Encoder, PairIndexer, Charges, Energies
 from .inputs import PositionsNode, SpeciesNode
 
 from ..indextypes import IdxType, index_type_coercion, elementwise_compare_reduce
@@ -120,20 +120,21 @@ def expansion1(self, charges, species, *, purpose, **kwargs):
     @_parent_expander.match(Charges, _BaseNode, SpeciesNode)
     def expansion2(self, charges, pos_or_pair, species, *, purpose, **kwargs):
         encoder, pidxer = acquire_encoding_padding(species, species_set=None, purpose=purpose)
-        return charges, pos_or_pair, encoder, pidxer
+        return charges, pos_or_pair, pidxer
 
-    @_parent_expander.match(Charges, PositionsNode, Encoder, PaddingIndexer)
-    def expansion3(self, charges, positions, encoder, pidxer, *, cutoff_distance, **kwargs):
+    @_parent_expander.match(Charges, PositionsNode, PaddingIndexer)
+    def expansion3(self, charges, positions, pidxer, *, cutoff_distance, **kwargs):
         try:
-            pairfinder = find_unique_relative((charges, positions, encoder, pidxer), PairIndexer)
+            pairfinder = find_unique_relative((charges, positions, pidxer), PairIndexer)
         except NodeNotFound:
             warnings.warn("Boundary conditions not specified, Building open boundary conditions.")
+            encoder = find_unique_relative(pidxer, Encoder)
             pairfinder = OpenPairIndexer("PairIndexer", (positions, encoder, pidxer), dist_hard_max=cutoff_distance)
         return charges, pairfinder, pidxer
 
     @_parent_expander.match(Charges, PairIndexer, AtomIndexer)
-    def expansion4(self, charges, pairfinder, pidxer, **kwargs):
-        self._validate_pairfinder(pairfinder, None)
+    def expansion4(self, charges, pairfinder, pidxer, *, cutoff_distance, **kwargs):
+        self._validate_pairfinder(pairfinder, cutoff_distance)
         pf = pairfinder
         return charges, pf.pair_dist, pf.pair_first, pf.pair_second, pidxer.mol_index, pidxer.n_molecules
 
@@ -142,10 +143,10 @@ def expansion4(self, charges, pairfinder, pidxer, **kwargs):
     _parent_expander.require_idx_states(IdxType.Atoms, *(None,) * 5)
 
 
-class CoulombEnergyNode(ChargePairSetup, AutoKw, MultiNode):
+class CoulombEnergyNode(ChargePairSetup, Energies, AutoKw, MultiNode):
     _input_names = "charges", "pair_dist", "pair_first", "pair_second", "mol_index", "n_molecules"
-    _output_names = "mol_energies", "atom_voltages"
-    _output_index_states = IdxType.Molecules, IdxType.Atoms
+    _output_names = "mol_energies", "atom_energies", "atom_voltages"
+    _output_index_states = IdxType.Molecules, IdxType.Atoms, IdxType.Atoms
     _main_output = "mol_energies"
     _auto_module_class = physics_layers.CoulombEnergy
 
@@ -159,7 +160,7 @@ def _validate_pairfinder(pairfinder, cutoff_distance):
 
         if pairfinder.torch_module.hard_dist_cutoff is not None:
             raise ValueError(
-                "dist_hard_max is set to a finite value,\n"
+                "hard_dist_cutoff is set to a finite value,\n"
                 "coulomb energy requires summing over the entire set of pairs"
             )
 
@@ -174,10 +175,11 @@ def __init__(self, name, parents, energy_conversion, module="auto"):
         super().__init__(name, parents, module=module)
 
 
-class ScreenedCoulombEnergyNode(ChargePairSetup, AutoKw, SingleNode):
+class ScreenedCoulombEnergyNode(ChargePairSetup, Energies, AutoKw, MultiNode):
     _input_names = "charges", "pair_dist", "pair_first", "pair_second", "mol_index", "n_molecules"
-    _output_names = "mol_energies", "atom_voltages"
-    _index_state = IdxType.Molecules
+    _output_names = "mol_energies", "atom_energies", "atom_voltages"
+    _output_index_states = IdxType.Molecules, IdxType.Atoms, IdxType.Atoms
+    _main_output = "mol_energies"
     _auto_module_class = physics_layers.ScreenedCoulombEnergy
 
     @staticmethod
@@ -306,3 +308,41 @@ def expansion1(self, features, species, **kwargs):
     def __init__(self, name, parents, module="auto", **kwargs):
         parents = self.expand_parents(parents)
         super().__init__(name, parents, module=module, **kwargs)
+
+
+
+class CombineEnergyNode(Energies, AutoKw, ExpandParents, MultiNode):
+    """
+    Combines Local atom energies from different Energy Nodes. 
+    """
+    _input_names = "input_atom_energy_1", "input_atom_energy_2", "mol_index", "n_molecules"
+    _output_names = "mol_energy", "atom_energies"
+    _main_output = "mol_energy"
+    _output_index_states = IdxType.Molecules, IdxType.Atoms,
+    _auto_module_class = physics_layers.CombineEnergy
+
+    @_parent_expander.match(_BaseNode, Energies)
+    def expansion0(self, energy_1, energy_2, **kwargs):
+        return energy_1, energy_2.atom_energies
+
+    @_parent_expander.match(Energies, _BaseNode)
+    def expansion0(self, energy_1, energy_2, **kwargs):
+        return energy_1.atom_energies, energy_2
+
+    @_parent_expander.match(_BaseNode, _BaseNode)
+    def expansion1(self, energy_1, energy_2, **kwargs):
+        pdindexer = find_unique_relative([energy_1, energy_2], AtomIndexer, why_desc="Generating CombineEnergies")
+        return energy_1, energy_2, pdindexer
+
+    @_parent_expander.match(_BaseNode, _BaseNode, PaddingIndexer)
+    def expansion2(self, energy_1, energy_2, pdindexer, **kwargs):
+        return energy_1, energy_2, pdindexer.mol_index, pdindexer.n_molecules
+
+    _parent_expander.assertlen(4)
+    _parent_expander.require_idx_states(IdxType.Atoms, IdxType.Atoms, None, None)
+
+    def __init__(self, name, parents, module="auto", module_kwargs=None, **kwargs):
+        self.module_kwargs = {} if module_kwargs is None else module_kwargs
+        parents = self.expand_parents(parents, **kwargs)
+        super().__init__(name, parents=parents, module=module, **kwargs)
+

hippynn/graphs/nodes/targets.py

@@ -4,6 +4,7 @@
 from .base import MultiNode, AutoKw, ExpandParents, find_unique_relative, _BaseNode
 from .indexers import PaddingIndexer
 from .tags import AtomIndexer, Network, PairIndexer, HAtomRegressor, Charges, Energies
+from .indexers import PaddingIndexer
 from ..indextypes import IdxType, index_type_coercion
 from ...layers import targets as target_modules
 
@@ -85,7 +86,7 @@ class HBondNode(ExpandParents, AutoKw, MultiNode):
     _output_index_states = IdxType.Pair, IdxType.Pair
     _input_names = "features", "pair_first", "pair_second", "pair_dist"
     _main_output = "bonds"
-
+ 
     @_parent_expander.matchlen(1)
     def expand0(self, features, *, purpose, **kwargs):
         pairfinder = find_unique_relative(features, PairIndexer, why_desc=purpose)

hippynn/interfaces/lammps_interface/mliap_interface.py

@@ -19,6 +19,7 @@
 from hippynn.graphs.nodes.tags import Encoder, PairIndexer
 from hippynn.graphs.nodes.physics import GradientNode, VecMag
 from hippynn.graphs.nodes.inputs import SpeciesNode
+from hippynn.graphs.nodes.pairs import PairFilter
 
 class MLIAPInterface(MLIAPUnified):
     """
@@ -125,6 +126,7 @@ def setup_LAMMPS_graph(energy):
     species_set = torch.as_tensor(encoder.species_set).to(torch.int64)
     min_radius = max(p.dist_hard_max for p in pair_indexers)
 
+
     ###############################################################
     # Set up graph to accept external pair indices and shifts
 
@@ -139,13 +141,27 @@ def setup_LAMMPS_graph(energy):
     pair_dist = VecMag("(LAMMPS)pair_dist", in_pair_coord)
 
     new_inputs = [species,in_pair_first,in_pair_second,in_pair_coord,in_nlocal]
-
+
+    # Construct Filters and replace the existing pair indexers with the 
+    # corresponding new (filtered) node that accepts external pairs of atoms
     for pi in pair_indexers:
-        replace_node(pi.pair_first, in_pair_first, disconnect_old=False)
-        replace_node(pi.pair_second, in_pair_second, disconnect_old=False)
-        replace_node(pi.pair_coord, in_pair_coord, disconnect_old=False)
-        replace_node(pi.pair_dist, pair_dist, disconnect_old=False)
-        pi.disconnect()
+        if pi.dist_hard_max == min_radius:
+            replace_node(pi.pair_first, in_pair_first, disconnect_old=False)
+            replace_node(pi.pair_second, in_pair_second, disconnect_old=False)
+            replace_node(pi.pair_coord, in_pair_coord, disconnect_old=False)
+            replace_node(pi.pair_dist, pair_dist, disconnect_old=False)
+            pi.disconnect()
+        else:
+            mapped_node = PairFilter(
+                "DistanceFilter-LAMMPS",
+                (pair_dist, in_pair_first, in_pair_second, in_pair_coord),
+                dist_hard_max=pi.dist_hard_max, 
+            )
+            replace_node(pi.pair_first, mapped_node.pair_first, disconnect_old=False)
+            replace_node(pi.pair_second, mapped_node.pair_second, disconnect_old=False)
+            replace_node(pi.pair_coord, mapped_node.pair_coord, disconnect_old=False)
+            replace_node(pi.pair_dist, mapped_node.pair_dist, disconnect_old=False)
+            pi.disconnect()
 
     energy, *new_required = new_required
     try:

hippynn/layers/physics.py

@@ -74,6 +74,13 @@ def forward(self, charges, positions, mol_index, n_molecules):
 
 
 class CoulombEnergy(torch.nn.Module):
+    """ Computes the Coulomb Energy of the molecule/configuration. 
+    
+    Coulomb energies is defined for pairs of atoms. Here, we adopt the 
+    convention that the Coulomby energy for a pair of atoms is evenly
+    partitioned to both atoms as the 'per-atom energies'. Therefore, the 
+    atom energies sum to the molecular energy; similar to the HEnergy. 
+    """
     def __init__(self, energy_conversion_factor):
         super().__init__()
         self.register_buffer("energy_conversion_factor", torch.tensor(energy_conversion_factor))
@@ -84,12 +91,18 @@ def forward(self, charges, pair_dist, pair_first, pair_second, mol_index, n_mole
         n_atoms, _ = charges.shape
         voltage_atom = torch.zeros((n_atoms, 1), device=charges.device, dtype=charges.dtype)
         voltage_atom.index_add_(0, pair_first, voltage_pairs)
-        coulomb_atom = voltage_atom * charges
-        coulomb_molecules = 0.5 * self.summer(coulomb_atom, mol_index, n_molecules)
-        return coulomb_molecules, voltage_atom
+        coulomb_atoms = 0.5*voltage_atom * charges
+        coulomb_molecule = self.summer(coulomb_atoms, mol_index, n_molecules)
+        return coulomb_molecule, coulomb_atoms, voltage_atom
 
 
 class ScreenedCoulombEnergy(CoulombEnergy):
+    """ Computes the Coulomb Energy of the molecule/configuration. 
+    
+    The convention for the atom energies is the same as CoulombEnergy
+    and the HEnergy. 
+    """
+
     def __init__(self, energy_conversion_factor, screening, radius=None):
         super().__init__(energy_conversion_factor)
         if screening is None:
@@ -106,37 +119,36 @@ def __init__(self, energy_conversion_factor, screening, radius=None):
         self.bond_summer = pairs.MolPairSummer()
 
     def forward(self, charges, pair_dist, pair_first, pair_second, mol_index, n_molecules):
-
-        # Calculation of pair terms
-        base_coulomb = charges[pair_first] * charges[pair_second] / pair_dist.unsqueeze(1)
         screening = self.screening(pair_dist, self.radius).unsqueeze(1)
         screening = torch.where((pair_dist < self.radius).unsqueeze(1), screening, torch.zeros_like(screening))
-        coulomb_pairs = base_coulomb * screening
-
-        # Add pair contributions to system
-        coulomb_molecule = self.bond_summer(coulomb_pairs, mol_index, n_molecules, pair_first)
 
-        # Finally, account for symmetry pairs and energy conversion factor.
-        coulomb_molecule = 0.5 * self.energy_conversion_factor * coulomb_molecule
+        # Voltage pairs for per-atom energy
+        voltage_pairs = self.energy_conversion_factor * (charges[pair_second] / pair_dist.unsqueeze(1)) 
+        voltage_pairs = voltage_pairs * screening 
+        n_atoms, _ = charges.shape
+        voltage_atom = torch.zeros((n_atoms, 1), device=charges.device, dtype=charges.dtype)
+        voltage_atom.index_add_(0, pair_first, voltage_pairs) 
+        coulomb_atoms = 0.5 * voltage_atom * charges
+        coulomb_molecule = self.summer(coulomb_atoms, mol_index, n_molecules)
 
-        return coulomb_molecule
+        return coulomb_molecule, coulomb_atoms, voltage_atom
 
 
 class CombineScreenings(torch.nn.Module):
-    """ Returns products of different screenings for Screened Coulomb Interactions. 
+    """ Returns products of different screenings for Screened Coulomb Interactions.
     """
     def __init__(self, screening_list):
         super().__init__()
         self.SL = torch.nn.ModuleList(screening_list)
 
     def forward(self, pair_dist, radius):
         """ Product of different screenings applied to pair_dist upto radius.
-        
+
         :param pair_dist: torch.tensor, dtype=float64: 'Neighborlist' distances for coulomb energies.
-        :param radius: Maximum radius that Screened-Coulomb is evaluated upto. 
+        :param radius: Maximum radius that Screened-Coulomb is evaluated upto.
         :return screening: Weights for screening for all pair_dist.
         """
-        screening = None 
+        screening = None
 
         for s in self.SL:
             if screening is None:
@@ -236,4 +248,29 @@ def forward(self, features, species):
 
 class VecMag(torch.nn.Module):
     def forward(self, vector_feature):
-        return torch.norm(vector_feature, dim=1).unsqueeze(1)
+        return torch.norm(vector_feature, dim=1) 
+
+
+
+
+class CombineEnergy(torch.nn.Module):
+    """
+    Combines the energies (molecular and atom energies) from two different 
+    nodes, e.g. HEnergy, Coulomb, or ScreenedCoulomb Energy Nodes. 
+    """
+    def __init__(self):
+        super().__init__()
+        self.summer = indexers.MolSummer()
+
+    def forward(self, atom_energy_1, atom_energy_2, mol_index, n_molecules):
+        """
+        :param: atom_energy_1 per-atom energy from first node. 
+        :param: atom_energy_2 per atom energy from second node. 
+        :param: mol_index the molecular index for atoms in the batch
+        :param: total number of molecules in the batch
+        :return: Total Energy
+        """
+        total_atom_energy = atom_energy_1 + atom_energy_2
+        mol_energy = self.summer(total_atom_energy, mol_index, n_molecules)
+
+        return mol_energy, total_atom_energy