Repository: staircase-dev/staircase (language: python)

Issue:
ENH: Allow Stairs.fillna to take a Stairs parameter
#### Is your feature request related to a problem?

We want the value parameter in Stairs.fillna to be able to take a staircase.Stairs parameter 

current definition https://github.com/staircase-dev/staircase/blob/cb0265d1b182ad3a8452f544c197ce54aef6feea/staircase/core/ops/masking.py#L160

docstring https://github.com/staircase-dev/staircase/blob/cb0265d1b182ad3a8452f544c197ce54aef6feea/staircase/core/ops/docstrings.py#L459

#### Describe the solution you'd like

See https://www.staircase.dev/en/latest/user_guide/cookbook.html#fill-undefined-intervals-of-one-step-function-with-another for an implementation
Separate docstrings into separate file



Discussion / hints:
take
Hi - whilst implementing the suggested approach is straightforward - it has the side effect of leaving zeros in the mutated SC object if the replacing SC also has nans that overlap.

I would expect this case to leave those nans as nans, replacing only those nans where the replacing SC has values

The logic for this is straightforward - but I'm still to fully understand the data structure to implement it.  A few attempts aren't yet working.

aka - I'm still plugging away
Hi @ChristopherShort, good catch!

Will certainly have to be one of the tests!

I have a solution which builds on the original one liner, without having to delve into the internals of staircase.  Don't want to ruin it for you if you're enjoying figuring out so I'm happy to either give you the fix, give you a hint, or neither.  It's up to you.


Fix patch:
diff --git a/docs/release_notes/changelog.rst b/docs/release_notes/changelog.rst
index 16854e4..506f53a 100644
--- a/docs/release_notes/changelog.rst
+++ b/docs/release_notes/changelog.rst
@@ -7,7 +7,8 @@ Changelog
 
 UNRELEASED
 
