Merge pull request #4950 from hakonhagland/pybind5

Get primary variables and fluid state from Python

opm/simulators/flow/python/PyBlackOilSimulator.hpp

@@ -23,7 +23,9 @@
 #include <opm/simulators/flow/Main.hpp>
 #include <opm/simulators/flow/FlowMain.hpp>
 #include <opm/models/utils/propertysystem.hh>
+#include <opm/models/utils/parametersystem.hh>
 #include <opm/simulators/flow/python/Pybind11Exporter.hpp>
+#include <opm/simulators/flow/python/PyFluidState.hpp>
 #include <opm/simulators/flow/python/PyMaterialState.hpp>
 #include <opm/input/eclipse/EclipseState/EclipseState.hpp>
 #include <opm/input/eclipse/Schedule/Schedule.hpp>
@@ -46,18 +48,27 @@ public:
     void advance(int report_step);
     bool checkSimulationFinished();
     int currentStep();
+    py::array_t<double> getFluidStateVariable(const std::string &name) const;
     py::array_t<double> getCellVolumes();
     double getDT();
     py::array_t<double> getPorosity();
+    py::array_t<double> getPrimaryVariable(const std::string &variable) const;
+    py::array_t<int> getPrimaryVarMeaning(const std::string &variable) const;
+    std::map<std::string, int> getPrimaryVarMeaningMap(const std::string &variable) const;
     int run();
     void setPorosity(
          py::array_t<double, py::array::c_style | py::array::forcecast> array);
+    void setPrimaryVariable(
+        const std::string &idx_name,
+        py::array_t<double,
+        py::array::c_style | py::array::forcecast> array);
     int step();
     int stepCleanup();
     int stepInit();
 
 private:
     Opm::FlowMain<TypeTag>& getFlowMain() const;
+    PyFluidState<TypeTag>& getFluidState() const;
     PyMaterialState<TypeTag>& getMaterialState() const;
 
     const std::string deck_filename_;
@@ -71,6 +82,7 @@ private:
 
     std::unique_ptr<Opm::FlowMain<TypeTag>> main_ebos_;
     Simulator *ebos_simulator_;
+    std::unique_ptr<PyFluidState<TypeTag>> fluid_state_;
     std::unique_ptr<PyMaterialState<TypeTag>> material_state_;
     std::shared_ptr<Opm::Deck> deck_;
     std::shared_ptr<Opm::EclipseState> eclipse_state_;

opm/simulators/flow/python/PyFluidState.hpp

