Add Python wrapper to set the lens strength (#3748)

* Add Python wrapper to set the lens strength

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Correct missing `c_real`

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add assert message

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

1 files changed, 22 insertions, 0 deletions

diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py
index ebd18e5ee..2a297b908 100755
--- a/Python/pywarpx/_libwarpx.py
+++ b/Python/pywarpx/_libwarpx.py
@@ -311,6 +311,7 @@ class LibWarpX():
         self.libwarpx_so.warpx_sett_new.argtypes = [ctypes.c_int, c_real]
         self.libwarpx_so.warpx_getdt.argtypes = [ctypes.c_int]
         self.libwarpx_so.warpx_setPotentialEB.argtypes = [ctypes.c_char_p]
+        self.libwarpx_so.warpx_setPlasmaLensStrength.argtypes = [ctypes.c_int, c_real, c_real]
 
     def _get_boundary_number(self, boundary):
         '''
@@ -1231,6 +1232,27 @@ class LibWarpX():
         """
         self.libwarpx_so.warpx_setPotentialEB(ctypes.c_char_p(potential.encode('utf-8')))
 
+    def set_plasma_lens_strength( self, i_lens, strength_E, strength_B ):
+        """
+        Set the strength of the `i_lens`-th lens
+
+        Parameters
+        ----------
+        i_lens: int
+            Index of the lens to be modified
+
+        strength_E, strength_B: floats
+            The electric and magnetic focusing strength of the lens
+        """
+        if self._numpy_real_dtype == 'f8':
+            c_real = ctypes.c_double
+        else:
+            c_real = ctypes.c_float
+
+        self.libwarpx_so.warpx_setPlasmaLensStrength(
+            ctypes.c_int(i_lens), c_real(strength_E), c_real(strength_B) )
+
+
     def _get_mesh_field_list(self, warpx_func, level, direction, include_ghosts):
         """
         Generic routine to fetch the list of field data arrays.