Merge branch 'development' into variable_optimizer

.github/workflows/build.yml

@@ -0,0 +1,24 @@
+# Check that package builds
+name: Build Checks
+
+on:
+    push:
+    pull_request:
+    workflow_dispatch:
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - uses: actions/setup-python@v2
+      with:
+        python-version: "3.10"
+
+    - name: Install dependencies
+      run: >-
+        python -m pip install --user --upgrade setuptools wheel
+    - name: Build
+      run: >-
+        python setup.py sdist bdist_wheel

.github/workflows/deploy.yml

@@ -5,8 +5,10 @@ name: PyPi Release
 
 on:
     push:
-    pull_request:
-    workflow_dispatch:
+      tags:
+        - hippynn-*
+    release:
+      types: [published]
 
 jobs:
   build:
@@ -31,4 +33,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/databases/database.py

@@ -191,7 +191,6 @@ def make_generator(self, split_type, evaluation_mode, batch_size=None, subsample
         if not self.splitting_completed:
             raise ValueError("Database has not yet been split.")
 
-
         if split_type not in self.splits:
             raise ValueError(f"Split {split_type} Invalid. Current splits:{list(self.splits.keys())}")
 
@@ -226,6 +225,51 @@ def make_generator(self, split_type, evaluation_mode, batch_size=None, subsample
         )
 
         return generator
+
+    def trim_all_arrays(self,index):
+        """
+        To be used in conjuction with remove_high_property
+        """
+        for key in self.arr_dict:
+            self.arr_dict[key] = self.arr_dict[key][index]
+
+    def remove_high_property(self,key,perAtom,species_key=None,cut=None,std_factor=10):
+        """
+        This function removes outlier data from the dataset
+        Must be called before splitting
+        "key": the property key in the dataset to check for high values
+        "perAtom": True if the property is defined per atom in axis 1, otherwise property is treated as full system
+        "std_factor": systems with values larger than this multiplier time the standard deviation of all data will be reomved. None to skip this step
+        "cut_factor": systems with values larger than this number are reomved. None to skip this step. This step is done first. 
+        """
+        if perAtom:
+            if species_key==None:
+                raise RuntimeError("species_key must be defined to trim a per atom quantity")
+            atom_ind = self.arr_dict[species_key] > 0
+        ndim = len(self.arr_dict[key].shape)
+        if cut!=None:
+            if perAtom:
+                Kmean = np.mean(self.arr_dict[key][atom_ind])
+            else:
+                Kmean = np.mean(self.arr_dict[key])
+            failArr = np.abs(self.arr_dict[key]-Kmean)>cut
+            #This does nothing with ndim=1
+            trimArr = np.sum(failArr,axis=tuple(range(1,ndim)))==0
+            self.trim_all_arrays(trimArr)
+
+        if std_factor!=None:
+            if perAtom:
+                atom_ind = self.arr_dict[species_key] > 0
+                Kmean = np.mean(self.arr_dict[key][atom_ind])
+                std_cut = np.std(self.arr_dict[key][atom_ind]) * std_factor
+            else: 
+                Kmean = np.mean(self.arr_dict[key])
+                std_cut = np.std(self.arr_dict[key]) * std_factor
+            failArr = np.abs(self.arr_dict[key]-Kmean)>std_cut
+            #This does nothing with ndim=1
+            trimArr = np.sum(failArr,axis=tuple(range(1,ndim)))==0
+            self.trim_all_arrays(trimArr)
+
 
 
 def compute_index_mask(indices, index_pool):

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/ase_interface/calculator.py

@@ -4,7 +4,6 @@
 import warnings
 import torch
 
-from ase.calculators import interface
 from ase.calculators.calculator import compare_atoms, PropertyNotImplementedError, Calculator # Calculator is required to allow HIPNN to be used with ASE Mixing Calculators
 
 from hippynn.graphs import find_relatives, find_unique_relative, get_subgraph, copy_subgraph, replace_node, GraphModule

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):
     """
@@ -40,7 +41,7 @@ def __init__(self, energy_node, element_types, ndescriptors=1,
         # Build the calculator
         self.rcutfac, self.species_set, self.graph = setup_LAMMPS_graph(energy_node)
         self.nparams = sum(p.nelement() for p in self.graph.parameters())
-        self.graph.to(torch.float64)
+        self.graph.to(torch.float32)
 
     def compute_gradients(self, data):
         pass
@@ -61,7 +62,7 @@ def compute_forces(self, data):
         z_vals = self.species_set[elems+1]
         pair_i = self.as_tensor(data.pair_i).type(torch.int64)
         pair_j = self.as_tensor(data.pair_j).type(torch.int64)
-        rij = self.as_tensor(data.rij).type(torch.float64)
+        rij = self.as_tensor(data.rij).type(torch.float32)
         nlocal = self.as_tensor(data.nlistatoms) 
 
         # note your sign for rij might need to be +1 or -1, depending on how your implementation works
@@ -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: