diff --git a/pyswiftsim/structure.py b/pyswiftsim/structure.py
index 38da5364c39507b431ee0726bdf4f673bc61f3d7..db7fbcd7aa6aa37a738ed9a579d0e89e8db60e42 100644
--- a/pyswiftsim/structure.py
+++ b/pyswiftsim/structure.py
@@ -1,5 +1,6 @@
 import struct
 import numpy
+from ctypes import *
 
 
 ######################################################################
@@ -129,6 +130,7 @@ class SwiftStruct(struct.Struct):
 
     def __str__(self):
         txt = "%s:\n" % type(self)
+        print(len(self.data))
         data = self.unpack(self.data)
         for name in self.struct_name:
             i, form, n = self._getInfoFromName(name)
@@ -144,7 +146,6 @@ class SwiftStruct(struct.Struct):
         
         # case where the attribute is in the structure
         else:
-            print("get", type(self))
             i, form, n = self._getInfoFromName(name)
 
             data = self.unpack(self.data)
@@ -168,7 +169,6 @@ class SwiftStruct(struct.Struct):
             object.__setattr__(self, name, value)
         # case where the attribute is in the structure
         else:
-            print("set", type(self))
             data = list(self.unpack(self.data))
 
             i, form, n = self._getInfoFromName(name)
@@ -192,12 +192,10 @@ class ArrayStruct(SwiftStruct):
         self._format = struct_format
 
     def __getitem__(self, ii):
-        print("getitem")
         data = self.unpack(self.data)
         return data[ii]
     
     def __setitem__(self, ii, value):
-        print("setitem")
         data = list(self.unpack(self.data))
         data[ii] = value
         setattr(self.parent, self._name, data)
@@ -253,13 +251,25 @@ class Part(SwiftStruct):
 
     @property
     def struct_format(self):
-        return "3d3ffP"
+        print("ERROR, need to fix density/time_bin")
+        return "qP3d3f3ffffffN7f4c"
+
 
     @property
     def struct_name(self):
         return  [
+            "id",
+            "gpart",
             "pos",
             "vel",
+            "a_hydro",
+            "h",
+            "mass",
+            "rho",
             "entropy",
-            "test"
+            "entropy_dt",
+            "last_offset",
+            "density",
+            "time_bin"
         ]
+
diff --git a/setup.py b/setup.py
index a71aec17a3963ccf73f73e52f9a8dc6f50022e7b..825e3e4fe04941ee458b5d44733e0ed6cd00a4cb 100644
--- a/setup.py
+++ b/setup.py
@@ -6,14 +6,17 @@ Wrapper around the SPH cosmological simulation code SWIFT
 
 from setuptools import setup, find_packages, Extension
 import sys
+import os
 from glob import glob
 import numpy
 
+os.environ["CC"] = "mpicc"
+
 # deal with arguments
 swift_path = None
 swift_arg = "--with-swift"
 if swift_arg in sys.argv:
-    i = sys.argv.find(swift_arg)
+    i = sys.argv.index(swift_arg)
     swift_path = sys.argv[i+1]
     sys.argv.remove(swift_arg)
     sys.argv.remove(swift_path)
@@ -21,15 +24,55 @@ if swift_arg in sys.argv:
 # python lib dependency
 install_requires=["numpy"]
 
+
+def getValueFromMakefile(swift_root, value):
+    makefile = swift_root + "/src/Makefile"
+    N = len(value)
+
+    with open(makefile, "r") as f:
+        for line in f.readlines():
+            if value == line[:N]:
+                return line[N:-1] # remove \n
+
+    raise ValueError("Value %s not found in Makefile" % value)
+
+
+if swift_path is not None:
+    hdf5_root = getValueFromMakefile(swift_path, "H5CC = ")
+
 # link to non pypi dependencies
 dependency_links = []
 
 # C includes
-include = [".", numpy.get_include()]
+include = [numpy.get_include()]
+
+if swift_path is not None:
+    # swift
+    include.extend([
+        swift_path + "/src",
+        swift_path
+    ])
+
+    # hdf5
+    include.append(hdf5_root + "/include")
 
 # C libraries
-lib = ["m"]
+lib = ["m",
+]
+
+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")
+
+
+    # mpi
+    
 #  src files
 c_src = []
 
