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 <limits>
#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 != std::numeric_limits<unsigned>::max();

public:
    explicit BlackOilHybridNewton(Simulator& simulator)
        : simulator_(simulator)
        , configsLoaded_(false)
    {}

    void tryApplyHybridNewton()
    {
        // Check if flag activated
        if (!Parameters::Get<Parameters::UseHybridNewton>())
            return;

        validateFluidSystem();

        if (!configsLoaded_) {
            std::string config_file = Parameters::Get<Parameters::HybridNewtonConfigFile>();
            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