From 6ca9cc657c9a36c1ebf568150b07af8ccda5cdda Mon Sep 17 00:00:00 2001
From: lhausamm <loic_hausammann@hotmail.com>
Date: Fri, 1 Dec 2017 14:40:46 +0100
Subject: [PATCH] Wrapper is working for unit_system, swift_params and
 phys_const.

---
 pyswiftsim/structure.py | 208 ++++++++++++++++++++++++++++++++++++----
 setup.py                |  19 ++--
 src/cooling_wrapper.c   |  39 ++++++++
 src/cooling_wrapper.h   |  11 +++
 src/parser_wrapper.c    |  20 ++++
 src/parser_wrapper.h    |   9 ++
 src/part_wrapper.c      |   6 +-
 src/part_wrapper.h      |   2 +-
 src/pyswiftsim_tools.c  | 142 +++++++++++++++++++++------
 src/pyswiftsim_tools.h  |  23 ++++-
 src/units_wrapper.c     |  40 +++++++-
 src/units_wrapper.h     |   4 +-
 src/wrapper.c           |  21 +++-
 test/test_cooling.py    |  18 ++++
 14 files changed, 491 insertions(+), 71 deletions(-)
 create mode 100644 src/cooling_wrapper.c
 create mode 100644 src/cooling_wrapper.h
 create mode 100644 src/parser_wrapper.c
 create mode 100644 src/parser_wrapper.h
 create mode 100644 test/test_cooling.py

diff --git a/pyswiftsim/structure.py b/pyswiftsim/structure.py
index db7fbcd..c183a04 100644
--- a/pyswiftsim/structure.py
+++ b/pyswiftsim/structure.py
@@ -2,6 +2,9 @@ import struct
 import numpy
 from ctypes import *
 
+PARSER_MAX_LINE_SIZE = 256
+PARSER_MAX_NO_OF_PARAMS = 256
+PARSER_MAX_NO_OF_SECTIONS = 64
 
 ######################################################################
 #                                                                    #
@@ -58,6 +61,14 @@ class SwiftStruct(struct.Struct):
         """
         raise NotImplementedError("SwiftStruct should not be used")
 
+    @property
+    def struct_substruct(self):
+        """
+        Dictionary containing the class of each substructure.
+        See for example SwiftParams
+        """
+        return {}
+
 
     def _getInfoFromName(self, name):
         """
@@ -115,28 +126,37 @@ class SwiftStruct(struct.Struct):
         ii = 0
         while ii < N:
             v = form[ii]
-            if v.isdigit():
-                out_nber.append(int(v))
-                out_form.append(form[ii+1])
-                # next value is the type => skip
+            count = ""
+            while v.isdigit():
+                count += v
                 ii += 1
-            else:
-                out_nber.append(1)
-                out_form.append(v)
+                v = form[ii]
+
+            if count == "":
+                count = 1
+
+            if v == "s":
+                count = 1
+                
+            count = int(count)
+            out_nber.append(count)
+            out_form.append(v)
             ii += 1
 
         return out_nber, out_form
 
 
     def __str__(self):
-        txt = "%s:\n" % type(self)
-        print(len(self.data))
-        data = self.unpack(self.data)
+        tab = ""
+        parent = self.parent
+        while parent is not None:
+            tab += "\t"
+            parent = parent.parent
+             
+        txt = tab + "%s:\n" % type(self)
         for name in self.struct_name:
-            i, form, n = self._getInfoFromName(name)
-            d = data[i]
-            txt += "\t%s: %s\n" % (name, d)
-
+            d = getattr(self, name)
+            txt += tab + "\t%s: %s\n" % (name, d)
         return txt
 
     def __getattr__(self, name):
@@ -150,12 +170,24 @@ class SwiftStruct(struct.Struct):
 
             data = self.unpack(self.data)
             if n == 1:
-                return data[i]
-
+                # if substruct
+                if name in self.struct_substruct:
+                    d = self.struct_substruct[name]
+                    cl = d["class"]
+                    tmp = []
+                    size = struct.calcsize(cl._format)
+                    for j in range(d["size"]):
+                        data_tmp = data[i][size*j:(j+1)*size]
+                        tmp.append(cl(data_tmp, parent=self))
+                    return tmp
+                # other case => array
+                else:
+                    return data[i]
+                
             else:
                 # transform scalar -> vector
                 nform = str(n) + form
