def test_varargs():
def sumall(a, *args, b=0, **kwargs):
return a + sum(args) + b + sum(kwargs.values())
op = operation(
sumall,
name="t",
needs=[
"a",
vararg("arg1"),
vararg("arg2"),
varargs("args"),
optional("b"),
optional("c"),
],
provides="sum",
)
exp = sum(range(8))
assert op.compute(dict(a=1, arg1=2, arg2=3, args=[4, 5], b=6,
c=7))["sum"] == exp
assert op.compute(dict(a=1, arg1=2, arg2=3, args=[4, 5],
c=7))["sum"] == exp - 6
assert op.compute(dict(a=1, arg1=2, arg2=3, args=[4, 5],
b=6))["sum"] == exp - 7
assert op.compute(dict(a=1, arg2=3, args=[4, 5], b=6,
c=7))["sum"] == exp - 2
assert op.compute(dict(a=1, arg1=2, arg2=3, b=6,
c=7))["sum"] == exp - 4 - 5
with pytest.raises(ValueError, match="Missing compulsory needs.+'a'"):
assert op.compute(dict(arg1=2, arg2=3, b=6, c=7))
def test_narrow_and_optionality(reverse):
op1 = operation(name="op1",
needs=[optional("a"), optional("bb")],
provides="sum1")(addall)
op2 = operation(name="op2", needs=["a", optional("bb")],
provides="sum2")(addall)
ops = [op1, op2]
provides = "'sum1', 'sum2'"
if reverse:
ops = list(reversed(ops))
provides = "'sum2', 'sum1'"
pipeline_str = f"Pipeline('t', needs=['a', 'bb'(?)], provides=[{provides}], x2 ops"
pipeline = compose("t", *ops)
assert repr(pipeline).startswith(pipeline_str)
## IO & predicate do not affect network, but solution.
## Compose with `inputs`
#
pipeline = compose("t", *ops)
assert repr(pipeline).startswith(pipeline_str)
assert repr(pipeline.compile("a")).startswith(
f"ExecutionPlan(needs=['a'], provides=[{provides}], x2 steps:")
#
pipeline = compose("t", *ops)
assert repr(pipeline).startswith(pipeline_str)
assert repr(pipeline.compile(["bb"])).startswith(
"ExecutionPlan(needs=['bb'(?)], provides=['sum1'], x1 steps:")
## Narrow by `provides`
#
pipeline = compose("t", *ops, outputs="sum1")
assert repr(pipeline).startswith(pipeline_str)
assert repr(pipeline.compile("bb")).startswith(
"ExecutionPlan(needs=['bb'(?)], provides=['sum1'], x3 steps:")
assert repr(pipeline.compile("bb")) == repr(
pipeline.compute({
"bb": 1
}).plan)
pipeline = compose("t", *ops, outputs=["sum2"])
assert repr(pipeline).startswith(pipeline_str)
assert not pipeline.compile("bb").steps
assert len(pipeline.compile("a").steps) == 3
assert repr(pipeline.compile("a")).startswith(
"ExecutionPlan(needs=['a'], provides=['sum2'], x3 steps:")
## Narrow by BOTH
#
pipeline = compose("t", *ops, outputs=["sum1"])
assert repr(pipeline.compile("a")).startswith(
"ExecutionPlan(needs=['a'(?)], provides=['sum1'], x3 steps:")
pipeline = compose("t", *ops, outputs=["sum2"])
with pytest.raises(ValueError, match="Unsolvable graph:"):
pipeline.compute({"bb": 11})
def test_node_predicate_based_prune():
pipeline = compose(
"N",
operation(name="A",
needs=["a"],
provides=["aa"],
node_props={"color": "red"})(identity),
operation(name="B",
needs=["b"],
provides=["bb"],
node_props={"color": "green"})(identity),
operation(name="C", needs=["c"], provides=["cc"])(identity),
operation(
name="SUM",
needs=[optional(i) for i in ("aa", "bb", "cc")],
provides=["sum"],
)(addall),
)
inp = {"a": 1, "b": 2, "c": 3}
assert pipeline(**inp)["sum"] == 6
assert len(pipeline.net.graph.nodes) == 11
pred = lambda n, d: d.get("color", None) != "red"
assert pipeline.withset(predicate=pred)(**inp)["sum"] == 5
assert len(pipeline.withset(predicate=pred).compile().dag.nodes) == 9
pred = lambda n, d: "color" not in d
assert pipeline.withset(predicate=pred)(**inp)["sum"] == 3
assert len(pipeline.withset(predicate=pred).compile().dag.nodes) == 7
def param_to_modifier(name: str, param: inspect.Parameter) -> str:
return (
optional(name)
# is optional?
if param.default is not inspect._empty # type: ignore
else keyword(name)
if param.kind == Parameter.KEYWORD_ONLY
else name
)
def test_parallel_execution(exemethod):
if not exemethod:
return
delay = 0.5
def fn(x):
sleep(delay)
print("fn %s" % (time() - t0))
return 1 + x
def fn2(a, b):
sleep(delay)
print("fn2 %s" % (time() - t0))
return a + b
def fn3(z, k=1):
sleep(delay)
print("fn3 %s" % (time() - t0))
return z + k
pipeline = compose(
"l",
# the following should execute in parallel under threaded execution mode
operation(name="a", needs="x", provides="ao")(fn),
operation(name="b", needs="x", provides="bo")(fn),
# this should execute after a and b have finished
operation(name="c", needs=["ao", "bo"], provides="co")(fn2),
operation(name="d", needs=["ao", optional("k")], provides="do")(fn3),
operation(name="e", needs=["ao", "bo"], provides="eo")(fn2),
operation(name="f", needs="eo", provides="fo")(fn),
operation(name="g", needs="fo", provides="go")(fn),
nest=False,
)
t0 = time()
result_threaded = pipeline.withset(parallel=True).compute(
{"x": 10}, ["co", "go", "do"]
)
# print("threaded result")
# print(result_threaded)
t0 = time()
pipeline = pipeline.withset(parallel=False)
result_sequential = pipeline.compute({"x": 10}, ["co", "go", "do"])
# print("sequential result")
# print(result_sequential)
# make sure results are the same using either method
assert result_sequential == result_threaded
def test_evict_instructions_vary_with_inputs():
# Check #21: eviction-steps positions vary when inputs change.
def count_evictions(steps):
return sum(isinstance(n, str) for n in steps)
pipeline = compose(
"pipeline",
operation(name="a free without b", needs=["a"],
provides=["aa"])(identity),
operation(name="satisfiable", needs=["a", "b"], provides=["ab"])(add),
operation(name="optional ab",
needs=["aa", optional("ab")],
provides=["asked"])(lambda a, ab=10: a + ab),
)
inp = {"a": 2, "b": 3}
exp = inp.copy()
exp.update({"aa": 2, "ab": 5, "asked": 7})
res = pipeline(**inp)
assert res == exp # ok
steps11 = pipeline.net.compile(inp).steps
res = pipeline.compute(inp, ["asked"])
assert res == filtdict(exp, "asked") # ok
steps12 = pipeline.net.compile(inp, ["asked"]).steps
inp = {"a": 2}
exp = inp.copy()
exp.update({"aa": 2, "asked": 12})
res = pipeline(**inp)
assert res == exp # ok
steps21 = pipeline.net.compile(inp).steps
res = pipeline.compute(inp, ["asked"])
assert res == filtdict(exp, "asked") # ok
steps22 = pipeline.net.compile(inp, ["asked"]).steps
# When no outs, no evict-instructions.
assert steps11 != steps12
assert count_evictions(steps11) == 0
assert steps21 != steps22
assert count_evictions(steps21) == 0
# Check steps vary with inputs
#
# FAILs in v1.2.4 + #18, PASS in #26
assert steps11 != steps21
# Check evicts vary with inputs
#
# FAILs in v1.2.4 + #18, PASS in #26
assert count_evictions(steps12) != count_evictions(steps22)
def test_cwd_fnop():
op = operation(
str,
None,
needs=[
"a",
"a/b",
"/r/b",
optional("o"),
keyword("k"),
implicit("i"),
vararg("v1"),
varargs("v2"),
sfx("s1"),
sfxed("s2", "s22"),
vcat("vc"),
],
provides=["A/B", "C", "/R"],
aliases=[("A/B", "aa"), ("C", "CC"), ("/R", "RR")],
cwd="root",
)
exp = """
FnOp(name='str',
needs=['root/a'($),
'root/a/b'($),
'/r/b'($),
'root/o'($?'o'),
'root/k'($>'k'),
'root/i'($),
'root/v1'($*),
'root/v2'($+),
sfx('s1'),
sfxed('root/s2'($),
's22'),
'root/vc'($)],
provides=['root/A/B'($),
'root/C'($),
'/R'($),
'root/aa'($),
'root/CC'($),
'root/RR'($)],
aliases=[('root/A/B'($), 'root/aa'($)),
('root/C'($), 'root/CC'($)),
('/R'($), 'root/RR'($))],
fn='str')
"""
assert oneliner(op) == oneliner(exp)
def test_pipe_rename():
pipe = compose(
"t",
operation(str, name="op1", needs=sfx("a")),
operation(
str,
name="op2",
needs=sfx("a"),
provides=["a", sfx("b")],
aliases=[("a", "b")],
),
)
def renamer(na):
assert na.op
assert not na.parent
return dep_renamed(na.name, lambda name: f"PP.{name}")
ren = pipe.withset(renamer=renamer)
got = str(ren)
assert got == ("""
Pipeline('t', needs=[sfx('PP.a')], provides=['PP.a', sfx('PP.b'), 'PP.b'], x2 ops: PP.op1, PP.op2)
""".strip())
got = str(ren.ops)
assert got == oneliner("""
[FnOp(name='PP.op1', needs=[sfx('PP.a')], fn='str'),
FnOp(name='PP.op2', needs=[sfx('PP.a')], provides=['PP.a', sfx('PP.b'), 'PP.b'],
aliases=[('PP.a', 'PP.b')], fn='str')]
""")
## Check dictionary with callables
#
ren = pipe.withset(renamer={
"op1": lambda n: "OP1",
"op2": False,
"a": optional("a"),
"b": "B",
})
got = str(ren.ops)
assert got == oneliner("""
[FnOp(name='OP1', needs=[sfx('a')], fn='str'),
FnOp(name='op2', needs=[sfx('a')], provides=['a'(?), sfx('b'), 'B'],
aliases=[('a'(?), 'B')], fn='str')]
""")
def test_evict_optional():
# Test that evictions included for optionals do not raise
# exceptions when the corresponding input is not provided.
# Function to add two values plus an optional third value.
def addplusplus(a, b, c=0):
return a + b + c
# Here, an eviction-dependency will be inserted for the optional need 'c'.
sum_op1 = operation(name="sum_op1",
needs=["a", "b", optional("c")],
provides="sum1")(addplusplus)
sum_op2 = operation(name="sum_op2",
needs=["sum1", "sum1"],
provides="sum2")(add)
net = compose("test_net", sum_op1, sum_op2)
# Evictions happen only when a subset of outputs are requested.
results = net.compute({"a": 4, "b": 3}, ["sum2"])
assert "sum2" in results
def test_optional_per_function_with_same_output(exemethod):
# Test that the same need can be both optional and not on different operations.
#
## ATTENTION, the selected function is NOT the one with more inputs
# but the 1st satisfiable function added in the network.
add_op = operation(name="add", needs=["a", "b"], provides="a+-b")(add)
sub_op_optional = operation(name="sub_opt",
needs=["a", optional("b")],
provides="a+-b")(lambda a, b=10: a - b)
# Normal order
#
pipeline = compose("partial_optionals",
add_op,
sub_op_optional,
parallel=exemethod)
#
named_inputs = {"a": 1, "b": 2}
assert pipeline(**named_inputs) == {"a": 1, "a+-b": -1, "b": 2}
assert pipeline.compute(named_inputs, ["a+-b"]) == {"a+-b": -1}
#
named_inputs = {"a": 1}
assert pipeline.compute(named_inputs) == {"a": 1, "a+-b": -9}
assert pipeline.compute(named_inputs, ["a+-b"]) == {"a+-b": -9}
# Inverse op order
#
pipeline = compose("partial_optionals",
sub_op_optional,
add_op,
parallel=exemethod)
#
named_inputs = {"a": 1, "b": 2}
assert pipeline(**named_inputs) == {"a": 1, "a+-b": 3, "b": 2}
assert pipeline.compute(named_inputs, ["a+-b"]) == {"a+-b": 3}
#
named_inputs = {"a": 1}
assert pipeline(**named_inputs) == {"a": 1, "a+-b": -9}
assert pipeline.compute(named_inputs, ["a+-b"]) == {"a+-b": -9}
def test_rescheduling(exemethod, resched, rescheduled):
canc = operation(lambda: None, name="canc", needs=["b"], provides="cc")
op = compose(
"pipeline",
operation(lambda: [1], name="op1", provides=["a", "b"], rescheduled=1),
canc,
operation(
lambda C=1: C and NO_RESULT,
name="op2",
needs=optional("C"),
provides=["c"],
rescheduled=1,
),
operation(
lambda *args: sum(args),
name="op3",
needs=["a", vararg("b"), vararg("c")],
provides=["d"],
),
parallel=exemethod,
)
sol = op.compute({})
assert sol == {"a": 1, "d": 1}
assert list(sol.canceled) == [canc]
dot = str(sol.plot())
assert "#a9a9a9" in dot # Canceled
assert 'BORDER="4"' in dot # Rescheduled
assert "x2 partial-ops" in str(sol.check_if_incomplete())
with pytest.raises(IncompleteExecutionError, match="x2 partial-ops"):
assert sol.scream_if_incomplete()
## Check if modified state fails the 2nd time.
assert op.compute({}) == {"a": 1, "d": 1}
## Tell op to cancel just 1 of 2 provides
# (the 2n one, 'b').
#
sol = op.compute({"C": False})
assert sol == {"C": False, "a": 1, "c": False, "d": 1}
def test_conveyor_identity_fn():
op = operation(name="copy values", needs="a")()
assert not op.fn
op = operation(None, needs="a", provides="A")
assert not op.fn
op = operation(None, name="a", needs="a", provides="A")
assert op.fn
assert op(a=5) == {"A": 5}
op = operation(name="a", needs=["a", "b"], provides=["A", "B"])()
assert op.compute({"a": 5, "b": 6}) == {"A": 5, "B": 6}
op = operation(name="a", needs=["a", keyword("b")], provides=["A", "B"])()
assert op(a=55, b=66) == {"A": 55, "B": 66}
op = operation(
fn=None,
name="a",
needs=[optional("a"), vararg("b"), "c"],
# positional, vararg, keyword, optional
provides=["C", "B", "A"],
)
assert op(c=7, a=5, b=6) == {"A": 5, "B": 6, "C": 7}
def test_optional():
# Test that optional() needs work as expected.
# Function to add two values plus an optional third value.
def addplusplus(a, b, c=0):
return a + b + c
sum_op = operation(name="sum_op1",
needs=["a", "b", optional("c")],
provides="sum")(addplusplus)
net = compose("test_net", sum_op)
# Make sure output with optional arg is as expected.
named_inputs = {"a": 4, "b": 3, "c": 2}
results = net(**named_inputs)
assert "sum" in results
assert results["sum"] == sum(named_inputs.values())
# Make sure output without optional arg is as expected.
named_inputs = {"a": 4, "b": 3}
results = net(**named_inputs)
assert "sum" in results
assert results["sum"] == sum(named_inputs.values())
class GearMultiIndexer:
"""
Multi-indexer for 2-level df columns like ``(item, gear)`` with 1-based & closed-bracket `gear`.
Example 2-level *grid_wots* columns::
p_avail p_avail ... n_foo n_foo
g1 g2 ... g1 g2
df.columns = gidx[:]
... Warning::
negative indices might not work as expected if :attr:`gnames` do not start from ``g1``
(e.g. when constructed with :meth:`from_df()` static method)
**Examples:**
- Without `items` you get simple gear-indexing:
>>> G = GearMultiIndexer.from_ngears(5)
>>> G.gnames
1 g1
2 g2
3 g3
4 g4
5 g5
dtype: object
>>> G[1:3]
['g1', 'g2', 'g3']
>>> G[::-1]
['g5', 'g4', 'g3', 'g2', 'g1']
>>> G[3:2:-1]
['g3', 'g2']
>>> G[3:]
['g3', 'g4', 'g5']
>>> G[3:-1]
['g3', 'g4', 'g5']
>>> G[-1:-2:-1]
['g5', 'g4']
>>> G[[1, 3, 2]]
['g1', 'g3', 'g2']
>>> G[-1]
'g5'
- When `items` are given, you get a "product" MultiIndex:
>>> G = G.with_item("foo")
>>> G[1:3]
MultiIndex([('foo', 'g1'),
('foo', 'g2'),
('foo', 'g3')],
names=['item', 'gear'])
>>> len(G)
5
>>> G.shape
(5, 2)
>>> G.size
10
Use no `items` to reset them:
>>> G = G.with_item()
>>> G[:]
['g1', 'g2', 'g3', 'g4', 'g5']
>>> G.shape
(5,)
- Notice that **G0** changes "negative" indices:
>>> G[[-5, -6, -7]]
['g1', 'g5', 'g4']
>>> G = GearMultiIndexer.from_ngears(5, gear0=True)
>>> G[:]
['g0', 'g1', 'g2', 'g3', 'g4', 'g5']
>>> G[[-5, -6, -7]]
['g1', 'g0', 'g5']
"""
#: 1st level column(s)
items: Optional[Sequence[str]]
#: 2-level columns; use a gear_namer like :func:`gear_names()` (default)
#:
#: to make a :class:`pd.Series` like::
#:
#: {1: 'g1', 2: 'g2', ...}
gnames: pd.Series
#: Setting it to a gear not in :attr:`gnames`, indexing with negatives
#: may not always work.
top_gear: int
#: a function returns the string representation of a gear, like ``1 --> 'g1'``
gear_namer: GearGenerator
level_names: Sequence[str] = dataclasses.field(default=("item", "gear"))
@classmethod
def from_ngears(
cls,
ngears: int,
items: Sequence[str] = None,
gear_namer: GearGenerator = gear_name,
gear0=False,
):
return GearMultiIndexer(
items,
pd.Series(
{i: gear_namer(i)
for i in range(int(not gear0), ngears + 1)}),
ngears,
gear_namer,
)
@classmethod
def from_gids(
cls,
gids: Iterable[int],
items: Sequence[str] = None,
gear_namer: GearGenerator = gear_name,
):
gids = sorted(gids)
gnames = pd.Series({i: gear_namer(i) for i in gids})
return GearMultiIndexer(items, gnames, gids[-1], gear_namer)
@classmethod
@autog.autographed(
name="make_gwots_multi_indexer",
needs=["gwots", optional("gear_namer")],
provides="gidx",
)
@autog.autographed(
name="make_cycle_multi_indexer",
needs=["cycle/OK_gear", optional("gear_namer")],
provides="gidx2",
)
def from_df(cls,
df,
items: Sequence[str] = None,
gear_namer: GearGenerator = gear_name):
"""
Derive gears from the deepest level columns, sorted, and the last one becomes `ngear`
:param df:
the regular or multilevel-level df, not stored, just to get gear-names.
... Warning::
Negative indices might not work as expected if :attr:`gnames`
does not start from ``g1``.
"""
gears = df.columns
if hasattr(gears, "levels"):
gears = gears.levels[-1]
gears = [g for g in gears if g]
gids = [int(i) for i in re.sub("[^0-9 ]", "", " ".join(gears)).split()]
gnames = pd.Series(gears, index=gids).sort_index()
return cls(items, gnames, gids[-1], gear_namer)
def with_item(self, *items: str):
"""
Makes a gear-indexer producing tuple of (items x gears).
Example:
>>> GearMultiIndexer.from_ngears(2, gear0=True).with_item("foo", "bar")[:]
MultiIndex([('foo', 'g0'),
('foo', 'g1'),
('foo', 'g2'),
('bar', 'g0'),
('bar', 'g1'),
('bar', 'g2')], )
"""
return type(self)(items or None, self.gnames, self.top_gear,
self.gear_namer) # type: ignore
def _level_names(self, items=None) -> Optional[Sequence[str]]:
if items is None:
items = self.items
n_levels = 1 if not items else 1 + len(items)
return (self.level_names[:n_levels]
if n_levels <= len(self.level_names) else None)
def __getitem__(self, key):
"""
1-based & closed-bracket indexing, like Series but with `-1` for the top-gear.
"""
top_gear = self.ng
# Support partial gears or G0!
offset = int(top_gear == self.top_gear)
def from_top_gear(i):
return offset + (i % top_gear) if isinstance(i,
int) and i < 0 else i
if isinstance(key, slice):
key = slice(from_top_gear(key.start), from_top_gear(key.stop),
key.step)
elif isinstance(key, int):
key = from_top_gear(key)
else: # assume Iterable[int]
key = [from_top_gear(g) for g in key]
gnames = self.gnames.loc[key]
## If no items, return just a list of gears.
#
if self.items is None:
if isinstance(gnames, pd.Series):
gnames = list(gnames)
return gnames
## Otherwise, return a product multi-index.
#
if not isinstance(gnames, pd.Series):
gnames = (gnames, )
return pd.MultiIndex.from_tuples(itt.product(self.items, gnames),
names=self._level_names())
def colidx_pairs(self,
item: Union[str, Iterable[str]],
gnames: Iterable[str] = None):
if gnames is None:
gnames = self.gnames
assert gnames, locals()
if isinstance(item, str):
item = (item, )
return pd.MultiIndex.from_tuples(itt.product(item, gnames),
names=self._level_names(item))
def __len__(self):
"""
The number of gears extracted from 2-level dataframe.
It equals :attr:`top_gear` if :attr:`gnames` are from 1-->top_gear.
"""
return len(self.gnames)
ng = property(__len__)
@property
def shape(self):
y = (1 + len(self.items), ) if self.items else ()
return (self.ng, *y)
@property
def size(self):
return self.ng * ((len(self.items) + 1) if self.items else 1)
def fill_insufficient_power(cycle):
c = wio.pstep_factory.get().cycle
idx_miss_gear = cycle[c.g_max0] < 0
ok_n = cycle.loc[:, c.ok_n]
p_remain = cycle.loc[:, c.p_remain]
cycle.loc[idx_miss_gear]
#################
## Graphtik code
## ...
#################
@autog.autographed(needs=optional("cycle"))
def get_forced_cycle(cycle: pd.DataFrame = None) -> Optional[pd.DataFrame]:
"""Extract any forced `cycle` from model. """
c = wio.pstep_factory.get().cycle
if cycle is None or getattr(cycle, "size", 0) == 0:
cycle = None
else:
if not isinstance(cycle, pd.DataFrame):
cycle = pd.DataFrame(cycle)
log.info("Found forced `cycle-run` table(%ix%i).", cycle.shape[0],
cycle.shape[1])
cycle.index.name = c.t
cycle.reset_index() ## Ensure Time-steps start from 0 (not 1!).