diff --git a/src/part_wrapper.c b/src/part_wrapper.c
new file mode 100644
index 0000000000000000000000000000000000000000..791efddad9108da63497d642311045642a579479
--- /dev/null
+++ b/src/part_wrapper.c
@@ -0,0 +1,36 @@
+#include "pyswiftsim_tools.h"
+
+#include <part.h>
+#include <hydro.h>
+
+#include <Python.h>
+#include <stdlib.h>
+#include <string.h>
+
+PyObject* part_test_struct(PyObject *self, PyObject *args)
+{
+
+  size_t N = sizeof(struct part);
+
+  struct part *p = malloc(N);
+  p->id = 0;
+  for(size_t k = 0; k < DIM; k++)
+    {
+      p->x[k] = 1. + k;
+      p->v[k] = 4. + k;
+      p->a_hydro[k] = 7. + k;
+    }
+
+  p->h = 10;
+  p->mass = 11.;
+  p->rho = 12.;
+  p->entropy = 13;
+  p->entropy_dt = 14;
+  p->last_offset = 15;
+  
+  hydro_init_part(p, NULL);
+
+  PyObject *object = tools_return(p, class_part);
+  
+  return object;
+}
diff --git a/src/part_wrapper.h b/src/part_wrapper.h
new file mode 100644
index 0000000000000000000000000000000000000000..b8fb132ef97efe5e45d7e0b2c230a24f621cdd3b
--- /dev/null
+++ b/src/part_wrapper.h
@@ -0,0 +1,9 @@
+#ifndef __PYSWIFTSIM_PART_H__
+#define __PYSWIFTSIM_PART_H__
+
+#include <Python.h>
+#include <part.h>
+
+PyObject* part_test_struct(PyObject *self, PyObject *args);
+
+#endif // __PYSWIFTSIM_PART_H__
diff --git a/src/pyswiftsim_tools.c b/src/pyswiftsim_tools.c
new file mode 100644
index 0000000000000000000000000000000000000000..3fd4274370a09730b724902e9d4c1a6fe401acf7
--- /dev/null
+++ b/src/pyswiftsim_tools.c
@@ -0,0 +1,78 @@
+#include "pyswiftsim_tools.h"
+
+/* include swift */
+#include <part.h>
+#include <units.h>
+
+#include <Python.h>
+
+PyObject* tools_return(void *p, int class)
+{
+
+  /* load module */
+  PyObject *module;
+  
+  module = PyImport_ImportModule("pyswiftsim.structure");
+
+  if (module == NULL)
+    {
+      PyErr_Print();
+      error("Failed to import module.");
+    }
+
+  /* get module dictionary */
+  PyObject *dict;
+
+  dict = PyModule_GetDict(module);
+  Py_DECREF(module);
+
+  if (dict == NULL)
+    {
+      PyErr_Print();
+      error("Failed to get the module dictionary");
+    }
+
+  /* get right 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);
+
+  if (!PyCallable_Check(python_class))
+    {
+      Py_DECREF(python_class);
+      error("Unable to create the return object");
+    }
+
+  /* create object */
+  PyObject *object, *args;
+  
+  args = PyTuple_Pack(1, PyBytes_FromStringAndSize((char *) p, nber_bytes));
+
+  object = PyObject_CallObject(python_class, args);
+
+  Py_DECREF(args);
+  Py_DECREF(python_class);
+
+  return object;
+  
+}
+
diff --git a/src/pyswiftsim_tools.h b/src/pyswiftsim_tools.h
new file mode 100644
index 0000000000000000000000000000000000000000..96a62fa6aaf9770506c0a416d77276007dff0a60
--- /dev/null
+++ b/src/pyswiftsim_tools.h
@@ -0,0 +1,35 @@
+#ifndef __PYSWIFTSIM_TOOLS_H__
+#define __PYSWIFTSIM_TOOLS_H__
+
+#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+
+#include <Python.h>
+
+#define DIM 3
+#define STRING_SIZE 200
+
+/* Set the error message for python (still need to return NULL) */
+#define error(s, ...)						\
+  ({								\
+    char error_msg[STRING_SIZE];				\
+    sprintf(error_msg, "%s:%s():%i: " s, __FILE__,		\
+	    __FUNCTION__, __LINE__, ##__VA_ARGS__);		\
+    PyErr_SetString(PyExc_RuntimeError, error_msg);		\
+    return FAIL;						\
+  })
+
+
+enum class {
+  class_units,
+  class_part,
+};
+
+enum error_code {
+  FAIL = 0,
+  SUCCESS,
+};
+
+PyObject* tools_return(void* p, int class);
+
+
+#endif // __PYSWIFTSIM_TOOLS_H__
diff --git a/src/units.c b/src/units.c
deleted file mode 100644
index d467e4e916e741dc9c741a4ab5c48f9efdfb7192..0000000000000000000000000000000000000000
--- a/src/units.c
+++ /dev/null
@@ -1,43 +0,0 @@
-#include "units.h"
-#include "tools.h"
-#include "part.h"
-
-#include <Python.h>
-#include <stdlib.h>
-#include <string.h>
-
-PyObject* test(PyObject *self, PyObject *args)
-{
-
-  /*
-  size_t N = sizeof(struct unit_system);
-
-  struct unit_system *us = malloc(N);
-  us->UnitMass_in_cgs = 1.;
-  us->UnitLength_in_cgs = 2.;
-  us->UnitTime_in_cgs = 3.;
-  us->UnitCurrent_in_cgs = 4.;
-  us->UnitTemperature_in_cgs = 5.;
-
-  //char *tmp = malloc(N);
-  //memcpy(tmp, us, N);
-
-  PyObject *object = tools_return(us, class_units);
-  */
-
-  size_t N = sizeof(struct part);
-  
-  struct part *p = malloc(N);
-  p->pos[0] = 1.;
-  p->pos[1] = 2.;
-  p->pos[2] = 3.;
-  p->vel[0] = 5.;
-  p->vel[1] = 6.;
-  p->vel[2] = 7.;
-  p->entropy = 8.;
-  p->test = p;
-
-  PyObject *object = tools_return(p, class_part);
-  
-  return object;
-}
diff --git a/src/units.h b/src/units.h
deleted file mode 100644
index 65a22e4850a2c166552c8dc9a80a42a75b507107..0000000000000000000000000000000000000000
--- a/src/units.h
+++ /dev/null
@@ -1,27 +0,0 @@
-#ifndef __PYSWIFTSIM_UNITS_H__
-#define __PYSWIFTSIM_UNITS_H__
-
-#include <Python.h>
-
-struct unit_system {
-
-  /*! Conversion factor from grams to internal mass units */
-  double UnitMass_in_cgs;
-
-  /*! Conversion factor from centimeters to internal length unit */
-  double UnitLength_in_cgs;
-
-  /*! Conversion factor from seconds to internal time units */
-  double UnitTime_in_cgs;
-
-  /*! Conversion factor from Ampere to internal current units */
-  double UnitCurrent_in_cgs;
-
-  /*! Conversion factor from Kelvins to internal temperature units. */
-  double UnitTemperature_in_cgs;
-};
-
-
-PyObject* test(PyObject *self, PyObject *args);
-
-#endif // __PYSWIFTSIM_UNITS_H__
diff --git a/src/units_wrapper.c b/src/units_wrapper.c
new file mode 100644
index 0000000000000000000000000000000000000000..d17d5c89763498f3f0ca0745fc7e5c668970e8d7
--- /dev/null
+++ b/src/units_wrapper.c
@@ -0,0 +1,24 @@
+#include "pyswiftsim_tools.h"
+
+#include <units.h>
+
+#include <Python.h>
+#include <stdlib.h>
+#include <string.h>
+
+PyObject* unit_system_test_struct(PyObject *self, PyObject *args)
+{
+
+  size_t N = sizeof(struct unit_system);
+
+  struct unit_system *us = malloc(N);
+  us->UnitMass_in_cgs = 1.;
+  us->UnitLength_in_cgs = 2.;
+  us->UnitTime_in_cgs = 3.;
+  us->UnitCurrent_in_cgs = 4.;
+  us->UnitTemperature_in_cgs = 5.;
+
+  PyObject *object = tools_return(us, class_units);
+
+  return object;
+}
diff --git a/src/units_wrapper.h b/src/units_wrapper.h
new file mode 100644
index 0000000000000000000000000000000000000000..f0f498aa2c4e65c086e8366b4575fcee39a38f95
--- /dev/null
+++ b/src/units_wrapper.h
@@ -0,0 +1,8 @@
+#ifndef __PYSWIFTSIM_UNITS_H__
+#define __PYSWIFTSIM_UNITS_H__
+
+#include <Python.h>
+
+PyObject* unit_system_test_struct(PyObject *self, PyObject *args);
+
+#endif // __PYSWIFTSIM_UNITS_H__
diff --git a/src/wrapper.c b/src/wrapper.c
index e858aba306060625817090c02ca9357afcdb039d..57951e50fdd674183b89a79b2211efe06d21b5a3 100644
--- a/src/wrapper.c
+++ b/src/wrapper.c
@@ -1,6 +1,7 @@
-#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+#include "units_wrapper.h"
+#include "part_wrapper.h"
+#include "pyswiftsim_tools.h"
 
-#include "units.h"
 
 #include <Python.h>
 #include <math.h>
@@ -11,11 +12,14 @@
       
 static PyMethodDef wrapper_methods[] = {
 
-          {"test",  test, METH_VARARGS,
-           "Test."},	   	         
-	   	   
-          {NULL, NULL, 0, NULL}        /* Sentinel */
-      };      
+  {"partTestStruct",  part_test_struct, METH_VARARGS,
+   "Construct a part object and return it."},	   	         
+
+  {"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_struct.py b/test/test_struct.py
index 8188b14136e688a6b9edd05147683dede1775904..7c228d849979d2229cfd9d741862f54535f63744 100644
--- a/test/test_struct.py
+++ b/test/test_struct.py
@@ -3,7 +3,13 @@
 from pyswiftsim import wrapper
 from pyswiftsim import structure
 
-us = wrapper.test()
+us = wrapper.unitSystemTestStruct()
+
+print(us)
+
+part = wrapper.partTestStruct()
+
+print(part)
 
 #us = structure.Units(t)
 
@@ -13,6 +19,7 @@ us = wrapper.test()
 #for i in us.struct_name:
 #    print(i, getattr(us, i))
 
+"""
 print(us.pos)
 us.pos = [12., 16., 11.]
 #us.pos[0] = 10.
@@ -36,3 +43,4 @@ d["UnitTemperature_in_cgs"] = 5.
 us = structure.UnitSystem(d)
 
 print(us)
+"""