Fixes for the Python RZ version

2 files changed, 13 insertions, 7 deletions

diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py
index 5044b1d67..4c3283b97 100755
--- a/Python/pywarpx/_libwarpx.py
+++ b/Python/pywarpx/_libwarpx.py
@@ -30,18 +30,24 @@ def _get_package_root():
 # --- Default to 3D if geometry is not setup yet.
 try:
     _prob_lo = geometry.prob_lo
+    _coord_sys = geometry.coord_sys
 except AttributeError:
-    geometry_dim = 3
+    geometry_dim = '3d'
 else:
-    geometry_dim = len(_prob_lo)
-    del _prob_lo
+    if _coord_sys == 0:
+        geometry_dim = '%dd'%len(_prob_lo)
+    elif _coord_sys == 1:
+        geometry_dim = 'rz'
+    else:
+        raise Exception('Undefined coordinate system %d'%_coord_sys)
+    del _prob_lo, _coord_sys
 
 _libc = ctypes.CDLL(_find_library('c'))
 
 try:
-    libwarpx = ctypes.CDLL(os.path.join(_get_package_root(), "libwarpx%dd.so"%geometry_dim))
+    libwarpx = ctypes.CDLL(os.path.join(_get_package_root(), "libwarpx%s.so"%geometry_dim))
 except OSError:
-    raise Exception('libwarpx%dd.so was not installed. It can be installed by running "make" in the Python directory of WarpX'%geometry_dim)
+    raise Exception('libwarpx%s.so was not installed. It can be installed by running "make" in the Python directory of WarpX'%geometry_dim)
 
 dim = libwarpx.warpx_SpaceDim()
 
diff --git a/Python/setup.py b/Python/setup.py
index f4917e7e2..ecb87190c 100644
--- a/Python/setup.py
+++ b/Python/setup.py
@@ -10,12 +10,12 @@ import argparse
 from setuptools import setup
 
 argparser = argparse.ArgumentParser(add_help=False)
-argparser.add_argument('--with-libwarpx', type=int, default=None, help='Install libwarpx with the given value as DIM. This option is only used by the makefile.')
+argparser.add_argument('--with-libwarpx', type=str, default=None, help='Install libwarpx with the given value as DIM. This option is only used by the makefile.')
 args, unknown = argparser.parse_known_args()
 sys.argv = [sys.argv[0]] + unknown
 
 if args.with_libwarpx:
-    package_data = {'pywarpx' : ['libwarpx%dd.so'%args.with_libwarpx]}
+    package_data = {'pywarpx' : ['libwarpx%s.so'%args.with_libwarpx]}
 else:
     package_data = {}