Implemented vector-style variables in C, Python, and Fortran APIs; unit test and more for Fortran/extract_variable

fortran/lammps.f90

@@ -66,8 +66,9 @@ MODULE LIBLAMMPS
     LMP_ERROR_WORLD = 4, &    ! error on comm->world
     LMP_ERROR_UNIVERSE = 8, & ! error on comm->universe
     LMP_VAR_EQUAL = 0, &      ! equal-style variables (and compatible)
-    LMP_VAR_ATOM = 1, &       ! atom-style variables (and compatible)
+    LMP_VAR_ATOM = 1, &       ! atom-style variables
-    LMP_VAR_STRING = 2        ! string variables (and compatible)
+    LMP_VAR_VECTOR = 2, &     ! vector variables
+    LMP_VAR_STRING = 3        ! string variables (everything else)
 
   ! "Constants" to use with extract_compute and friends
   TYPE lammps_style
@@ -1078,12 +1079,13 @@ CONTAINS
     CHARACTER(LEN=*), INTENT(IN), OPTIONAL :: group
     TYPE(lammps_variable_data) :: variable_data
 
-    TYPE(c_ptr) :: Cptr, Cname, Cgroup
+    TYPE(c_ptr) :: Cptr, Cname, Cgroup, Cveclength
     INTEGER :: length, i
     CHARACTER(KIND=c_char, LEN=1), DIMENSION(:), POINTER :: Cstring
     INTEGER(c_int) :: datatype
-    REAL(c_double), POINTER :: double
+    REAL(c_double), POINTER :: double => NULL()
-    REAL(c_double), DIMENSION(:), POINTER :: double_vec
+    REAL(c_double), DIMENSION(:), POINTER :: double_vec => NULL()
+    INTEGER(c_int), POINTER :: Clength => NULL()
 
     Cname = f2c_string(name)
     IF ( PRESENT(group) ) THEN
@@ -1107,8 +1109,27 @@ CONTAINS
         variable_data%datatype = DATA_DOUBLE_1D
         length = lmp_extract_setting(self, 'nlocal')
         CALL C_F_POINTER(Cptr, double_vec, [length])
+        IF ( ALLOCATED(variable_data%r64_vec) ) &
+          DEALLOCATE(variable_data%r64_vec)
+        ALLOCATE( variable_data%r64_vec(length) )
         variable_data%r64_vec = double_vec
         CALL lammps_free(Cptr)
+      CASE (LMP_VAR_VECTOR)
+        variable_data%datatype = DATA_DOUBLE_1D
+        Cgroup = f2c_string('LMP_SIZE_VECTOR') ! must match library.cpp
+        Cname = f2c_string(name)
+        Cveclength = lammps_extract_variable(self%handle, Cname, Cgroup)
+        CALL C_F_POINTER(Cveclength, Clength)
+        length = Clength
+        CALL lammps_free(Cgroup)
+        CALL lammps_free(Cname)
+        CALL lammps_free(Cveclength)
+        CALL C_F_POINTER(Cptr, double_vec, [length])
+        IF ( ALLOCATED(variable_data%r64_vec) ) &
+          DEALLOCATE(variable_data%r64_vec)
+        ALLOCATE( variable_data%r64_vec(length) )
+        variable_data%r64_vec = double_vec
+        ! DO NOT deallocate the C pointer
       CASE (LMP_VAR_STRING)
         variable_data%datatype = DATA_STRING
         length = c_strlen(Cptr)
@@ -1117,6 +1138,11 @@ CONTAINS
         FORALL ( i=1:length )
           variable_data%str(i:i) = Cstring(i)
         END FORALL
