fixed Python support

CFD-GO · Nov 6, 2023 · 1eff4d3 · 1eff4d3
1 parent 0985da4
commit 1eff4d3
Show file tree

Hide file tree

Showing 3 changed files with 115 additions and 7 deletions.
diff --git a/src/Handlers/cbPythonCall.cpp b/src/Handlers/cbPythonCall.cpp
@@ -84,7 +84,9 @@ int cbPythonCall::DoIt () {
 
 	    pModule = PyImport_ImportModule(module.value());
 	    //Py_DECREF(pName);
-
+
+            const auto lattice = solver->getCartLattice();
+
 	    if (pModule != NULL) {
             int buff_id = 0; 
 	        pFunc = PyObject_GetAttrString( pModule, function.value() );
@@ -104,7 +106,7 @@ int cbPythonCall::DoIt () {
                     real_t * buffer;
                     long int dims[3];
 
-                    long int size = solver->getComponentIntoBuffer(component, buffer, dims, offsets );
+                    long int size = lattice->getComponentIntoBuffer(component, buffer, dims, offsets );
 
                     if (sizeof(real_t) == sizeof(float) ) {
                         pInputData = PyArray_SimpleNewFromData(3, dims, NPY_FLOAT, buffer);
@@ -128,7 +130,7 @@ int cbPythonCall::DoIt () {
                     real_t * buffer;
                     long int dims[4];
 
-                    long int size = solver->getQuantityIntoBuffer(quantity, buffer, dims, offsets );
+                    long int size = lattice->getQuantityIntoBuffer(quantity, buffer, dims, offsets );
 
                     if (sizeof(real_t) == sizeof(float) ) {
                         pInputData = PyArray_SimpleNewFromData(4, dims, NPY_FLOAT, buffer);
@@ -153,9 +155,9 @@ int cbPythonCall::DoIt () {
                 }
 
                 pGlobalSize = PyTuple_New(3);    
-                PyTuple_SetItem(pGlobalSize, 0, PyLong_FromLong(solver->info.region.nx));
-                PyTuple_SetItem(pGlobalSize, 1, PyLong_FromLong(solver->info.region.ny));
-                PyTuple_SetItem(pGlobalSize, 2, PyLong_FromLong(solver->info.region.nz));
+                PyTuple_SetItem(pGlobalSize, 0, PyLong_FromLong(lattice->connectivity.global_region.nx));
+                PyTuple_SetItem(pGlobalSize, 1, PyLong_FromLong(lattice->connectivity.global_region.ny));
+                PyTuple_SetItem(pGlobalSize, 2, PyLong_FromLong(lattice->connectivity.global_region.nz));
 
                 //first one defines number of extra arguments used
                 PyTuple_SetItem(pArgs, 0, PyLong_FromLong( extra_args ));
@@ -199,7 +201,7 @@ int cbPythonCall::DoIt () {
                 for (int k =0; k < 10; k++){
                     debug1("PythonCall after,comp %s,  buffer %d, value %d: %f\n",component,buff_id,k, buffers[buff_id][k]);
                 }
-                int status = solver->loadComponentFromBuffer(component, buffers[buff_id]);
+                int status = lattice->loadComponentFromBuffer(component, buffers[buff_id]);
                 buff_id++;
             }           
             for (; buff_id <= all_buff; buff_id++){

diff --git a/src/Lattice.cpp.Rt b/src/Lattice.cpp.Rt
@@ -1308,6 +1308,108 @@ int Lattice::loadComp(const char* filename, const char* comp) const {
     return 0;
 }
 
+
+int Lattice::getComponentIntoBuffer(const char* comp, real_t* &buf, long int*  dim, long int*  offsets ) {
+    const auto& region = getLocalRegion();
+    int n = region.size();
+
+    dim[0] = region.nx;
+    dim[1] = region.ny;
+    dim[2] = region.nz;
+
+    offsets[0] = region.dx;
+    offsets[1] = region.dy;
+    offsets[2] = region.dz;
+
+    buf = new real_t[n];
+
+    output("Providing component %s to buffer..\n", comp);
+    bool somethingWritten = false;
+<?R
+    for (d in rows(DensityAll)) if (d$parameter){
+?>
+	if (std::string_view(comp) == "<?%s d$name ?>") {
+        Get_<?%s d$nicename ?>(buf);
+        somethingWritten = true;
+    }
+<?R
+}
+?>
+    if (!somethingWritten){
+        output("Possible densities:\n");
+    <?R
+        for (d in rows(DensityAll)) if (d$parameter){
+            ?> output("-> <?%s d$name ?>\n") ; <?R
+        }
+    ?>
+    }
+    assert(somethingWritten);
+    if (somethingWritten) {
+        output("...provided %s\n", comp);
+    } else {
+     	output("...not saved %s\n", comp);
+    }
+	return n;
+}
+
+int Lattice::getQuantityIntoBuffer(const char* comp, real_t* &buf, long int*  dim, long int*  offsets ) {
+    const auto& region = getLocalRegion();
+    int n = region.size();
+
+    dim[0] = region.nx;
+    dim[1] = region.ny;
+    dim[2] = region.nz;
+    dim[3] = 1;
+
+    offsets[0] = region.dx;
+    offsets[1] = region.dy;
+    offsets[2] = region.dz;
+
+	output("Providing quantity %s to buffer..\n", comp);
+    bool somethingWritten = false;
+<?R
+for (q in rows(Quantities)){ ifdef(q$adjoint);
+?>
+    if (std::string_view(comp) == "<?%s q$name ?>") {
+        <?%s q$type ?>* tmp_<?%s q$name ?> = new <?%s q$type ?>[n];
+        <?R if (q$type == 'vector_t') { ?>
+            dim[3] = 3;
+<?R } ?>
+        Get<?%s q$name ?>(region, tmp_<?%s q$name ?>, 1);
+        buf = (real_t*)tmp_<?%s q$name ?>;
+        somethingWritten = true;
+    }
+<?R
+}
+ifdef();
+?>
+    if (!somethingWritten){
+        output("Possible quantities:\n");
+<?R
+for (d in rows(Quantities)) {
+     ?> output("-> <?%s d$name ?>\n"); <?R
+}
+?>
+    }
+    assert(somethingWritten);
+    if (somethingWritten) {
+        output("...provided %s\n", comp);
+    } else {
+        output("...not saved %s\n", comp);
+    }
+    return n;
+}
+
+
+int Lattice::loadComponentFromBuffer(const char* comp,  real_t* buf) {
+    int n = getLocalRegion().size();
+<?R for (d in rows(DensityAll)) if (d$parameter) { ?>
+    if (std::string_view(comp) == "<?%s d$name ?>") Set_<?%s d$nicename ?>(buf); <?R
+} ?>
+    delete [] buf;
+    return 0;
+}
+
 void Lattice::resetAverage() {
 	data.reset_iter = data.iter;
         <?R for (f in rows(Fields))  if (f$average) { ?>

diff --git a/src/Lattice.h.Rt b/src/Lattice.h.Rt
@@ -165,6 +165,10 @@ void GetQuantity(int quant, const lbRegion& over, real_t * tab, real_t scale);
   int getPar(const ParStruct& par_struct, double * wb);
   int setPar(const ParStruct& par_struct, double * w);
 
+  int getComponentIntoBuffer(const char*, real_t *&, long int* , long int* );
+  int loadComponentFromBuffer(const char*, real_t*);
+  int getQuantityIntoBuffer(const char*, real_t*&, long int*, long int*);
+
   void updateAllSamples();
   void resetAverage();
 };