diff --git a/src/ramanchada2/misc/spectrum_deco/spectrum_filter.py b/src/ramanchada2/misc/spectrum_deco/spectrum_filter.py index 4b538548..008d36ac 100644 --- a/src/ramanchada2/misc/spectrum_deco/spectrum_filter.py +++ b/src/ramanchada2/misc/spectrum_deco/spectrum_filter.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 from functools import wraps -from copy import copy +from copy import copy, deepcopy from ramanchada2.spectrum.spectrum import Spectrum from .dynamically_added import dynamically_added_filters import logging @@ -13,7 +13,9 @@ def add_spectrum_filter(fun): @wraps(fun) def retf(old_spe: Spectrum, *args, **kwargs) -> Spectrum: new_spe = copy(old_spe) - new_spe._applied_processings.append(proc=fun.__name__, args=args, kwargs=kwargs) + new_spe._applied_processings.append(proc=fun.__name__, + args=deepcopy(args), + kwargs=deepcopy(kwargs)) fun(old_spe, new_spe, *args, **kwargs) new_spe.write_cache() return new_spe diff --git a/src/ramanchada2/spectrum/baseline/add_baseline.py b/src/ramanchada2/spectrum/baseline/add_baseline.py index 1d467aac..02d0f667 100644 --- a/src/ramanchada2/spectrum/baseline/add_baseline.py +++ b/src/ramanchada2/spectrum/baseline/add_baseline.py @@ -14,8 +14,14 @@ def generate_baseline( n_freq: int = Field(..., gt=2), size: int = Field(..., gt=2), - rng_seed: Union[int, None] = None): - rng = np.random.default_rng(rng_seed) + # validation for rng_seed is removed because + # it makes in-place modification impossible + rng_seed=None): + if isinstance(rng_seed, dict): + rng = np.random.default_rng() + rng.__setstate__(rng_seed) + else: + rng = np.random.default_rng(rng_seed) k = rng.normal(0, size, size=(2, n_freq)) k[1][0] = 0 z = k[0] + k[1]*1j @@ -26,6 +32,8 @@ def generate_baseline( base = base[:size] base -= base.min() base /= base.max() + if isinstance(rng_seed, dict): + rng_seed.update(rng.__getstate__()) return base diff --git a/src/ramanchada2/spectrum/filters/add_gaussian_noise.py b/src/ramanchada2/spectrum/filters/add_gaussian_noise.py index 775f3173..27e1976a 100644 --- a/src/ramanchada2/spectrum/filters/add_gaussian_noise.py +++ b/src/ramanchada2/spectrum/filters/add_gaussian_noise.py @@ -1,7 +1,5 @@ #!/usr/bin/env python3 -from typing import Union - import numpy as np from pydantic import validate_arguments, PositiveFloat @@ -15,7 +13,9 @@ def add_gaussian_noise( old_spe: Spectrum, new_spe: Spectrum, /, sigma: PositiveFloat, - rng_seed: Union[int, None] = None): + # validation for rng_seed is removed because + # it makes in-place modification impossible + rng_seed=None): r""" Add gaussian noise to the spectrum. Random number i.i.d. $N(0, \sigma)$ is added to every sample @@ -24,11 +24,18 @@ def add_gaussian_noise( ---------- sigma : float sigma of the gaussian noise - rng_seed : int, optional - seed for the random generator + rng_seed : int or rng state, optional + seed for the random generator. + If a state is provided it is updated in-place """ - rng = np.random.default_rng(rng_seed) + if isinstance(rng_seed, dict): + rng = np.random.default_rng() + rng.__setstate__(rng_seed) + else: + rng = np.random.default_rng(rng_seed) dat = old_spe.y + rng.normal(0., sigma, size=len(old_spe.y)) if any(dat < 0): dat += abs(dat.min()) + if isinstance(rng_seed, dict): + rng_seed.update(rng.__getstate__()) new_spe.y = np.array(dat) diff --git a/src/ramanchada2/spectrum/filters/add_gaussian_noise_drift.py b/src/ramanchada2/spectrum/filters/add_gaussian_noise_drift.py new file mode 100644 index 00000000..b2004203 --- /dev/null +++ b/src/ramanchada2/spectrum/filters/add_gaussian_noise_drift.py @@ -0,0 +1,74 @@ +import numpy as np +from pydantic import PositiveFloat, confloat, validate_arguments + +from ramanchada2.misc.spectrum_deco import add_spectrum_filter + +from ..spectrum import Spectrum + + +@validate_arguments(config=dict(arbitrary_types_allowed=True)) +def generate_add_gaussian_noise_drift(y, /, + sigma: PositiveFloat, + coef: confloat(ge=0, le=1), # type: ignore [valid-type] + # validation for rng_seed is removed because + # it makes in-place modification impossible + rng_seed=None): + if isinstance(rng_seed, dict): + rng = np.random.default_rng() + rng.__setstate__(rng_seed) + else: + rng = np.random.default_rng(rng_seed) + gaus = rng.normal(0., sigma+coef/np.sqrt(2), size=len(y)) + cs = np.cumsum(gaus) + # coef*sum(cs[:i]) + (1-coef)*gaus is identical to + # coef*sum(cs[:i-1]) + gaus + noise = coef*cs + gaus*(1-coef) + noise -= np.std(noise) + dat = y + noise + if any(dat < 0): + dat += abs(dat.min()) + if isinstance(rng_seed, dict): + rng_seed.update(rng.__getstate__()) + return np.array(dat) + + +@add_spectrum_filter +@validate_arguments(config=dict(arbitrary_types_allowed=True)) +def add_gaussian_noise_drift( + old_spe: Spectrum, + new_spe: Spectrum, /, + sigma: PositiveFloat, + coef: confloat(ge=0, le=1), # type: ignore [valid-type] + # validation for rng_seed is removed because + # it makes in-place modification impossible + rng_seed=None): + r""" + Add cumulative gaussian noise to the spectrum. + Exponential-moving-average-like gaussian noise is added + to each sample. The goal is to mimic the low-frequency noise + (or random substructures in spectra). + The additive noise is + .. math:: + a_i = coef*\sum_{j=0}^{i-1}g_j + g_i, + where + .. math:: + g_i = \mathcal{N}(0, 1+\frac{coef}{\sqrt 2}). + This way drifting is possible while keeping the + .. math:: + \sigma(\Delta(a)) \approx 1. + + Parameters + ---------- + sigma : float + sigma of the gaussian noise + coef : float in [0, 1] + drifting coefficient. if `coef == 0` the result is + identical to `add_gaussian_noise()`. + rng_seed : int or rng state, optional + seed for the random generator. + If a state is provided it is updated in-place + """ + new_spe.y = generate_add_gaussian_noise_drift(old_spe.y, + sigma=sigma, + coef=coef, + rng_seed=rng_seed) diff --git a/src/ramanchada2/spectrum/filters/add_poisson_noise.py b/src/ramanchada2/spectrum/filters/add_poisson_noise.py index b947b2ae..6e8f3102 100644 --- a/src/ramanchada2/spectrum/filters/add_poisson_noise.py +++ b/src/ramanchada2/spectrum/filters/add_poisson_noise.py @@ -24,10 +24,17 @@ def add_poisson_noise( ---------- scale : float, optional scale the amplitude of the noise, by default 1 - rng_seed : int, optional + rng_seed : int or rng state, optional seed for the random generator + If a state is provided it is updated in-place """ - rng = np.random.default_rng(rng_seed) + if isinstance(rng_seed, dict): + rng = np.random.default_rng() + rng.__setstate__(rng_seed) + else: + rng = np.random.default_rng(rng_seed) dat = old_spe.y + [rng.normal(0., np.sqrt(i*scale)) for i in old_spe.y] dat[dat < 0] = 0 + if isinstance(rng_seed, dict): + rng_seed.update(rng.__getstate__()) new_spe.y = np.array(dat)