Merge pull request #6424 from lechevaa/HybridNewton_flag

Hybrid newton flag

CMakeLists_files.cmake

@@ -934,6 +934,7 @@ list (APPEND PUBLIC_HEADER_FILES
   opm/simulators/flow/GenericThresholdPressure_impl.hpp
   opm/simulators/flow/GenericTracerModel.hpp
   opm/simulators/flow/GenericTracerModel_impl.hpp
+  opm/simulators/flow/HybridNewton.hpp
   opm/simulators/flow/InterRegFlows.hpp
   opm/simulators/flow/KeywordValidation.hpp
   opm/simulators/flow/LogOutputHelper.hpp

opm/models/blackoil/blackoilenergymodules.hh

@@ -264,7 +264,7 @@ public:
                                   const FluidState& fluidState)
     {
         if constexpr (enableEnergy) {
-            priVars[temperatureIdx] = fluidState.temperature(/*phaseIdx=*/0);
+            priVars[temperatureIdx] = getValue(fluidState.temperature(/*phaseIdx=*/0));
         }
     }
 

opm/simulators/flow/FlowProblem.hpp

@@ -349,7 +349,7 @@ public:
     /*!
      * \brief Called by the simulator before each time integration.
      */
-    void beginTimeStep()
+    virtual void beginTimeStep()
     {
         OPM_TIMEBLOCK(beginTimeStep);
         const int episodeIdx = this->episodeIndex();
@@ -376,7 +376,6 @@ public:
         wellModel_.beginTimeStep();
         aquiferModel_.beginTimeStep();
         tracerModel_.beginTimeStep();
-
     }
 
     /*!

opm/simulators/flow/FlowProblemBlackoil.hpp

@@ -52,6 +52,8 @@
 #include <opm/simulators/flow/MixingRateControls.hpp>
 #include <opm/simulators/flow/OutputBlackoilModule.hpp>
 #include <opm/simulators/flow/VtkTracerModule.hpp>
+#include <opm/simulators/flow/HybridNewton.hpp> 
+#include <opm/simulators/flow/HybridNewtonConfig.hpp> 
 
 #include <opm/simulators/utils/satfunc/SatfuncConsistencyCheckManager.hpp>
 
@@ -148,6 +150,7 @@ private:
     using DiffusionModule = BlackOilDiffusionModule<TypeTag, enableDiffusion>;
     using ConvectiveMixingModule = BlackOilConvectiveMixingModule<TypeTag, enableConvectiveMixing>;
     using ModuleParams = typename BlackOilLocalResidualTPFA<TypeTag>::ModuleParams;
+    using HybridNewton = BlackOilHybridNewton<TypeTag>;
 
     using InitialFluidState = typename EquilInitializer<TypeTag>::ScalarFluidState;
     using EclWriterType = EclWriter<TypeTag, OutputBlackOilModule<TypeTag> >;
@@ -188,6 +191,7 @@ public:
                          simulator.vanguard().summaryState(),
                          this->wellModel_,
                          simulator.vanguard().grid().comm())
+        , hybridNewton_(simulator)
     {
         this->model().addOutputModule(std::make_unique<VtkTracerModule<TypeTag>>(simulator));
 
@@ -264,6 +268,15 @@ public:
                                              this->moduleParams_.convectiveMixingModuleParam);
     }
 
+    /*!
+     * \brief Called by the simulator before each time integration.
+     */
+    void beginTimeStep() override
+    {
+        FlowProblemType::beginTimeStep();
+        hybridNewton_.tryApplyHybridNewton();
+    }
+
     /*!
      * \copydoc FvBaseProblem::finishInit
      */
@@ -1695,6 +1708,8 @@ protected:
 
     ModuleParams moduleParams_;
 
+    HybridNewton hybridNewton_;
+
 private:
     /// Whether or not the current epsiode will end at the end of the
     /// current time step.

opm/simulators/flow/FlowProblemParameters.cpp

@@ -68,6 +68,12 @@ void registerFlowProblemParameters()
 
     Parameters::Register<Parameters::NumSatfuncConsistencySamplePoints>
         ("Maximum number of reported failures for each individual saturation function consistency check");
+    
+    Parameters::Register<Parameters::HyNeConfigFile>
+        ("Use config files for Hybrid Newton");
+
+    Parameters::Register<Parameters::UseHyNe>
+        ("Wheter or not to use Hybrid Newton");
 
     // By default, stop it after the universe will probably have stopped
     // to exist. (the ECL problem will finish the simulation explicitly

opm/simulators/flow/FlowProblemParameters.hpp

@@ -56,6 +56,11 @@ struct OutputMode { static constexpr auto value = "all"; };
 // between writing restart files
 struct RestartWritingInterval { static constexpr int value = 0xffffff; }; // disable
 
