Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_generate_dataframes(self, simple_linear_model, tmpdir):
"""Test TablesRecorder.generate_dataframes """
from pywr.parameters import ConstantScenarioParameter
model = simple_linear_model
scA = Scenario(model, name='A', size=4)
scB = Scenario(model, name='B', size=2)
otpt = model.nodes['Output']
inpt = model.nodes['Input']
inpt.max_flow = ConstantScenarioParameter(model, scA, [10, 20, 30, 40])
otpt.max_flow = ConstantScenarioParameter(model, scB, [20, 40])
otpt.cost = -2.0
h5file = tmpdir.join('output.h5')
TablesRecorder(model, h5file)
model.run()
dfs = {}
for node, df in TablesRecorder.generate_dataframes(h5file):
dfs[node] = df
def test_parameter_df_upsampling_multiple_columns(model):
""" Test that the `DataFrameParameter` works with multiple columns that map to a `Scenario`
"""
scA = Scenario(model, 'A', size=20)
scB = Scenario(model, 'B', size=2)
# scenario indices (not used for this test)
# Use a 7 day timestep for this test and run 2015
model.timestepper.delta = datetime.timedelta(7)
model.timestepper.start = pd.to_datetime('2015-01-01')
model.timestepper.end = pd.to_datetime('2015-12-31')
model.timestepper.setup()
# Daily time-step
index = pd.date_range('2015-01-01', periods=365, freq='D')
df = pd.DataFrame(np.random.rand(365, 20), index=index)
p = DataFrameParameter(model, df, scenario=scA)
p.setup()
A = df.resample('7D', axis=0).mean()
def test_annual_count_index_threshold_recorder(simple_storage_model, params):
"""
The test sets uses a simple reservoir model with different inputs that
trigger a control curve failure after different numbers of years.
"""
from pywr.parameters import ConstantScenarioParameter, ConstantParameter
from pywr.parameters.control_curves import ControlCurveIndexParameter
model = simple_storage_model
scenario = Scenario(model, 'A', size=2)
# Simulate 5 years
model.timestepper.start = '2015-01-01'
model.timestepper.end = '2019-12-31'
# Control curve parameter
param = ControlCurveIndexParameter(model, model.nodes['Storage'], ConstantParameter(model, 0.25))
# Storage model has a capacity of 20, but starts at 10 Ml
# Demand is roughly 2 Ml/d per year
# First ensemble balances the demand
# Second ensemble should fail during 3rd year
demand = 2.0 / 365
model.nodes['Input'].max_flow = ConstantScenarioParameter(model, scenario, [demand, 0])
model.nodes['Output'].max_flow = demand
# Create the recorder with a threshold of 1
rec = AnnualCountIndexThresholdRecorder(model, [param] * params, 'TestRec', 1)
def test_dirty_scenario(simple_linear_model):
"""Adding a scenario to a model makes it dirty"""
model = simple_linear_model
model.setup()
assert(not model.dirty)
scenario = Scenario(model, "test", size=42)
assert(model.dirty)
def test_user_scenarios(self, simple_linear_model, tmpdir):
"""
Test the TablesRecorder with user defined scenario subset
"""
from pywr.parameters import ConstantScenarioParameter
model = simple_linear_model
scA = Scenario(model, name='A', size=4)
scB = Scenario(model, name='B', size=2)
# Use first and last combinations
model.scenarios.user_combinations = [[0, 0], [3, 1]]
otpt = model.nodes['Output']
inpt = model.nodes['Input']
inpt.max_flow = ConstantScenarioParameter(model, scA, [10, 20, 30, 40])
otpt.max_flow = ConstantScenarioParameter(model, scB, [20, 40])
otpt.cost = -2.0
h5file = tmpdir.join('output.h5')
import tables
with tables.open_file(str(h5file), 'w') as h5f:
rec = TablesRecorder(model, h5f)
def test_agg(self, simple_linear_model, agg_func):
model = simple_linear_model
model.timestepper.delta = 15
scenarioA = Scenario(model, "Scenario A", size=2)
scenarioB = Scenario(model, "Scenario B", size=5)
values = np.arange(366, dtype=np.float64)
p1 = DailyProfileParameter(model, values)
p2 = ConstantScenarioParameter(model, scenarioB, np.arange(scenarioB.size, dtype=np.float64))
p = AggregatedParameter(model, [p1, p2], agg_func=agg_func)
func = TestAggregatedParameter.funcs[agg_func]
@assert_rec(model, p)
def expected_func(timestep, scenario_index):
x = p1.get_value(scenario_index)
y = p2.get_value(scenario_index)
return func(np.array([x,y]))
model.run()
def test_parameter_mean_recorder(daily_profile_model):
model = daily_profile_model
node = model.nodes["Input"]
scenario = Scenario(model, "dummy", size=3)
timesteps = 3
rec_mean = RollingWindowParameterRecorder(model, node.max_flow, timesteps,
temporal_agg_func="mean", name="rec_mean")
rec_sum = RollingWindowParameterRecorder(model, node.max_flow, timesteps,
temporal_agg_func="sum", name="rec_sum")
rec_min = RollingWindowParameterRecorder(model, node.max_flow, timesteps,
temporal_agg_func="min", name="rec_min")
rec_max = RollingWindowParameterRecorder(model, node.max_flow, timesteps,
temporal_agg_func="max", name="rec_max")
model.run()
assert_allclose(rec_mean.data[[0, 1, 2, 3, 364], 0], [0, 0.5, 1, 2, 363])
assert_allclose(rec_max.data[[0, 1, 2, 3, 364], 0], [0, 1, 2, 3, 364])
assert_allclose(rec_min.data[[0, 1, 2, 3, 364], 0], [0, 0, 0, 1, 362])
def test_parameter_array_indexed_scenario_monthly_factors_json(model):
model.path = os.path.join(TEST_DIR, "models")
scA = Scenario(model, 'Scenario A', size=2)
scB = Scenario(model, 'Scenario B', size=3)
p1 = ArrayIndexedScenarioMonthlyFactorsParameter.load(model, {
"scenario": "Scenario A",
"values": list(range(32)),
"factors": [list(range(1, 13)),list(range(13, 25))],
})
p2 = ArrayIndexedScenarioMonthlyFactorsParameter.load(model, {
"scenario": "Scenario B",
"values": {
"url": "timeseries1.csv",
"index_col": "Timestamp",
"column": "Data",
},
"factors": {
def test_scenario(simple_linear_model, ):
"""Basic test of Scenario functionality"""
model = simple_linear_model # Convenience renaming
scenario = Scenario(model, 'Inflow', size=2)
model.nodes["Input"].max_flow = ConstantScenarioParameter(model, scenario, [5.0, 10.0])
model.nodes["Output"].max_flow = 5.0
model.nodes["Output"].cost = -2.0
expected_node_results = {
"Input": [5.0, 5.0],
"Link": [5.0, 5.0],
"Output": [5.0, 5.0],
}
assert_model(model, expected_node_results)
def test_scenario_slices(simple_linear_model):
"""Test slicing of scenarios"""
model = simple_linear_model
# create two scenarios
s1 = Scenario(model=model, name="A", size=20)
s2 = Scenario(model=model, name="B", size=3)
combinations = model.scenarios.get_combinations()
assert(len(combinations) == 20 * 3)
s1.slice = slice(0, None, 2)
combinations = model.scenarios.get_combinations()
assert(len(combinations) == 10 * 3)
# check multiindex respects scenario slices
index = model.scenarios.multiindex
assert(len(index.levels) == 2)
assert(len(index.levels[0]) == 10)
assert(len(index.levels[1]) == 3)
assert(len(index.labels) == 2)
assert(len(index.labels[0]) == 10 * 3)