-- warnings generated if using :meth:`staircase.StairsArray.layer` with non-standard argument types (#GH132)
+- :meth:`staircase.Stairs.fillna` can be used with :class:`staircase.Stairs` argument (#GH82)
+- warnings generated if using :meth:`staircase.Stairs.layer` with non-standard argument types (#GH132)
 
 Please list new changes above this comment
 
diff --git a/docs/user_guide/cookbook.rst b/docs/user_guide/cookbook.rst
index c280f2a..7101b61 100644
--- a/docs/user_guide/cookbook.rst
+++ b/docs/user_guide/cookbook.rst
@@ -126,37 +126,6 @@ A variant of the above problem, suppose a collection of events is defined by sta
     9  38.199949   43.50357
 
 
-Fill undefined intervals of one step function with another
------------------------------------------------------------
-
-:meth:`staircase.Stairs.fillna` allows undefined intervals in a step function to be redefined (i.e. "filled") with a number.  This recipe shows a simple one-liner which fills the undefined values of step function *a* with the values of step function *b*.
-
-.. plot::
-    :context: close-figs
-
-    >>> # test data
-    >>> def gen_test_step_function(seed):
-    ...     return (
-    ...         sc.make_test_data(dates=False, seed=seed)
-    ...         .pipe(sc.Stairs, "start", "end")
-    ...     )
-    ...
-    >>> a = gen_test_step_function(0).mask((20,30)).mask((80,90))
-    >>> b = gen_test_step_function(1)
-
-    >>> # recipe
-    >>> result = a.fillna(0) + b.where(a.isna()).fillna(0)
-
-    >>> # plot
-    >>> fig, axes = plt.subplots(ncols=3, figsize=(8,3), sharex=True, sharey=True)
-    >>> a.plot(axes[0])
-    >>> axes[0].set_title("a")
-    >>> b.plot(axes[1])
-    >>> axes[1].set_title("b")
-    >>> result.plot(axes[2])
-    >>> axes[2].set_title("result")
-
-
 Stitch two step functions together at a point
 ----------------------------------------------
 
diff --git a/staircase/core/ops/masking.py b/staircase/core/ops/masking.py
index ff25fe3..13bb9cf 100644
--- a/staircase/core/ops/masking.py
+++ b/staircase/core/ops/masking.py
@@ -7,6 +7,7 @@
 import staircase as sc
 from staircase.constants import inf
 from staircase.core.ops import docstrings
+from staircase.core.ops.common import requires_closed_match
 from staircase.util import _replace_none_with_infs
 from staircase.util._decorators import Appender
 
@@ -159,28 +160,53 @@ def func(self):
 notna = _make_null_comparison_func(docstrings.notna_docstring, lambda x: ~np.isnan(x))
 
 
-# TODO: test
-@Appender(docstrings.fillna_docstring, join="\n", indents=1)
-def fillna(self, value):
-
-    if not isinstance(value, str) and np.isnan(self.initial_value):
-        initial_value = value
-    else:
-        initial_value = self.initial_value
-
+def _make_data_fillna_method(self, value):
+    # value is a string
+    initial_value = self.initial_value
     if self._data is None:
         data = None
     else:
         values = self._get_values().copy()
         if value in ("pad", "ffill") and np.isnan(values.iloc[0]):
             values.iloc[0] = self.initial_value
-        if isinstance(value, str):
-            values = values.fillna(method=value)
-        else:
-            values = values.fillna(value=value)
+        values = values.fillna(method=value)
         if value in ("backfill", "bfill") and np.isnan(self.initial_value):
             initial_value = values.iloc[0]
         data = pd.DataFrame({"value": values})
+    return initial_value, data
+
+
+def _make_data_fillna_scalar(self, value):
+    if np.isnan(self.initial_value):
+        initial_value = value
+    else:
+        initial_value = self.initial_value
+
+    if self._data is None:
+        data = None
+    else:
+        values = self._get_values().copy().fillna(value=value)
+        data = pd.DataFrame({"value": values})
+    return initial_value, data
+
+
+@requires_closed_match
+def _fillna_with_stairs(self, value):
+    # value is Stairs
+    return self.fillna(0) + value * self.isna()
+
+
+# TODO: test
+@Appender(docstrings.fillna_docstring, join="\n", indents=1)
+def fillna(self, value):
+
+    if isinstance(value, sc.Stairs):
+        return _fillna_with_stairs(self, value)
+
+    if isinstance(value, str):
+        initial_value, data = _make_data_fillna_method(self, value)
+    else:
+        initial_value, data = _make_data_fillna_scalar(self, value)
 
     new_instance = sc.Stairs._new(
         initial_value=initial_value, data=data, closed=self.closed


Test patch:
diff --git a/tests/test_masking/test_fillna.py b/tests/test_masking/test_fillna.py
new file mode 100644
index 0000000..82d3d43
--- /dev/null
+++ b/tests/test_masking/test_fillna.py
@@ -0,0 +1,110 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+import staircase as sc
+from staircase.core.ops.common import ClosedMismatchError
+
+
+def s1(initial_value=0):
+    int_seq1 = sc.Stairs(initial_value=initial_value)
+    int_seq1.layer(1, 10, 2)
+    int_seq1.layer(-4, 5, -1.75)
+    int_seq1.layer(3, 5, 2.5)
+    int_seq1.layer(6, 7, -2.5)
+    int_seq1.layer(7, 10, -2.5)
+    return int_seq1
+
+
+def s2(initial_value=0):
+    int_seq2 = sc.Stairs(initial_value=initial_value)
+    int_seq2.layer(1, 7, -2.5)
+    int_seq2.layer(8, 10, 5)
+    int_seq2.layer(2, 5, 4.5)
+    int_seq2.layer(2.5, 4, -2.5)
+    int_seq2.layer(-2, 1, -1.75)
+    return int_seq2
+
+
+@pytest.fixture
+def s1_fix():
+    return s1(0)
+
+
+@pytest.fixture
+def s2_fix():
+    return s2(0)
+
+
+def test_fillna(s1_fix):
+    s = s1_fix.clip(None, 10).mask((-2, 0)).mask((4, 7)).fillna(1)
+    pd.testing.assert_series_equal(
+        s.step_values,
+        pd.Series(
+            [-1.75, 1, -1.75, 0.25, 2.75, 1, -0.5, 1],
+            index=[-4, -2, 0, 1, 3, 4, 7, 10],
+        ),
+        check_names=False,
+        check_index_type=False,
+    )
+    assert s.initial_value == s1_fix.initial_value
+
+
+@pytest.mark.parametrize(
+    "method",
+    ["ffill", "pad"],
+)
+def test_fillna_ffill(s1_fix, method):
+    s = s1_fix.mask((None, -2)).mask((0, 2)).mask((8, None)).fillna(method)
+    pd.testing.assert_series_equal(
+        s.step_values,
+        pd.Series(
+            [-1.75, 0.25, 2.75, 2.0, -0.5],
+            index=[-2, 2, 3, 5, 6],
+        ),
+        check_names=False,
+        check_index_type=False,
+    )
+    assert np.isnan(s.initial_value)
+
+
+@pytest.mark.parametrize(
+    "method",
+    ["bfill", "backfill"],
+)
+def test_fillna_bfill(s1_fix, method):
+    s = s1_fix.mask((None, -2)).mask((0, 2)).mask((8, None)).fillna(method)
+    pd.testing.assert_series_equal(
+        s.step_values,
+        pd.Series(
+            [0.25, 2.75, 2.0, -0.5, np.nan],
+            index=[0, 3, 5, 6, 8],
+        ),
+        check_names=False,
+        check_index_type=False,
+    )
+    assert s.initial_value == -1.75
+
+
+def test_fillna_stairs(s1_fix, s2_fix):
+    s = (
+        s1_fix.mask((None, -2))
+        .mask((0, 2))
+        .mask((8, None))
+        .fillna(s2_fix.mask((11, None)).mask((0.5, 1.5)))
+    )
+    pd.testing.assert_series_equal(
+        s.step_values,
+        pd.Series(
+            [-1.75, np.nan, -2.5, 0.25, 2.75, 2.0, -0.5, 5.0, 0.0, np.nan],
+            index=[-2.0, 0.5, 1.5, 2.0, 3.0, 5.0, 6.0, 8.0, 10.0, 11.0],
+        ),
+        check_names=False,
+        check_index_type=False,
+    )
+    assert s.initial_value == 0
+
+
+def test_fillna_stairs_error(s1_fix):
+    with pytest.raises(ClosedMismatchError):
+        s1_fix.fillna(sc.Stairs(start=1, end=2, closed="right"))
diff --git a/tests/test_masking/test_mask.py b/tests/test_masking/test_mask.py
index 3c35fa0..1b6d587 100644
--- a/tests/test_masking/test_mask.py
+++ b/tests/test_masking/test_mask.py
@@ -313,17 +313,3 @@ def test_notna(s1_fix):
         check_dtype=False,
     )
     assert s.initial_value == 1
-
-
-def test_fillna(s1_fix):
-    s = s1_fix.clip(None, 10).mask((-2, 0)).mask((4, 7)).fillna(1)
-    pd.testing.assert_series_equal(
-        s.step_values,
-        pd.Series(
-            [-1.75, 1, -1.75, 0.25, 2.75, 1, -0.5, 1],
-            index=[-4, -2, 0, 1, 3, 4, 7, 10],
-        ),
-        check_names=False,
-        check_index_type=False,
-    )
-    assert s.initial_value == s1_fix.initial_value