+// Path to the config file containing all Hybrid Newton parameters
+struct HyNeConfigFile { static constexpr auto value = "hybridNewtonConfig.json"; };
+// Wheter or not to use Hybrid Newton nonlinear preconditioning
+struct UseHyNe { static constexpr bool value = false; }; 
+
 } // namespace Opm::Parameters
 
 namespace Opm {

opm/simulators/flow/HybridNewton.hpp

@@ -0,0 +1,581 @@
+// -*- mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*-
+// vi: set et ts=4 sw=4 sts=4:
+/*
+  Copyright 2025 NORCE Research AS
+
+  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 HYBRID_NEWTON_HPP
+#define HYBRID_NEWTON_HPP
+
+
+#include <opm/simulators/flow/FlowBaseProblemProperties.hpp>
+#include <opm/simulators/flow/HybridNewtonConfig.hpp>
+#include <opm/input/eclipse/Units/UnitSystem.hpp>
+
+#include <opm/ml/ml_model.hpp>
+
+#include <string>         
+#include <vector>
+
+namespace Opm {
+
+/*!
+ * \brief Hybrid Newton solver extension for the black-oil model.
+ *
+ * This class integrates machine learning–based corrections into the
+ * standard Newton solver of the black-oil simulator. It uses one or
+ * more trained Hybrid Newton models to adjust the initial guess of the
+ * nonlinear solver at specific timesteps.
+ *
+ * Each model is described by a HybridNewtonConfig, which specifies:
+ *   - the path to the model file,
+ *   - the cells to which the model applies,
+ *   - the time points when it should be applied,
+ *   - the input and output features with their transformations and scalings.
+ *
+ * At runtime, the class:
+ *   1. Loads the relevant configuration(s),
+ *   2. Constructs an input tensor [n_cells x n_input_features],
+ *   3. Runs the ML model to produce an output tensor [n_cells x n_output_features],
+ *   4. Updates the nonlinear solver’s initial guess accordingly.
+ */
+template <typename TypeTag>
+class BlackOilHybridNewton
+{
+protected:
+    using Simulator  = GetPropType<TypeTag, Properties::Simulator>;
+    using Evaluation = GetPropType<TypeTag, Properties::Evaluation>;
+    using FluidSystem = GetPropType<TypeTag, Properties::FluidSystem>;
+    using MaterialLaw = GetPropType<TypeTag, Properties::MaterialLaw>;
+    using Indices = GetPropType<TypeTag, Properties::Indices>;
+    using Scalar  = GetPropType<TypeTag, Properties::Scalar>;
+
+    enum { numPhases = FluidSystem::numPhases };
+
+    enum { gasPhaseIdx = FluidSystem::gasPhaseIdx };
+    enum { oilPhaseIdx = FluidSystem::oilPhaseIdx };
+    enum { waterPhaseIdx = FluidSystem::waterPhaseIdx };
+
+    static constexpr bool compositionSwitchEnabled = Indices::compositionSwitchIdx >= 0;
+
+public:
+    explicit BlackOilHybridNewton(Simulator& simulator)
+        : simulator_(simulator)
+        , configsLoaded_(false)
+    {}
+
+    /*!
+    * \brief Attempt to apply the Hybrid Newton correction at the current timestep.
+    *
+    * This function acts as the entry point for Hybrid Newton corrections.
+    * It first checks whether the Hybrid Newton mechanism is enabled
+    * (via `Parameters::UseHyNe`). If enabled, it validates the fluid system
+    * and loads model configurations from the parameter-specified
+    * configuration file (if not already loaded).
+    *
+    * At each timestep, it iterates over all loaded configurations and applies
+    * those that match the current simulation time.
+    *
+    * Typical call site: at the beginning of a nonlinear solve for each timestep.
+    *
+    * \throws std::runtime_error if configuration validation fails or if an
+    *         invalid fluid system is encountered.
+    */
+    void tryApplyHybridNewton()
+    {   
+        // Check if flag activated 
+        if (!Parameters::Get<Parameters::UseHyNe>())
+            return;
+        
+        validateFluidSystem();
+
+        if (!configsLoaded_) {
+            std::string config_file = Parameters::Get<Parameters::HyNeConfigFile>();
+            PropertyTree pt(config_file);
+            for (const auto& model_key : pt.get_child_keys()) {
+                HybridNewtonConfig config(pt.get_child(model_key));
+                config.validateConfig(compositionSwitchEnabled);
+                configs_.push_back(std::move(config));
+            }
+            configsLoaded_ = true;
+        }
+
+        Scalar current_time = simulator_.time();
+        // Find and apply all models that should run at this time
+        for (const auto& config : configs_) {
+            if (shouldApplyHybridNewton(current_time, config)) {
+                runHybridNewton(config);
+            }
+        }
+    }
+
+private:
+    /*!
+    * \brief Apply the Hybrid Newton method using a given model configuration.
+    *
+    * This function loads the cell indices on demand from the configuration's
+    * `cell_indices_file`, constructs the input tensor, evaluates the model to
+    * produce the output tensor, and applies the result to update the initial
+    * guess of the nonlinear solver.
+    *
+    * \param config The HybridNewtonConfig object containing model path,
+    *        feature specifications, and cell indices file reference.
+    *
+    * \throws std::runtime_error if cell indices cannot be loaded or if
+    *         feature extraction or application fails.
+    */
+    void runHybridNewton(const HybridNewtonConfig& config)
+    {
+        auto input  = constructInputTensor(config);
+        auto output = constructOutputTensor(input, config);
+        updateInitialGuess(output, config);
+    }
+
+protected:
+
+    void validateFluidSystem()
+    {
+        const auto& eclState = simulator_.vanguard().eclState();
+        const auto& phases = eclState.runspec().phases();
+        
+        bool hasWater = phases.active(Phase::WATER);
+        bool hasGas = phases.active(Phase::GAS);
+        bool hasOil = phases.active(Phase::OIL);
+        bool hasSolvent = phases.active(Phase::SOLVENT);
+
+        // Check for three-phase black oil system
+        if (!(hasWater && hasOil && hasGas && !hasSolvent)) {
+            OPM_THROW(std::runtime_error, 
+                "HybridNewton: Unsupported fluid system. Only three-phase black oil is supported.");
+        }
+    }
+
+    /*!
+    * \brief Check whether the Hybrid Newton method should be applied at the given time.
+    *
+    * This function evaluates the current simulation time against the application times
+    * specified in the HybridNewtonConfig. The behavior depends on how many entries are
+    * provided in `config.apply_times`:
+    *
+    * - If one entry is provided, the method is applied exactly at that time (within
+    *   a tolerance of 1e-6).
+    * - If two entries are provided, the method is applied during the inclusive interval
+    *   between the start and end times.
+    * - For any other number of entries, the method is not applied.
+    *
+    * \param current_time The current simulation time.
+    * \param config The HybridNewtonConfig specifying when the Hybrid Newton method
+    *        should be applied.
+    *
+    * \return True if the Hybrid Newton method should be applied at \p current_time,
+    *         false otherwise.
+    */
+    bool shouldApplyHybridNewton(Scalar current_time, const HybridNewtonConfig& config) const
+    {
+        if (config.apply_times.size() == 1) {
+            // Apply exactly at one point in time (with optional tolerance)
+            constexpr Scalar tolerance = 1e-6;
+            bool apply = std::abs(current_time - config.apply_times[0]) < tolerance;
+            return apply;
+        }
+
+        if (config.apply_times.size() == 2) {
+            Scalar start_time = config.apply_times[0];
+            Scalar end_time = config.apply_times[1];
+            bool apply = (current_time >= start_time) && (current_time <= end_time);
+            return apply;
+        }
+
+        // If apply_times has an unexpected number of entries
+        return false;
+    }
+
+    
+    /*!
+    * \brief Construct the input feature tensor for the Hybrid Newton model.
+    *
+    * The tensor is constructed in the exact order specified by
+    * `config.input_features`:
+    *   - Scalar features (e.g., TIMESTEP) contribute one entry each.
+    *   - Per-cell features (e.g., PRESSURE, SGAS) contribute a block of
+    *     `n_cells` entries, one per cell index in `config.cell_indices`.
+    *
+    * Example:
+    *   Suppose `config.input_features = [SGAS, TIMESTEP, PRESSURE]`
+    *   and `config.n_cells = 3`. The resulting tensor layout is:
+    *
+    *     [ SGAS(cell0), SGAS(cell1), SGAS(cell2),
+    *       TIMESTEP,
+    *       PRESSURE(cell0), PRESSURE(cell1), PRESSURE(cell2) ]
+    *
+    * Thus, the tensor length equals:
+    *
+    *     (# of scalar features) + (# of per-cell features × n_cells).
+    *
+    * \param config  HybridNewtonConfig containing feature definitions and cell indices.
+    * \return        A 1D tensor of Evaluation values with the computed layout.
+    */
+    ML::Tensor<Evaluation> constructInputTensor(
+        const HybridNewtonConfig& config
+    )
+    {
+        const auto& features = config.input_features;
+
+        // compute total length directly
+        std::size_t input_tensor_length = 0;
+        for (const auto& feature : features) {
+            const FeatureSpec& spec = feature.second;
+            if (spec.actual_name == "TIMESTEP") {
+                input_tensor_length += 1;
+            } else {
+                input_tensor_length += config.n_cells;
+            }
+        }
+
+        ML::Tensor<Evaluation> input(input_tensor_length);
+        std::size_t offset = 0;
+
+        // fill in exact feature order from config
+        for (const auto& feature: features) {
+            const FeatureSpec& spec = feature.second;
+
+            if (spec.actual_name == "TIMESTEP") {
+                auto value = static_cast<Evaluation>(getScalarFeatureValue(spec));
+                input(offset++) = value;
+            } else {
+                for (std::size_t i = 0; i < config.n_cells; ++i) {
+                    auto value = static_cast<Evaluation>(
+                        getPerCellFeatureValue(spec, config.cell_indices[i])
+                    );
+                    input(offset + i) = value;
+                }
+                offset += config.n_cells;
+            }
+        }
+
+        return input;
+    }
+
+    /*!
+    * \brief Retrieve and transform a scalar feature (global across the domain).
+    *
+    * Supported scalar feature:
+    *   - "TIMESTEP": current timestep size.
+    *
+    * The raw value is passed through the feature's transformation
+    * and scaling functions before being returned.
+    *
+    * \param spec The feature specification, including transform and scaling.
+    * \return The transformed and scaled feature value.
+    *
+    * \throws std::runtime_error if the feature is unknown.
+    */
+    Scalar getScalarFeatureValue(const FeatureSpec& spec)
+    {
+        Scalar value = 0.0;
+
+        if (spec.actual_name == "TIMESTEP") {
+            value = simulator_.timeStepSize();
+        } else {
+            OPM_THROW(std::runtime_error, "Unknown scalar feature: " + spec.actual_name);
+        }
+
+        value = spec.transform.apply(value);
+        return spec.scaler.scale(value);
+    }
+
+    /*!
+    * \brief Retrieve and transform a per-cell feature value.
+    *
+    * Supported per-cell features include:
+    *   - PRESSURE, SWAT, SOIL, SGAS, RS, RV, PERMX.
+    *
+    * The raw value is taken from the simulator state,
+    * converted into the configured unit system,
+    * then passed through the feature's transformation
+    * and scaling functions.
+    *
+    * \param spec The feature specification, including transform and scaling.
+    * \param cell_index The cell index for which to retrieve the value.
+    * \return The transformed and scaled feature value.
+    *
+    * \throws std::runtime_error if the feature is unknown.
+    */
+    Scalar getPerCellFeatureValue(const FeatureSpec& spec, int cell_index)
+    {
+        const auto& intQuants = simulator_.model().intensiveQuantities(cell_index, 0);
+        const auto& fs = intQuants.fluidState();
+        const auto& unitSyst = simulator_.vanguard().schedule().getUnits();
+
+        Scalar value = 0.0;
+
+        if (spec.actual_name == "PRESSURE") {
+            value = getValue(fs.pressure(oilPhaseIdx));
+            value = unitSyst.from_si(UnitSystem::measure::pressure, value);
+        } else if (spec.actual_name == "SWAT") {
+            value = getValue(fs.saturation(waterPhaseIdx));
+        } else if (spec.actual_name == "SGAS") {
+            value = getValue(fs.saturation(gasPhaseIdx));
+        } else if (spec.actual_name == "SOIL") {
+            value = getValue(fs.saturation(oilPhaseIdx));
+        } else if (spec.actual_name == "RS") {
+            value = getValue(fs.Rs());
+            value = unitSyst.from_si(UnitSystem::measure::gas_oil_ratio, value);
+        } else if (spec.actual_name == "RV") {
+            value = getValue(fs.Rv());
+            value = unitSyst.from_si(UnitSystem::measure::oil_gas_ratio, value);
+        } else if (spec.actual_name == "PERMX") {
+            const auto& eclState = simulator_.vanguard().eclState();
+            const auto& fp = eclState.fieldProps();
+            auto permX = fp.get_double("PERMX");
+            value = permX[cell_index];
+            value = unitSyst.from_si(UnitSystem::measure::permeability, value);
+        } else {
+            OPM_THROW(std::runtime_error, "Unknown per-cell feature: " + spec.actual_name);
+        }
+
+        Scalar transformed = spec.transform.apply(value);
+        Scalar scaled = spec.scaler.scale(transformed);
+        return scaled;
+    }
+
+
+    /*!
+    * \brief Run the Hybrid Newton model to produce output predictions.
+    *
+    * Uses the input tensor (prepared by constructInputTensor) together
+    * with the model specified in \p config to compute the corresponding
+    * output tensor. Scaling and transformations defined in the configuration
+    * are applied automatically.
+    *
+    * \param input The input tensor of shape 
+    *              [(# of scalar features) + (# of per-cell features × n_cells)].
+    * \param config The HybridNewtonConfig specifying model path and output
+    *        feature definitions.
+    * \return A tensor of shape [n_cells x n_output_features], where rows
+    *         correspond to cells and columns correspond to output features.
+    *
+    * \throws std::runtime_error if model inference fails or if the
+    *         output tensor does not match the expected feature layout.
+    */
+    ML::Tensor<Evaluation>
+    constructOutputTensor(const ML::Tensor<Evaluation>& input, const HybridNewtonConfig& config)
+    {
+        const auto& features = config.output_features;
+        const int n_features = features.size();
+        
+        ML::NNModel<Evaluation> model;
+        model.loadModel(config.model_path);
+
+        ML::Tensor<Evaluation> output(1, config.n_cells * n_features);
+        model.apply(input, output);
+
+        return output; 
+    }
+
+    /*!
+    * \brief Update the nonlinear solver's initial guess using ML predictions.
+    *
+    * For each target cell, this function:
+    *   1. Reads predicted output features from the ML tensor,
+    *   2. Applies inverse scaling and transformations,
+    *   3. If marked as delta values (`spec.is_delta == true`),
+    *      adds the predicted delta to the current simulator state,
+    *   4. Updates fluid state variables accordingly:
+    *        - PRESSURE: replaces oil-phase pressure and reconstructs
+    *          other phase pressures using capillary pressure relations,
+    *        - SWAT, SOIL, SGAS: updates saturations with closure rules
+    *          if fewer than 3 are predicted,
+    *        - RS, RV: updates dissolved/volatilized gas-oil ratios
+    *          if composition-switching is enabled.
+    *
+    * Special handling:
+    *   - Saturations are clamped to [0,1] and normalized to ensure
+    *     their sum is unity.
+    *   - Invalid states (e.g., zero saturation sum) raise exceptions.
+    *
+    * \param output The ML output tensor [n_cells x n_output_features].
+    * \param config The model configuration specifying output feature mappings.
+    *
+    * \throws std::runtime_error if an unknown output feature is encountered
+    *         or if state consistency cannot be enforced.
+    */
+    void updateInitialGuess(
+        ML::Tensor<Evaluation>& output, 
+        const HybridNewtonConfig& config 
+    )
+    {   
+        const auto& features = config.output_features;
+
+        FeatureFlags flags = flagFeatures(features);
+
+        const auto& unitSyst = simulator_.vanguard().schedule().getUnits();
+
+        for (std::size_t i = 0; i < config.n_cells; ++i) {
+            const int cell_idx = config.cell_indices[i];
+            const auto& intQuants = simulator_.model().intensiveQuantities(cell_idx, /*timeIdx*/0);
+            auto fs = intQuants.fluidState();
+
+            int feature_idx = 0;
+
+            // Temp variables per cell
+            Scalar sw_val = -1.0;
+            Scalar so_val = -1.0;
+            Scalar sg_val = -1.0; 
+            Scalar po_val = -1.0;
+
+            for (const auto& [name, spec] : features) {
+
+                auto scaled_value = getValue(output(feature_idx * config.n_cells + i));
+                
+                // Inverse scaling
+                Scalar raw_value = spec.scaler.unscale(scaled_value);
+
+                // Inverse transform
+                raw_value = spec.transform.applyInverse(raw_value);
+
+                if (spec.is_delta) { 
+                    if (spec.actual_name == "PRESSURE") { 
+                        raw_value = raw_value + unitSyst.from_si(UnitSystem::measure::pressure, getValue(fs.pressure(oilPhaseIdx))); 
+                    } else if (spec.actual_name == "SWAT") { 
+                        raw_value += getValue(fs.saturation(waterPhaseIdx)); 
+                    } else if (spec.actual_name == "SOIL") { 
+                        raw_value += getValue(fs.saturation(oilPhaseIdx)); 
+                    } else if (spec.actual_name == "SGAS") { 
+                        raw_value += getValue(fs.saturation(gasPhaseIdx)); 
+                    } else if (spec.actual_name == "RS") { 
+                        raw_value = raw_value + unitSyst.from_si(UnitSystem::measure::gas_oil_ratio, getValue(fs.Rs())); 
+                    } else if (spec.actual_name == "RV") { 
+                        raw_value = raw_value + unitSyst.from_si(UnitSystem::measure::oil_gas_ratio, getValue(fs.Rv())); 
+                    } else { 
+                        OPM_THROW(std::runtime_error, "Unknown delta feature: " + name); 
+                    } 
+                }
+
+                if (spec.actual_name == "PRESSURE") {
+                    po_val = unitSyst.to_si(UnitSystem::measure::pressure, raw_value);
+                } else if (spec.actual_name == "SWAT") {
+                    sw_val = raw_value;
+                } else if (spec.actual_name == "SOIL") {
+                    so_val = raw_value;
+                } else if (spec.actual_name == "SGAS") {
+                    sg_val = raw_value;
+                } else if (spec.actual_name == "RS") {
+                    if constexpr (compositionSwitchEnabled) {
+                        raw_value = unitSyst.to_si(UnitSystem::measure::gas_oil_ratio, raw_value); 
+                        fs.setRs(raw_value);
+                    }
+                } else if (spec.actual_name == "RV") {
+                    if constexpr (compositionSwitchEnabled) {
+                        raw_value = unitSyst.to_si(UnitSystem::measure::oil_gas_ratio, raw_value);
+                        fs.setRv(raw_value);
+                    }
+                } else {
+                    OPM_THROW(std::runtime_error, "Unknown output feature: " + name);
+                }
+
+                ++feature_idx;
+            }
+
+            int sat_count = static_cast<int>(flags.has_SWAT) + static_cast<int>(flags.has_SOIL) + static_cast<int>(flags.has_SGAS);
+
+            if (sat_count >= 2) {
+                Scalar sw = sw_val;
+                Scalar so = so_val;
+                Scalar sg = sg_val;
+
+                if (!flags.has_SWAT) {
+                    sw = 1.0 - so - sg;
+                } else if (!flags.has_SOIL) {
+                    so = 1.0 - sw - sg;
+                } else if (!flags.has_SGAS) {
+                    sg = 1.0 - sw - so;
+                }
+
+                sw = max(0.0, min(sw, 1.0));
+                so = max(0.0, min(so, 1.0));
+                sg = max(0.0, min(sg, 1.0));
+
+                Scalar sum = sw + so + sg;
+                if (sum <= 1e-12) {
+                    OPM_THROW(std::runtime_error, "Saturation sum is zero in cell " + std::to_string(cell_idx));
+                }
+
+                fs.setSaturation(waterPhaseIdx, sw / sum);
+                fs.setSaturation(oilPhaseIdx,   so / sum);
+                fs.setSaturation(gasPhaseIdx,   sg / sum);
+            }
+
+            if (flags.has_PRESSURE) {
+                std::array<Evaluation, numPhases> pC;
+                const auto& materialParams = simulator_.problem().materialLawParams(cell_idx);
+                MaterialLaw::capillaryPressures(pC, materialParams, fs);
+
+                for (unsigned phaseIdx = 0; phaseIdx < numPhases; ++phaseIdx) {
+                    if (!FluidSystem::phaseIsActive(phaseIdx)) continue;
+                    if (phaseIdx == oilPhaseIdx) {
+                        fs.setPressure(phaseIdx, po_val);
+                    } else {
+                        fs.setPressure(phaseIdx, po_val - pC[phaseIdx]);
+                    }
+                }
+            }
+
+            auto& primaryVars = simulator_.model().solution(/*timeIdx*/0)[cell_idx];
+            primaryVars.assignNaive(fs);
+        }
+
+        simulator_.model().invalidateAndUpdateIntensiveQuantities(/*timeIdx*/0);
+    }
+    
+    struct FeatureFlags {
+    bool has_SWAT = false;
+    bool has_SOIL = false;
+    bool has_SGAS = false;
+    bool has_PRESSURE = false;
+    };
+
+    FeatureFlags flagFeatures(const std::vector<std::pair<std::string, FeatureSpec>>& features) 
+    {
+        FeatureFlags flags;
+
+        for (const auto& [name, spec] : features) {
+
+            if (spec.actual_name == "SWAT") {
+                flags.has_SWAT = true;
+            } else if (spec.actual_name == "SOIL") {
+                flags.has_SOIL = true;
+            } else if (spec.actual_name == "SGAS") {
+                flags.has_SGAS = true;
+            } else if (spec.actual_name == "PRESSURE") {
+                flags.has_PRESSURE = true;
+            }
+        }
+
+        return flags;
+    }
+
+    protected:
+        Simulator& simulator_;
+        std::vector<HybridNewtonConfig> configs_;
+        bool configsLoaded_;
+    };
+
+} // namespace Opm
+
+#endif // HYBRID_NEWTON_CLASS_HPP

