made collapse_tree() method work correctly in Variable class for Python function wrappers, also added support for vector-style variables to use Python function wrappers

doc/src/python.rst

@@ -98,20 +98,21 @@ function.
 
 Later execution can be triggered in one of two ways.  One is to use
 the python command again with its *invoke* keyword.  The other is to
-trigger the evaluation of a python-style, equal-style, or atom-style
+trigger the evaluation of a python-style, equal-style, vector-style,
-variable.  A python-style variable invokes its associated Python
+or atom-style variable.  A python-style variable invokes its
-function; its return value becomes the value of the python-style
+associated Python function; its return value becomes the value of the
-variable.  Equal- and atom-style variables can use a Python function
+python-style variable.  Equal-, vector-, and atom-style variables can
-wrapper in their formulas which encodes the Python function name, and
+use a Python function wrapper in their formulas which encodes the
-specifies arguments to pass to the function.
+Python function name, and specifies arguments to pass to the function.
 
-Note python-style, equal-style, and atom-style variables can be used
+Note that python-style, equal-style, vectir-style, and atom-style
-in many different ways within LAMMPS.  They can be evaulated directly
+variables can be used in many different ways within LAMMPS.  They can
-in an input script, effectively replacing the variable with its value.
+be evaulated directly in an input script, effectively replacing the
-Or they can be passed to various commands as arguments, so that the
+variable with its value.  Or they can be passed to various commands as
-variable is evaluated multiple times during a simulation run.  See the
+arguments, so that the variable is evaluated multiple times during a
-:doc:`variable <variable>` command doc page for more details on
+simulation run.  See the :doc:`variable <variable>` command doc page
-variable styles which enable Python function evaluation.
+for more details on variable styles which enable Python function
+evaluation.
 
 The Python code for the function can be included directly in the input
 script or in a separate Python file.  The function can be standard
@@ -199,7 +200,9 @@ command.  The style of variable must be consistent with the *format*
 keyword specification for the type of data that is passed to Python.
 Each time the Python function is invoked, the LAMMPS variable is
 evaluated and its value is passed as an argument to the Python
-function.
+function.  Note that a python-style variable can be used as an
+argument, which means that the a Python function can use arguments
+which invoke other Python functions.
 
 A LAMMPS internal-style variable can also be used as an *input*
 argument, specified as iv_name, where "name" is the name of the
@@ -208,26 +211,35 @@ be defined in the input script (though it can be); if it is not
 defined, this command creates an :doc:`internal-style variable
 <variable>` with the specified name.
 
-An internal-style variable must be used when an equal-style or
+An internal-style variable must be used when an equal-style,
-atom-style variable triggers the invocation of the Python function
+vector-style, or atom-style variable triggers the invocation of the
-defined by this command, by including a Python function wrapper in its
+Python function defined by this command, by including a Python
-formula, with one or more arguments also included in the formula.
+function wrapper with arguments in its formula.  Each of the arguments
+must be specified as an internal-style variable via the *input*
+keyword.
 
 In brief, the syntax for a Python function wrapper in a variable
 formula is py_varname(arg1,arg2,...argN), where "varname" is the name
 of a python-style variable associated with a Python function defined
 by this command.  One or more arguments to the function wrapper can
-themselves be formulas which the variable command will evaluate and
+themselves be sub-formulas which the variable command will evaluate
-pass as arguments to the Python function.  This is done by assigning
+and pass as arguments to the Python function.  This is done by
-the numeric result for each argument to an internal-style variable;
+assigning the numeric result for each argument to an internal-style
-this the *input* keyword must specify the arguments as internal-style
+variable; thus the *input* keyword must specify the arguments as
-variables and their format (see below) as "f" for floating point.
+internal-style variables and their format (see below) as "f" for
-This is because LAMMPS variable formulas are calculated with floating
+floating point.  This is because LAMMPS variable formulas are
-point arithmetic (any integer values are converted to floating point).
+calculated with floating point arithmetic (any integer values are
+converted to floating point).  Note that the Python function can also
+have additional inputs, also specified by the *input* keyword, which
+are NOT arguments in the Python function wrapper.  See the example
+below for the "mixedargs" Python function.
 
 See the :doc:`variable <variable>` command doc page for full details
 on formula syntax including for Python function wrappers.  Examples
-using Python function wrappers are shown below.
+using Python function wrappers are shown below.  Note that as
+explained above with python-style variables, Python function wrappers
+can be nested; a sub-formula for an argument can contain its own
+Python function wrapper which invokes another Python function.
 
 The *return* keyword is only needed if the Python function returns a
 value.  The specified *varReturn* is of the form v_name, where "name"
