def test_piecewise_linear():
from larch.util.data_expansion import piecewise_linear
func = piecewise_linear(X.DataName, P.ParamName, [3, 5, 7])
assert func[0] == P('ParamName ① up to 3') * X('piece(DataName,None,3)')
assert func[1] == P('ParamName ② 3 to 5') * X('piece(DataName,3,5)')
assert func[2] == P('ParamName ③ 5 to 7') * X('piece(DataName,5,7)')
assert func[3] == P('ParamName ④ over 7') * X('piece(DataName,7,None)')
assert len(func) == 4
func = piecewise_linear(X.DataName, breaks=[3, 5, 7])
assert func[0] == P('DataName ① up to 3') * X('piece(DataName,None,3)')
assert func[1] == P('DataName ② 3 to 5') * X('piece(DataName,3,5)')
assert func[2] == P('DataName ③ 5 to 7') * X('piece(DataName,5,7)')
assert func[3] == P('DataName ④ over 7') * X('piece(DataName,7,None)')
assert len(func) == 4
func = piecewise_linear('GenName', breaks=[3, 5, 7])
assert func[0] == P('GenName ① up to 3') * X('piece(GenName,None,3)')
assert func[1] == P('GenName ② 3 to 5') * X('piece(GenName,3,5)')
assert func[2] == P('GenName ③ 5 to 7') * X('piece(GenName,5,7)')
assert func[3] == P('GenName ④ over 7') * X('piece(GenName,7,None)')
assert len(func) == 4
with pytest.raises(ValueError):
func = piecewise_linear('GenName', [3, 5, 7])
def test_piecewise_expansion():
import pandas, io, numpy
from larch.util.data_expansion import piecewise_expansion, piecewise_linear
df = pandas.DataFrame(
numpy.linspace(0, 10, 25),
columns=['DataName'],
)
expanded = piecewise_expansion(df, [3, 5, 7])
s = '''
piece(DataName,None,3) piece(DataName,3,5) piece(DataName,5,7) piece(DataName,7,None)
0 0.000000 0.000000 0.000000 0.000000
1 0.416667 0.000000 0.000000 0.000000
2 0.833333 0.000000 0.000000 0.000000
3 1.250000 0.000000 0.000000 0.000000
4 1.666667 0.000000 0.000000 0.000000
5 2.083333 0.000000 0.000000 0.000000
6 2.500000 0.000000 0.000000 0.000000
7 2.916667 0.000000 0.000000 0.000000
8 3.000000 0.333333 0.000000 0.000000
9 3.000000 0.750000 0.000000 0.000000
10 3.000000 1.166667 0.000000 0.000000
11 3.000000 1.583333 0.000000 0.000000
12 3.000000 2.000000 0.000000 0.000000
13 3.000000 2.000000 0.416667 0.000000
14 3.000000 2.000000 0.833333 0.000000
15 3.000000 2.000000 1.250000 0.000000
16 3.000000 2.000000 1.666667 0.000000
17 3.000000 2.000000 2.000000 0.083333
18 3.000000 2.000000 2.000000 0.500000
19 3.000000 2.000000 2.000000 0.916667
20 3.000000 2.000000 2.000000 1.333333
21 3.000000 2.000000 2.000000 1.750000
22 3.000000 2.000000 2.000000 2.166667
23 3.000000 2.000000 2.000000 2.583333
24 3.000000 2.000000 2.000000 3.000000
'''
correct = pandas.read_csv(io.StringIO(s), sep='\s+')
pandas.testing.assert_frame_equal(expanded, correct)
func = piecewise_linear(X.DataName, P.ParamName, [3, 5, 7])
expanded2 = piecewise_expansion(df, func)
pandas.testing.assert_frame_equal(expanded2, correct)
