initial implementation of python functions in variable formulas, including for atom-style vars

doc/src/variable.rst

@@ -45,7 +45,8 @@ Syntax
        *universe* args = one or more strings
        *world* args = one string for each partition of processors
 
-       *equal* or *vector* or *atom* args = one formula containing numbers, thermo keywords, math operations, built-in functions, atom values and vectors, compute/fix/variable references
+       *equal* or *vector* or *atom* args = one formula containing numbers, thermo keywords,
+           math operations, built-in functions, atom values and vectors, compute/fix/variable references
          numbers = 0.0, 100, -5.4, 2.8e-4, etc
          constants = PI, version, on, off, true, false, yes, no
          thermo keywords = vol, ke, press, etc from :doc:`thermo_style <thermo_style>`
@@ -67,8 +68,12 @@ Syntax
                            bound(group,dir,region), gyration(group,region), ke(group,reigon),
                            angmom(group,dim,region), torque(group,dim,region),
                            inertia(group,dimdim,region), omega(group,dim,region)
-         special functions = sum(x), min(x), max(x), ave(x), trap(x), slope(x), sort(x), rsort(x), gmask(x), rmask(x), grmask(x,y), next(x), is_file(name), is_os(name), extract_setting(name), label2type(kind,label), is_typelabel(kind,label), is_timeout()
+         special functions = sum(x), min(x), max(x), ave(x), trap(x), slope(x), sort(x), rsort(x), \                                  gmask(x), rmask(x), grmask(x,y), next(x), is_file(name), is_os(name),
-         feature functions = is_available(category,feature), is_active(category,feature), is_defined(category,id)
+                             extract_setting(name), label2type(kind,label),
+                             is_typelabel(kind,label), is_timeout()
+         feature functions = is_available(category,feature), is_active(category,feature),
+                             is_defined(category,id)
+         python function = py_varname(x,y,z,...)
          atom value = id[i], mass[i], type[i], mol[i], x[i], y[i], z[i], vx[i], vy[i], vz[i], fx[i], fy[i], fz[i], q[i]
          atom vector = id, mass, type, mol, radius, q, x, y, z, vx, vy, vz, fx, fy, fz
          custom atom property = i_name, d_name, i_name[i], d_name[i], i2_name[i], d2_name[i], i2_name[i][j], d2_name[i][j]
@@ -127,18 +132,21 @@ command), or used as input to an averaging fix (see the :doc:`fix
 ave/time <fix_ave_time>` command).  Variables of style *vector* store
 a formula which produces a vector of such values which can be used as
 input to various averaging fixes, or elements of which can be part of
-thermodynamic output.  Variables of style *atom* store a formula which
+thermodynamic output.
-when evaluated produces one numeric value per atom which can be output
+
-to a dump file (see the :doc:`dump custom <dump>` command) or used as
+Variables of style *atom* store a formula which when evaluated
-input to an averaging fix (see the :doc:`fix ave/chunk
+produces one numeric value per atom which can be output to a dump file
-<fix_ave_chunk>` and :doc:`fix ave/atom <fix_ave_atom>` commands).
+(see the :doc:`dump custom <dump>` command) or used as input to an
-Variables of style *atomfile* can be used anywhere in an input script
+averaging fix (see the :doc:`fix ave/chunk <fix_ave_chunk>` and
-that atom-style variables are used; they get their per-atom values
+:doc:`fix ave/atom <fix_ave_atom>` commands).  Variables of style
-from a file rather than from a formula.  Variables of style *python*
+*atomfile* can be used anywhere in an input script that atom-style
-can be hooked to Python functions using code you provide, so that the
+variables are used; they get their per-atom values from a file rather
-variable gets its value from the evaluation of the Python code.
+than from a formula.
-Variables of style *internal* are used by a few commands which set
+
-their value directly.
+Variables of style *python* can be hooked to Python functions using
+Python code you provide, so that the variable gets its value from the
+evaluation of the Python code.  Variables of style *internal* are used
+by a few commands which set their value directly.
 
 .. note::
 
@@ -166,15 +174,16 @@ simulation.
 
 .. note::
 
