2 files changed, 17 insertions, 4 deletions
diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py
index 57caf5d66..77a7a2f1c 100755
--- a/Python/pywarpx/_libwarpx.py
+++ b/Python/pywarpx/_libwarpx.py
@@ -6,6 +6,8 @@ from ctypes.util import find_library as _find_library
 import numpy as np
 from numpy.ctypeslib import ndpointer as _ndpointer
 
+from .Geometry import geometry
+
 # --- Is there a better way of handling constants?
 clight = 2.99792458e+8 # m/s
 
@@ -22,11 +24,22 @@ def _get_package_root():
             return ''
         cur = os.path.dirname(cur)
 
+# --- Use geometry to determine whether to import the 2D or 3D version.
+# --- This assumes that the input is setup before this module is imported,
+# --- which should normally be the case.
+# --- Default to 3D if geometry is not setup yet.
+try:
+    _prob_lo = geometry.prob_lo
+except AttributeError:
+    geometry_dim = 3
+else:
+    geometry_dim = len(_prob_lo)
+    del _prob_lo
 
 try:
-    libwarpx = ctypes.CDLL(os.path.join(_get_package_root(), "libwarpx.so"))
+    libwarpx = ctypes.CDLL(os.path.join(_get_package_root(), "libwarpx%dd.so"%geometry_dim))
 except OSError:
-    raise Exception('libwarpx.so was not installed. It can be installed by running "make" in the Python directory of WarpX')
+    raise Exception('libwarpx%dd.so was not installed. It can be installed by running "make" in the Python directory of WarpX'%geometry_dim)
 
 _libc = ctypes.CDLL(_find_library('c'))
 
diff --git a/Python/setup.py b/Python/setup.py
index 5f762dabd..e68870809 100644
--- a/Python/setup.py
+++ b/Python/setup.py
@@ -15,12 +15,12 @@ except ImportError:
     from distutils.command.build_py import build_py
 
 argparser = argparse.ArgumentParser(add_help=False)
-argparser.add_argument('--with-libwarpx', action='store_true', help='Install libwarpx. This option is only used by the makefile.')
+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.')
 args, unknown = argparser.parse_known_args()
 sys.argv = [sys.argv[0]] + unknown
 
 if args.with_libwarpx:
-    package_data = {'pywarpx' : ['libwarpx.so']}
+    package_data = {'pywarpx' : ['libwarpx%dd.so'%args.with_libwarpx]}
 else:
     package_data = {}