diff --git a/stormpy/src/core/model.cpp b/stormpy/src/core/model.cpp index 79f9dca4e..25b78a634 100644 --- a/stormpy/src/core/model.cpp +++ b/stormpy/src/core/model.cpp @@ -6,6 +6,16 @@ std::shared_ptr buildModel(storm::prism::Program const return storm::buildSymbolicModel(program, std::vector>(1,formula)).model; } +// Thin wrapper for getting initial states +template +std::vector getInitialStates(storm::models::sparse::Model const& model) { + std::vector initialStates; + for (auto entry : model.getInitialStates()) { + initialStates.push_back(entry); + } + return initialStates; +} + // Define python bindings void define_model(py::module& m) { @@ -39,6 +49,11 @@ void define_model(py::module& m) { // Models py::class_, std::shared_ptr>>(m, "SparseModel", "A probabilistic model where transitions are represented by doubles and saved in a sparse matrix", py::base()) + .def("labels", [](storm::models::sparse::Model const& model) { + return model.getStateLabeling().getLabels(); + }, "Get labels") + .def("labels_state", &storm::models::sparse::Model::getLabelsOfState, "Get labels") + .def("initial_states", &getInitialStates, "Get initial states") ; py::class_, std::shared_ptr>>(m, "SparseDtmc", "DTMC in sparse representation", py::base>()) ; @@ -46,6 +61,11 @@ void define_model(py::module& m) { ; py::class_, std::shared_ptr>>(m, "SparseParametricModel", "A probabilistic model where transitions are represented by rational functions and saved in a sparse matrix", py::base()) .def("collect_probability_parameters", &storm::models::sparse::getProbabilityParameters, "Collect parameters") + .def("labels", [](storm::models::sparse::Model const& model) { + return model.getStateLabeling().getLabels(); + }, "Get labels") + .def("labels_state", &storm::models::sparse::Model::getLabelsOfState, "Get labels") + .def("initial_states", &getInitialStates, "Get initial states") ; py::class_, std::shared_ptr>>(m, "SparseParametricDtmc", "pDTMC in sparse representation", py::base>()) ; diff --git a/stormpy/tests/core/test_model.py b/stormpy/tests/core/test_model.py index cd91abccc..8958080b7 100644 --- a/stormpy/tests/core/test_model.py +++ b/stormpy/tests/core/test_model.py @@ -62,3 +62,37 @@ class TestModel: assert model.supports_parameters() assert not model.has_parameters() assert type(model) is stormpy.SparseParametricDtmc + + def test_label(self): + program = stormpy.parse_program("../examples/dtmc/die/die.pm") + formulas = stormpy.parse_formulas_for_program("P=? [ F \"one\" ]", program) + model = stormpy.build_model(program, formulas[0]) + labels = model.labels() + assert len(labels) == 2 + assert "init" in labels + assert "one" in labels + assert "init" in model.labels_state(0) + assert "one" in model.labels_state(7) + + def test_initial_states(self): + program = stormpy.parse_program("../examples/dtmc/die/die.pm") + formulas = stormpy.parse_formulas_for_program("P=? [ F \"one\" ]", program) + model = stormpy.build_model(program, formulas[0]) + initial_states = model.initial_states() + assert len(initial_states) == 1 + assert 0 in initial_states + + def test_label_parametric(self): + program = stormpy.parse_program("../examples/pdtmc/brp/brp_16_2.pm") + formulas = stormpy.parse_formulas_for_program("P=? [ F \"target\" ]", program) + model = stormpy.build_parametric_model(program, formulas[0]) + labels = model.labels() + assert len(labels) == 2 + assert "init" in labels + assert "target" in labels + assert "init" in model.labels_state(0) + assert "target" in model.labels_state(28) + assert "target" in model.labels_state(611) + initial_states = model.initial_states() + assert len(initial_states) == 1 + assert 0 in initial_states