-                i = slice(i.start, n)
+                i = slice(i.start, i.start+n)
                 data = data[i]
 
                 # compress data and create return struct
@@ -199,6 +231,10 @@ class ArrayStruct(SwiftStruct):
         data = list(self.unpack(self.data))
         data[ii] = value
         setattr(self.parent, self._name, data)
+
+    def __str__(self):
+        data = self.unpack(self.data)
+        return str(data)
         
         
     def getArray(self):
@@ -273,3 +309,141 @@ class Part(SwiftStruct):
             "time_bin"
         ]
 
+######################################################################
+#                                                                    #
+#                       Parameter                                    #
+#                                                                    #
+######################################################################
+class Parameter(SwiftStruct):
+    _format = "{line_size}c{line_size}c".format(
+            line_size=PARSER_MAX_LINE_SIZE
+        )
+
+    _name = [
+            "name",
+            "value"
+        ]
+
+
+    def __init__(self, data, parent=None):
+        super().__init__(self.struct_format, data, parent)
+
+    @property
+    def struct_format(self):
+        return Parameter._format
+
+    @property
+    def struct_name(self):
+        return Parameter._name
+    
+######################################################################
+#                                                                    #
+#                       Section                                      #
+#                                                                    #
+######################################################################
+class Section(SwiftStruct):
+    _format = "{line_size}c".format(
+            line_size=PARSER_MAX_LINE_SIZE
+        )
+    _name = [
+            "name"
+        ]
+
+    def __init__(self, data, parent=None):
+        super().__init__(self.struct_format, data, parent)
+
+    @property
+    def struct_format(self):
+        return Section._format
+
+
+    @property
+    def struct_name(self):
+        return Section._name 
+
+######################################################################
+#                                                                    #
+#                       SwiftParams                                  #
+#                                                                    #
+######################################################################
+class SwiftParams(SwiftStruct):
+    _format = "{sec}s{data}sii{line_size}c".format(
+            sec=struct.calcsize(Section._format)*PARSER_MAX_NO_OF_SECTIONS,
+            data=struct.calcsize(Parameter._format)*PARSER_MAX_NO_OF_PARAMS,
+            line_size=PARSER_MAX_LINE_SIZE
+        )
+
+    _name = [
+            "section",
+            "data_params",
+            "sectionCount",
+            "paramCount",
+            "filename"
+        ]
+
+    def __init__(self, data, parent=None):
+        super().__init__(self.struct_format, data, parent)
+
+    @property
+    def struct_format(self):
+        return SwiftParams._format
+
+
+    @property
+    def struct_name(self):
+        return SwiftParams._name
+
+    @property
+    def struct_substruct(self):
+        sec = {
+            "class": Section,
+            "size": PARSER_MAX_NO_OF_SECTIONS
+        }
+
+        param = {
+            "class": Parameter,
+            "size": PARSER_MAX_NO_OF_PARAMS
+        }
+        return {
+            "section": sec,
+            "data_params": param
+        }
+
+######################################################################
+#                                                                    #
+#                        PhysConst                                   #
+#                                                                    #
+######################################################################
+class PhysConst(SwiftStruct):
+    _format = "dddddddddddddddd"
+    _name = [
+        "const_newton_G",
+        "const_speed_light_c",
+        "const_planck_h",
+        "const_planck_hbar",
+        "const_boltzmann_k",
+        "const_thomson_cross_section",
+        "const_electron_charge",
+        "const_electron_volt",
+        "const_electron_mass",
+        "const_proton_mass",
+        "const_year",
+        "const_astronomical_unit",
+        "const_parsec",
+        "const_light_year",
+        "const_solar_mass",
+        "const_earth_mass",
+    ]
+
+
+    def __init__(self, data, parent=None):
+        super().__init__(self.struct_format, data, parent)
+
+    @property
+    def struct_format(self):
+        return PhysConst._format
+
+
+    @property
+    def struct_name(self):
+        return PhysConst._name
diff --git a/setup.py b/setup.py
index 825e3e4..285ff2f 100644
--- a/setup.py
+++ b/setup.py
@@ -58,20 +58,16 @@ if swift_path is not None:
 
 # C libraries
 lib = ["m",
+       "swiftsim",
+       "hdf5"
 ]
 