opm/simulators/flow/HybridNewtonConfig.hpp

@@ -0,0 +1,291 @@
+// -*- mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*-
+// vi: set et ts=4 sw=4 sts=4:
+/*
+  Copyright 2025 NORCE Research AS
+
+  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 HYBRID_NEWTON_CONFIG_HPP
+#define HYBRID_NEWTON_CONFIG_HPP
+
+#include <algorithm>
+#include <cmath>
+#include <fstream>
+#include <string>
+#include <stdexcept>
+#include <vector>
+
+#include <opm/simulators/linalg/PropertyTree.hpp>
+
+namespace Opm {
+
+/*!
+ * \brief Represents scaling information for a feature.
+ *
+ * Supports standard (mean/std) and min-max scaling.
+ */
+struct Scaler {
+    enum class Type { None, Standard, MinMax } type = Type::None;
+    double mean = 0.0;
+    double std  = 1.0;
+    double min  = 0.0;
+    double max  = 1.0;
+
+    double scale(double raw_value) const {
+        switch (type) {
+            case Type::Standard:
+                return (std == 0.0) ? mean : (raw_value - mean) / std;
+            case Type::MinMax: {
+                double denom = max - min;
+                return (denom == 0.0) ? min : (raw_value - min) / denom;
+            }
+            case Type::None:
+            default:
+                return raw_value;
+        }
+    }
+
+    double unscale(double scaled_value) const {
+        switch (type) {
+            case Type::Standard: return scaled_value * std + mean;
+            case Type::MinMax:   return scaled_value * (max - min) + min;
+            case Type::None:
+            default:             return scaled_value;
+        }
+    }
+};
+
+/*!
+ * \brief Represents a transformation applied to a feature.
+ *
+ * Supports log, log10, log1p, and no transform. Provides forward
+ * and inverse methods.
+ */
+struct Transform {
+    enum class Type { None, Log, Log10, Log1p } type = Type::None;
+
+    Transform() = default;
+    explicit Transform(const std::string& name) {
+        if      (name == "log10") type = Type::Log10;
+        else if (name == "log")   type = Type::Log;
+        else if (name == "log1p") type = Type::Log1p;
+        else                       type = Type::None;
+    }
+
+    double apply(double raw_value) const {
+        switch (type) {
+            case Type::Log10: return std::log10(raw_value);
+            case Type::Log:   return std::log(raw_value);
+            case Type::Log1p: return std::log1p(raw_value);
+            case Type::None:
+            default:          return raw_value;
+        }
+    }
+
+    double applyInverse(double transformed_value) const {
+        switch (type) {
+            case Type::Log10: return std::pow(10.0, transformed_value);
+            case Type::Log:   return std::exp(transformed_value);
+            case Type::Log1p: return std::expm1(transformed_value);
+            case Type::None:
+            default:          return transformed_value;
+        }
+    }
+};
+
+/*!
+ * \brief Metadata for a single feature (input or output).
+ *
+ * Includes transformation, scaling, delta flag, and actual feature name.
+ */
+struct FeatureSpec {
+    Transform transform;
+    Scaler scaler;
+    bool is_delta = false;
+    std::string actual_name;
+
+    FeatureSpec() = default;
+};
+
+/*!
+ * \brief Configuration for a Hybrid Newton ML model.
+ *
+ * Encapsulates model path, cell indices, apply times, input/output
+ * features, and validation. Can be constructed from a PropertyTree.
+ */
+class HybridNewtonConfig {
+public:
+    std::string model_path;
+    std::string cell_indices_file;
+    std::vector<int> cell_indices;
+    std::size_t n_cells = 0;
+    std::vector<double> apply_times;
+    std::vector<std::pair<std::string, FeatureSpec>> input_features;
+    std::vector<std::pair<std::string, FeatureSpec>> output_features;
+
+    // Default constructor
+    HybridNewtonConfig() = default;
+
+    /*!
+    * \brief Construct configuration from a PropertyTree.
+    *
+    * Loads model path, cell indices, apply times, and features from
+    * the provided PropertyTree. Throws on missing or invalid entries.
+    */
+    explicit HybridNewtonConfig(const PropertyTree& model_config)
+    {
+        model_path = model_config.get<std::string>("model_path", "");
+        if (model_path.empty())
+            throw std::runtime_error("Missing 'model_path' in HybridNewton config");
+
+        cell_indices_file = model_config.get<std::string>("cell_indices_file", "");
+        if (cell_indices_file.empty())
+            throw std::runtime_error("Missing 'cell_indices_file' in HybridNewton config");
+
+        // Load cell indices immediately
+        cell_indices = loadCellIndicesFromFile(cell_indices_file);
+        n_cells = cell_indices.size();
+
+        // Load apply times
+        auto applyTimesOpt = model_config.get_child_optional("apply_times");
+        if (!applyTimesOpt)
+            throw std::runtime_error("Missing 'apply_times' in HybridNewton config");
+
+        for (const auto& key : applyTimesOpt->get_child_keys())
+            apply_times.push_back(applyTimesOpt->get_child(key).get<double>(""));
+
+        if (apply_times.empty())
+            throw std::runtime_error("'apply_times' must contain at least one value");
+
+        // Parse features
+        parseFeatures(model_config, "features.inputs", input_features);
+        parseFeatures(model_config, "features.outputs", output_features);
+    }
+
+    bool hasInputFeature(const std::string& name) const {
+        return std::any_of(input_features.begin(), input_features.end(),
+                           [&](const auto& p) { return p.first == name; });
+    }
+
+    bool hasOutputFeature(const std::string& name) const {
+        return std::any_of(output_features.begin(), output_features.end(),
+                           [&](const auto& p) { return p.first == name; });
+    }
+
+    /*!
+    * \brief Validate feature compatibility with simulator settings.
+    *
+    * Ensures RS/RV features are only used if composition support is enabled.
+    * Throws std::runtime_error otherwise.
+    */
+    void validateConfig(bool compositionSwitchEnabled) const
+    {
+        bool hasRsFeature = hasInputFeature("RS") ||
+                            hasOutputFeature("RS") ||
+                            hasOutputFeature("DELTA_RS");
+
+        bool hasRvFeature = hasInputFeature("RV") ||
+                            hasOutputFeature("RV") ||
+                            hasOutputFeature("DELTA_RV");
+
+        if ((hasRsFeature || hasRvFeature) && !compositionSwitchEnabled) {
+            OPM_THROW(std::runtime_error,
+                "HybridNewton: RS or RV features detected but composition support is disabled. "
+                "CompositionSwitch must be enabled for RS/RV features."
+            );
+        }
+    }
+    
+private:
+    /*!
+    * \brief Load cell indices from a plain text file.
+    *
+    * Each line must contain one integer index. Lines starting with '#' or
+    * empty lines are ignored. Throws if the file is missing, invalid, or empty.
+    */
+    std::vector<int> loadCellIndicesFromFile(const std::string& filename) const {
+        std::vector<int> indices;
+        std::ifstream cellFile(filename);
+        if (!cellFile.is_open())
+            throw std::runtime_error("Cannot open cell indices file: " + filename);
+
+        std::string line;
+        int lineNumber = 0;
+        while (std::getline(cellFile, line)) {
+            ++lineNumber;
+            if (line.empty() || line[0] == '#') continue;
+            try { indices.push_back(std::stoi(line)); }
+            catch (...) {
+                throw std::runtime_error("Invalid cell index at line " + std::to_string(lineNumber) +
+                                         " in file " + filename + ": " + line);
+            }
+        }
+
+        if (indices.empty())
+            throw std::runtime_error("No valid cell indices found in file: " + filename);
+
+        return indices;
+    }
+
+    /*!
+    * \brief Parse feature specifications from a PropertyTree.
+    *
+    * Reads transform, scaling parameters, and delta flag.
+    * Stores the results in the given `features` vector.
+    *
+    * \param pt       PropertyTree containing the model configuration.
+    * \param path     Path to the feature subtree ("features.inputs" or "features.outputs").
+    * \param features Destination vector of (name, FeatureSpec) pairs.
+    */
+    void parseFeatures(const PropertyTree& pt, const std::string& path,
+                       std::vector<std::pair<std::string, FeatureSpec>>& features) {
+        auto subtreeOpt = pt.get_child_optional(path);
+        if (!subtreeOpt) return;
+
+        for (const auto& name : subtreeOpt->get_child_keys()) {
+            const PropertyTree& ft = subtreeOpt->get_child(name);
+            FeatureSpec spec;
+            spec.transform = Transform(ft.get<std::string>("feature_engineering", "none"));
+
+            if (auto sOpt = ft.get_child_optional("scaling_params")) {
+                const PropertyTree& s = *sOpt;
+                if (s.get_child_optional("mean") && s.get_child_optional("std")) {
+                    spec.scaler.type = Scaler::Type::Standard;
+                    spec.scaler.mean = s.get<double>("mean", 0.0);
+                    spec.scaler.std  = s.get<double>("std", 1.0);
+                } else if (s.get_child_optional("min") && s.get_child_optional("max")) {
+                    spec.scaler.type = Scaler::Type::MinMax;
+                    spec.scaler.min  = s.get<double>("min", 0.0);
+                    spec.scaler.max  = s.get<double>("max", 1.0);
+                } else {
+                    spec.scaler.type = Scaler::Type::None;
+                }
+            } else {
+                spec.scaler.type = Scaler::Type::None;
+            }
+
+            spec.is_delta = name.compare(0, 6, "DELTA_") == 0;
+            spec.actual_name = spec.is_delta ? name.substr(6) : name;
+
+            features.emplace_back(name, std::move(spec));
+        }
+    }
+};
+
+} // namespace Opm
+
+#endif

