change how internal vars are defined by various commands, enable Python class to return a numeric value directly (for speed)

src/PYTHON/python_impl.cpp

@@ -62,6 +62,7 @@
 using namespace LAMMPS_NS;
 
 enum { NONE, INT, DOUBLE, STRING, PTR };
+enum { VALUE, VARIABLE, INTERNALVAR };
 
 /* ---------------------------------------------------------------------- */
 
@@ -174,7 +175,17 @@ void PythonImpl::command(int narg, char **arg)
                    arg[0], pfuncs[ifunc].ovarname, pfuncs[ifunc].name);
     }
 
-    invoke_function(ifunc, str);
+    // NOTE to Richard - if this function returns a value,
+    //   I don't believe it will be accessible to the input script
+    //   b/c that only occurs if Variable::retreive() or Variable::compute_equal() is called
+    //   so if a Python function with a return is invoked this way,
+    //     it might be better to issue an error or warning ?
+    //     i.e. it only make sense to call a setup-style kind of Python function this way
+    //          one with no return value
+    //   it also means there is no need to call Variable::pythonstyle() here
+    //     or even define the pythonstyle() method in Variable ?
+    
+    invoke_function(ifunc, str, nullptr);
     return;
   }
 
@@ -317,7 +328,7 @@ void PythonImpl::command(int narg, char **arg)
 
 /* ------------------------------------------------------------------ */
 