-   When an input script line is encountered that defines a variable
+   When an input script line is encountered that defines a variable of
-   of style *equal* or *vector* or *atom* or *python* that contains a
+   style *equal* or *vector* or *atom* or *python* that contains a
-   formula or Python code, the formula is NOT immediately evaluated.  It
+   formula or links to Python code, the formula or Python code is NOT
-   will be evaluated every time when the variable is **used** instead.  If
+   immediately evaluated.  Instead, it is evaulated aech time the
-   you simply want to evaluate a formula in place you can use as
+   variable is **used**.  If you simply want to evaluate a formula in
-   so-called. See the section below about "Immediate Evaluation of
+   place you can use a so-called immediate variable. as described in
-   Variables" for more details on the topic.  This is also true of a
+   the preceding note.  Or see the section below about "Immediate
-   *format* style variable since it evaluates another variable when it is
+   Evaluation of Variables" for more details on the topic.  This is
-   invoked.
+   also true of a *format* style variable since it evaluates another
+   variable when it is invoked.
 
 Variables of style *equal* and *vector* and *atom* can be used as
 inputs to various other commands which evaluate their formulas as
@@ -183,12 +192,12 @@ this context, variables of style *timer* or *internal* or *python* can
 be used in place of an equal-style variable, with the following two
 caveats.
 
-First, internal-style variables can be used except by commands that
+First, internal-style variables require their values be set by code
-set the value stored by the internal variable.  When the LAMMPS
+elsewhere in LAMMPS.  When a LAMMPS input script or command evaluates
-command evaluates the internal-style variable, it will use the value
+an internal-style variable, it must have a current value set
-set (internally) by another command.  Second, python-style variables
+(internally) via that mechanism.  Second, python-style variables can
-can be used so long as the associated Python function, as defined by
+be used so long as the associated Python function, as defined by the
-the :doc:`python <python>` command, returns a numeric value.  When the
+:doc:`python <python>` command, returns a numeric value.  When the
 LAMMPS command evaluates the python-style variable, the Python
 function will be executed.
 
@@ -439,6 +448,14 @@ python-style variable can be used in place of an equal-style variable
 anywhere in an input script, e.g. as an argument to another command
 that allows for equal-style variables.
 
+A python-style variable can also be used within the formula for an
+equal-style or atom-style formula with the syntax
+py_varname(arg1,arg2,...) as explained below for variable formulas.
+When used in an atom-style formula, it can the variable can be invoked
+once per atom using arguments specific to each atom.  The resulting
+values in the atom-style variable can thus be calculated by Python
+code.
+
 ----------
 
 For the *string* style, a single string is assigned to the variable.
@@ -528,9 +545,9 @@ is a valid (though strange) variable formula:
 
 Specifically, a formula can contain numbers, constants, thermo
 keywords, math operators, math functions, group functions, region
-functions, special functions, feature functions, atom values, atom
+functions, special functions, feature functions, the python function,
-vectors, custom atom properties, compute references, fix references, and references to other
+atom values, atom vectors, custom atom properties, compute references,
-variables.
+fix references, and references to other variables.
 
 +------------------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | Number                 | 0.2, 100, 1.0e20, -15.4, etc                                                                                                                                                                                                                                                                                                                               |
@@ -550,6 +567,7 @@ variables.
 | Special functions      | sum(x), min(x), max(x), ave(x), trap(x), slope(x), sort(x), rsort(x), gmask(x), rmask(x), grmask(x,y), next(x), is_file(name), is_os(name), extract_setting(name), label2type(kind,label), is_typelabel(kind,label), is_timeout()                                                                                                                          |
 +------------------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | Feature functions      | is_available(category,feature), is_active(category,feature), is_defined(category,id)                                                                                                                                                                                                                                                                       |
++------------------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| Python function        | py_varname(x,y,z,...)                                                                                                                                                                                                                                                                                                                                      |
 +------------------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | Atom values            | id[i], mass[i], type[i], mol[i], x[i], y[i], z[i], vx[i], vy[i], vz[i], fx[i], fy[i], fz[i], q[i]                                                                                                                                                                                                                                                          |
 +------------------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
@@ -1161,6 +1179,23 @@ variable name.
 
 ----------
 
+Python Function
+-----------------
+
+NOTE: this needs work
+explain why use this versus just reference a python variable
+(b/c can pass args)
+(b/c can use it in an atom-style varible)
+
+Math operators are written in the usual way, where the "x" and "y" in
+the examples can themselves be arbitrarily complex formulas, as in the
+examples above.  In this syntax, "x" and "y" can be scalar values or
+per-atom vectors.  For example, "ke/natoms" is the division of two
+scalars, where "vy+vz" is the element-by-element sum of two per-atom
+vectors of y and z velocities.
+
+----------
+
 Atom Values and Vectors
 -----------------------
 

src/variable.cpp