opm/simulators/linalg/PropertyTree.cpp

@@ -96,6 +96,15 @@ PropertyTree::get_child_optional(const std::string& key) const
   return PropertyTree(pt.get());
 }
 
+std::vector<std::string> PropertyTree::get_child_keys() const
+{
+    std::vector<std::string> keys;
+    for (auto it = tree_->begin(); it != tree_->end(); ++it) {
+        keys.push_back(it->first);
+    }
+    return keys;
+}
+
 template <typename T>
 std::optional<std::vector<T>>
 PropertyTree::get_child_items_as_vector(const std::string& child) const

opm/simulators/linalg/PropertyTree.hpp

@@ -129,6 +129,11 @@ public:
     std::optional<PropertyTree>
     get_child_optional(const std::string& key) const;
 
+    /// Retrieve all child keys of this property tree node
+    ///
+    /// \return Vector of strings containing the names of all immediate children
+    std::vector<std::string> get_child_keys() const;
+
     /// Retrieve node items as linearised vector.
     ///
     /// Assumes that the node's child is an array type of homongeneous

python/test/ml/README.md

@@ -0,0 +1,13 @@
+requirements (same as opm-common/python/opm/ml/ml_tools/requirements.txt)
+
+# Python v.3.9.0 and above
+python>=3.9.0<=3.12.0
+
+# Numpy 1.23.0 and above
+numpy>=1.23.0
+
+# TensorFlow v.2.1 and above include both CPU and GPU versions.
+tensorflow>=2.1.0
+
+# Keras v.2.12.0 and above
+keras>=2.12.0