-void PythonImpl::invoke_function(int ifunc, char *result)
+void PythonImpl::invoke_function(int ifunc, char *result, double *dvalue)
 {
   PyUtils::GIL lock;
   PyObject *pValue;
@@ -336,27 +347,33 @@ void PythonImpl::invoke_function(int ifunc, char *result)
   for (int i = 0; i < ninput; i++) {
     int itype = pfuncs[ifunc].itype[i];
     if (itype == INT) {
-      if (pfuncs[ifunc].ivarflag[i]) {
+      if (pfuncs[ifunc].ivarflag[i] == VARIABLE) {
         str = input->variable->retrieve(pfuncs[ifunc].svalue[i]);
         if (!str)
           error->all(FLERR, "Could not evaluate Python function {} input variable: {}",
                      pfuncs[ifunc].name, pfuncs[ifunc].svalue[i]);
         pValue = PY_INT_FROM_LONG(PY_LONG_FROM_STRING(str));
+      } else if (pfuncs[ifunc].ivarflag[i] == INTERNALVAR) {
+        double value = input->variable->compute_equal(pfuncs[ifunc].internal_var[i]);
+        pValue = PyLong_FromDouble(value);
       } else {
         pValue = PY_INT_FROM_LONG(pfuncs[ifunc].ivalue[i]);
       }
     } else if (itype == DOUBLE) {
-      if (pfuncs[ifunc].ivarflag[i]) {
+      if (pfuncs[ifunc].ivarflag[i] == VARIABLE) {
         str = input->variable->retrieve(pfuncs[ifunc].svalue[i]);
         if (!str)
           error->all(FLERR, "Could not evaluate Python function {} input variable: {}",
                      pfuncs[ifunc].name, pfuncs[ifunc].svalue[i]);
         pValue = PyFloat_FromDouble(std::stod(str));
+      } else if (pfuncs[ifunc].ivarflag[i] == INTERNALVAR) {
+        double value = input->variable->compute_equal(pfuncs[ifunc].internal_var[i]);
+        pValue = PyFloat_FromDouble(value);
       } else {
         pValue = PyFloat_FromDouble(pfuncs[ifunc].dvalue[i]);
       }
     } else if (itype == STRING) {
-      if (pfuncs[ifunc].ivarflag[i]) {
+      if (pfuncs[ifunc].ivarflag[i] == VARIABLE) {
         str = input->variable->retrieve(pfuncs[ifunc].svalue[i]);
         if (!str)
           error->all(FLERR, "Could not evaluate Python function {} input variable: {}",
@@ -385,17 +402,24 @@ void PythonImpl::invoke_function(int ifunc, char *result)
   }
 
   // function returned a value
-  // assign it to result string stored by python-style variable
+  // if result is non-NULL, assign to result string stored by python-style variable
-  // or if user specified a length, assign it to longstr
+  //   or if value is string and user specified a length, assign it to longstr
+  // if dvalue is non-NULL, assign numeric value directly to dvalue
 
   if (pfuncs[ifunc].noutput) {
     int otype = pfuncs[ifunc].otype;
     if (otype == INT) {
+      if (dvalue) *dvalue = (double) PY_INT_AS_LONG(pValue);
+      else {
         auto value = fmt::format("{}", PY_INT_AS_LONG(pValue));
         strncpy(result, value.c_str(), Variable::VALUELENGTH - 1);
+      }
     } else if (otype == DOUBLE) {
+      if (*dvalue) *dvalue = PyFloat_AsDouble(pValue);
+      else {
         auto value = fmt::format("{:.15g}", PyFloat_AsDouble(pValue));
         strncpy(result, value.c_str(), Variable::VALUELENGTH - 1);
+      }
     } else if (otype == STRING) {
       const char *pystr = PyUnicode_AsUTF8(pValue);
       if (pfuncs[ifunc].longstr)
@@ -416,10 +440,22 @@ int PythonImpl::find(const char *name)
   return -1;
 }
 
-/* ------------------------------------------------------------------ */
+/* ---------------------------------------------------------------------
+   called by Variable class when a python-style variable is evaluated
+     this will call invoke_function() in this class
+   either via Variable::retrieve() or Variable::equalstyle
+     retrieve calls with numeric = 0, equalstyle with numeric = 1
+   ensure name matches a Python function
+   ensure the Python function produces an output
+   ensure the Python function outputs to the matching python-style variable
+   ensure a string is returned only if retrieve() is the caller
+--------------------------------------------------------------------- */
 
-int PythonImpl::variable_match(const char *name, const char *varname, int numeric)
+int PythonImpl::function_match(const char *name, const char *varname, int numeric)
 {
+  // NOTE to Richard - any reason not to put error messages here instead of in Variable class ?
+  //   error messages appear 3x in Variable
+  
   int ifunc = find(name);
   if (ifunc < 0) return -1;
   if (pfuncs[ifunc].noutput == 0) return -2;
@@ -428,6 +464,50 @@ int PythonImpl::variable_match(const char *name, const char *varname, int numeri
   return ifunc;
 }
 
+/* ---------------------------------------------------------------------
+   called by Variable class when evaluating a Python wrapper function
+     which will call invoke_function()
+   either via equal-style or atom-style variable formula
+     the latter calls invoke_function() once per atom
+   same error checks as function_match() plus 2 new ones
+   ensure match of number of Python function args mapped to internal-style variables
+   ensure each internal-style variable still exists
+     must check now in case user input script deleted variables between runs
+       which could invalidate indices set in create_event()
+     other classes avoid this issue by setting variable indices in their init() method
+--------------------------------------------------------------------- */
+
+int PythonImpl::wrapper_match(const char *name, const char *varname, int narg, int *argvars)
+{
+  // NOTE to Richard - any reason not to put 2 extra error messages here instead of in Variable class ?
+  //   only this class knows the name of the missing internal var, so can generate better error message
+  
+  int ifunc = function_match(name, varname, 1);
+  if (ifunc < 0) return ifunc;
+
+  int internal_count = 0;
+  for (int i = 0; i < pfuncs[ifunc].ninput; i++)
+    if (pfuncs[ifunc].ivarflag[i] == INTERNALVAR) internal_count++;
+  if (internal_count != narg) return -5;
+
+  // set argvars of internal-style variables for use by Variable class
+  //   in Python wrapper functions
+  // also set internal_var for use by invoke_function()
+  //   so that invoke_function() is as fast as possible for args which are internal-style vars
+  
+  int j = 0;
+  for (int i = 0; i < pfuncs[ifunc].ninput; i++)
+    if (pfuncs[ifunc].ivarflag[i] == INTERNALVAR) {
+      int ivar = input->variable->find(pfuncs[ifunc].svalue[i]);
+      if (ivar < 0) return -6;
+      pfuncs[ifunc].internal_var[i] = ivar;
+      argvars[j++] = ivar;
+    }
+  
+  return ifunc;
+}
+
+
 /* ------------------------------------------------------------------ */
 
 char *PythonImpl::long_string(int ifunc)
@@ -469,6 +549,7 @@ int PythonImpl::create_entry(char *name, int ninput, int noutput, int length_lon
   pfuncs[ifunc].ivalue = new int[ninput];
   pfuncs[ifunc].dvalue = new double[ninput];
   pfuncs[ifunc].svalue = new char *[ninput];
+  pfuncs[ifunc].internal_var = new int[ninput];
 
   for (int i = 0; i < ninput; i++) {
     pfuncs[ifunc].svalue[i] = nullptr;
@@ -476,32 +557,56 @@ int PythonImpl::create_entry(char *name, int ninput, int noutput, int length_lon
     if (type == 'i') {
       pfuncs[ifunc].itype[i] = INT;
       if (utils::strmatch(istr[i], "^v_")) {
-        pfuncs[ifunc].ivarflag[i] = 1;
+        pfuncs[ifunc].ivarflag[i] = VARIABLE;
         pfuncs[ifunc].svalue[i] = utils::strdup(istr[i] + 2);
+      } else if (utils::strmatch(istr[i], "^iv_")) {
+        pfuncs[ifunc].ivarflag[i] = INTERNALVAR;
+        pfuncs[ifunc].svalue[i] = utils::strdup(istr[i] + 3);
+        char *vname = pfuncs[ifunc].svalue[i];
+        int ivar = input->variable->find(vname);
+        if (ivar < 0) {   // create internal variable if does not exist
+          input->variable->internal_create(vname,0.0); 
+          ivar = input->variable->find(vname);
+        }
+        if (!input->variable->internalstyle(ivar))
+          error->all(FLERR, Error::NOLASTLINE, "Variable {} for python command is invalid style", vname);
       } else {
-        pfuncs[ifunc].ivarflag[i] = 0;
+        pfuncs[ifunc].ivarflag[i] = VALUE;
         pfuncs[ifunc].ivalue[i] = utils::inumeric(FLERR, istr[i], false, lmp);
       }
     } else if (type == 'f') {
       pfuncs[ifunc].itype[i] = DOUBLE;
       if (utils::strmatch(istr[i], "^v_")) {
-        pfuncs[ifunc].ivarflag[i] = 1;
+        pfuncs[ifunc].ivarflag[i] = VARIABLE;
         pfuncs[ifunc].svalue[i] = utils::strdup(istr[i] + 2);
+      } else if (utils::strmatch(istr[i], "^iv_")) {
+        pfuncs[ifunc].ivarflag[i] = INTERNALVAR;
+        pfuncs[ifunc].svalue[i] = utils::strdup(istr[i] + 3);
+        char *vname = pfuncs[ifunc].svalue[i];
+        int ivar = input->variable->find(vname);
+        if (ivar < 0) {   // create internal variable if does not exist
+          input->variable->internal_create(vname,0.0); 
+          ivar = input->variable->find(vname);
+        }
+        if (!input->variable->internalstyle(ivar))
+          error->all(FLERR, Error::NOLASTLINE, "Variable {} for python command is invalid style", vname);
       } else {
-        pfuncs[ifunc].ivarflag[i] = 0;
+        pfuncs[ifunc].ivarflag[i] = VALUE;
         pfuncs[ifunc].dvalue[i] = utils::numeric(FLERR, istr[i], false, lmp);
       }
     } else if (type == 's') {
       pfuncs[ifunc].itype[i] = STRING;
       if (utils::strmatch(istr[i], "^v_")) {
-        pfuncs[ifunc].ivarflag[i] = 1;
+        pfuncs[ifunc].ivarflag[i] = VARIABLE;
         pfuncs[ifunc].svalue[i] = utils::strdup(istr[i] + 2);
+      } else if (utils::strmatch(istr[i], "^iv_")) {
+        error->all(FLERR, "Input argument {} cannot be internal variable with string format", istr[i]);
       } else {
-        pfuncs[ifunc].ivarflag[i] = 0;
+        pfuncs[ifunc].ivarflag[i] = VALUE;
         pfuncs[ifunc].svalue[i] = utils::strdup(istr[i]);
       }
     } else if (type == 'p') {
-      pfuncs[ifunc].ivarflag[i] = 0;
+      pfuncs[ifunc].ivarflag[i] = VALUE;
       pfuncs[ifunc].itype[i] = PTR;
       if (strcmp(istr[i], "SELF") != 0) error->all(FLERR, "Invalid python command");
 
@@ -574,6 +679,7 @@ void PythonImpl::deallocate(int i)
   delete[] pfuncs[i].dvalue;
   for (int j = 0; j < pfuncs[i].ninput; j++) delete[] pfuncs[i].svalue[j];
   delete[] pfuncs[i].svalue;
+  delete[] pfuncs[i].internal_var;
   delete[] pfuncs[i].ovarname;
   delete[] pfuncs[i].longstr;
 }

src/PYTHON/python_impl.h

@@ -25,9 +25,10 @@ class PythonImpl : protected Pointers, public PythonInterface {
   PythonImpl(class LAMMPS *);
   ~PythonImpl() override;
   void command(int, char **) override;
-  void invoke_function(int, char *) override;
+  void invoke_function(int, char *, double *) override;
   int find(const char *) override;
-  int variable_match(const char *, const char *, int) override;
+  int function_match(const char *, const char *, int) override;
+  int wrapper_match(const char *, const char *, int, int *) override;
   char *long_string(int) override;
   int execute_string(char *) override;
   int execute_file(char *) override;
@@ -44,6 +45,7 @@ class PythonImpl : protected Pointers, public PythonInterface {
     int *ivalue;
     double *dvalue;
     char **svalue;
+    int *internal_var;        // stores per-arg index of internal variable
     int otype;
     char *ovarname;
     char *longstr;

src/compute_angle_local.cpp

@@ -108,7 +108,10 @@ ComputeAngleLocal::ComputeAngleLocal(LAMMPS *lmp, int narg, char **arg) :
 
     if (tstr) {
       tvar = input->variable->find(tstr);
-      if (tvar < 0) error->all(FLERR, "Variable name for compute angle/local does not exist");
+      if (tvar < 0) {
+        input->variable->internal_create(tstr,0.0); 
+        tvar = input->variable->find(tstr);
+      }
       if (!input->variable->internalstyle(tvar))
         error->all(FLERR, "Variable for compute angle/local is invalid style");
     }

src/compute_bond_local.cpp

@@ -154,7 +154,10 @@ ComputeBondLocal::ComputeBondLocal(LAMMPS *lmp, int narg, char **arg) :
 
     if (dstr) {
       dvar = input->variable->find(dstr);
-      if (dvar < 0) error->all(FLERR, "Variable name for compute bond/local does not exist");
+      if (dvar < 0) {
+        input->variable->internal_create(dstr,0.0); 
+        dvar = input->variable->find(dstr);
+      }
       if (!input->variable->internalstyle(dvar))
         error->all(FLERR, "Variable for compute bond/local is invalid style");
     }

src/compute_dihedral_local.cpp

@@ -102,7 +102,10 @@ ComputeDihedralLocal::ComputeDihedralLocal(LAMMPS *lmp, int narg, char **arg) :
 
     if (pstr) {
       pvar = input->variable->find(pstr);
-      if (pvar < 0) error->all(FLERR, "Variable name for compute dihedral/local does not exist");
+      if (pvar < 0) {
+        input->variable->internal_create(pstr,0.0); 
+        pvar = input->variable->find(pstr);
+      }
       if (!input->variable->internalstyle(pvar))
         error->all(FLERR, "Variable for compute dihedral/local is invalid style");
     }

src/create_atoms.cpp

@@ -392,22 +392,28 @@ void CreateAtoms::command(int narg, char **arg)
 
     if (xstr) {
       xvar = input->variable->find(xstr);
-      if (xvar < 0)
+      if (xvar < 0) {
-        error->all(FLERR, Error::NOLASTLINE, "Variable {} for create_atoms does not exist", xstr);
+        input->variable->internal_create(xstr,0.0); 
+        xvar = input->variable->find(xstr);
+      }
       if (!input->variable->internalstyle(xvar))
         error->all(FLERR, Error::NOLASTLINE, "Variable {} for create_atoms is invalid style", xstr);
     }
     if (ystr) {
       yvar = input->variable->find(ystr);
-      if (yvar < 0)
+      if (yvar < 0) {
-        error->all(FLERR, Error::NOLASTLINE, "Variable {} for create_atoms does not exist", ystr);
+        input->variable->internal_create(ystr,0.0); 
+        yvar = input->variable->find(ystr);
+      }
       if (!input->variable->internalstyle(yvar))
         error->all(FLERR, Error::NOLASTLINE, "Variable {} for create_atoms is invalid style", ystr);
     }
     if (zstr) {
       zvar = input->variable->find(zstr);
-      if (zvar < 0)
+      if (zvar < 0) {
-        error->all(FLERR, Error::NOLASTLINE, "Variable {} for create_atoms does not exist", zstr);
+        input->variable->internal_create(zstr,0.0); 
+        zvar = input->variable->find(zstr);
+      }
       if (!input->variable->internalstyle(zvar))
         error->all(FLERR, Error::NOLASTLINE, "Variable {} for create_atoms is invalid style", zstr);
     }

src/fix_deposit.cpp

@@ -868,19 +868,28 @@ void FixDeposit::options(int narg, char **arg)
 
     if (xstr) {
       xvar = input->variable->find(xstr);
-      if (xvar < 0) error->all(FLERR, "Variable {} for fix deposit does not exist", xstr);
+      if (xvar < 0) {
+        input->variable->internal_create(xstr,0.0); 
+        xvar = input->variable->find(xstr);
+      }
       if (!input->variable->internalstyle(xvar))
         error->all(FLERR, "Variable for fix deposit is invalid style");
     }
     if (ystr) {
       yvar = input->variable->find(ystr);
-      if (yvar < 0) error->all(FLERR, "Variable {} for fix deposit does not exist", ystr);
+      if (yvar < 0) {
+        input->variable->internal_create(ystr,0.0); 
+        yvar = input->variable->find(ystr);
+      }
       if (!input->variable->internalstyle(yvar))
         error->all(FLERR, "Variable for fix deposit is invalid style");
     }
     if (zstr) {
       zvar = input->variable->find(zstr);
-      if (zvar < 0) error->all(FLERR, "Variable {} for fix deposit does not exist", zstr);
+      if (zvar < 0) {
+        input->variable->internal_create(zstr,0.0); 
+        zvar = input->variable->find(zstr);
+      }
       if (!input->variable->internalstyle(zvar))
         error->all(FLERR, "Variable for fix deposit is invalid style");
     }

src/lmppython.cpp

@@ -67,10 +67,10 @@ void Python::command(int narg, char **arg)
 
 /* ------------------------------------------------------------------ */
 
-void Python::invoke_function(int ifunc, char *result)
+void Python::invoke_function(int ifunc, char *result, double *dvalue)
 {
   init();
-  impl->invoke_function(ifunc, result);
+  impl->invoke_function(ifunc, result, dvalue);
 }
 
 /* ------------------------------------------------------------------ */
@@ -83,10 +83,19 @@ int Python::find(const char *name)
 
 /* ------------------------------------------------------------------ */
 
-int Python::variable_match(const char *name, const char *varname, int numeric)
+int Python::function_match(const char *name, const char *varname, int numeric)
 {
   init();
-  return impl->variable_match(name, varname, numeric);
+  return impl->function_match(name, varname, numeric);
+}
+
+/* ------------------------------------------------------------------ */
+
+int Python::wrapper_match(const char *name, const char *varname,
+                          int narg, int *argvars)
+{
+  init();
+  return impl->wrapper_match(name, varname, narg, argvars);
 }
 
 /* ------------------------------------------------------------------ */

src/lmppython.h

@@ -22,9 +22,10 @@ class PythonInterface {
  public:
   virtual ~PythonInterface() noexcept(false) {}
   virtual void command(int, char **) = 0;
-  virtual void invoke_function(int, char *) = 0;
+  virtual void invoke_function(int, char *, double *) = 0;
   virtual int find(const char *) = 0;
-  virtual int variable_match(const char *, const char *, int) = 0;
+  virtual int function_match(const char *, const char *, int) = 0;
+  virtual int wrapper_match(const char *, const char *, int, int *) = 0;
   virtual char *long_string(int ifunc) = 0;
   virtual int execute_string(char *) = 0;
   virtual int execute_file(char *) = 0;
@@ -37,9 +38,10 @@ class Python : protected Pointers {
   ~Python() override;
 
   void command(int, char **);
-  void invoke_function(int, char *);
+  void invoke_function(int, char *, double *);
   int find(const char *);
-  int variable_match(const char *, const char *, int);
+  int function_match(const char *, const char *, int);
+  int wrapper_match(const char *, const char *, int, int *);
   char *long_string(int ifunc);
   int execute_string(char *);
   int execute_file(char *);

src/variable.cpp

@@ -117,6 +117,7 @@ Variable::Variable(LAMMPS *lmp) : Pointers(lmp)
   num = nullptr;
   which = nullptr;
   pad = nullptr;
+  pyindex = nullptr;
   reader = nullptr;
   data = nullptr;
   dvalue = nullptr;
@@ -163,6 +164,7 @@ Variable::~Variable()
   memory->destroy(num);
   memory->destroy(which);
   memory->destroy(pad);
+  memory->destroy(pyindex);
   memory->sfree(reader);
   memory->sfree(data);
   memory->sfree(dvalue);
@@ -960,9 +962,27 @@ int Variable::equalstyle(int ivar)
   if (style[ivar] == EQUAL || style[ivar] == TIMER ||
       style[ivar] == INTERNAL) return 1;
   if (style[ivar] == PYTHON) {
-    int ifunc = python->variable_match(data[ivar][0],names[ivar],1);
+    pyindex[ivar] = python->function_match(data[ivar][0],names[ivar],1);
-    if (ifunc < 0) return 0;
+    if (pyindex[ivar] < 0) {
-    else return 1;
+      int ierror = pyindex[ivar];
+      if (ierror == -1) {
+        error->all(FLERR, "Python function {} specified by variable {} not found",
+                   data[ivar][0], names[ivar]);
+      } else if (ierror == -2) {
+        error->all(FLERR, "Python function {} for variable {} does not return a value",
+                   data[ivar][0], names[ivar]);
+      } else if (ierror == -3) {
+        error->all(FLERR, "Python function {} and variable {} do not not link to each other",
+                   data[ivar][0], names[ivar]);
+      } else if (ierror == -4) {
+        error->all(FLERR, "Python function {} for variable {} returns a string",
+                   data[ivar][0], names[ivar]);
+      } else {
+        error->all(FLERR, "Unknown error linking Python function {} to variable {}",
+                   data[ivar][0],names[ivar]);
+      }
+      return 0;
+    } else return 1;
   }
   return 0;
 }
@@ -990,7 +1010,7 @@ int Variable::vectorstyle(int ivar)
 }
 
 /* ----------------------------------------------------------------------
-   check if variable with name is PYTHON and matches funcname
+   check if variable with name is PYTHON style and matches funcname
    called by Python class before it invokes a Python function
    return data storage so Python function can return a value for this variable
    return nullptr if not a match
@@ -1007,7 +1027,7 @@ char *Variable::pythonstyle(char *name, char *funcname)
 
 /* ----------------------------------------------------------------------
    return 1 if variable is INTERNAL style, 0 if not
-   this is checked before call to internal_set() to assure it can be set
+   this is checked before call to internal_set() to ensure it can be set
 ------------------------------------------------------------------------- */
 
 int Variable::internalstyle(int ivar)
@@ -1083,23 +1103,23 @@ char *Variable::retrieve(const char *name)
     str = data[ivar][1] = utils::strdup(result);
 
   } else if (style[ivar] == PYTHON) {
-    int ifunc = python->variable_match(data[ivar][0],name,0);
+    int ifunc = python->function_match(data[ivar][0],name,0);
     if (ifunc < 0) {
       if (ifunc == -1) {
-        error->all(FLERR, "Could not find Python function {} linked to variable {}",
+        error->all(FLERR, "Python function {} specified by variable {} not found",
                    data[ivar][0], name);
       } else if (ifunc == -2) {
-        error->all(FLERR, "Python function {} for variable {} does not have a return value",
+        error->all(FLERR, "Python function {} for variable {} does not return a value",
                    data[ivar][0], name);
       } else if (ifunc == -3) {
-        error->all(FLERR,"Python variable {} does not match variable name registered with "
+        error->all(FLERR, "Python function {} and variable {} do not not link to each other",
-                   "Python function {}", name, data[ivar][0]);
+                   data[ivar][0], name);
       } else {
-        error->all(FLERR, "Unknown error verifying function {} linked to python style variable {}",
+        error->all(FLERR, "Unknown error linking Python function {} to variable {}",
                    data[ivar][0],name);
       }
     }
-    python->invoke_function(ifunc,data[ivar][1]);
+    python->invoke_function(ifunc,data[ivar][1],nullptr);
     str = data[ivar][1];
 
     // if Python func returns a string longer than VALUELENGTH
@@ -1158,17 +1178,7 @@ double Variable::compute_equal(int ivar)
   if (style[ivar] == EQUAL) value = evaluate(data[ivar][0],nullptr,ivar);
   else if (style[ivar] == TIMER) value = dvalue[ivar];
   else if (style[ivar] == INTERNAL) value = dvalue[ivar];
-  else if (style[ivar] == PYTHON) {
+  else if (style[ivar] == PYTHON) python->invoke_function(pyindex[ivar],nullptr,&value);
-    int ifunc = python->find(data[ivar][0]);
-    if (ifunc < 0)
-      print_var_error(FLERR,fmt::format("cannot find python function {}",data[ivar][0]),ivar);
-    python->invoke_function(ifunc,data[ivar][1]);
-    try {
-      value = std::stod(data[ivar][1]);
-    } catch (std::exception &e) {
-      print_var_error(FLERR,"has an invalid value", ivar);
-    }
-  }
 
   // round to zero on underflow
   if (fabs(value) < std::numeric_limits<double>::min()) value = 0.0;
@@ -1335,6 +1345,30 @@ void Variable::internal_set(int ivar, double value)
   dvalue[ivar] = value;
 }
 
+/* ----------------------------------------------------------------------
+   create an INTERNAL style variable with name, set to value
+------------------------------------------------------------------------- */
+
+void Variable::internal_create(char *name, double value)
+{
+  if (find(name) >= 0)
+    error->all(FLERR,"Creation of internal-style variable {} which already exists", name);
+
+  if (nvar == maxvar) grow();
+  style[nvar] = INTERNAL;
+  num[nvar] = 1;
+  which[nvar] = 0;
+  pad[nvar] = 0;
+  data[nvar] = new char *[num[nvar]];
+  data[nvar][0] = new char[VALUELENGTH];
+  dvalue[nvar] = value;
+
+  if (!utils::is_id(name))
+    error->all(FLERR, "Variable name '{}' must have only letters, numbers, or underscores", name);
+  names[nvar] = utils::strdup(name);
+  nvar++;
+}
+
 /* ----------------------------------------------------------------------
    remove Nth variable from list and compact list
    delete reader explicitly if it exists
@@ -1382,6 +1416,7 @@ void Variable::grow()
   memory->grow(num,maxvar,"var:num");
   memory->grow(which,maxvar,"var:which");
   memory->grow(pad,maxvar,"var:pad");
+  memory->grow(pyindex,maxvar,"var:pyindex");
 
   reader = (VarReader **)
     memory->srealloc(reader,maxvar*sizeof(VarReader *),"var:reader");
@@ -4146,31 +4181,47 @@ int Variable::math_function(char *word, char *contents, Tree **tree, Tree **tree
     if (tree) newtree->type = SIGN;
     else argstack[nargstack++] = (value1 >= 0.0) ? 1.0 : -1.0; // sign(value1);
 
-  // Python function tied to varname
+  // Python wrapper function tied to python-style variable
-  // narg arguments tied to internal variables pyarg1, pyarg2, ... pyargN
+  // text following py_ = python-style variable name tied to Python function
+  // narg arguments are tied to internal variables defined by python command
 
   } else if (strstr(word,"py_") == word) {
 
-    // text following py_ = python-style variable name
+    // pyvar = index of python-style variable which invokes Python function
 
     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
+    // check that wrapper matches Python function
-    // store their indices in jvars
+    // jvars = returned indices of narg internal variables used by Python function
 
     int *jvars = new int[narg];
-    char *internal_varname;
+    pyindex[pyvar] = python->wrapper_match(data[pyvar][0],names[pyvar],narg,jvars);
-
+    if (pyindex[pyvar] < 0) {
-    for (int iarg = 0; iarg < narg; iarg++) {
+      int ierror = pyindex[pyvar];
-      internal_varname = utils::strdup(fmt::format("pyarg{}", iarg+1));
+      if (ierror == -1) {
-      jvars[iarg] = find(internal_varname);
+        error->all(FLERR, "Python function {} specified by variable {} not found",
-      if (jvars[iarg] < 0)
+                   data[ivar][0], names[ivar]);
-        print_var_error(FLERR,"Invalid python function arg in variable formula",ivar);
+      } else if (ierror == -2) {
-      if (!internalstyle(jvars[iarg]))
+        error->all(FLERR, "Python function {} for variable {} does not return a value",
-        print_var_error(FLERR,"Invalid python function arg in variable formula",ivar);
+                   data[ivar][0], names[ivar]);
-      delete[] internal_varname;
+      } else if (ierror == -3) {
+        error->all(FLERR, "Python function {} and variable {} do not not link to each other",
+                   data[ivar][0], names[ivar]);
+      } else if (ierror == -4) {
+        error->all(FLERR, "Python function {} for variable {} returns a string",
+                   data[ivar][0], names[ivar]);
+      } else if (ierror == -5) {
+        error->all(FLERR, "Python function {} does not use {} internal variable args",
+                   data[ivar][0], narg);
+      } else if (ierror == -6) {
+        error->all(FLERR,"Python function {} cannot find an internal variable",
+                   data[ivar][0]);
+      } else {
+        error->all(FLERR, "Unknown error linking Python function {} to variable {}",
+                   data[ivar][0],names[ivar]);
+      }
     }
       
     // if tree: store python variable and arg info in tree for later eval

src/variable.h

@@ -49,6 +49,7 @@ class Variable : protected Pointers {
   void compute_atom(int, int, double *, int, int);
   int compute_vector(int, double **);
   void internal_set(int, double);
+  void internal_create(char *, double);
 
   tagint int_between_brackets(char *&, int);
   double evaluate_boolean(char *);
@@ -87,6 +88,7 @@ class Variable : protected Pointers {
   int *num;                    // # of values for each variable
   int *which;                  // next available value for each variable
   int *pad;                    // 1 = pad loop/uloop variables with 0s, 0 = no pad
+  int *pyindex;                // indices to Python funcs for python-style vars
   class VarReader **reader;    // variable that reads from file
   char ***data;                // str value of each variable's values
   double *dvalue;              // single numeric value for internal variables

src/write_dump.cpp

@@ -85,7 +85,7 @@ void WriteDump::command(int narg, char **arg)
   if (strcmp(arg[1], "image") == 0) (dynamic_cast<DumpImage *>(dump))->multifile_override = 1;
   if (strcmp(arg[1], "cfg") == 0) (dynamic_cast<DumpCFG *>(dump))->multifile_override = 1;
   if ((update->first_update == 0) && (comm->me == 0) && (noinitwarn == 0))
-    error->warning(FLERR, "Calling write_dump before a full system init.");
+    error->warning(FLERR, "Calling write_dump before a full system init");
 
   dump->init();
   dump->write();