@@ -81,6 +81,7 @@ enum{DONE,ADD,SUBTRACT,MULTIPLY,DIVIDE,CARAT,MODULO,UNARY,
      RAMP,STAGGER,LOGFREQ,LOGFREQ2,LOGFREQ3,STRIDE,STRIDE2,
      VDISPLACE,SWIGGLE,CWIGGLE,SIGN,GMASK,RMASK,
      GRMASK,IS_ACTIVE,IS_DEFINED,IS_AVAILABLE,IS_FILE,EXTRACT_SETTING,
+     PYFUNCTION,
      VALUE,ATOMARRAY,TYPEARRAY,INTARRAY,BIGINTARRAY,VECTORARRAY};
 
 // customize by adding a special function
@@ -1006,7 +1007,7 @@ char *Variable::pythonstyle(char *name, char *funcname)
 
 /* ----------------------------------------------------------------------
    return 1 if variable is INTERNAL style, 0 if not
-   this is checked before call to set_internal() to assure it can be set
+   this is checked before call to internal_set() to assure it can be set
 ------------------------------------------------------------------------- */
 
 int Variable::internalstyle(int ivar)
@@ -2361,13 +2362,13 @@ double Variable::evaluate(char *str, Tree **tree, int ivar)
         }
 
       // ----------------