-if swift_path is not None:
-    lib.insert(0, swift_path + "/src/.libs/libswiftsim")
-
-    # hdf5
-    lib.insert(0, hdf5_root + "/lib/libhdf5")
-
-else:
-    lib.append("swiftsim")
-    lib.append("hdf5")
+lib_dir = []
 
+if swift_path is not None:
+    lib_dir.append(swift_path + "/src/.libs")
+    lib_dir.append(hdf5_root + "/lib")
 
-    # mpi
     
 #  src files
 c_src = []
@@ -88,7 +84,8 @@ c_src = glob("src/*.c")
 ext_modules = Extension("pyswiftsim.wrapper",
                         c_src,
                         include_dirs=include,
-                        libraries=lib)
+                        libraries=lib,
+                        library_dirs=lib_dir)
 
 ext_modules = [ext_modules]
     
diff --git a/src/cooling_wrapper.c b/src/cooling_wrapper.c
new file mode 100644
index 0000000..ce09541
--- /dev/null
+++ b/src/cooling_wrapper.c
@@ -0,0 +1,39 @@
+#include "pyswiftsim_tools.h"
+#include "cooling_wrapper.h"
+
+#include <cooling.h>
+#include <cooling_struct.h>
+
+PyObject* pycooling_init(PyObject* self, PyObject* args) {
+  PyObject* pyparams;
+  PyObject* pyus;
+  PyObject* pypconst;
+
+  if (!PyArg_ParseTuple(args, "OOO", &pyparams, &pyus, &pypconst))
+      return NULL;
+
+  struct swift_params *params = pytools_construct(pyparams, class_swift_params);
+  if (params == NULL)
+    return NULL;
+
+  struct unit_system *us = pytools_construct(pyus, class_unit_system);
+  if (us == NULL)
+    return NULL;
+  
+  struct phys_const *pconst = pytools_construct(pypconst, class_phys_const);
+  if (pconst == NULL)
+    return NULL;
+
+
+  struct cooling_function_data cooling;
+
+  cooling_init_backend(params, us, pconst, &cooling);
+
+  PyObject *pycooling = pytools_return(&cooling, class_cooling_function_data);
+
+  return pycooling;
+}
+
+PyObject* pycooling_rate(PyObject* self, PyObject* args) {
+  pyerror("Not implemented");
+}
diff --git a/src/cooling_wrapper.h b/src/cooling_wrapper.h
new file mode 100644
index 0000000..55dcec6
--- /dev/null
+++ b/src/cooling_wrapper.h
@@ -0,0 +1,11 @@
+#ifndef __PYSWIFTSIM_COOLING_H__
+#define __PYSWIFTSIM_COOLING_H__
+
+#include "pyswiftsim_tools.h"
+
+PyObject* pycooling_init(PyObject* self, PyObject* args);
+
+PyObject* pycooling_rate(PyObject* self, PyObject* args);
+
+#endif // __PYSWIFTSIM_COOLING_H__
+
diff --git a/src/parser_wrapper.c b/src/parser_wrapper.c
new file mode 100644
index 0000000..a6ecc0b
--- /dev/null
+++ b/src/parser_wrapper.c
@@ -0,0 +1,20 @@
+#include "parser_wrapper.h"
+#include "pyswiftsim_tools.h"
+
+#include <parser.h>
+
+PyObject* pyparser_read_file(PyObject *self, PyObject *args)
+{
+
+  char *filename;
+  if (!PyArg_ParseTuple(args, "s", &filename))
+    return NULL;
+
+  struct swift_params params;
+
+  parser_read_file(filename, &params);
+
+  PyObject* obj = pytools_return(&params, class_swift_params);
+
+  return obj;
+}
diff --git a/src/parser_wrapper.h b/src/parser_wrapper.h
new file mode 100644
index 0000000..b94ca16
--- /dev/null
+++ b/src/parser_wrapper.h
@@ -0,0 +1,9 @@
+#ifndef __PYSWIFTSIM_PARSER_H__
+#define __PYSWIFTSIM_PARSER_H__
+
+#include <Python.h>
+
+PyObject* pyparser_read_file(PyObject *self, PyObject *args);
+
+
+#endif // __PYSWIFTSIM_PARSER_H__
diff --git a/src/part_wrapper.c b/src/part_wrapper.c
index 791efdd..742d1a7 100644
--- a/src/part_wrapper.c
+++ b/src/part_wrapper.c
@@ -7,7 +7,7 @@
 #include <stdlib.h>
 #include <string.h>
 
