Merge pull request #12 from h2020charisma/devel-gg

central moments and arbitrary function as baseline

examples/synthetic-spectrum-generation.ipynb

@@ -24,8 +24,8 @@
     "def convolve(lineshape, sigma):\n",
     "    spe_tmp['convolved'] = spe_tmp['orig'].convolve(lineshape=lineshape, sigma=sigma)\n",
     "    spe_tmp['convolved'].plot()\n",
-    "def add_baseline(n_freq, amplitude, pedestal, rng_seed):\n",
-    "    spe_tmp['baseline'] = spe_tmp['convolved'].add_baseline(n_freq=n_freq, amplitude=amplitude, pedestal=pedestal, rng_seed=rng_seed)\n",
+    "def add_baseline(n_freq, amplitude, pedestal, rng_seed, a1, a2):\n",
+    "    spe_tmp['baseline'] = spe_tmp['convolved'].add_baseline(n_freq=n_freq, amplitude=amplitude, pedestal=pedestal, rng_seed=rng_seed, func=lambda x: x*a1 + x**2*a2)\n",
     "    spe_tmp['baseline'].plot()\n",
     "def add_noise(scale, rng_seed):\n",
     "    spe_tmp['noise'] = spe_tmp['baseline'].add_poisson_noise(scale=scale, rng_seed=rng_seed)\n",
@@ -49,7 +49,7 @@
    "source": [
     "spe = rc2.spectrum.from_delta_lines(deltas={500:1e3, 700:1.5e3})\n",
     "spe = spe.convolve(lineshape='voigt', sigma=5, gamma=3)\n",
-    "spe = spe.add_baseline(n_freq=50, amplitude=10, pedestal=10)\n",
+    "spe = spe.add_baseline(n_freq=50, amplitude=10, pedestal=10, func=lambda x: x*.006 + x**2*.00001)\n",
     "spe = spe.add_poisson_noise(scale=.01)\n",
     "spe.plot()"
    ]
@@ -137,6 +137,8 @@
     "ipyw.interact(\n",
     "    add_baseline,\n",
     "    n_freq=ipyw.widgets.IntSlider(min=3, max=500, value=50),\n",
+    "    a1=ipyw.widgets.FloatSlider(min=-.2, max=.2, value=.0, step=.00001),\n",
+    "    a2=ipyw.widgets.FloatSlider(min=-.0002, max=.0002, value=.0, step=.00001,  readout_format='.5f'),\n",
     "    amplitude=ipyw.widgets.FloatSlider(min=0, max=50, value=10),\n",
     "    pedestal=ipyw.widgets.FloatSlider(min=1, max=50, value=10),\n",
     "    rng_seed=ipyw.widgets.IntSlider(min=0, max=999999, value=0)\n",

src/ramanchada2/spectrum/__init__.py

@@ -3,6 +3,7 @@
 from .spectrum import Spectrum  # noqa
 from .arithmetics import *  # noqa
 from .baseline import *  # noqa
+from .calc import *  # noqa
 from .calibration import *  # noqa
 from .filters import *  # noqa
 from .peaks import *  # noqa

src/ramanchada2/spectrum/baseline/add_baseline.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-from typing import Union
+from typing import Union, Callable
 
 from pydantic import validate_arguments, Field
 import numpy as np
@@ -32,12 +32,14 @@ def generate_baseline(
 @add_spectrum_filter
 @validate_arguments(config=dict(arbitrary_types_allowed=True))
 def add_baseline(old_spe: Spectrum, new_spe: Spectrum, /,
-                 n_freq, amplitude, pedestal, rng_seed=None):
+                 n_freq, amplitude, pedestal, func: Callable = None, rng_seed=None):
     """
     Add artificial baseline to the spectrum.
     A random baseline is generated in frequency domain using uniform random numbers.
     The baseline in frequency domain is tapered with bohman window to reduce the bandwidth
     of the baseline to first `n_freq` frequencies and is transformed to "time" domain.
+    Additionaly by using `func` parameter the user can define arbitrary function
+    to be added as baseline.
 
     Parameters
     ----------
@@ -47,9 +49,15 @@ def add_baseline(old_spe: Spectrum, new_spe: Spectrum, /,
         upper boundary for the uniform random generator
     pedestal : float
         additive constant pedestal to the spectrum
+    func : callable
+        user defined function to be added as baseline.
+        Example: func=lambda x: x*.01 + x**2*.0001
     rng_seed : int, optional
         seed for the random generator
     """
     size = len(old_spe.y)
     base = generate_baseline(n_freq=n_freq, size=size, rng_seed=rng_seed)
-    new_spe.y = old_spe.y + amplitude*base + pedestal
+    y = old_spe.y + amplitude*base + pedestal
+    if func is not None:
+        y += func(old_spe.x) + old_spe.y
+    new_spe.y = y

src/ramanchada2/spectrum/calc/__init__.py

@@ -0,0 +1,8 @@
+import os
+import glob
+
+__all__ = [
+    os.path.basename(f)[:-3]
+    for f in glob.glob(os.path.dirname(__file__)+"/*.py")
+    if os.path.isfile(f) and not os.path.basename(f).startswith('_')
+]

src/ramanchada2/spectrum/calc/central_moments.py

@@ -0,0 +1,26 @@
+import numpy as np
+from pydantic import validate_arguments
+
+from ..spectrum import Spectrum
+from ramanchada2.misc.spectrum_deco import add_spectrum_method
+
+
+@add_spectrum_method
+@validate_arguments(config=dict(arbitrary_types_allowed=True))
+def central_moments(spe: Spectrum, /,
+                    left_idx=0, right_idx=-1, moments=[1, 2, 3, 4], normalize=False
+                    ):
+    mom = dict()
+    x = spe.x[left_idx:right_idx]
+    p = spe.y[left_idx:right_idx]
+    p -= p.min()
+    p /= p.sum()
+    mom[1] = np.sum(x*p)
+    mom[2] = np.sum((x - mom[1])**2 * p)
+    for i in moments:
+        if i <= 2:
+            continue
+        mom[i] = np.sum((x - mom[1])**i * p)
+        if normalize and i > 2:
+            mom[i] /= mom[2] ** i/2
+    return [mom[i] for i in moments]