laserfields 0.4.2__tar.gz → 0.4.3__tar.gz

{laserfields-0.4.2 → laserfields-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: laserfields
-Version: 0.4.2
+Version: 0.4.3
 Summary: Python library for describing time-dependent laserfields by Johannes Feist.
 Author-email: Johannes Feist <johannes.feist@gmail.com>
 Requires-Python: >=3.7

{laserfields-0.4.2 → laserfields-0.4.3}/laserfields/__init__.py

@@ -1,6 +1,6 @@
 """Python library for describing time-dependent laserfields by Johannes Feist."""
 
-__version__ = "0.4.2"
+__version__ = "0.4.3"
 
 __all__ = [
     "intensity_Wcm2_to_Eau",

{laserfields-0.4.2 → laserfields-0.4.3}/laserfields/laserfields.py

@@ -208,31 +208,36 @@ class SinExpLaserField(LaserField):
         return self.T * gamma(0.5 + n_photon * self.exponent) / (sqrt(π) * gamma(1 + n_photon * self.exponent))
 
 
+def _make_flattop_jit_envelope(ramponfunc, ramponfuncpr):
+    ramponfunc_jit = njit(ramponfunc)
+    ramponfuncpr_jit = njit(ramponfuncpr)
+
+    @vectorize(nopython=True)
+    def _jit_envelope(tr, E0, Tflat, Tramp, getprime):
+        if abs(tr) > Tflat / 2 + Tramp:
+            return 0.0
+        elif abs(tr) > Tflat / 2:
+            trel = (Tramp + Tflat / 2 - abs(tr)) / Tramp
+            if getprime:
+                return -E0 * np.sign(tr) * ramponfuncpr_jit(trel) / Tramp
+            else:
+                return E0 * ramponfunc_jit(trel)
+        else:
+            return 0.0 if getprime else E0
+
+    return _jit_envelope
+
+
 @dataclass
 class FlatTopLaserField(LaserField):
     Tflat: float
     Tramp: float
 
-    def _envelope(self, tr):
-        # we cannot pass functions to numba, so make a closure
-        if not hasattr(self, "_jit_envelope"):
-            ramponfunc = njit(self.ramponfunc)
-            ramponfuncpr = njit(self.ramponfuncpr)
-
-            @vectorize(nopython=True)
-            def _jit_envelope(tr, E0, Tflat, Tramp, getprime):
-                if abs(tr) > Tflat / 2 + Tramp:
-                    return 0.0
-                elif abs(tr) > Tflat / 2:
-                    trel = (Tramp + Tflat / 2 - abs(tr)) / Tramp
-                    if getprime:
-                        return -E0 * np.sign(tr) * ramponfuncpr(trel) / Tramp
-                    else:
-                        return E0 * ramponfunc(trel)
-                else:
-                    return 0.0 if getprime else E0
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls._jit_envelope = _make_flattop_jit_envelope(cls.ramponfunc, cls.ramponfuncpr)
 
-            self._jit_envelope = _jit_envelope
+    def _envelope(self, tr):
         env = self._jit_envelope(tr, self.E0, self.Tflat, self.Tramp, False)
         envpr = self._jit_envelope(tr, self.E0, self.Tflat, self.Tramp, True)
         return env, envpr