-PyObject* part_test_struct(PyObject *self, PyObject *args)
+PyObject* pypart_test_struct(PyObject *self, PyObject *args)
 {
 
   size_t N = sizeof(struct part);
@@ -30,7 +30,9 @@ PyObject* part_test_struct(PyObject *self, PyObject *args)
   
   hydro_init_part(p, NULL);
 
-  PyObject *object = tools_return(p, class_part);
+  PyObject *object = pytools_return(p, class_part);
+
+  free(p);
   
   return object;
 }
diff --git a/src/part_wrapper.h b/src/part_wrapper.h
index b8fb132..9f34c85 100644
--- a/src/part_wrapper.h
+++ b/src/part_wrapper.h
@@ -4,6 +4,6 @@
 #include <Python.h>
 #include <part.h>
 
-PyObject* part_test_struct(PyObject *self, PyObject *args);
+PyObject* pypart_test_struct(PyObject *self, PyObject *args);
 
 #endif // __PYSWIFTSIM_PART_H__
diff --git a/src/pyswiftsim_tools.c b/src/pyswiftsim_tools.c
index 3fd4274..a01ede6 100644
--- a/src/pyswiftsim_tools.c
+++ b/src/pyswiftsim_tools.c
@@ -3,21 +3,39 @@
 /* include swift */
 #include <part.h>
 #include <units.h>
+#include <parser.h>
+#include <physical_constants.h>
+#include <cooling_struct.h>
 
 #include <Python.h>
 
