Add STL files support in pywarpx (#3089)

* Added STL files support in pywarpx

* Add new EB bucket

* Added name to file headers

* Fixing year in  file headers

* Checking that stl file and imp function are not both specified

* Renamed EB bucket to EB2

* Adding STL documentation

* Implement suggestions from code review

Co-authored-by: lgiacome <lorenzo.giacome@cern.ch>

4 files changed, 53 insertions, 11 deletions

diff --git a/Python/pywarpx/EB2.py b/Python/pywarpx/EB2.py
new file mode 100644
index 000000000..4b74aafb0
--- /dev/null
+++ b/Python/pywarpx/EB2.py
@@ -0,0 +1,9 @@
+# Copyright 2022 Lorenzo Giacomel
+#
+# This file is part of WarpX.
+#
+# License: BSD-3-Clause-LBNL
+
+from .Bucket import Bucket
+
+eb2 = Bucket('eb2')
diff --git a/Python/pywarpx/WarpX.py b/Python/pywarpx/WarpX.py
index 50e34a72a..4701fcfb0 100644
--- a/Python/pywarpx/WarpX.py
+++ b/Python/pywarpx/WarpX.py
@@ -1,5 +1,5 @@
-# Copyright 2016-2020 Andrew Myers, David Grote, Maxence Thevenet
-# Remi Lehe
+# Copyright 2016-2022 Andrew Myers, David Grote, Maxence Thevenet
+# Remi Lehe, Lorenzo Giacomel
 #
 # This file is part of WarpX.
 #
@@ -13,6 +13,7 @@ from .Bucket import Bucket
 from .Collisions import collisions, collisions_list
 from .Constants import my_constants
 from .Diagnostics import diagnostics
+from .EB2 import eb2
 from .Geometry import geometry
 from .Interpolation import interpolation
 from .Langmuirwave import langmuirwave
@@ -38,6 +39,7 @@ class WarpX(Bucket):
         argv += langmuirwave.attrlist()
         argv += interpolation.attrlist()
         argv += psatd.attrlist()
+        argv += eb2.attrlist()
 
         # --- Search through species_names and add any predefined particle objects in the list.
         particles_list_names = [p.instancename for p in particles_list]
diff --git a/Python/pywarpx/__init__.py b/Python/pywarpx/__init__.py
index a41b50dc9..5590c7d6f 100644
--- a/Python/pywarpx/__init__.py
+++ b/Python/pywarpx/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2016-2019 Andrew Myers, David Grote
+# Copyright 2016-2022 Andrew Myers, David Grote, Lorenzo Giacomel
 #
 # This file is part of WarpX.
 #
@@ -10,6 +10,7 @@ from .Boundary import boundary
 from .Collisions import collisions
 from .Constants import my_constants
 from .Diagnostics import diagnostics
+from .EB2 import eb2
 from .Geometry import geometry
 from .Interpolation import interpolation
 from .Langmuirwave import langmuirwave
diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py
index ae5fc8bd4..bb50ed42d 100644
--- a/Python/pywarpx/picmi.py
+++ b/Python/pywarpx/picmi.py
@@ -1,5 +1,5 @@
-# Copyright 2018-2020 Andrew Myers, David Grote, Ligia Diana Amorim
-# Maxence Thevenet, Remi Lehe, Revathi Jambunathan
+# Copyright 2018-2022 Andrew Myers, David Grote, Ligia Diana Amorim
+# Maxence Thevenet, Remi Lehe, Revathi Jambunathan, Lorenzo Giacomel
 #
 #
 # This file is part of WarpX.
@@ -917,21 +917,43 @@ class EmbeddedBoundary(picmistandard.base._ClassWithInit):
     """
     Custom class to handle set up of embedded boundaries specific to WarpX.
     If embedded boundary initialization is added to picmistandard this can be
-    changed to inherit that functionality.
+    changed to inherit that functionality. The geometry can be specified either as
+    an implicit function or as an STL file (ASCII or binary). In the latter case the
+    geometry specified in the STL file can be scaled, translated and inverted.
     - implicit_function: Analytic expression describing the embedded boundary
+    - stl_file: STL file path (string), file contains the embedded boundary geometry
+    - stl_scale: factor by which the STL geometry is scaled (pure number)
+    - stl_center: vector by which the STL geometry is translated (in meters)
+    - stl_reverse_normal: if True inverts the orientation of the STL geometry
     - potential: Analytic expression defining the potential. Can only be specified
                  when the solver is electrostatic. Optional, defaults to 0.
      Parameters used in the expressions should be given as additional keyword arguments.
     """
-    def __init__(self, implicit_function, potential=None, **kw):
+    def __init__(self, implicit_function=None, stl_file=None, stl_scale=None, stl_center=None, stl_reverse_normal=False,
+                 potential=None, **kw):
+
+        assert stl_file is None or implicit_function is None, Exception('Only one between implicit_function and '
+                                                                            'stl_file can be specified')
+
         self.implicit_function = implicit_function
+        self.stl_file = stl_file
+
+        if stl_file is None:
+            assert stl_scale is None, Exception('EB can only be scaled only when using an stl file')
+            assert stl_center is None, Exception('EB can only be translated only when using an stl file')
+            assert stl_reverse_normal is False, Exception('EB can only be reversed only when using an stl file')
+
+        self.stl_scale = stl_scale
+        self.stl_center = stl_center
+        self.stl_reverse_normal = stl_reverse_normal
+
         self.potential = potential
 
         # Handle keyword arguments used in expressions
         self.user_defined_kw = {}
         for k in list(kw.keys()):
-            if (re.search(r'\b%s\b'%k, implicit_function) or
-                (potential is not None and re.search(r'\b%s\b'%k, potential))):
+            if (implicit_function is not None and re.search(r'\b%s\b'%k, implicit_function) or
+               (potential is not None and re.search(r'\b%s\b'%k, potential))):
                 self.user_defined_kw[k] = kw[k]
                 del kw[k]
 
@@ -944,8 +966,16 @@ class EmbeddedBoundary(picmistandard.base._ClassWithInit):
         # defined in my_constants with the same name but different value.
         self.mangle_dict = pywarpx.my_constants.add_keywords(self.user_defined_kw)
 
-        expression = pywarpx.my_constants.mangle_expression(self.implicit_function, self.mangle_dict)
-        pywarpx.warpx.eb_implicit_function = expression
+        if self.implicit_function is not None:
+            expression = pywarpx.my_constants.mangle_expression(self.implicit_function, self.mangle_dict)
+            pywarpx.warpx.eb_implicit_function = expression
+
+        if self.stl_file is not None:
+            pywarpx.eb2.geom_type = "stl"
+            pywarpx.eb2.stl_file = self.stl_file
+            pywarpx.eb2.stl_scale = self.stl_scale
+            pywarpx.eb2.stl_center = self.stl_center
+            pywarpx.eb2.stl_reverse_normal = self.stl_reverse_normal
 
         if self.potential is not None:
             assert isinstance(solver, ElectrostaticSolver), Exception('The potential is only supported with the ElectrostaticSolver')