ICODE-ADC/lammps
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Steve Plimpton
2025-05-01 07:59:41 -06:00
parent 7ca493917a
commit 19d25a3654
3 changed files with 163 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
doc/src/variable.rst
View File

@ -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()
feature functions = is_available(category,feature), is_active(category,feature), is_defined(category,id)
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 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
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
input to an averaging fix (see the :doc:`fix ave/chunk
<fix_ave_chunk>` and :doc:`fix ave/atom <fix_ave_atom>` commands).
Variables of style *atomfile* can be used anywhere in an input script
that atom-style variables are used; they get their per-atom values
from a file rather than from a formula. Variables of style *python*
can be hooked to Python functions using 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.
thermodynamic output.
Variables of style *atom* store a formula which 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 input to an
averaging fix (see the :doc:`fix ave/chunk <fix_ave_chunk>` and
:doc:`fix ave/atom <fix_ave_atom>` commands). Variables of style
*atomfile* can be used anywhere in an input script that atom-style
variables are used; they get their per-atom values from a file rather
than from a formula.
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
of style *equal* or *vector* or *atom* or *python* that contains a
formula or Python code, the formula is NOT immediately evaluated. It
will be evaluated every time when the variable is **used** instead. If
you simply want to evaluate a formula in place you can use as
so-called. See the section below about "Immediate Evaluation of
Variables" for more details on the topic. This is also true of a
*format* style variable since it evaluates another variable when it is
invoked.
When an input script line is encountered that defines a variable of
style *equal* or *vector* or *atom* or *python* that contains a
formula or links to Python code, the formula or Python code is NOT
immediately evaluated. Instead, it is evaulated aech time the
variable is **used**. If you simply want to evaluate a formula in
place you can use a so-called immediate variable. as described in
the preceding note. Or see the section below about "Immediate
Evaluation of Variables" for more details on the topic. This is
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
set the value stored by the internal variable. When the LAMMPS
command evaluates the internal-style variable, it will use the value
set (internally) by another command. Second, python-style variables
can be used so long as the associated Python function, as defined by
the :doc:`python <python>` command, returns a numeric value. When the
First, internal-style variables require their values be set by code
elsewhere in LAMMPS. When a LAMMPS input script or command evaluates
an internal-style variable, it must have a current value set
(internally) via that mechanism. Second, python-style variables can
be used so long as the associated Python function, as defined by 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
vectors, custom atom properties, compute references, fix references, and references to other
variables.
functions, special functions, feature functions, the python function,
atom values, atom vectors, custom atom properties, compute references,
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
-----------------------

94
src/variable.cpp
View File

@ -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()
------------------------------------------------------------------------- */

6
src/variable.h
View File

@ -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)
{
}
};