@@ -545,12 +557,14 @@ simulation run at each time step using the :doc:`fix python/invoke
 ----------
 
 A Python function can also be invoked within the formula for an
-equal-style or atom-style varaible.  This means the Python function
+equal-style, vector-style, or atom-style varaible.  This means the
-will be invoked whenever the variable is invoked.  In the case of an
+Python function will be invoked whenever the variable is invoked.  In
-atom-style varaible, the Python function can be invoked once per atom.
+the case of a vector-style variable, the Python function can be
+invoked once per element of the global vector.  In the case of an
+atom-style variable, the Python function can be invoked once per atom.
 
-Here are two simple examples using equal- and atom-style variables to
+Here are three simple examples using equal-, vector-, and atom-style
-trigger execution of a Python function:
+variables to trigger execution of a Python function:
 
 .. code-block:: LAMMPS
 
@@ -571,8 +585,24 @@ keyword for the *python* command, specifies an internal-style variable
 named "arg" as iv_arg which is required to invoke the Python function
 from a Python function wrapper.
 
-The last 2 lines can be replaced by these to define atom-style
+The last 2 lines can be replaced by these to define a vector-style
-varaibles which invoke the same Python "truncate" function:
+variable which invokes the same Python "truncate" function:
+
+.. code-block:: LAMMPS
+
+  compute         ke all temp
+  variable        ke vector c_ke
+  variable        ketrunc vector py_foo(v_ke)
+  thermo_style    custom step temp epair v_ketrunc[*]
+
+The vector-style variable "ketrunc" invokes the Python "truncate"
+function on each of the 6 components of the global kinetic energy
+tensor calculated by the :doc:`compute ke <compute_ke>` command.  The
+6 truncated values will be printed with thermo output to the screen
+and log file.
+
+Alternatively, the last 2 lines can be replaced by these to define
+atom-style variables which invoke the same Python "truncate" function:
 
 .. code-block:: LAMMPS
 
@@ -581,7 +611,7 @@ varaibles which invoke the same Python "truncate" function:
    variable        ztrunc atom py_foo(z)
    dump            1 all custom 100 tmp.dump id x y z v_xtrunc v_ytrunc v_ztrunc
 
-Now when the dump command invokes the 3 atom-style variables, their
+When the dump command invokes the 3 atom-style variables, their
 arguments x,y,z to the Python function wrapper are the current
 per-atom coordinates of each atom.  The Python "truncate" function is
 thus invoked 3 times for each atom, and the truncated coordinate

src/variable.cpp

@@ -1460,7 +1460,7 @@ void Variable::copy(int narg, char **from, char **to)
                       sin(x),cos(x),tan(x),asin(x),atan2(y,x),...
      group function = count(group), mass(group), xcm(group,x), ...
      special function = sum(x),min(x), ...
-     python function wrapper = py_varname(x,y,z,...)
+     python function wrapper = py_varname(x,y,z,...) (up to MAXFUNCARG)
      atom value = x[i], y[i], vx[i], ...
      atom vector = x, y, vx, ...
      custom atom property = i/d_name, i/d_name[i], i/d2_name[i], i/d2_name[i][j]
@@ -2663,7 +2663,8 @@ double Variable::evaluate(char *str, Tree **tree, int ivar)
      atan2(y,x),random(x,y,z),normal(x,y,z),ceil(),floor(),round(),ternary(x,y,z),
      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),gmask(x),rmask(x),grmask(x,y)
+     cwiggle(x,y,z),sign(x),py_varname(x,y,z,...),
+     gmask(x),rmask(x),grmask(x,y)
 ---------------------------------------------------------------------- */
 
 double Variable::collapse_tree(Tree *tree)
@@ -3220,7 +3221,7 @@ double Variable::collapse_tree(Tree *tree)
     tree->value = (arg1 >= 0.0) ? 1.0 : -1.0; // sign(arg1);
     return tree->value;
   }
-
+    
   if (tree->type == PYWRAPPER) {
     int narg = tree->argcount;
     int *argvars = tree->argvars;
@@ -3231,6 +3232,16 @@ double Variable::collapse_tree(Tree *tree)
       else arg = collapse_tree(tree->extra[iarg-2]);
       internal_set(argvars[iarg],arg);
     }
+    for (int iarg = 0; iarg < narg; iarg++) {
+      if (iarg == 0) {
+        if (tree->first->type != VALUE) return 0.0;
+      } else if (iarg == 1) {
+        if (tree->second->type != VALUE) return 0.0;
+      } else {
+        if (tree->extra[iarg-2]->type != VALUE) return 0.0;
+      }
+    }
+    tree->type = VALUE;
     tree->value = compute_equal(tree->pyvar);
     return tree->value;
   }
@@ -3248,12 +3259,14 @@ double Variable::collapse_tree(Tree *tree)
    evaluate an atom-style or vector-style variable parse tree
    index I = atom I or vector index I
    tree was created by one-time parsing of formula string via evaluate()
+     followed by collapse_tree() operation to streamline tree as much as possible
    customize by adding a function:
      sqrt(),exp(),ln(),log(),sin(),cos(),tan(),asin(),acos(),atan(),
      atan2(y,x),random(x,y,z),normal(x,y,z),ceil(),floor(),round(),ternary(x,y,z),
      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),gmask(x),rmask(x),grmask(x,y)
+     swiggle(x,y,z),cwiggle(x,y,z),sign(x),py_varname(x,y,z,...),
+     gmask(x),rmask(x),grmask(x,y)
 ---------------------------------------------------------------------- */
 
 double Variable::eval_tree(Tree *tree, int i)
@@ -3749,7 +3762,6 @@ 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 name
    contents = str between parentheses with comma-separated args
@@ -3759,8 +3771,7 @@ 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(x,y,z,...)
-     py_varname(arg1,arg2,...) (up to MAXFUNCARG)
 ------------------------------------------------------------------------- */
 
 int Variable::math_function(char *word, char *contents, Tree **tree, Tree **treestack,
@@ -5429,7 +5440,11 @@ void Variable::print_var_error(const std::string &srcfile, const int lineno,
 
 void Variable::print_tree(Tree *tree, int level)
 {
-  printf("TREE %d: %d %g\n",level,tree->type,tree->value);
+  if (tree->type == VALUE) {
+    printf("TREE %d: %d %g\n",level,tree->type,tree->value);
+    return;
+  }
+  printf("TREE %d: %d\n",level,tree->type);
   if (tree->first) print_tree(tree->first,level+1);
   if (tree->second) print_tree(tree->second,level+1);
   if (tree->nextra)