@@ -0,0 +1,71 @@
+/*
+  Copyright 2023 Equinor ASA.
+
+  This file is part of the Open Porous Media project (OPM).
+
+  OPM is free software: you can redistribute it and/or modify
+  it under the terms of the GNU General Public License as published by
+  the Free Software Foundation, either version 3 of the License, or
+  (at your option) any later version.
+
+  OPM is distributed in the hope that it will be useful,
+  but WITHOUT ANY WARRANTY; without even the implied warranty of
+  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+  GNU General Public License for more details.
+
+  You should have received a copy of the GNU General Public License
+  along with OPM.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+#ifndef OPM_PY_FLUID_STATE_HEADER_INCLUDED
+#define OPM_PY_FLUID_STATE_HEADER_INCLUDED
+
+#include <opm/models/utils/propertysystem.hh>
+
+#include <exception>
+#include <iostream>
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace Opm::Pybind
+{
+    template <class TypeTag>
+    class PyFluidState {
+        using Simulator = GetPropType<TypeTag, Opm::Properties::Simulator>;
+        using Problem = GetPropType<TypeTag, Opm::Properties::Problem>;
+        using Model = GetPropType<TypeTag, Opm::Properties::Model>;
+        using ElementContext = GetPropType<TypeTag, Opm::Properties::ElementContext>;
+        using FluidSystem = GetPropType<TypeTag, Opm::Properties::FluidSystem>;
+        using Indices = GetPropType<TypeTag, Opm::Properties::Indices>;
+        using GridView = GetPropType<TypeTag, Opm::Properties::GridView>;
+        using PrimaryVariables = GetPropType<TypeTag, Properties::PrimaryVariables>;
+
+        enum class VariableType {
+            // Primary variables: Sw, Sg, po, pg, Rs, Rv
+            Sw, Sg, So, pw, pg, po, Rs, Rv, rho_w, rho_g, rho_o, T
+        };
+    public:
+        PyFluidState(Simulator *ebos_simulator);
+        std::vector<double> getFluidStateVariable(const std::string &name) const;
+        std::vector<int> getPrimaryVarMeaning(const std::string &variable) const;
+        std::map<std::string, int> getPrimaryVarMeaningMap(const std::string &variable) const;
+        std::vector<double> getPrimaryVariable(const std::string &idx_name) const;
+        void setPrimaryVariable(const std::string &idx_name, const double *data, std::size_t size);
+
+    private:
+        std::size_t getPrimaryVarIndex_(const std::string &idx_name) const;
+        int getVariableMeaning_(PrimaryVariables &primary_vars, const std::string &variable) const;
+        VariableType getVariableType_(const std::string &name) const;
+        template <class FluidState> double getVariableValue_(
+            FluidState &fs, VariableType var_type, const std::string &name) const;
+        void variableNotFoundError_(const std::string &name) const;
+
+        Simulator *ebos_simulator_;
+    };
+}
+#include "PyFluidState_impl.hpp"
+
+#endif // OPM_PY_FLUID_STATE_HEADER_INCLUDED
+

opm/simulators/flow/python/PyFluidState_impl.hpp

@@ -0,0 +1,316 @@
+/*
+  Copyright 2023 Equinor ASA.
+
+  This file is part of the Open Porous Media project (OPM).
+
+  OPM is free software: you can redistribute it and/or modify
+  it under the terms of the GNU General Public License as published by
+  the Free Software Foundation, either version 3 of the License, or
+  (at your option) any later version.
+
+  OPM is distributed in the hope that it will be useful,
+  but WITHOUT ANY WARRANTY; without even the implied warranty of
+  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+  GNU General Public License for more details.
+
+  You should have received a copy of the GNU General Public License
+  along with OPM.  If not, see <http://www.gnu.org/licenses/>.
+*/
+#include <fmt/format.h>
+
+namespace Opm::Pybind {
+
+template <class TypeTag>
+PyFluidState<TypeTag>::
+PyFluidState(Simulator* ebos_simulator) : ebos_simulator_(ebos_simulator)
+{
+
+}
+
+// Public methods alphabetically sorted
+// ------------------------------------
+
+template <class TypeTag>
+std::vector<int>
+PyFluidState<TypeTag>::
+getPrimaryVarMeaning(const std::string &variable) const {
+    Model &model = this->ebos_simulator_->model();
+    auto &sol = model.solution(/*timeIdx*/0);
+    auto size = model.numGridDof();
+    std::vector<int> array(size);
+    for (unsigned dof_idx = 0; dof_idx < size; ++dof_idx) {
+        auto primary_vars = sol[dof_idx];
+        array[dof_idx] = getVariableMeaning_(primary_vars, variable);
+    }
+    return array;
+}
+
+template <class TypeTag>
+std::map<std::string, int>
+PyFluidState<TypeTag>::
+getPrimaryVarMeaningMap(const std::string &variable) const
+{
+    if (variable.compare("pressure") == 0) {
+        return {{ "Po", static_cast<int>(PrimaryVariables::PressureMeaning::Po) },
+                { "Pw", static_cast<int>(PrimaryVariables::PressureMeaning::Pw) },
+                { "Pg", static_cast<int>(PrimaryVariables::PressureMeaning::Pg) }};
+    }
+    else if (variable.compare("water") == 0) {
+        return {{ "Sw", static_cast<int>(PrimaryVariables::WaterMeaning::Sw) },
+                { "Rvw", static_cast<int>(PrimaryVariables::WaterMeaning::Rvw) },
+                { "Rsw", static_cast<int>(PrimaryVariables::WaterMeaning::Rsw) },
+                { "Disabled", static_cast<int>(PrimaryVariables::WaterMeaning::Disabled) }};
+    }
+    else if (variable.compare("gas") == 0) {
+        return {{ "Sg", static_cast<int>(PrimaryVariables::GasMeaning::Sg) },
+                { "Rs", static_cast<int>(PrimaryVariables::GasMeaning::Rs) },
+                { "Rv", static_cast<int>(PrimaryVariables::GasMeaning::Rv) },
+                { "Disabled", static_cast<int>(PrimaryVariables::GasMeaning::Disabled) }};
+    }
+    else if (variable.compare("brine") == 0) {
+        return {{ "Cs", static_cast<int>(PrimaryVariables::BrineMeaning::Cs) },
+                { "Sp", static_cast<int>(PrimaryVariables::BrineMeaning::Sp) },
+                { "Disabled", static_cast<int>(PrimaryVariables::BrineMeaning::Disabled) }};
+    }
+    else {
+        const std::string msg = fmt::format(
+            "Unknown variable meaning '{}': Expected pressure, water, gas, or brine", variable);
+        throw std::runtime_error(msg);
+    }
+}
+
+/* Meaning of the primary variables: Sw, Sg, po, pg, Rs, Rv
+ * 1. Sw_po_Sg -> threephase case
+ * 2. Sw_po_Rs -> water + oil case
+ * 3. Sw_pg_Rv -> water + gas case
+ */
+
+/* Variables:
+   Sw = Water saturation,
+   So = Oil saturation,
+   Sg = Gas saturation,
+   pw = Water pressure,
+   po = Oil pressure,
+   pg = Gas pressure,
+   Rs = The solution gas oil ratio: The amount of gas dissolved in the oil
+   Rv = The oil vaporization factor of the gas phase
+   invB = The inverse formation volume factor of a fluid phase
+   rho_w = Water density,
+   rho_o = Oil density,
+   rho_g = Gas density,
+   mu_w = Water viscosity,
+   mu_o = Oil viscosity,
+   mu_g = Gas viscosity,
+   kr_w = Water relperm,
+   kr_o = Oil relperm,
+   kr_g = Gas relperm,
+ */
+template <class TypeTag>
+std::vector<double>
+PyFluidState<TypeTag>::
+getFluidStateVariable(const std::string &name) const
+{
+    Model &model = this->ebos_simulator_->model();
+    auto size = model.numGridDof();
+    std::vector<double> array(size);
+    const auto& grid_view = this->ebos_simulator_->vanguard().gridView();
+    /* NOTE: grid_view.size(0) should give the same value as
+     *  model.numGridDof()
+     */
+    ElementContext elem_ctx(*this->ebos_simulator_);
+    auto var_type = getVariableType_(name);
+    for (const auto& elem : elements(grid_view, Dune::Partitions::interior)) {
+        elem_ctx.updatePrimaryStencil(elem);
+        elem_ctx.updatePrimaryIntensiveQuantities(/*timeIdx=*/0);
+        for (unsigned dof_idx = 0; dof_idx < elem_ctx.numPrimaryDof(/*timeIdx=*/0); ++dof_idx) {
+            const auto& int_quants = elem_ctx.intensiveQuantities(dof_idx, /*timeIdx=*/0);
+            const auto& fs = int_quants.fluidState();
+            unsigned global_dof_idx = elem_ctx.globalSpaceIndex(dof_idx, /*timeIdx=*/0);
+            array[global_dof_idx] = getVariableValue_(fs, var_type, name);
+        }
+    }
+    return array;
+}
+
+template <class TypeTag>
+std::vector<double>
+PyFluidState<TypeTag>::
+getPrimaryVariable(const std::string &idx_name) const
+{
+    std::size_t primary_var_idx = getPrimaryVarIndex_(idx_name);
+    Model &model = this->ebos_simulator_->model();
+    auto &sol = model.solution(/*timeIdx*/0);
+    auto size = model.numGridDof();
+    std::vector<double> array(size);
+    for (unsigned dof_idx = 0; dof_idx < size; ++dof_idx) {
+        auto primary_vars = sol[dof_idx];
+        array[dof_idx] = primary_vars[primary_var_idx];
+    }
+    return array;
+}
+
+template <class TypeTag>
+void
+PyFluidState<TypeTag>::
+setPrimaryVariable(const std::string &idx_name, const double *data, std::size_t size)
+{
+    std::size_t primary_var_idx = getPrimaryVarIndex_(idx_name);
+    Model &model = this->ebos_simulator_->model();
+    auto &sol = model.solution(/*timeIdx*/0);
+    auto model_size = model.numGridDof();
+    if (model_size != size) {
+        const std::string msg = fmt::format(
+            "Cannot set primary variable. Expected array of size {} but got array of size: {}",
+            model_size, size);
+        throw std::runtime_error(msg);
+    }
+    for (unsigned dof_idx = 0; dof_idx < size; ++dof_idx) {
+        auto &primary_vars = sol[dof_idx];
+        primary_vars[primary_var_idx] = data[dof_idx];
+    }
+}
+
+// Private methods alphabetically sorted
+// -------------------------------------
+
+template <class TypeTag>
+std::size_t
+PyFluidState<TypeTag>::
+getPrimaryVarIndex_(const std::string &idx_name) const
+{
+    if (idx_name.compare("pressure") == 0) {
+        return Indices::pressureSwitchIdx;
+    }
+    else if (idx_name.compare("water_saturation") == 0) {
+        return Indices::waterSwitchIdx;
+    }
+    else if (idx_name.compare("composition") == 0) {
+        return Indices::compositionSwitchIdx;
+    }
+    else {
+        const std::string msg = fmt::format("Unknown primary variable index name: {}", idx_name);
+        throw std::runtime_error(msg);
+    }
+}
+
+template <class TypeTag>
+int
+PyFluidState<TypeTag>::
+getVariableMeaning_(PrimaryVariables &primary_vars, const std::string &variable) const
+{
+    if (variable.compare("pressure") == 0) {
+        return static_cast<int>(primary_vars.primaryVarsMeaningPressure());
+    }
+    else if(variable.compare("water") == 0) {
+        return static_cast<int>(primary_vars.primaryVarsMeaningWater());
+    }
+    else if (variable.compare("gas") == 0) {
+        return static_cast<int>(primary_vars.primaryVarsMeaningGas());
+    }
+    else if (variable.compare("brine") == 0) {
+        return static_cast<int>(primary_vars.primaryVarsMeaningBrine());
+    }
+    else {
+        const std::string msg = fmt::format(
+            "Unknown variable meaning '{}': Expected pressure, water, gas, or brine", variable);
+        throw std::runtime_error(msg);
+    }
+}
+
+template <class TypeTag>
+typename PyFluidState<TypeTag>::VariableType
+PyFluidState<TypeTag>::
+getVariableType_(const std::string &name) const
+{
+    static std::map<std::string, VariableType> variable_type_map =
+       {
+           {"Sw", VariableType::Sw},
+           {"Sg", VariableType::Sg},
+           {"So", VariableType::So},
+           {"pw", VariableType::pw},
+           {"pg", VariableType::pg},
+           {"po", VariableType::po},
+           {"Rs", VariableType::Rs},
+           {"Rv", VariableType::Rv},
+           {"rho_w", VariableType::rho_w},
+           {"rho_g", VariableType::rho_g},
+           {"rho_o", VariableType::rho_o},
+           {"T", VariableType::T}
+       };
+
+    if (variable_type_map.count(name) == 0) {
+        variableNotFoundError_(name);
+    }
+    return variable_type_map.at(name);
+}
+
+template <class TypeTag>
+template <class FluidState>
+double
+PyFluidState<TypeTag>::
+getVariableValue_(FluidState &fs, VariableType var_type, const std::string &name) const
+{
+    double value;
+    switch(var_type) {
+    case VariableType::pw :
+        value = Opm::getValue(
+            fs.pressure(FluidSystem::waterPhaseIdx));
+        break;
+    case VariableType::pg :
+        value = Opm::getValue(
+            fs.pressure(FluidSystem::gasPhaseIdx));
+        break;
+    case VariableType::po :
+        value = Opm::getValue(
+            fs.pressure(FluidSystem::oilPhaseIdx));
+        break;
+    case VariableType::rho_w :
+        value = Opm::getValue(
+            fs.density(FluidSystem::waterPhaseIdx));
+        break;
+    case VariableType::rho_g :
+        value = Opm::getValue(
+            fs.density(FluidSystem::gasPhaseIdx));
+        break;
+    case VariableType::rho_o :
+        value = Opm::getValue(
+            fs.density(FluidSystem::oilPhaseIdx));
+        break;
+    case VariableType::Rs :
+        value = Opm::getValue(fs.Rs());
+        break;
+    case VariableType::Rv :
+        value = Opm::getValue(fs.Rv());
+        break;
+    case VariableType::Sw :
+        value = Opm::getValue(
+            fs.saturation(FluidSystem::waterPhaseIdx));
+        break;
+    case VariableType::Sg :
+        value = Opm::getValue(
+            fs.saturation(FluidSystem::gasPhaseIdx));
+        break;
+    case VariableType::So :
+        value = Opm::getValue(
+            fs.saturation(FluidSystem::oilPhaseIdx));
+        break;
+    case VariableType::T :
+        value = Opm::getValue(
+            fs.temperature(FluidSystem::waterPhaseIdx));
+        break;
+    default:
+        variableNotFoundError_(name);
+    }
+    return value;
+}
+
+template <class TypeTag>
+void
+PyFluidState<TypeTag>::
+variableNotFoundError_(const std::string &name) const
+{
+    const std::string msg = fmt::format("Access to variable '{}' is not implemented yet!", name);
+    throw std::runtime_error(msg);
+}
+
+} // namespace Opm::Pybind

