Merge pull request mala-project#649 from RandomDefaultUser/ace_docs

ACE documentation updates

Author: RandomDefaultUser

mala/descriptors/ace.py

@@ -1,9 +1,9 @@
 """
 ACE potential class.
 
-Computes the coupling coefficients. ACE_DOCS_MISSING: Why is this called
-"potential" - would this work as an actual ACE potential or is this just the
-expansion coefficients?
+It writes a .yace file (ACE potential file) that is readable by LAMMPS. When only the
+coupling coefficients are saved to the .yace file, this acts as a parameter
+file with which one can evaluate ACE descriptors.
 """
 
 import itertools
@@ -16,49 +16,64 @@
 
 class ACEPotential:
     """
-    Class for calculation of ACE coupling coefficients.
+    Class to manage interface between ACE descriptor enumeration library and LAMMPS. By default, it will enumerate ACE descriptors according to the Permutation-adapted approach described in https://doi.org/10.1016/j.jcp.2024.113073.
 
-    ACE_DOCS_MISSING: Explain scope of class. I think we always assume PA
-    basis, but I don't know what that is.
+    However, there are options for assigning descriptor labels if desired, manually for example.
 
+    After enumerating descriptors, it assigns all relevant hyperparamters needed to evaluate the ACE descriptors in LAMMPS. It saves a .yace file needed to evaluate ACE descriptors in LAMMPS (containing coupling coefficients). It does this by writing an ACE potential file, readable by LAMNMPS, that contains only information to evaluate descriptors, not energy models.
+    
     Parameters
     ----------
     elements : List
-        List of elements (symbols)
+        List of elements (symbols), including 'G' for the grid points. In ACE, all possible
+        combinations of these elements determine the "bond types" spanned by the ACE chemical
+        basis (the chemical basis is the delta function basis used in Drautz 2019). For
+        example, the "bond types" resulting from all possible combinations of ['Al','G'] are
+        determined with :code:`itertools.product()` in python. They are (Al,Al)(Al,G)(G,Al)(G,G).
+        For mala, only those of type (G,X) are kept, (grid-atom interactions) and only a placeholder
+        is kept for other interaction types.
 
     reference_ens : List
-        List of floats, has same dimensions as elements. ACE_DOCS_MISSING:
-        What does it do?
+        List of reference energies. (To be applied only for linear models) with a constant
+        shift to the energy per element type. Values other than 0 not be necessary in MALA.
 
     ranks : List
-        ACE_DOCS_MISSING
+        Orders of the expansion, referred to as `N` in Drautz 2019, of the
+        descriptors to be enumerated
 
     nmax: List
-        ACE_DOCS_MISSING
+        Maximum radial basis function index per descriptor rank
 
     lmax: List
-        ACE_DOCS_MISSING
+        Maximum angular momentum number per descriptor rank (maximum angular function index)
 
     nradbase : int
-        ACE_DOCS_MISSING
+        Maximum radial basis function index OVERALL: max(nmax) - in the future, may be used
+        to define the number of g_k(r) comprising R_nl from Drautz 2019 radial basis.
 
-    rcut : float
-        ACE_DOCS_MISSING
+    rcut : float/list
+        radial basis function cutoff per bond type. For example, if elements are ['Al','G']
+        then rcut must be supplied for each:(Al,Al)(Al,G)(G,Al)(G,G)
 
-    lmbda : float
-        ACE_DOCS_MISSING
+    lmbda : float/list
+        Exponential factor for scaled distance in Drautz 2019 used in the radial basis
+        functions. As with the radial cutoff, lambda must be supplied per bond type. For
+        example, if elements are ['Al','G'] then lambda must be supplied for 
+        each:(Al,Al)(Al,G)(G,Al)(G,G)
 
     css : dict
-        ACE_DOCS_MISSING
+        Dictionary of coupling coefficients of the format: {rank:{l:{m:coupling_coefficient}}
 
     rcutinner : List
-        ACE_DOCS_MISSING
+        Inner cutoff to turn on soft-core repulsions. This parameter should be 0.0 (OFF) for
+        each bond type in MALA. 
 
     drcutinner : List
-        ACE_DOCS_MISSING
+        Parameter for soft-core repulsions. This parameter should not matter if rcutinner is
+         0.0 (OFF) for each bond type in MALA. 
 
-    lmin : int
-        ACE_DOCS_MISSING
+    lmin : int/list
+        Lower bound on angular momentum quantum number per rank.
 
     manual_labels : str
         File for loading labels. If not None, then labels will be loaded from
@@ -67,12 +82,17 @@ class ACEPotential:
 
     **kwarg : dict
         Additional keyword arguments.
+
+    Returns
+    -------
+    ACEPotential : class
+        Class containing ACE descriptor and hyperparamter info.
     """
 
     def __init__(
         self,
         elements,
-        reference_ensemble,
+        reference_energy,
         ranks,
         nmax,
         lmax,
@@ -89,15 +109,22 @@ def __init__(
         if kwargs is not None:
             self.__dict__.update(kwargs)
 
-        # ACE_DOCS_MISSING: Explain what all these variables do, at least
-        # briefly, and give them more meaningful names.
+        #NOTE Our unused variable names here match exactly what is in LAMMPS
+        # I'm hesitant to change these to something else. We wont use them in MALA
+        # and if people are really interested, they can find them in the lammps
+        # source code directly
+        #coupling coefficients (generalized Wigner symbols or Generalized Clebsch-Gordan coefficients)
         self.__global_ccs = css
         self.__global_ccs[1] = {"0": {tuple([]): {"0": 1.0}}}
-        self.__E0 = reference_ensemble
+        #0th-order expansion term in ACE (e.g., constant energy shift)
+        self.__E0 = reference_energy
         self.__ranks = ranks
         self.__elements = elements
+        #linear model coefficients
         self.__betas = None
+        #descriptor labels in FitSNAP/LAMMPS format - as described elsewhere
         self.__nus = None
+        #hyperparameters for ACE basis (relevant for density embeddings)
         self.__deltaSplineBins = 0.001
         self.__global_ndensity = 1
         self.__global_rhocut = 100000
@@ -195,17 +222,17 @@ def __init__(
 
     def __set_embeddings(self, npoti="FinnisSinclair", FSparams=[1.0, 1.0]):
         """
-        Set embeddings.
+        Set 'embeddings' in the .yace file for lammps. Some terms here control if a square-root embedding term is added.
 
-        ACE_DOCS_MISSING
+        This should always be OFF for descriptor calculations in MALA.
 
         Parameters
         ----------
         npoti : str
-            ACE_DOCS_MISSING
+            paramter to specify descriptor type in LAMMPS
 
         FSparams : List
-            ACE_DOCS_MISSING
+            parameters to specify embedding terms in LAMMPS
         """
         # embeddings =dict()#OrderedDict() #{ind:None for ind in range(len(self.elements))}
         embeddings = {ind: None for ind in range(len(self.__elements))}
@@ -227,9 +254,12 @@ def __set_bonds(self):
 
     def __set_bond_base(self):
         """
-        Set bond base.
+        Set the per-bond radial basis parameters. The 'bonds' are determined based on the elements.
+
+        If the elements are ['Al','G'], the bonds may be determined with :code:`itertools.product()` in python. They are (Al,Al)(Al,G)(G,Al)(G,G).
 
-        ACE_DOCS_MISSING - what does that mean, what does this function do?
+        Some hyperparameters must be specified per bond label shown here.
+        
         """
         bondstrs = ["[%d, %d]" % (b[0], b[1]) for b in self.__bondlsts]
         bonds = {bondstr: None for bondstr in bondstrs}
@@ -312,17 +342,18 @@ def __set_bond_base(self):
 
     def set_funcs(self, nulst=None, print_0s=True):
         """
-        Set functions.
-
-        ACE_DOCS_MISSING - what does this function do?
+        Set ctilde 'functions' in the .yace file, used to calculate descriptors in lammps.
 
         Parameters
         ----------
         nulst : List
-            List of nus ACE_DOCS_MISSING - what are those?
+            List of nus, a.k.a. ACE descriptor labels to write to 
+            the .yace file.
 
         print_0s : bool
-            ACE_DOCS_MISSING - what does this do?
+            Logical to include 0-valued descriptors, this should always
+            be True in MALA as 0-valued descriptors only arise after 
+            training linear models with sparsifying solvers like LASSO
         """
         if nulst is None:
             if self.__nus is not None:

mala/descriptors/acelib/ace_potential.py

@@ -8,32 +8,35 @@ def clebsch_gordan_coupling(
     clebsch_gordan_coefficients, ldict, L_R=0, use_permutations=True
 ):
     """
-    Compute Clebsch-Gordan couplings for a given L_R.
+    Compute generalized Clebsch-Gordan coefficients for a given L_R.
 
-    ACE_DOCS_MISSING: Clarify what this means, and also maybe the name,
-    because I am not so sure about it (see ace.py).
+    These are all defined according to the pairwise coupling scheme in Goff et al. 2024. 
+
+    The approach to construct generalized Clebsch-Gordan coefficients is manually defined for each rank for pedagogical and clarity purposes. 
 
     Parameters
     ----------
     clebsch_gordan_coefficients : dict
         Precomputed Clebsch-Gordan coefficients.
 
     ldict : dict
-        Dictionary of ranks and their corresponding L values.
+        Dictionary of ranks and their corresponding l values.
 
     L_R : int
-        ACE_DOCS_MISSING
+        Resultant angular momentum quantum number. This determines the equivariant
+        character of the rank N descriptor after reduction. L_R=0 corresponds to
+        a rotationally invariant feature, L_R=1 corresponds to a feature that
+        transforms like a vector, L_R=2 a tensor, etc.
 
     use_permutations : bool
-        ACE_DOCS_MISSING
+        Flag to consider permutations of l in generation of coupling coefficients.
 
     Returns
     -------
     coupling : dict
-        Clebsch-Gordan couplings.
+        Generalized Clebsch-Gordan coefficients, {M_R:{rank:{l:{m:coupling_coefficient}}}
     """
     # Define the functions for the individual ranks.
-    # ACE_DOCS_MISSING Maybe we can briefly summarize what the differences
     # between the functions for the individual ranks are.
 
     def _rank_1_tree(_clebsch_gordan_coefficients, _l, _L_R=0, _M_R=0):
@@ -46,18 +49,18 @@ def _rank_1_tree(_clebsch_gordan_coefficients, _l, _L_R=0, _M_R=0):
             Precomputed Clebsch-Gordan coefficients.
 
         _l : List
-            ACE_DOCS_MISSING
+            list of angular momentum quantum numbers [l1]
 
         _L_R : int
-            ACE_DOCS_MISSING
+            Resultant angular momentum quantum number. This determines the equivariant character of the rank N descriptor after reduction. L_R=0 corresponds to a rotationally invariant feature, L_R=1 corresponds to a feature that transforms like a vector, L_R=2 a tensor, etc.
 
         _M_R : int
-            ACE_DOCS_MISSING
+            Resultant projection quantum number. This also determines the equivariant character of the rank N descriptor after reduction. M_R must obey -L_R <= M_R <= L_R.
 
         Returns
         -------
         decomposed : dict
-            Clebsch-Gordan couplings for rank 1.
+            Generalized Clebsch-Gordan coefficient for rank 1
         """
         mstrs = mala.descriptors.acelib.common_utils.get_ms(_l, _M_R)
         full_inter_tuples = [()]
@@ -94,18 +97,24 @@ def _rank_2_tree(_clebsch_gordan_coefficients, _l, _L_R=0, _M_R=0):
             Precomputed Clebsch-Gordan coefficients.
 
         _l : List
-            ACE_DOCS_MISSING
+            list of angular momentum quantum numbers [l1,l2]
 
         _L_R : int
-            ACE_DOCS_MISSING
+            Resultant angular momentum quantum number. This determines the equivariant
+            character of the rank N descriptor after reduction. L_R=0 corresponds to
+            a rotationally invariant feature, L_R=1 corresponds to a feature that
+            transforms like a vector, L_R=2 a tensor, etc.
 
         _M_R : int
-            ACE_DOCS_MISSING
+            Resultant projection quantum number. This also determines the equivariant
+            character of the rank N descriptor after reduction. M_R must obey
+            -L_R <= M_R <= L_R
 
         Returns
         -------
         decomposed : dict
-            Clebsch-Gordan couplings for rank 2.
+            Generalized Clebsch-Gordan coefficients for rank 2
+
         """
         nodes, remainder = ace_coupling_utils.build_quick_tree(_l)
         node = nodes[0]
@@ -147,18 +156,23 @@ def _rank_3_tree(_clebsch_gordan_coefficients, _l, _L_R=0, _M_R=0):
             Precomputed Clebsch-Gordan coefficients.
 
         _l : List
-            ACE_DOCS_MISSING
+            list of angular momentum quantum numbers [l1,l2,l3]
 
         _L_R : int
-            ACE_DOCS_MISSING
+            Resultant angular momentum quantum number. This determines the equivariant
+            character of the rank N descriptor after reduction. L_R=0 corresponds to
+            a rotationally invariant feature, L_R=1 corresponds to a feature that
+            transforms like a vector, L_R=2 a tensor, etc.
 
         _M_R : int
-            ACE_DOCS_MISSING
+            Resultant projection quantum number. This also determines the equivariant
+            character of the rank N descriptor after reduction. M_R must obey
+            -L_R <= M_R <= L_R
 
         Returns
         -------
         decomposed : dict
-            Clebsch-Gordan couplings for rank 3.
+            Generalized Clebsch-Gordan coefficients for rank 3
         """
         full_inter_tuples = ace_coupling_utils.build_tree_for_l_intermediates(
             _l, L_R=_L_R
@@ -203,18 +217,23 @@ def _rank_4_tree(_clebsch_gordan_coefficients, _l, _L_R=0, _M_R=0):
             Precomputed Clebsch-Gordan coefficients.
 
         _l : List
-            ACE_DOCS_MISSING
+            list of angular momentum quantum numbers [l1,l2,l3,l4]
 
         _L_R : int
-            ACE_DOCS_MISSING
+            Resultant angular momentum quantum number. This determines the equivariant
+            character of the rank N descriptor after reduction. L_R=0 corresponds to
+            a rotationally invariant feature, L_R=1 corresponds to a feature that
+            transforms like a vector, L_R=2 a tensor, etc.
 
         _M_R : int
-            ACE_DOCS_MISSING
+            Resultant projection quantum number. This also determines the equivariant
+            character of the rank N descriptor after reduction. M_R must obey
+            -L_R <= M_R <= L_R
 
         Returns
         -------
         decomposed : dict
-            Clebsch-Gordan couplings for rank 4.
+            Generalized Clebsch-Gordan coefficients for rank 4
         """
         nodes, remainder = ace_coupling_utils.build_quick_tree(_l)
         mstrs = mala.descriptors.acelib.common_utils.get_ms(_l, _M_R)
@@ -266,6 +285,8 @@ def _rank_4_tree(_clebsch_gordan_coefficients, _l, _L_R=0, _M_R=0):
     ranks = list(ldict.keys())
     coupling = {M_R: {rank: {} for rank in ranks} for M_R in M_Rs}
 
+    #NOTE higher-rank couplings (to enable higher rank descriptors)
+    # are available in FitSNAP
     for M_R in M_Rs:
         for rank in ranks:
             rnk = rank