-PyObject* tools_return(void *p, int class)
-{
+const size_t class_size[class_count] = {
+  sizeof(struct unit_system),
+  sizeof(struct part),
+  sizeof(struct swift_params),
+  sizeof(struct phys_const),
+  sizeof(struct cooling_function_data)
+};
+  
+const char *class_name[class_count] = {
+  "UnitSystem",
+  "Part",
+  "SwiftParams",
+  "PhysConst",
+  "CoolingFunctionData"
+};
 
+
+PyObject* pytools_import(char* module_name, char* object_name)
+{
   /* load module */
   PyObject *module;
   
-  module = PyImport_ImportModule("pyswiftsim.structure");
+  module = PyImport_ImportModule(module_name);
 
   if (module == NULL)
     {
-      PyErr_Print();
-      error("Failed to import module.");
+      pyerror("Failed to import module '%s'.", module_name);
     }
 
   /* get module dictionary */
@@ -28,38 +46,44 @@ PyObject* tools_return(void *p, int class)
 
   if (dict == NULL)
     {
-      PyErr_Print();
-      error("Failed to get the module dictionary");
+      pyerror("Failed to get module '%s' dictionary", module_name);
     }
 
   /* get right class */
+  PyObject *python_obj = PyDict_GetItemString(dict, object_name);
+  Py_DECREF(dict);
+
+  if (python_obj == NULL)
+    pyerror("Object %s does not exist in module %s", object_name, module_name);
+
+  return python_obj;
+}
+
+
+PyObject* pytools_return(void *p, int class)
+{
+
   PyObject *python_class;
   size_t nber_bytes;
-  
-  switch(class)
-    {
-    case class_units:
-      python_class = PyDict_GetItemString(dict, "UnitSystem");
-      nber_bytes = sizeof(struct unit_system);
-      break;
-
-    case class_part:
-      python_class = PyDict_GetItemString(dict, "Part");
-      nber_bytes = sizeof(struct part);
-      break;
-      
-    default:
-      Py_DECREF(dict);
-      error("Class not implemented");
-      break;
-    }
 
-  Py_DECREF(dict);
+  char module_name[STRING_SIZE] = "pyswiftsim.structure";
+  char *class_pyname;
+
+  if (class >= class_count)
+    pyerror("Class %i does not exists", class);
+
+  nber_bytes = class_size[class];
+  class_pyname = class_name[class];
 
+  python_class = pytools_import(module_name, class_pyname);
+
+  if (python_class == NULL)
+    return NULL;
+      
   if (!PyCallable_Check(python_class))
     {
       Py_DECREF(python_class);
-      error("Unable to create the return object");
+      pyerror("Unable to import class %s from %s", class_pyname, module_name);
     }
 
   /* create object */
@@ -76,3 +100,67 @@ PyObject* tools_return(void *p, int class)
   
 }
 
+char* pytools_get_type_name(PyObject *obj)
+{
+  PyObject *type = PyObject_Type(obj);
+  if (type == NULL)
+    {
+      Py_DECREF(type);
+      pyerror("Unable to get type");
+    }
+  
+  PyObject* recv = PyObject_Str(type);
+  Py_DECREF(type);
+
+  if (recv == NULL)
+    {
+      Py_DECREF(recv);
+      pyerror("Unable to get string representation");
+    }
+  
+  size_t size;
+  char *name = PyUnicode_AsUTF8AndSize(recv, size);
+  Py_DECREF(recv);
+
+  if (name == NULL)
+    {
+      pyerror("Unable to convert string to char");
+    }
+
+  return name;
+}
+
+
+char* pytools_construct(PyObject* obj, int class)
+{
+  char *module_name = "pyswiftsim.structure";
+  char *class_pyname;
+
+  if (class >= class_count)
+    pyerror("Class %i does not exists", class);
+
+  class_pyname = class_name[class];
+
+  PyObject *pyclass = pytools_import(module_name, class_pyname);
+
+  int test = !PyObject_IsInstance(obj, pyclass);
+  Py_DECREF(pyclass);
+  if (test)
+    {
+      char *recv = pytools_get_type_name(obj);
+      if (recv == NULL)
+	return NULL;
+      pyerror("Expecting class %s, received %s", class_pyname, recv);
+    }
+
+  
+  PyObject* data = PyObject_GetAttrString(obj, "data");
+
+  if (data == NULL)
+    pyerror("Unable to get the attribute 'data'");
+
+  char *ret = PyBytes_AsString(data);
+
+  Py_DECREF(data);
+  return ret;
+}
diff --git a/src/pyswiftsim_tools.h b/src/pyswiftsim_tools.h
index 96a62fa..fb2a366 100644
--- a/src/pyswiftsim_tools.h
+++ b/src/pyswiftsim_tools.h
@@ -9,10 +9,11 @@
 #define STRING_SIZE 200
 
 /* Set the error message for python (still need to return NULL) */
-#define error(s, ...)						\
+#define pyerror(s, ...)						\
   ({								\
     char error_msg[STRING_SIZE];				\
-    sprintf(error_msg, "%s:%s():%i: " s, __FILE__,		\
+    PyErr_Print();						\
+    sprintf(error_msg, "\n%s:%s():%i: " s, __FILE__,		\
 	    __FUNCTION__, __LINE__, ##__VA_ARGS__);		\
     PyErr_SetString(PyExc_RuntimeError, error_msg);		\
     return FAIL;						\
@@ -20,16 +21,28 @@
 
 
 enum class {
-  class_units,
+  class_unit_system,
   class_part,
+  class_swift_params,
+  class_phys_const,
+  class_cooling_function_data,
+  class_count /* should always be last! */
 };
 
+extern const size_t class_size[];
+extern const char *class_name[];
+
 enum error_code {
-  FAIL = 0,
+  FAIL = 0, // ensure NULL == FAIL
   SUCCESS,
 };
 
-PyObject* tools_return(void* p, int class);
+PyObject* pytools_return(void* p, int class);
+
+char* pytools_construct(PyObject* obj, int class);
+
+PyObject* pytools_import(char* module, char* object_name);
 
+char* pytools_get_type_name(PyObject *obj);
 
 #endif // __PYSWIFTSIM_TOOLS_H__
diff --git a/src/units_wrapper.c b/src/units_wrapper.c
index d17d5c8..47ada02 100644
--- a/src/units_wrapper.c
+++ b/src/units_wrapper.c
@@ -1,12 +1,14 @@
 #include "pyswiftsim_tools.h"
 
 #include <units.h>
+#include <parser.h>
+#include <physical_constants.h>
 
 #include <Python.h>
 #include <stdlib.h>
 #include <string.h>
 
-PyObject* unit_system_test_struct(PyObject *self, PyObject *args)
+PyObject* pyunit_system_test_struct(PyObject *self, PyObject *args)
 {
 
   size_t N = sizeof(struct unit_system);
@@ -18,7 +20,41 @@ PyObject* unit_system_test_struct(PyObject *self, PyObject *args)
   us->UnitCurrent_in_cgs = 4.;
   us->UnitTemperature_in_cgs = 5.;
 
-  PyObject *object = tools_return(us, class_units);
+  PyObject *object = pytools_return(us, class_unit_system);
+  free(us);
+  if (object == NULL)
+    return NULL;
 
   return object;
 }
+
+PyObject* pyunit_system_init(PyObject *self, PyObject *args)
+{
+  PyObject* parser;
+
+  if (!PyArg_ParseTuple(args, "O", &parser))
+    return NULL;
+
+  struct swift_params *params = pytools_construct(parser, class_swift_params);
+
+  if (params == NULL)
+    return NULL;
+
+  struct unit_system us;
+  units_init(&us, params, "InternalUnitSystem");
+
+  struct phys_const pconst;
+  phys_const_init(&us, &pconst);
+
+  PyObject *pyus = pytools_return(&us, class_unit_system);
+  if (pyus == NULL)
+    return NULL;
+  
+  PyObject *pypconst = pytools_return(&pconst, class_phys_const);
+  if (pypconst == NULL)
+    return NULL;
+
+  return PyTuple_Pack(2, pyus, pypconst);
+}
+
+
diff --git a/src/units_wrapper.h b/src/units_wrapper.h
index f0f498a..0d5d6ff 100644
--- a/src/units_wrapper.h
+++ b/src/units_wrapper.h
@@ -3,6 +3,8 @@
 
 #include <Python.h>
 
-PyObject* unit_system_test_struct(PyObject *self, PyObject *args);
+PyObject* pyunit_system_test_struct(PyObject *self, PyObject *args);
+
+PyObject* pyunit_system_init(PyObject *self, PyObject *args);
 
 #endif // __PYSWIFTSIM_UNITS_H__
diff --git a/src/wrapper.c b/src/wrapper.c
index 57951e5..8c4b0ba 100644
--- a/src/wrapper.c
+++ b/src/wrapper.c
@@ -1,5 +1,7 @@
 #include "units_wrapper.h"
 #include "part_wrapper.h"
+#include "parser_wrapper.h"
+#include "cooling_wrapper.h"
 #include "pyswiftsim_tools.h"
 
 
@@ -12,12 +14,21 @@
       
 static PyMethodDef wrapper_methods[] = {
 
-  {"partTestStruct",  part_test_struct, METH_VARARGS,
-   "Construct a part object and return it."},	   	         
+  {"partTestStruct", pypart_test_struct, METH_VARARGS,
+   "Construct a part object and return it."},
+
+  {"unitSystemTestStruct", pyunit_system_test_struct, METH_VARARGS,
+   "Construct a unit_system object and return it."},
+
+  {"parserReadFile", pyparser_read_file, METH_VARARGS,
+   "Read a swift params file."},
+
+  {"unitSystemInit", pyunit_system_init, METH_VARARGS,
+   "Construct a unit_system object and return it."},
+
+  {"coolingInit", pycooling_init, METH_VARARGS,
+   "Initialize cooling."},
 
-  {"unitSystemTestStruct",  unit_system_test_struct, METH_VARARGS,
-   "Construct a unit_system object and return it."},	   	         
-  
   {NULL, NULL, 0, NULL}        /* Sentinel */
 };      
       
diff --git a/test/test_cooling.py b/test/test_cooling.py
new file mode 100644
index 0000000..616b6af
--- /dev/null
+++ b/test/test_cooling.py
@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+
+from pyswiftsim import wrapper
+from pyswiftsim import structure
+
+filename = "/home/loikki/swift_test/cooling_sedov/sedov.yml"
+
+params = wrapper.parserReadFile(filename)
+
+us, pconst = wrapper.unitSystemInit(params)
+
+#print(us)
+#print(pconst)
+
+print(type(params), type(us), type(pconst))
+cooling = wrapper.coolingInit(params, us, pconst)
+
+print(cooling)
-- 
GitLab