+        ! DO NOT deallocate the C pointer
+      CASE (-1)
+        CALL lmp_error(self, LMP_ERROR_ALL + LMP_ERROR_WORLD, &
+          'Variable "' // TRIM(name) // &
+          '" not found [Fortran/extract_variable]')
       CASE DEFAULT
         CALL lmp_error(self, LMP_ERROR_ALL + LMP_ERROR_WORLD, &
           'Unknown variable type returned from &
@@ -1453,10 +1479,12 @@ CONTAINS
   END SUBROUTINE assign_double_to_lammps_variable_data
 
   SUBROUTINE assign_doublevec_to_lammps_variable_data (lhs, rhs)
-    REAL(c_double), DIMENSION(:), INTENT(OUT), POINTER :: lhs
+    REAL(c_double), DIMENSION(:), ALLOCATABLE, INTENT(OUT) :: lhs
     CLASS(lammps_variable_data), INTENT(IN) :: rhs
 
     IF ( rhs%datatype == DATA_DOUBLE_1D ) THEN
+      IF ( ALLOCATED(lhs) ) DEALLOCATE(lhs)
+      ALLOCATE( lhs(SIZE(rhs%r64_vec)) )
       lhs = rhs%r64_vec
     ELSE
       CALL assignment_error(rhs, 'vector of doubles')

python/lammps/constants.py

@@ -42,7 +42,8 @@ LMP_ERROR_UNIVERSE = 8
 
 LMP_VAR_EQUAL      = 0
 LMP_VAR_ATOM       = 1
-LMP_VAR_STRING     = 2
+LMP_VAR_VECTOR     = 2
+LMP_VAR_STRING     = 3
 
 # -------------------------------------------------------------------------
 

python/lammps/core.py

@@ -1136,6 +1136,26 @@ class lammps(object):
         self.lib.lammps_free(ptr)
       else: return None
       return result
+    elif vartype == LMP_VAR_VECTOR :
+      nvector = 0
+      self.lib.lammps_extract_variable.restype = POINTER(c_int)
+      ptr = self.lib.lammps_extract_variable(self.lmp,name,
+              'LMP_SIZE_VECTOR'.encode())
+      if ptr :
+        nvector = ptr[0]
+        self.lib.lammps_free(ptr)
+      else :
+        return None
+      self.lib.lammps_extract_variable.restype = POINTER(c_double)
+      result = (c_double*nvector)()
+      values = self.lib.lammps_extract_variable(self.lmp,name,group)
+      if values :
+        for i in range(nvector) :
+          result[i] = values[i]
+        # do NOT free the values pointer (points to internal vector data)
+        return result
+      else :
+        return None
     elif vartype == LMP_VAR_STRING :
       self.lib.lammps_extract_variable.restype = c_char_p
       with ExceptionCheck(self) :

src/library.cpp

@@ -2043,16 +2043,19 @@ void *lammps_extract_fix(void *handle, const char *id, int style, int type,
 
 This function returns a pointer to data from a LAMMPS :doc:`variable`
 identified by its name.  When the variable is either an *equal*\ -style
-compatible or an *atom*\ -style variable the variable is evaluated and
+compatible variable, a *vector*\ -style variable, or an *atom*\ -style
-the corresponding value(s) returned.  Variables of style *internal*
+variable, the variable is evaluated and the corresponding value(s) returned.
-are compatible with *equal*\ -style variables and so are *python*\
+Variables of style *internal* are compatible with *equal*\ -style variables and
--style variables, if they return a numeric value.  For other
+so are *python*\ -style variables, if they return a numeric value.  For other
-variable styles their string value is returned.  The function returns
+variable styles, their string value is returned.  The function returns
 ``NULL`` when a variable of the provided *name* is not found or of an
 incompatible style.  The *group* argument is only used for *atom*\
--style variables and ignored otherwise.  If set to ``NULL`` when
+-style variables and ignored otherwise, with one exception: for style *vector*,
-extracting data from and *atom*\ -style variable, the group is assumed
+if *group* is "GET_VECTOR_SIZE", the returned pointer will yield the length
-to be "all".
+of the vector to be returned when dereferenced. This pointer must be
+deallocated after the value is read to avoid a memory leak.
+If *group* is set to ``NULL`` when extracting data from an *atom*\ -style
+variable, the group is assumed to be "all".
 
 When requesting data from an *equal*\ -style or compatible variable
 this function allocates storage for a single double value, copies the
@@ -2066,15 +2069,23 @@ use to avoid a memory leak. Example:
    double value = *dptr;
    lammps_free((void *)dptr);
 
-For *atom*\ -style variables the data returned is a pointer to an
+For *atom*\ -style variables, the return value is a pointer to an
 allocated block of storage of double of the length ``atom->nlocal``.
 Since the data returned are a copy, the location will persist, but its
 content will not be updated in case the variable is re-evaluated.
 To avoid a memory leak, this pointer needs to be freed after use in
 the calling program.
 
+For *vector*\ -style variables, the returned pointer is to actual LAMMPS data.
+The pointer should not be deallocated. Its length depends on the variable,
+compute, or fix data used to construct the *vector*\ -style variable.
+This length can be fetched by calling this function with *group* set to the
+constant "LMP_SIZE_VECTOR", which returns a ``void\*`` pointer that can be
+dereferenced to an integer that is the length of the vector. This pointer
+needs to be deallocated when finished with it to avoid memory leaks.
+
 For other variable styles the returned pointer needs to be cast to
-a char pointer.
+a char pointer. It should not be deallocated.
 
 .. code-block:: c
 
@@ -2084,7 +2095,7 @@ a char pointer.
 .. note::
 
    LAMMPS cannot easily check if it is valid to access the data
-   referenced by the variables (e.g., computes or fixes or thermodynamic
+   referenced by the variables (e.g., computes, fixes, or thermodynamic
    info), so it may fail with an error.  The caller has to make certain
    that the data are extracted only when it safe to evaluate the variable
    and thus an error or crash are avoided.
@@ -2118,6 +2129,15 @@ void *lammps_extract_variable(void *handle, const char *name, const char *group)
       auto vector = (double *) malloc(nlocal*sizeof(double));
       lmp->input->variable->compute_atom(ivar,igroup,vector,1,0);
       return (void *) vector;
+    } else if (lmp->input->variable->vectorstyle(ivar)) {
+      double *values = nullptr;
+      int nvector = lmp->input->variable->compute_vector(ivar, &values);
+      if ( group != nullptr && strcmp(group,"LMP_SIZE_VECTOR") == 0 ) {
+          int* nvecptr = (int *) malloc(sizeof(int));
+          *nvecptr = nvector;
+          return (void *) nvecptr;
+      } else
+        return (void *) values;
     } else {
       return lmp->input->variable->retrieve(name);
     }
@@ -2162,6 +2182,8 @@ int lammps_extract_variable_datatype(void *handle, const char *name)
       return LMP_VAR_EQUAL;
     else if (lmp->input->variable->atomstyle(ivar))
       return LMP_VAR_ATOM;
+    else if (lmp->input->variable->vectorstyle(ivar))
+      return LMP_VAR_VECTOR;
     else
       return LMP_VAR_STRING;
   }

src/library.h

@@ -99,7 +99,8 @@ enum _LMP_ERROR_CONST {
 enum _LMP_VAR_CONST {
   LMP_VAR_EQUAL = 0,  /*!< compatible with equal-style variables */
   LMP_VAR_ATOM = 1,   /*!< compatible with atom-style variables */
-  LMP_VAR_STRING = 2 /*!< return value will be a string (catch-all) */
+  LMP_VAR_VECTOR = 2, /*!< compatible with vector-style variables */
+  LMP_VAR_STRING = 3  /*!< return value will be a string (catch-all) */
 };
 
 /* Ifdefs to allow this file to be included in C and C++ programs */

unittest/fortran/test_fortran_extract_variable.f90

@@ -17,12 +17,6 @@ MODULE keepvar
       'pair_style lj/cut 2.5',                              &
       'pair_coeff 1 1 1.0 1.0',                             &
       'mass 1 2.0' ]
-  CHARACTER(LEN=60), DIMENSION(4), PARAMETER :: py_input = &
-      [ CHARACTER(LEN=60) :: &
-      'python square_it input 1 v_lp return v_square here """', &
-      'def square_it(N) :', &
-      '  return N*N', &
-      '"""' ]
 
 CONTAINS
 
@@ -115,6 +109,14 @@ SUBROUTINE f_lammps_setup_extract_variable () BIND(C)
   USE keepvar, ONLY : lmp, demo_input, cont_input, pair_input, absolute_path
   IMPLICIT NONE
 
+  ! Had to do this one as one string because lammps_commands_list and
+  ! lammps_commands_string do not play well with triple quotes
+  CHARACTER(LEN=256), PARAMETER :: py_input = &
+      'python square_it input 1 v_lp return v_py format ff here """' &
+        // NEW_LINE(' ') // 'def square_it(N) :' &
+        // NEW_LINE(' ') // '  return N*N' &
+        // NEW_LINE(' ') // '"""'
+
   CALL lmp%command('atom_modify map array')
   CALL lmp%commands_list(demo_input)
   CALL lmp%commands_list(cont_input)
@@ -135,12 +137,17 @@ SUBROUTINE f_lammps_setup_extract_variable () BIND(C)
   CALL lmp%command('variable atfile atomfile ' &
     // absolute_path('atomdata.txt'))
   IF ( lmp%config_has_package('PYTHON') ) THEN
+    CALL lmp%command(py_input)
     CALL lmp%command('variable py python square_it')
   END IF
   CALL lmp%command('variable time timer')
   CALL lmp%command('variable int internal 4')
-  CALL lmp%command("variable nat equal count(all)")
+  CALL lmp%command('variable at_z atom z')
-  CALL lmp%command("variable ts equal step")
+  CALL lmp%command("compute COM all com") ! defines a global vector
+  CALL lmp%command("variable center vector c_COM")
+  ! make sure COM is computable...
+  CALL lmp%command("thermo_style custom step pe c_COM[1] v_center[1]")
+  CALL lmp%command("run 0") ! so c_COM and v_center have values
 END SUBROUTINE f_lammps_setup_extract_variable
 
 FUNCTION f_lammps_extract_variable_index_1 () BIND(C)
@@ -305,9 +312,76 @@ FUNCTION f_lammps_extract_variable_atomfile(i) BIND(C)
   IMPLICIT NONE
   INTEGER(c_int), INTENT(IN), VALUE :: i
   REAL(c_double) :: f_lammps_extract_variable_atomfile
-  REAL(c_double), DIMENSION(:), POINTER :: atom_data
+  REAL(c_double), DIMENSION(:), ALLOCATABLE :: atom_data
 
   atom_data = lmp%extract_variable('atfile')
-print*, 'TESTING: atom_data is', atom_data
   f_lammps_extract_variable_atomfile = atom_data(i)
 END FUNCTION f_lammps_extract_variable_atomfile
+
+FUNCTION f_lammps_extract_variable_python(i) BIND(C)
+  USE, INTRINSIC :: ISO_C_BINDING, ONLY : C_int, C_double
+  USE LIBLAMMPS
+  USE keepvar, ONLY : lmp
+  IMPLICIT NONE
+  INTEGER(c_int), INTENT(IN), VALUE :: i
+  REAL(c_double) :: f_lammps_extract_variable_python
+
+  f_lammps_extract_variable_python = lmp%extract_variable('py')
+END FUNCTION f_lammps_extract_variable_python
+
+FUNCTION f_lammps_extract_variable_timer() BIND(C)
+  USE, INTRINSIC :: ISO_C_BINDING, ONLY : C_double
+  USE LIBLAMMPS
+  USE keepvar, ONLY : lmp
+  IMPLICIT NONE
+  REAL(c_double) :: f_lammps_extract_variable_timer
+
+  f_lammps_extract_variable_timer = lmp%extract_variable('time')
+END FUNCTION f_lammps_extract_variable_timer
+
+FUNCTION f_lammps_extract_variable_internal() BIND(C)
+  USE, INTRINSIC :: ISO_C_BINDING, ONLY : C_double
+  USE LIBLAMMPS
+  USE keepvar, ONLY : lmp
+  IMPLICIT NONE
+  REAL(c_double) :: f_lammps_extract_variable_internal
+
+  f_lammps_extract_variable_internal = lmp%extract_variable('int')
+END FUNCTION f_lammps_extract_variable_internal
+
+FUNCTION f_lammps_extract_variable_equal() BIND(C)
+  USE, INTRINSIC :: ISO_C_BINDING, ONLY : C_double
+  USE LIBLAMMPS
+  USE keepvar, ONLY : lmp
+  IMPLICIT NONE
+  REAL(c_double) :: f_lammps_extract_variable_equal
+
+  f_lammps_extract_variable_equal = lmp%extract_variable('ex')
+END FUNCTION f_lammps_extract_variable_equal
+
+FUNCTION f_lammps_extract_variable_atom(i) BIND(C)
+  USE, INTRINSIC :: ISO_C_BINDING, ONLY : C_double, C_int
+  USE LIBLAMMPS
+  USE keepvar, ONLY : lmp
+  IMPLICIT NONE
+  INTEGER(c_int), INTENT(IN), VALUE :: i
+  REAL(c_double) :: f_lammps_extract_variable_atom
+  REAL(c_double), DIMENSION(:), ALLOCATABLE :: atom
+
+  atom = lmp%extract_variable('at_z') ! z-coordinates
+  f_lammps_extract_variable_atom = atom(i)
+END FUNCTION f_lammps_extract_variable_atom
+
+FUNCTION f_lammps_extract_variable_vector(i) BIND(C)
+  USE, INTRINSIC :: ISO_C_BINDING, ONLY : C_double, C_int
+  USE LIBLAMMPS
+  USE keepvar, ONLY : lmp
+  IMPLICIT NONE
+  INTEGER(c_int), INTENT(IN), VALUE :: i
+  REAL(c_double) :: f_lammps_extract_variable_vector
+  REAL(c_double), DIMENSION(:), ALLOCATABLE :: vector
+
+  vector = lmp%extract_variable('center') ! z-coordinates
+  f_lammps_extract_variable_vector = vector(i)
+END FUNCTION f_lammps_extract_variable_vector
+! vim: sts=2 ts=2 sw=2 et

unittest/fortran/wrap_extract_variable.cpp

@@ -9,6 +9,9 @@
 #include <cstdlib>
 #include <cstdint>
 #include <cmath>
+#include <chrono>
+#include <thread>
+#include <cmath>
 
 #include "gtest/gtest.h"
 
@@ -36,10 +39,9 @@ double f_lammps_extract_variable_atomfile(int);
 double f_lammps_extract_variable_python();
 double f_lammps_extract_variable_timer();
 double f_lammps_extract_variable_internal();
-double f_lammps_extract_variable_equal_natoms();
+double f_lammps_extract_variable_equal();
-double f_lammps_extract_variable_equal_dt();
-double f_lammps_extract_variable_vector(int);
 double f_lammps_extract_variable_atom(int);
+double f_lammps_extract_variable_vector(int);
 }
 
 class LAMMPS_extract_variable : public ::testing::Test {
@@ -196,8 +198,64 @@ TEST_F(LAMMPS_extract_variable, atomfile)
    EXPECT_DOUBLE_EQ(f_lammps_extract_variable_atomfile(1), 5.2);
    EXPECT_DOUBLE_EQ(f_lammps_extract_variable_atomfile(2), 1.6);
    EXPECT_DOUBLE_EQ(f_lammps_extract_variable_atomfile(3), -1.4);
-/*   lammps_command(lmp, "next atfile");
+   lammps_command(lmp, "next atfile");
    EXPECT_DOUBLE_EQ(f_lammps_extract_variable_atomfile(1), -1.1);
    EXPECT_DOUBLE_EQ(f_lammps_extract_variable_atomfile(2), 0.0);
-   EXPECT_DOUBLE_EQ(f_lammps_extract_variable_atomfile(3), 2.5); */
+   EXPECT_DOUBLE_EQ(f_lammps_extract_variable_atomfile(3), 2.5);
+};
+
+TEST_F(LAMMPS_extract_variable, python)
+{
+   if ( lammps_config_has_package("PYTHON") ) {
+      f_lammps_setup_extract_variable();
+      for (int i = 1; i <= 10; i++) {
+         EXPECT_DOUBLE_EQ(f_lammps_extract_variable_python(),
+            static_cast<double>(i*i));
+         lammps_command(lmp, "next lp");
+      }
+   }
+};
+
+TEST_F(LAMMPS_extract_variable, timer)
+{
+   f_lammps_setup_extract_variable();
+   double initial_t, final_t;
+   initial_t = f_lammps_extract_variable_timer();
+   std::this_thread::sleep_for(std::chrono::milliseconds(100));
+   lammps_command(lmp,"variable time timer"); // update the time
+   final_t = f_lammps_extract_variable_timer();
+   EXPECT_GT(final_t, initial_t + 0.1);
+};
+
+TEST_F(LAMMPS_extract_variable, internal)
+{
+   f_lammps_setup_extract_variable();
+   EXPECT_DOUBLE_EQ(f_lammps_extract_variable_internal(), 4.0);
+};
+
+TEST_F(LAMMPS_extract_variable, equal)
+{
+   f_lammps_setup_extract_variable();
+   int i;
+   for ( i = 1; i <= 10; i++ ) {
+      EXPECT_DOUBLE_EQ(f_lammps_extract_variable_equal(),
+         std::exp(static_cast<double>(i)));
+      lammps_command(lmp, "next lp");
+   }
+};
+
+TEST_F(LAMMPS_extract_variable, atom)
+{
+   f_lammps_setup_extract_variable();
+   EXPECT_DOUBLE_EQ(f_lammps_extract_variable_atom(1), 1.5);
+   EXPECT_DOUBLE_EQ(f_lammps_extract_variable_atom(2), 0.1);
+   EXPECT_DOUBLE_EQ(f_lammps_extract_variable_atom(3), 0.5);
+};
+
+TEST_F(LAMMPS_extract_variable, vector)
+{
+   f_lammps_setup_extract_variable();
+   EXPECT_DOUBLE_EQ(f_lammps_extract_variable_vector(1), (1+0.2+0.5)/3.0);
+   EXPECT_DOUBLE_EQ(f_lammps_extract_variable_vector(2), (1+0.1+0.5)/3.0);
+   EXPECT_DOUBLE_EQ(f_lammps_extract_variable_vector(3), (1.5+0.1+0.5)/3.0);
 };