More datagen classes

src/labcore/data/datagen.py

@@ -41,6 +41,7 @@
 @dataclass
 class DataGen(ABC):
     noise_std: float = 1.0
+    imaginary: bool = False
 
     @staticmethod
     @abstractmethod
@@ -53,12 +54,15 @@ def generate(self, coordinates: NDArray[Any], **kwargs: Any) -> NDArray[Any]:
         params = asdict(self)
         params.update(kwargs)
         noise_std = params.pop("noise_std")
+        imaginary = params.pop("imaginary")
 
         one_d = coordinates.ndim == 1
         coordinates = np.atleast_2d(coordinates)
         model_outputs = np.array(
             [
-                self.model(coords, **params) + self.noise(coords, noise_std)
+                self.model(coords, **params)
+                + self.noise(coords, noise_std)
+                + (1j * self.noise(coords, noise_std) if imaginary else 0)
                 for coords in coordinates
             ]
         )
@@ -69,11 +73,12 @@ def generate(self, coordinates: NDArray[Any], **kwargs: Any) -> NDArray[Any]:
 
     @staticmethod
     def noise(coordinates: NDArray[Any], std: float) -> NDArray[Any]:
-        return np.random.normal(scale=std, size=len(coordinates))
+        n = len(coordinates)
+        return np.random.normal(scale=std, size=n)
 
 
 @dataclass
-class ExponentialDataGen(DataGen):
+class Exponential(DataGen):
     base: float = np.e
 
     @staticmethod
@@ -82,7 +87,7 @@ def model(coordinates: NDArray[Any], base: float) -> NDArray[Any]:
 
 
 @dataclass
-class SineDataGen(DataGen):
+class Sine(DataGen):
     A: float = 1
     f: float = 1
     phi: float = 0
@@ -96,7 +101,7 @@ def model(
 
 
 @dataclass
-class GaussianDataGen(DataGen):
+class Gaussian(DataGen):
     x0: float = 0
     sigma: float = 1
     A: float = 1
@@ -107,3 +112,48 @@ def model(
         coordinates: NDArray[Any], x0: float, sigma: float, A: float, of: float
     ) -> NDArray[Any]:
         return A * np.exp(-((coordinates - x0) ** 2) / (2 * sigma**2)) + of
+
+
+@dataclass
+class ExponentialDecay(DataGen):
+    A: float = 1
+    tau: float = 1
+    of: float = 0
+
+    @staticmethod
+    def model(
+        coordinates: NDArray[Any], A: float, tau: float, of: float
+    ) -> NDArray[Any]:
+        return A * np.exp(-coordinates / tau) + of
+
+
+@dataclass
+class ExponentialDecayingSine(DataGen):
+    A: float = 1
+    f: float = 1
+    phi: float = 0
+    tau: float = 1
+    of: float = 0
+
+    @staticmethod
+    def model(
+        coordinates: NDArray[Any], A: float, f: float, phi: float, tau: float, of: float
+    ) -> NDArray[Any]:
+        return (
+            A * np.exp(-coordinates / tau) * np.sin(2 * np.pi * f * coordinates + phi)
+            + of
+        )
+
+
+@dataclass
+class Lorentzian(DataGen):
+    A: float = 1
+    x0: float = 0
+    gamma: float = 1
+    of: float = 0
+
+    @staticmethod
+    def model(
+        coordinates: NDArray[Any], A: float, x0: float, gamma: float, of: float
+    ) -> NDArray[Any]:
+        return A * (gamma**2) / ((coordinates - x0) ** 2 + gamma**2) + of