Merge pull request #171 from h2020charisma/drop_spectrum_nans

Drop NaNs from spectrum

src/ramanchada2/spectrum/filters/dropna.py

@@ -0,0 +1,32 @@
+import numpy as np
+from pydantic import validate_call
+
+from ramanchada2.misc.spectrum_deco import add_spectrum_filter
+
+from ..spectrum import Spectrum
+
+
+@add_spectrum_filter
+@validate_call(config=dict(arbitrary_types_allowed=True))
+def dropna(old_spe: Spectrum,
+           new_spe: Spectrum):
+    """
+    Remove non finite numbers on both axes
+
+    Args:
+        old_spe: internal use only
+        new_spe: internal use only
+
+    Returns: modified Spectrum
+    """
+
+    x = old_spe.x
+    y = old_spe.y
+    idx = np.isfinite(x)
+    x = x[idx]
+    y = y[idx]
+    idx = np.isfinite(y)
+    x = x[idx]
+    y = y[idx]
+    new_spe.x = x
+    new_spe.y = y

tests/spectrum/filters/test_dropna.py

@@ -0,0 +1,16 @@
+import numpy as np
+
+import ramanchada2 as rc2
+
+
+def test_dropna():
+    x = np.array([0, 1, 2, np.nan, np.nan, 3, 4, 5, 6, np.nan])
+    y = np.array([np.nan, 1, 2, 3, np.nan, np.nan, 4, 5, 6, np.nan])
+    spe = rc2.spectrum.Spectrum(x=np.array([1]), y=np.array([2]))
+    spe.x = x
+    spe.y = y
+
+    spe_new = spe.dropna()
+
+    assert np.all(spe_new.x == spe_new.y)
+    assert np.all(spe_new.x == np.array([1, 2, 4, 5, 6]))