python/test/ml/hybrid_newton/README.md

@@ -0,0 +1,104 @@
+# Hybrid Newton Test Suite
+
+This test suite validates the integration of **Hybrid Newton (HyNE)** with machine learning models inside the `flow` reservoir simulator.
+It ensures that adding ML-based approximations does not degrade solver performance and, in some cases, improves convergence.
+
+---
+
+## Dependencies
+
+In a similar way as for opm-common ML module, these tests require **TensorFlow** to be installed. 
+The CPU version is enough for this test and allows faster installation.
+```bash
+pip install tensorflow-cpu
+```
+
+---
+
+## Test Framework
+
+* Implemented with Python’s built-in [`unittest`](https://docs.python.org/3/library/unittest.html).
+* Each test case defines a specific **Hybrid Newton configuration** via JSON.
+* Simulations are run using `flow`(therefore assuming it is installed):
+
+  * **Baseline** → Standard Newton solver (no ML).
+  * **Hybrid Newton** → Newton solver with ML-based updates injected.
+
+---
+
+## Baseline Simulation
+
+Before running hybrid cases:
+
+* Runs a standard `flow` simulation with the given deck.
+* Extracts:
+
+  * **Baseline Newton iterations** (`baseline_iters`) from simulation logs.
+  * **Reservoir state variables** (`PRESSURE`, `SGAS`, `SWAT`, etc.) from UNRST files.
+* Stores `n_cells` and `times` for consistency across tests.
+
+---
+
+## Hybrid Newton Simulation
+
+For each test case:
+
+1. A **minimal Keras model** is built that maps input → output features.
+
+   * Weights are initialized to copy output variables directly.
+   * Models are exported using `kerasify.export_model`.
+2. A **JSON config** is generated with:
+
+   * Input/output features.
+   * Scaling and feature engineering options.
+   * Apply time and active cells.
+3. `flow` is run with:
+
+   ```bash
+   --use-hy-ne=true --hy-ne-config-file=config.json
+   ```
+4. Hybrid Newton iteration counts are extracted and compared to baseline.
+
+---
+
+## Assertions
+
+Each test validates:
+
+* **Convergence parity**
+  Hybrid and baseline must produce the same number of timesteps.
+* **Iteration bounds**
+  `hybrid_iters[t] <= baseline_iters[t]` for all timesteps.
+* **Zero Newton cases**
+  Some cases require the first two timesteps to have `0` Newton iterations.
+* **Config integrity**
+  Config JSONs must exist and be valid.
+
+---
+
+## Test Categories
+
+| Category                    | Description                                           |
+| --------------------------- | ----------------------------------------------------- |
+| `ABSOLUTE_CASES`            | Tests absolute features.                              |
+| `RELATIVE_CASES`            | Tests relative features.                              |
+| `FEATURE_ENGINEERING_CASES` | Tests features engineered inputs/outputs (e.g., log). |
+| `SCALING_CASES`             | Applies scaling strategies (e.g., min-max, standard). |
+| `MULTI_MODEL_CASES`         | Combines multiple models into one simulation.         |
+| `ZERO_NEWTON_CASES`         | Use all features, resulting in 0 Newton iterations.   |
+| `ALL_CASES`                 | Runs all cases together as a stress test.             |
+
+---
+
+## Utilities
+
+The suite uses helper functions from `utils.py`:
+
+* `collect_input_features` → Prepares input arrays from UNRST data.
+* `compute_output_vars` → Builds expected outputs from reservoir state.
+* `write_config` → Writes the HyNE JSON configuration.
+* `export_model` → Saves Keras models in a format usable by `flow`.
+* `extract_newtit_from_file` → Extracts iteration counts from solver logs.
+* `extract_unrst_variables` → Reads simulation outputs from UNRST files.
+
+---

python/test/ml/hybrid_newton/test_hybrid_newton.py

@@ -0,0 +1,201 @@
+import os
+import unittest
+from pathlib import Path
+import numpy as np
+import subprocess
+
+from utils import collect_input_features, compute_output_vars, write_config, extract_newtit_from_file, extract_unrst_variables
+from utils import (ABSOLUTE_CASES, RELATIVE_CASES, FEATURE_ENGINEERING_CASES, SCALING_CASES, MULTI_MODEL_CASES, ZERO_NEWTON_CASES, ALL_CASES)
+
+import tensorflow as tf
+from kerasify import export_model
+
+
+class TestHybridNewton(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        """Run baseline simulation via subprocess and collect timesteps & fluid state."""
+        test_dir = Path(__file__).parent
+        cls.data_dir = test_dir.parent.parent.parent / "test_data/SPE1CASE1b"
+        cls.deck_file = "SPE1CASE1.DATA"
+
+        cls.unrst = {}
+        cls.rd_init = {"PERMX": []}
+        cls.times = []
+
+        baseline_cmd = [
+            "flow",
+            str(cls.data_dir / cls.deck_file),
+            "--output-extra-convergence-info=steps,iterations",
+            "--newton-min-iterations=0",
+            "--full-time-step-initially=true",
+        ]
+
+        print("Running baseline simulation...")
+        subprocess.run(baseline_cmd, cwd=cls.data_dir, stdout=subprocess.DEVNULL, check=True)
+
+        # Extract Newton iterations
+        cls.baseline_iters = extract_newtit_from_file(cls.data_dir / Path(cls.deck_file).stem)
+        print(f"Baseline Newton iterations: {cls.baseline_iters}")
+
+        # Extract fluid state from UNRST for all timesteps
+        cls.unrst, cls.times = extract_unrst_variables(cls.data_dir, cls.deck_file)
+        cls.rd_init["PERMX"] = [10**2] * len(cls.unrst["PRESSURE"][0])
+        cls.n_cells = len(cls.unrst["PRESSURE"][0])
+  
+    def _run_cases(self, cases):
+        models_dir = self.data_dir / "models"
+        os.makedirs(models_dir, exist_ok=True)
+
+        for case in cases:
+            with self.subTest(case=case["label"]):
+                json_path = self._build_case_models(case, models_dir)
+                self.assertTrue(json_path and os.path.isfile(json_path),
+                                f"Config JSON was not created for case {case['label']}")
+
+                # Run hybrid Newton test
+                hybrid_iters = self._run_hybrid_newton_test(json_path)
+        return hybrid_iters
+    
+    def _build_case_models(self, case, models_dir):
+        models_to_write = []
+
+        if case.get("mode", "single") == "single":
+            model_entries = [case]
+            multi_mode = False
+        else:
+            model_entries = case.get("models", [])
+            multi_mode = True
+
+        for model_case in model_entries:
+            start_idx = model_case.get("start_idx", case.get("start_idx"))
+            end_idx = model_case.get("end_idx", case.get("end_idx"))
+            input_features = model_case["input_features"]
+            output_features = model_case["output_features"]
+            fe_input = model_case.get("feature_engineering_input")
+            fe_output = model_case.get("feature_engineering_output")
+            scaling_input = model_case.get("scaling_input")
+            scaling_output = model_case.get("scaling_output")
+
+            input_flat, input_scaling_params = collect_input_features(
+                input_features, start_idx, self.unrst, self.rd_init, self.times,
+                feat_eng=fe_input, scaling=scaling_input
+            )
+            output_flat, output_scaling_params = compute_output_vars(
+                output_features, start_idx, end_idx, self.unrst, self.rd_init,
+                feature_engineering=fe_output, scaling=scaling_output,
+            )
+
+            input_flat = input_flat.reshape(-1)
+            output_flat = output_flat.reshape(-1)
+
+            model = tf.keras.Sequential([
+                tf.keras.Input(shape=(input_flat.shape[0],)),
+                tf.keras.layers.Dense(output_flat.shape[0], use_bias=True)
+            ])
+            model.layers[0].set_weights([np.zeros((input_flat.shape[0], output_flat.shape[0])), output_flat])
+            apply_time = self.times[start_idx]
+            
+            active_cells = list(range(self.n_cells))
+
+            models_dir.mkdir(exist_ok=True)
+            model_name = "_".join(input_features) + "__TO__" + "_".join(output_features) + f"_{start_idx}_{end_idx}.model"
+            model_path = models_dir / model_name
+            export_model(model, model_path)
+    
+            models_to_write.append({
+                "input_features": input_features,
+                "output_features": output_features,
+                "model_path": model_path,
+                "active_cells": active_cells,
+                "apply_time": apply_time,
+                "feature_engineering_input": fe_input,
+                "feature_engineering_output": fe_output,
+                "scaling_input": scaling_input,
+                "scaling_output": scaling_output,
+                "input_scaling_params": input_scaling_params,
+                "output_scaling_params": output_scaling_params
+            })
+
+        if not models_to_write:
+            print(f"No valid models for {case['label']}")
+            return None
+
+        return write_config(models_to_write, models_dir, multi_mode=multi_mode, case_label=case['label'])
+
+    def _run_hybrid_newton_test(self, json_path):
+        """Run hybrid Newton via subprocess and extract Newton iterations."""
+        cmd = [
+            "flow",
+            str(self.data_dir / self.deck_file),
+            f"--hy-ne-config-file={json_path}",
+            "--use-hy-ne=true",
+            "--output-extra-convergence-info=steps,iterations",
+            "--newton-min-iterations=0",
+            "--full-time-step-initially=true"
+        ]
+
+        print(f"Running hybrid Newton simulation with config {json_path}...")
+        subprocess.run(cmd, cwd=self.data_dir, 
+                       stdout=subprocess.DEVNULL, 
+                       check=True)
+
+        hybrid_iters = extract_newtit_from_file(self.data_dir / Path(self.deck_file).stem) 
+        return hybrid_iters
+
+    def test_absolute_cases(self):
+        self._run_cases(ABSOLUTE_CASES)
+
+    def test_relative_cases(self):
+        self._run_cases(RELATIVE_CASES)
+
+    def test_feature_engineering_cases(self):
+        self._run_cases(FEATURE_ENGINEERING_CASES)
+
+    def test_scaling_cases(self):
+        self._run_cases(SCALING_CASES)
+
+    def test_multi_model_cases(self):
+        self._run_cases(MULTI_MODEL_CASES)
+    
+    def test_zero_newton_cases(self):
+        hybrid_iters = self._run_cases(ZERO_NEWTON_CASES)
+
+        # Length must match (failed newton convergence otherwise)
+        assert len(hybrid_iters) == len(self.baseline_iters), (
+            f"Hybrid iterations length {len(hybrid_iters)} "
+            f"does not match baseline length {len(self.baseline_iters)}"
+        )
+
+        # First two must be 0
+        assert hybrid_iters[0] == 0, f"Expected hybrid_iters[0] == 0, got {hybrid_iters[0]}"
+        assert hybrid_iters[1] == 0, f"Expected hybrid_iters[1] == 0, got {hybrid_iters[1]}"
+
+        # Hybrid must not exceed baseline
+        for i, (h, b) in enumerate(zip(hybrid_iters, self.baseline_iters)):
+            assert h <= b, (
+                f"Hybrid iterations exceeded baseline at index {i}: "
+                f"got {h}, expected <= {b}"
+            )
+
+    def test_all_cases(self):
+        hybrid_iters = self._run_cases(ALL_CASES)
+        # Length must match (failed newton convergence otherwise)
+        assert len(hybrid_iters) == len(self.baseline_iters), (
+            f"Hybrid iterations length {len(hybrid_iters)} "
+            f"does not match baseline length {len(self.baseline_iters)}"
+        )
+
+        # First two must be 0
+        assert hybrid_iters[0] == 0, f"Expected hybrid_iters[0] == 0, got {hybrid_iters[0]}"
+        assert hybrid_iters[1] == 0, f"Expected hybrid_iters[1] == 0, got {hybrid_iters[1]}"
+
+        # Hybrid must not exceed baseline
+        for i, (h, b) in enumerate(zip(hybrid_iters, self.baseline_iters)):
+            assert h <= b, (
+                f"Hybrid iterations exceeded baseline at index {i}: "
+                f"got {h}, expected <= {b}"
+            )
+
+if __name__ == "__main__":
+    unittest.main()

python/test/ml/hybrid_newton/utils.py

@@ -0,0 +1,414 @@
+import json
+import os
+import numpy as np
+from pathlib import Path
+
+from opm.io.ecl import ERst, ESmry
+
+
+def extract_unrst_variables(data_dir, deck_file):
+    """
+    Extracts fluid state variables and timesteps from the .UNRST file
+    corresponding to a deck.
+
+    Parameters
+    ----------
+    data_dir : Path or str
+        Directory containing the UNRST file.
+    deck_file : str
+        Name of the deck file (e.g., SPE1CASE1.DATA).
+
+    Returns
+    -------
+    unrst_vars : dict
+        Dictionary with keys "PRESSURE", "SWAT", "SOIL", "SGAS", "RS"
+        Each value is a list of np.ndarray for each report step.
+    dts : list
+        List of timestep durations (approximated from report steps if needed).
+    """
+    deck_stem = Path(deck_file).stem
+    unrst_file = str(Path(data_dir) / f"{deck_stem}.UNRST")  # convert to string
+    unsmry_file = str(Path(data_dir) / f"{deck_stem}.SMSPEC")
+    rst = ERst(unrst_file)
+    unsmry = ESmry(unsmry_file)
+    # Load all report steps
+    report_steps = sorted(rst.report_steps)
+    for step in report_steps:
+        rst.load_report_step(step)
+
+    variables_of_interest = ["PRESSURE", "SWAT", "SOIL", "SGAS", "RS"]
+    unrst_vars = {var: [] for var in variables_of_interest}
+
+    for step in report_steps:
+        for var in variables_of_interest:
+            if (var, step) in rst:
+                unrst_vars[var].append(np.array(rst[var, step], copy=True))
+            else:
+                # Fill with zeros if variable not in UNRST for this step
+                n_cells = len(rst["PRESSURE", step])  # assume same size as PRESSURE
+                unrst_vars[var].append(np.zeros(n_cells, dtype=np.float32))
+
+    # Compute approximate timesteps (differences between report steps)
+    dates = unsmry.dates()
+    times = [0.] + [(d - unsmry.start_date).total_seconds() for d in dates]
+
+    return unrst_vars, times
+
+
+def write_config(models, model_dir, multi_mode=False, case_label=None):
+    """
+    Write one or multiple model configs to JSON in OPM-supported format.
+    """
+
+    model_dir = Path(model_dir)
+    if multi_mode and not case_label:
+        raise ValueError("case_label must be provided for multi_mode")
+
+    os.makedirs(model_dir, exist_ok=True)
+    all_configs = {}
+
+    for idx, m in enumerate(models):
+        # Ensure input/output features are lists
+        input_features = m["input_features"] if isinstance(m["input_features"], list) else [m["input_features"]]
+        output_features = m["output_features"] if isinstance(m["output_features"], list) else [m["output_features"]]
+
+        fe_input = m.get("feature_engineering_input") or [None]*len(input_features)
+        fe_output = m.get("feature_engineering_output") or [None]*len(output_features)
+        scaling_params_input = m.get("input_scaling_params") or [None]*len(input_features)
+        scaling_params_output = m.get("output_scaling_params") or [None]*len(output_features)
+
+        # Build input/output blocks
+        input_block = {}
+        for i, f in enumerate(input_features):
+            d = {}
+            if fe_input[i] and fe_input[i].lower() != "none":
+                d["feature_engineering"] = fe_input[i].lower()
+            if scaling_params_input[i] is not None:
+                # convert numpy float32 to native float if needed
+                d["scaling_params"] = {k: float(v) for k, v in scaling_params_input[i].items()}
+            input_block[f] = d
+
+        output_block = {}
+        for i, f in enumerate(output_features):
+            d = {}
+            if fe_output[i] and fe_output[i].lower() != "none":
+                d["feature_engineering"] = fe_output[i].lower()
+            if scaling_params_output[i] is not None:
+                d["scaling_params"] = {k: float(v) for k, v in scaling_params_output[i].items()}
+            output_block[f] = d
+
+        # Active cells
+        active_cells = m.get("active_cells")
+        if active_cells is None:
+            n_cells = getattr(m, "n_cells", None)
+            if n_cells is None:
+                raise ValueError("Cannot determine active_cells; provide n_cells or active_cells")
+            active_cells = list(range(n_cells))
+
+        cells_file = model_dir / f"{os.path.splitext(os.path.basename(m['model_path']))[0]}_active_cells.txt"
+        with open(cells_file, "w") as f:
+            for cell in active_cells:
+                f.write(f"{cell}\n")
+
+        # Apply time(s)
+        apply_time = m["apply_time"]
+        if isinstance(apply_time, float):
+            apply_times_dict = {"0": float(apply_time)}
+        elif isinstance(apply_time, list):
+            apply_times_dict = {str(i): float(t) for i, t in enumerate(apply_time)}
+        else:
+            raise ValueError("apply_time must be float or list of floats")
+
+        model_name_rel = f"models/{os.path.basename(m['model_path'])}"
+        cells_file_rel = f"models/{cells_file.name}"
+
+        all_configs[str(idx)] = {
+            "model_path": model_name_rel,
+            "cell_indices_file": cells_file_rel,
+            "apply_times": apply_times_dict,
+            "features": {"inputs": input_block, "outputs": output_block}
+        }
+
+    # JSON filename
+    if multi_mode:
+        json_name = case_label + ".json"
+    else:
+        json_name = f"{os.path.splitext(os.path.basename(models[0]['model_path']))[0]}.json"
+
+    json_path = model_dir / json_name
+    with open(json_path, "w") as f:
+        json.dump(all_configs, f, indent=2)
+
+    return json_path
+
+
+def seconds_to_days(sec):
+    return sec / 86400.0
+
+
+def apply_feature_engineering(data, method=None):
+    if not method or method.lower() == "none":
+        return data
+    if method.lower() == "log":
+        return np.log(data)
+    if method.lower() == "log10":
+        return np.log10(data)
+    if method.lower() == "log1p":
+        return np.log1p(data)
+    raise ValueError(f"Unsupported feature engineering method: {method}")
+
+
+def apply_scaling_with_params(data, method):
+    data = np.array(data)
+    if method == "standard":
+        mean, std = np.mean(data), np.std(data)
+        return ((data - mean) / std if std != 0 else data - mean), {"mean": mean, "std": std}
+    if method == "minmax":
+        mn, mx = np.min(data), np.max(data)
+        return ((data - mn) / (mx - mn) if mx != mn else data - mn), {"min": mn, "max": mx}
+    return data, None
+
+
+def get_feature_data(name, timestep, unrst, rd_init, times=None, feat_eng=None):
+    if name == "PERMX":
+        val = np.array(rd_init["PERMX"])
+    elif name == "SOIL":
+        val = 1.0 - np.array(unrst["SGAS"][timestep]) - np.array(unrst["SWAT"][timestep])
+    elif name == "TIMESTEP":
+        if times is None:
+            raise ValueError("times must be provided for TIMESTEP")
+        val = times[timestep + 1] - times[timestep]
+    else:
+        val = np.array(unrst[name][timestep])
+    return apply_feature_engineering(val, feat_eng)
+
+
+def collect_input_features(input_features, start_idx, unrst, rd_init, times=None, feat_eng=None, scaling=None):
+    input_data = []
+    scaling_params = []
+
+    for i, f in enumerate(input_features):
+        fe = feat_eng[i] if feat_eng else None
+        sc = scaling[i] if scaling else None
+        val = get_feature_data(f, start_idx, unrst, rd_init, times=times, feat_eng=fe)
+        val, params = apply_scaling_with_params(val, sc) if sc else (val, None)
+        input_data.append(np.array(val).flatten())
+        scaling_params.append(params)
+
+    return np.concatenate(input_data), scaling_params
+
+
+def compute_output_vars(output_features, start_idx, end_idx, unrst, rd_init, feature_engineering=None, scaling=None):
+    output_data = []
+    scaling_params = []
+
+    for i, f in enumerate(output_features):
+        fe = feature_engineering[i] if feature_engineering else None
+        sc = scaling[i] if scaling else None
+
+        if f.startswith("DELTA_"):
+            base = f[6:]
+            val = get_feature_data(base, end_idx, unrst, rd_init) - get_feature_data(base, start_idx, unrst, rd_init)
+        else:
+            val = get_feature_data(f, end_idx, unrst, rd_init)
+
+        val = apply_feature_engineering(val, fe)
+        val, params = apply_scaling_with_params(val, sc) if sc else (val, None)
+        output_data.append(np.array(val).flatten())
+        scaling_params.append(params)
+
+    return np.stack(output_data), scaling_params
+
+
+def extract_newtit_from_file(filename: str | Path) -> list[int]:
+    """
+    Parse an INFOSTEP file and return the values of the 'NewtIt' column.
+
+    Parameters
+    ----------
+    filename : str or Path
+        Path to the file (with or without .INFOSTEP extension).
+
+    Returns
+    -------
+    list[int]
+        List of NewtIt values for each data row.
+    """
+    path = Path(filename)
+    if path.suffix.lower() != ".infostep":
+        path = path.with_suffix(".INFOSTEP")
+    if not path.exists():
+        raise FileNotFoundError(f"INFOSTEP file not found: {path}")
+
+    with path.open("r") as f:
+        lines = [line.strip() for line in f if line.strip()]
+
+    # Locate header with "NewtIt"
+    header_idx = next((i for i, line in enumerate(lines) if "NewtIt" in line), None)
+    if header_idx is None:
+        raise ValueError(f"'NewtIt' column not found in {path}")
+
+    header = lines[header_idx].split()
+    idx = header.index("NewtIt")
+
+    newt_values = []
+    for row in lines[header_idx + 1 :]:
+        cols = row.split()
+        if len(cols) > idx:
+            try:
+                newt_values.append(int(cols[idx]))
+            except ValueError:
+                continue  # skip malformed rows
+
+    return newt_values
+
+# -------------------------
+# Absolute (single timestep predictions)
+# -------------------------
+ABSOLUTE_CASES = [
+    {
+        "label": "single_RS",
+        "mode": "single",
+        "start_idx": 0, "end_idx": 1,
+        "input_features": ["TIMESTEP"],
+        "output_features": ["RS"],
+    },
+    {
+        "label": "single_PRESSURE",
+        "mode": "single",
+        "start_idx": 0, "end_idx": 1,
+        "input_features": ["PRESSURE"],
+        "output_features": ["PRESSURE"],
+    },
+    {
+        "label": "single_PERMX_to_SGAS",
+        "mode": "single",
+        "start_idx": 0, "end_idx": 1,
+        "input_features": ["PERMX"],
+        "output_features": ["SGAS"],
+    },
+    {
+        "label": "single_SGAS_to_SWAT",
+        "mode": "single",
+        "start_idx": 0, "end_idx": 1,
+        "input_features": ["SGAS"],
+        "output_features": ["SWAT"],
+    },
+]
+
+# -------------------------
+# Relative / delta predictions
+# -------------------------
+RELATIVE_CASES = [
+    {"label": "delta_RS", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["RS"], "output_features": ["DELTA_RS"]},
+    {"label": "delta_SGAS", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["SGAS"], "output_features": ["DELTA_SGAS"]},
+    {"label": "delta_SWAT", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["SWAT"], "output_features": ["DELTA_SWAT"]},
+    {"label": "delta_PRESSURE", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["PRESSURE"], "output_features": ["DELTA_PRESSURE"]},
+    {"label": "delta_PERMX_to_RS", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["PERMX"], "output_features": ["DELTA_RS"]},
+    {"label": "delta_PERMX_to_PRESSURE", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["PERMX"], "output_features": ["DELTA_PRESSURE"]},
+]
+
+# -------------------------
+# Feature engineering cases
+# -------------------------
+FEATURE_ENGINEERING_CASES = [
+    {"label": "feat_eng_RS_to_PRESSURE", "mode": "single", "start_idx": 0, "end_idx": 1,
+     "input_features": ["RS"], "feature_engineering_input": ["log"],
+     "output_features": ["PRESSURE"], "feature_engineering_output": ["log1p"]},
+    {"label": "feat_eng_PRESSURE", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["SGAS"], "feature_engineering_input": ["log1p"],
+     "output_features": ["PRESSURE"], "feature_engineering_output": ["log10"]},
+    {"label": "feat_eng_PERMX_to_PRESSURE", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["PERMX"], "feature_engineering_input": ["log"],
+     "output_features": ["PRESSURE"], "feature_engineering_output": ["log"]},
+]
+
+# -------------------------
+# Scaling cases
+# -------------------------
+SCALING_CASES = [
+    {"label": "scaling_RS_minmax", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["RS"], "scaling_input": ["minmax"],
+     "output_features": ["RS"], "scaling_output": ["minmax"]},
+    {"label": "scaling_PRESSURE_standard", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["PRESSURE"], "scaling_input": ["standard"],
+     "output_features": ["PRESSURE"], "scaling_output": ["standard"]},
+    {"label": "scaling_SGAS_SWAT_standard", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["SGAS", "SWAT"], "scaling_input": ["standard", "none"], 
+     "output_features": ["SGAS", "SWAT"], "scaling_output": ["none", "standard"]},
+    {"label": "scaling_PRESSURE_SWAT_minmax", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["PRESSURE", "SWAT"], "scaling_input": ["minmax", "none"],
+     "output_features": ["SWAT", "PRESSURE"], "scaling_output": ["none", "minmax"]},
+    {"label": "scaling_PERMX_TIMESTEP_standard", "mode": "single", "start_idx": 1, "end_idx": 2,
+     "input_features": ["PRESSURE", "TIMESTEP"], "scaling_input": ["minmax", "none"],
+     "output_features": ["SGAS", "SWAT"]},
+]
+
+# -------------------------
+# Multi-model cases
+# -------------------------
+MULTI_MODEL_CASES = [
+    {"label": "multi_same_index", "mode": "multi",
+     "models": [
+         {"input_features": ["PERMX"], "output_features": ["RS"], "start_idx": 1, "end_idx": 2},
+         {"input_features": ["SGAS"], "output_features": ["SGAS"], "start_idx": 1, "end_idx": 2},
+     ]},
+    {"label": "multi_diff_index", "mode": "multi",
+     "models": [
+         {"input_features": ["PERMX"], "output_features": ["RS"], "start_idx": 1, "end_idx": 2},
+         {"input_features": ["PRESSURE", "TIMESTEP", "PERMX"], "output_features": ["PRESSURE"], "start_idx": 8, "end_idx": 9},
+         {"input_features": ["SGAS"], "output_features": ["SGAS"], "start_idx": 10, "end_idx": 11},
+     ]},
+]
+
+ZERO_NEWTON_CASES = [
+    {
+        "label": "all_features",
+        "mode": "multi",
+        "models": [
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 0, "end_idx": 1},
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 1, "end_idx": 2},
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 2, "end_idx": 3},
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 3, "end_idx": 4},
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 4, "end_idx": 5},
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 5, "end_idx": 6},
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 6, "end_idx": 7},
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 7, "end_idx": 8},
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 8, "end_idx": 9},
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 9, "end_idx": 10},
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 10, "end_idx": 11},
+            {"input_features": ["PERMX"], "output_features": ["RS", "PRESSURE", "SGAS", "SOIL"], "start_idx": 11, "end_idx": 12},
+        ],
+
+    },
+]
+
+ALL_CASES = [
+    {
+        "label": "all_cases",
+        "mode": "multi",
+        "models": [
+            {"input_features": ["PERMX", "TIMESTEP", "RS", "SOIL", "SWAT", "PRESSURE"],
+             "scaling_input": ["none", "none", "standard", "standard", "standard","minmax"],
+             "feature_engineering_input": ["log", "log10", "log1p", "log1p", "log1p", "log10"],
+             "output_features": ["RS", "PRESSURE", "DELTA_SGAS", "DELTA_SOIL"],
+             "scaling_output": ["standard", "minmax", "standard", "standard"],
+             "feature_engineering_output": ["log", "log10", "log1p", "log1p"],
+             "start_idx": 0, "end_idx": 1},
+             {"input_features": ["PERMX", "TIMESTEP", "RS", "SOIL", "SWAT", "PRESSURE"],
+             "scaling_input": ["none", "none", "standard", "standard", "standard","minmax"],
+             "feature_engineering_input": ["log", "log10", "log1p", "log1p", "log1p", "log10"],
+             "output_features": ["DELTA_RS", "DELTA_PRESSURE", "DELTA_SGAS", "DELTA_SOIL"],
+             "scaling_output": ["standard", "minmax", "standard", "standard"],
+             "feature_engineering_output": ["none", "none", "log1p", "log1p"],
+             "start_idx": 1, "end_idx": 2},
+        ],
+
+    },
+]

python/test_data/SPE1CASE1b/SPE1CASE1.DATA

@@ -59,6 +59,9 @@ WELLDIMS
 -- 	   - there must be two wells in a group as there are two wells in total
    2 1 1 2 /
 
+UNIFIN
+UNIFOUT
+
 GRID
 
 -- The INIT keyword is used to request an .INIT file. The .INIT file