HybridNewton.hpp

// -*- 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 {

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)
    {}

    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:
    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.");
        }
    }

    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;
    }
 
    
    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;
    }

    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);
    }

    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;
    }
 

    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; 
    }

    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