-      // math/group/special/labelmap function or atom value/vector or constant or thermo keyword
+      // math/group/region/special/feature function or atom value/vector or constant or thermo keyword
       // ----------------
 
       } else {
 
         // ----------------
-        // math or group or special function
+        // math or group/region or special or feature function
         // ----------------
 
         if (str[i] == '(') {
@@ -3183,6 +3184,20 @@ double Variable::collapse_tree(Tree *tree)
     return tree->value;
   }
 
+  if (tree->type == PYFUNCTION) {
+    int narg = tree->argcount;
+    int *argvars = tree->argvars;
+    double arg;
+    for (int iarg = 0; iarg < narg; iarg++) {
+      if (iarg == 0) arg = collapse_tree(tree->first);
+      else if (iarg == 1) arg = collapse_tree(tree->second);
+      else arg = collapse_tree(tree->extra[iarg-2]);
+      internal_set(argvars[iarg],arg);
+    }
+    tree->value = compute_equal(tree->pyvar);
+    return tree->value;
+  }
+
   // mask functions do not become a single collapsed value
 
   if (tree->type == GMASK) return 0.0;
@@ -3520,6 +3535,18 @@ double Variable::eval_tree(Tree *tree, int i)
   if (tree->type == SIGN)
     return (eval_tree(tree->first,i) >= 0.0) ? 1.0 : -1.0;   // sign(eval_tree(tree->first,i));
   
+  if (tree->type == PYFUNCTION) {
+    int narg = tree->argcount;
+    for (int iarg = 0; iarg < narg; iarg++) {
+      if (iarg == 0) arg = eval_tree(tree->first,i);
+      else if (iarg == 1) arg = eval_tree(tree->second,i);
+      else arg = eval_tree(tree->extra[iarg-2],i);
+      internal_set(tree->argvars[iarg],arg);
+    }
+    arg = compute_equal(tree->pyvar);
+    return arg;
+  }
+  
   if (tree->type == GMASK) {
     if (atom->mask[i] & tree->ivalue) return 1.0;
     else return 0.0;
@@ -3583,6 +3610,7 @@ void Variable::free_tree(Tree *tree)
     for (int i = 0; i < tree->nextra; i++) free_tree(tree->extra[i]);
     delete[] tree->extra;
   }
+  if (tree->argvars) delete[] tree->argvars;
 
   if (tree->selfalloc) memory->destroy(tree->array);
   delete tree;
@@ -3684,8 +3712,9 @@ tagint Variable::int_between_brackets(char *&ptr, int varallow)
 
 /* ----------------------------------------------------------------------
    process a math function in formula
+   includes a Python function with syntax py_varname(arg1,arg2,...)
    push result onto tree or arg stack
-   word = math function
+   word = math function name
    contents = str between parentheses with comma-separated args
    return 0 if not a match, 1 if successfully processed
    customize by adding a math function:
@@ -3693,7 +3722,8 @@ tagint Variable::int_between_brackets(char *&ptr, int varallow)
      atan2(y,x),random(x,y,z),normal(x,y,z),ceil(),floor(),round(),ternary(),
      ramp(x,y),stagger(x,y),logfreq(x,y,z),logfreq2(x,y,z),
      logfreq3(x,y,z),stride(x,y,z),stride2(x,y,z,a,b,c),vdisplace(x,y),
-     swiggle(x,y,z),cwiggle(x,y,z),sign(x)
+     swiggle(x,y,z),cwiggle(x,y,z),sign(x),
+     py_varname(arg1,arg2,...) (up to MAXFUNCARG)
 ------------------------------------------------------------------------- */
 
 int Variable::math_function(char *word, char *contents, Tree **tree, Tree **treestack,
@@ -3711,7 +3741,8 @@ int Variable::math_function(char *word, char *contents, Tree **tree, Tree **tree
       strcmp(word,"logfreq") != 0 && strcmp(word,"logfreq2") != 0 &&
       strcmp(word,"logfreq3") != 0 && strcmp(word,"stride") != 0 &&
       strcmp(word,"stride2") != 0 && strcmp(word,"vdisplace") != 0 &&
-      strcmp(word,"swiggle") != 0 && strcmp(word,"cwiggle") != 0 && strcmp(word,"sign") != 0)
+      strcmp(word,"swiggle") != 0 && strcmp(word,"cwiggle") != 0 && strcmp(word,"sign") != 0 &&
+      strstr(word,"py_") != word)
     return 0;
 
   // parse contents for comma-separated args
@@ -4106,11 +4137,60 @@ int Variable::math_function(char *word, char *contents, Tree **tree, Tree **tree
       double value = value1 + value2*(1.0-cos(omega*delta*update->dt));
       argstack[nargstack++] = value;
     }
+    
   } else if (strcmp(word,"sign") == 0) {
     if (narg != 1)
       print_var_error(FLERR,"Invalid math function in variable formula",ivar);
     if (tree) newtree->type = SIGN;
     else argstack[nargstack++] = (value1 >= 0.0) ? 1.0 : -1.0; // sign(value1);
+
+  // Python function tied to varname
+  // narg arguments tied to internal variables pyarg1, pyarg2, ... pyargN
+
+  } else if (strstr(word,"py_") == word) {
+
+    // text following py_ = python-style variable name
+    
+    int pyvar = find(&word[3]);
+    if (style[pyvar] != PYTHON)
+      print_var_error(FLERR,"Invalid python function variable name",ivar);
+
+    // check for existence of narg internal variables with names pyarg1 to pyargN
+    // store their indices in jvars
+
+    int *jvars = new int[narg];
+    char *internal_varname;
+    
+    for (int iarg = 0; iarg < narg; iarg++) {
+      internal_varname = utils::strdup(fmt::format("pyarg{}", iarg+1));
+      jvars[iarg] = find(internal_varname);
+      if (jvars[iarg] < 0)
+        print_var_error(FLERR,"Invalid python function arg in variable formula",ivar);
+      if (!internalstyle(jvars[iarg]))
+        print_var_error(FLERR,"Invalid python function arg in variable formula",ivar);
+      delete[] internal_varname;
+    }
+
+    // if tree: store python variable and arg info in tree for later eval
+    // else: one-time eval of python function now
+    
+    if (tree) {
+      newtree->type = PYFUNCTION;
+      newtree->pyvar = pyvar;
+      newtree->argcount = narg;
+      newtree->argvars = new int[narg];
+      for (int iarg = 0; iarg < narg; iarg++)
+        newtree->argvars[iarg] = jvars[iarg];
+    } else {
+      for (int iarg = 0; iarg < narg; iarg++) {
+        if (iarg == 0) internal_set(jvars[iarg],value1);
+        else if (iarg == 1) internal_set(jvars[iarg],value2);
+        else internal_set(jvars[iarg],values[iarg-2]);
+      }
+      argstack[nargstack++] = compute_equal(pyvar);
+    }
+
+    delete[] jvars;
   }
 
   // delete stored args
@@ -4377,7 +4457,7 @@ Region *Variable::region_function(char *id, int ivar)
    customize by adding a special function:
      sum(x),min(x),max(x),ave(x),trap(x),slope(x),
      gmask(x),rmask(x),grmask(x,y),next(x),is_file(x),is_os(x),
-     extract_setting(x),label2type(x,y),is_tpelabel(x,y)
+     extract_setting(x),label2type(x,y),is_typelabel(x,y)
      is_timeout()
 ------------------------------------------------------------------------- */
 

src/variable.h

@@ -123,9 +123,13 @@ class Variable : protected Pointers {
     Tree *first, *second;    // ptrs further down tree for first 2 args
     Tree **extra;            // ptrs further down tree for nextra args
 
+    int pyvar;               // index of Python variable invoked as py_name()
+    int argcount;            // # of args to associated Python function
+    int *argvars;            // indices of internal variables for each arg
+
     Tree() :
         array(nullptr), iarray(nullptr), barray(nullptr), selfalloc(0), ivalue(0), nextra(0),
-        region(nullptr), first(nullptr), second(nullptr), extra(nullptr)
+        region(nullptr), first(nullptr), second(nullptr), extra(nullptr), argvars(nullptr)
     {
     }
   };