opm/simulators/flow/python/PyMaterialState.hpp

@@ -45,8 +45,8 @@ namespace Opm::Pybind
         PyMaterialState(Simulator *ebos_simulator)
             : ebos_simulator_(ebos_simulator) { }
 
-        std::unique_ptr<double []> getCellVolumes( std::size_t *size);
-        std::unique_ptr<double []> getPorosity( std::size_t *size);
+        std::vector<double> getCellVolumes();
+        std::vector<double> getPorosity();
         void setPorosity(const double *poro, std::size_t size);
     private:
         Simulator *ebos_simulator_;

opm/simulators/flow/python/PyMaterialState_impl.hpp

@@ -22,29 +22,29 @@
 namespace Opm::Pybind {
 
 template <class TypeTag>
-std::unique_ptr<double []>
+std::vector<double>
 PyMaterialState<TypeTag>::
-getCellVolumes( std::size_t *size)
+getCellVolumes()
 {
     Model &model = this->ebos_simulator_->model();
-    *size = model.numGridDof();
-    auto array = std::make_unique<double []>(*size);
-    for (unsigned dof_idx = 0; dof_idx < *size; ++dof_idx) {
+    auto size = model.numGridDof();
+    std::vector<double> array(size);
+    for (unsigned dof_idx = 0; dof_idx < size; ++dof_idx) {
         array[dof_idx] = model.dofTotalVolume(dof_idx);
     }
     return array;
 }
 
 template <class TypeTag>
-std::unique_ptr<double []>
+std::vector<double>
 PyMaterialState<TypeTag>::
-getPorosity( std::size_t *size)
+getPorosity()
 {
     Problem &problem = this->ebos_simulator_->problem();
     Model &model = this->ebos_simulator_->model();
-    *size = model.numGridDof();
-    auto array = std::make_unique<double []>(*size);
-    for (unsigned dof_idx = 0; dof_idx < *size; ++dof_idx) {
+    auto size = model.numGridDof();
+    std::vector<double> array(size);
+    for (unsigned dof_idx = 0; dof_idx < size; ++dof_idx) {
         array[dof_idx] = problem.referencePorosity(dof_idx, /*timeIdx*/0);
     }
     return array;

opm/simulators/flow/python/Pybind11Exporter.hpp

@@ -3,6 +3,7 @@
 
 #include <pybind11/pybind11.h>
 #include <pybind11/numpy.h>
+#include <pybind11/stl.h>
 //#include <pybind11/embed.h>
 
 namespace py = pybind11;

python/simulators/CMakeLists.txt

@@ -41,21 +41,13 @@ if(OPM_ENABLE_PYTHON_TESTS)
   #   splitting the python tests into multiple add_test() tests instead
   #   of having a single "python -m unittest" test call that will run all
   #   the tests in the "test" sub directory.
-  add_test(NAME python_basic
+  foreach(case_name IN ITEMS basic fluidstate_variables primary_variables schedule throw)
+    add_test(NAME python_${case_name}
         WORKING_DIRECTORY ${PROJECT_BINARY_DIR}/python
         COMMAND ${CMAKE_COMMAND}
         -E env PYTHONPATH=${PYTHON_PATH} ${PYTHON_EXECUTABLE}
-      -m unittest test/test_basic.py)
-  add_test(NAME python_schedule
-      WORKING_DIRECTORY ${PROJECT_BINARY_DIR}/python
-      COMMAND ${CMAKE_COMMAND}
-      -E env PYTHONPATH=${PYTHON_PATH} ${PYTHON_EXECUTABLE}
-      -m unittest test/test_schedule.py)
-  add_test(NAME python_throw
-      WORKING_DIRECTORY ${PROJECT_BINARY_DIR}/python
-      COMMAND ${CMAKE_COMMAND}
-      -E env PYTHONPATH=${PYTHON_PATH} ${PYTHON_EXECUTABLE}
-      -m unittest test/test_throw.py)
+        -m unittest test/test_${case_name}.py)
+  endforeach()
 endif()
 
 find_file(PYTHON_INSTALL_PY install.py

python/simulators/PyBlackOilSimulator.cpp

@@ -95,9 +95,8 @@ int PyBlackOilSimulator::currentStep()
 }
 
 py::array_t<double> PyBlackOilSimulator::getCellVolumes() {
-    std::size_t len;
-    auto array = getMaterialState().getCellVolumes(&len);
-    return py::array(len, array.get());
+    auto vector = getMaterialState().getCellVolumes();
+    return py::array(vector.size(), vector.data());
 }
 
 double PyBlackOilSimulator::getDT() {
@@ -106,9 +105,40 @@ double PyBlackOilSimulator::getDT() {
 
 py::array_t<double> PyBlackOilSimulator::getPorosity()
 {
-    std::size_t len;
-    auto array = getMaterialState().getPorosity(&len);
-    return py::array(len, array.get());
+    auto vector = getMaterialState().getPorosity();
+    return py::array(vector.size(), vector.data());
+}
+
+py::array_t<double>
+PyBlackOilSimulator::
+getFluidStateVariable(const std::string &name) const
+{
+    auto vector = getFluidState().getFluidStateVariable(name);
+    return py::array(vector.size(), vector.data());
+}
+
+py::array_t<double>
+PyBlackOilSimulator::
+getPrimaryVariable(const std::string &variable) const
+{
+    auto vector = getFluidState().getPrimaryVariable(variable);
+    return py::array(vector.size(), vector.data());
+}
+
+py::array_t<int>
+PyBlackOilSimulator::
+getPrimaryVarMeaning(const std::string &variable) const
+{
+    auto vector = getFluidState().getPrimaryVarMeaning(variable);
+    return py::array(vector.size(), vector.data());
+}
+
+std::map<std::string, int>
+PyBlackOilSimulator::
+getPrimaryVarMeaningMap(const std::string &variable) const
+{
+
+    return getFluidState().getPrimaryVarMeaningMap(variable);
 }
 
 int PyBlackOilSimulator::run()
@@ -125,6 +155,19 @@ void PyBlackOilSimulator::setPorosity( py::array_t<double,
     getMaterialState().setPorosity(poro, size_);
 }
 
+void
+PyBlackOilSimulator::
+setPrimaryVariable(
+    const std::string &idx_name,
+    py::array_t<double,
+    py::array::c_style | py::array::forcecast> array
+)
+{
+    std::size_t size_ = array.size();
+    const double *data = array.data();
+    getFluidState().setPrimaryVariable(idx_name, data, size_);
+}
+
 int PyBlackOilSimulator::step()
 {
     if (!this->has_run_init_) {
@@ -177,6 +220,7 @@ int PyBlackOilSimulator::stepInit()
         int result = this->main_ebos_->executeInitStep();
         this->has_run_init_ = true;
         this->ebos_simulator_ = this->main_ebos_->getSimulatorPtr();
+        this->fluid_state_ = std::make_unique<PyFluidState<TypeTag>>(this->ebos_simulator_);
         this->material_state_ = std::make_unique<PyMaterialState<TypeTag>>(this->ebos_simulator_);
         return result;
     }
@@ -200,7 +244,20 @@ Opm::FlowMain<typename Opm::Pybind::PyBlackOilSimulator::TypeTag>&
     }
 }
 
-PyMaterialState<typename Opm::Pybind::PyBlackOilSimulator::TypeTag>&
+PyFluidState<typename PyBlackOilSimulator::TypeTag>&
+PyBlackOilSimulator::
+getFluidState() const
+{
+    if (this->fluid_state_) {
+        return *this->fluid_state_;
+    }
+    else {
+        throw std::runtime_error("BlackOilSimulator not initialized: "
+            "Cannot get reference to FlowMainEbos object" );
+    }
+}
+
+PyMaterialState<typename PyBlackOilSimulator::TypeTag>&
 PyBlackOilSimulator::getMaterialState() const
 {
     if (this->material_state_) {
@@ -222,18 +279,29 @@ void export_PyBlackOilSimulator(py::module& m)
             std::shared_ptr<Opm::EclipseState>,
             std::shared_ptr<Opm::Schedule>,
             std::shared_ptr<Opm::SummaryConfig> >())
+        .def("advance", &PyBlackOilSimulator::advance, py::arg("report_step"))
+        .def("current_step", &PyBlackOilSimulator::currentStep)
         .def("get_cell_volumes", &PyBlackOilSimulator::getCellVolumes,
             py::return_value_policy::copy)
         .def("get_dt", &PyBlackOilSimulator::getDT)
+        .def("get_fluidstate_variable", &PyBlackOilSimulator::getFluidStateVariable,
+            py::return_value_policy::copy, py::arg("name"))
         .def("get_porosity", &PyBlackOilSimulator::getPorosity,
             py::return_value_policy::copy)
+        .def("get_primary_variable_meaning", &PyBlackOilSimulator::getPrimaryVarMeaning,
+            py::return_value_policy::copy, py::arg("variable"))
+        .def("get_primary_variable_meaning_map", &PyBlackOilSimulator::getPrimaryVarMeaningMap,
+            py::return_value_policy::copy, py::arg("variable"))
+        .def("get_primary_variable", &PyBlackOilSimulator::getPrimaryVariable,
+            py::return_value_policy::copy, py::arg("variable"))
         .def("run", &PyBlackOilSimulator::run)
         .def("set_porosity", &PyBlackOilSimulator::setPorosity)
+        .def("set_primary_variable", &PyBlackOilSimulator::setPrimaryVariable,
+            py::arg("idx_name"), py::arg("value"))
         .def("current_step", &PyBlackOilSimulator::currentStep)
         .def("step", &PyBlackOilSimulator::step)
-        .def("advance", &PyBlackOilSimulator::advance, py::arg("report_step"))
-        .def("step_init", &PyBlackOilSimulator::stepInit)
-        .def("step_cleanup", &PyBlackOilSimulator::stepCleanup);
+        .def("step_cleanup", &PyBlackOilSimulator::stepCleanup)
+        .def("step_init", &PyBlackOilSimulator::stepInit);
 }
 
 } // namespace Opm::Pybind

python/test/test_basic.py

@@ -1,6 +1,5 @@
 import os
 import unittest
-from contextlib import contextmanager
 from pathlib import Path
 from opm.simulators import BlackOilSimulator
 from .pytest_common import pushd

python/test/test_fluidstate_variables.py

@@ -0,0 +1,45 @@
+import os
+import unittest
+from pathlib import Path
+from opm.simulators import BlackOilSimulator
+from .pytest_common import pushd
+
+class TestBasic(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        # NOTE: See comment in test_basic.py for the reason why we are
+        #   only using a single test_all() function instead of splitting
+        #   it up in multiple test functions
+        test_dir = Path(os.path.dirname(__file__))
+        cls.data_dir = test_dir.parent.joinpath("test_data/SPE1CASE1a")
+
+
+    def test_all(self):
+        with pushd(self.data_dir):
+            sim = BlackOilSimulator("SPE1CASE1.DATA")
+            sim.step_init()
+            sim.step()
+            oil_pressure = sim.get_fluidstate_variable(name='po')
+            self.assertAlmostEqual(oil_pressure[0], 41729978.837, places=2, msg='value of oil pressure')
+            gas_pressure = sim.get_fluidstate_variable(name='pg')
+            self.assertAlmostEqual(gas_pressure[0], 41729978.837, places=2, msg='value of gas pressure')
+            water_pressure = sim.get_fluidstate_variable(name='pw')
+            self.assertAlmostEqual(water_pressure[0], 41729978.837, places=2, msg='value of water pressure')
+            rho_w = sim.get_fluidstate_variable(name='rho_w')
+            self.assertAlmostEqual(rho_w[0], 1001.7549054, places=6, msg='value of water density')
+            rho_g = sim.get_fluidstate_variable(name='rho_g')
+            self.assertAlmostEqual(rho_g[0], 275.72397867, places=7, msg='value of gas density')
+            rho_o = sim.get_fluidstate_variable(name='rho_o')
+            self.assertAlmostEqual(rho_o[0], 639.64061021, places=7, msg='value of oil density')
+            Rs = sim.get_fluidstate_variable(name='Rs')
+            self.assertAlmostEqual(Rs[0], 226.196660482, places=7, msg='value of solution gas-oil ratio')
+            Rv = sim.get_fluidstate_variable(name='Rv')
+            self.assertAlmostEqual(Rv[0], 0.0, places=7, msg='value of volatile gas-oil ratio')
+            Sw = sim.get_fluidstate_variable(name='Sw')
+            self.assertAlmostEqual(Sw[0], 0.11905577997, places=10, msg='value of water saturation')
+            So = sim.get_fluidstate_variable(name='So')
+            self.assertAlmostEqual(So[0], 0.56951652831, places=10, msg='value of oil saturation')
+            Sg = sim.get_fluidstate_variable(name='Sg')
+            self.assertAlmostEqual(Sg[0], 0.31142769170, places=10, msg='value of gas saturation')
+            T = sim.get_fluidstate_variable(name='T')
+            self.assertAlmostEqual(T[0], 288.705, places=3, msg='value of temperature')

python/test/test_primary_variables.py

@@ -0,0 +1,43 @@
+import os
+import unittest
+from pathlib import Path
+from opm.simulators import BlackOilSimulator
+from .pytest_common import pushd
+
+class TestBasic(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        # NOTE: See comment in test_basic.py for the reason why we are
+        #   only using a single test_all() function instead of splitting
+        #   it up in multiple test functions
+        test_dir = Path(os.path.dirname(__file__))
+        cls.data_dir = test_dir.parent.joinpath("test_data/SPE1CASE1a")
+
+
+    def test_all(self):
+        with pushd(self.data_dir):
+            sim = BlackOilSimulator("SPE1CASE1.DATA")
+            sim.step_init()
+            sim.step()
+            pressure = sim.get_primary_variable(variable='pressure')
+            self.assertAlmostEqual(pressure[0], 41729978.837, places=2, msg='value of pressure')
+            pressure_meaning = sim.get_primary_variable_meaning(
+                variable='pressure')
+            pressure_meaning_map = sim.get_primary_variable_meaning_map(
+                variable='pressure')
+            self.assertEqual(pressure_meaning[0], pressure_meaning_map["Po"])
+            water_meaning = sim.get_primary_variable_meaning(
+                variable='water')
+            water_meaning_map = sim.get_primary_variable_meaning_map(
+                variable='water')
+            self.assertEqual(water_meaning[0], water_meaning_map["Sw"])
+            gas_meaning = sim.get_primary_variable_meaning(
+                variable='gas')
+            gas_meaning_map = sim.get_primary_variable_meaning_map(
+                variable='gas')
+            self.assertEqual(gas_meaning[0], gas_meaning_map["Sg"])
+            brine_meaning = sim.get_primary_variable_meaning(
+                variable='brine')
+            brine_meaning_map = sim.get_primary_variable_meaning_map(
+                variable='brine')
+            self.assertEqual(brine_meaning[0], brine_meaning_map["Disabled"])

python/test/test_schedule.py

@@ -1,6 +1,5 @@
 import os
 import unittest
-from contextlib import contextmanager
 import datetime as dt
 from pathlib import Path
 import re