LgrOutputHelpers.hpp

/*
  Copyright 2025 Equinor ASA.

  This file is part of the Open Porous Media project (OPM).

  OPM is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  OPM is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with OPM.  If not, see <http://www.gnu.org/licenses/>.
*/

#ifndef OPM_GRID_CPGRID_LGROUTPUTHELPERS_HEADER_INCLUDED
#define OPM_GRID_CPGRID_LGROUTPUTHELPERS_HEADER_INCLUDED

#include <opm/grid/CpGrid.hpp>
#include <opm/grid/cpgrid/LevelCartesianIndexMapper.hpp>

#if HAVE_OPM_COMMON
#include <opm/input/eclipse/Units/UnitSystem.hpp>
#include <opm/output/data/Cells.hpp>
#include <opm/output/data/Solution.hpp>
#include <opm/output/eclipse/RestartValue.hpp>
#endif

#include <algorithm>    // for std::min/max
#include <cstddef>      // for std::size_t
#include <utility>      // for std::move
#include <type_traits>  // for std::is_same_v
#include <vector>

template<class ScalarType>
struct MinChildrenData{
    void operator()(int level,
                    int levelIdx,
                    const std::vector<std::vector<ScalarType>>& levelVectors)
    {
        min = std::min(min, levelVectors[ level ][ levelIdx ]);
    }

    ScalarType getValue()
    {
        return min;
    }

    ScalarType min{};
};

template<class ScalarType>
struct MaxChildrenData{
    void operator()(int level,
                    int levelIdx,
                    const std::vector<std::vector<ScalarType>>& levelVectors)
    {
        max = std::max(max, levelVectors[ level ][ levelIdx ]);
    }

    ScalarType getValue()
    {
        return max;
    }

    ScalarType max{};
};

template<class ScalarType>
struct AverageChildrenData{
    void operator()(int level,
                    int levelIdx,
                    const std::vector<std::vector<ScalarType>>& levelVectors)
    {
        partialSum+= levelVectors[ level ][ levelIdx ];
        ++count;
    }

    ScalarType getValue()
    {
        return partialSum/count;
    }

    std::size_t count{};
    ScalarType partialSum{};
};

namespace Opm
{
namespace Lgr
{

std::vector<int> mapLevelIndicesToCartesianOutputOrder(const Dune::CpGrid& grid,
                                                       const Opm::LevelCartesianIndexMapper<Dune::CpGrid>& levelCartMapp,
                                                       int level);

template <typename Container>
Container reorderForOutput(const Container& simulatorContainer,
                           const std::vector<int>& toOutput);

std::vector<std::unordered_map<int,int>>
levelCartesianToLevelCompressedMaps(const Dune::CpGrid& grid,
                                    const Opm::LevelCartesianIndexMapper<Dune::CpGrid>& levelCartMapp);

template <typename ScalarType>
ScalarType processChildrenData(const std::vector<std::vector<ScalarType>>& levelVectors,
                               const Dune::cpgrid::Entity<0>& element,
                               const Dune::CpGrid& grid);

template <typename ScalarType>
void populateDataVectorLevelGrids(const Dune::CpGrid& grid,
                                  int maxLevel,
                                  const std::vector<ScalarType>& leafVector,
                                  const std::vector<std::vector<int>>& toOutput_refinedLevels,
                                  std::vector<std::vector<ScalarType>>& levelVectors);

#if HAVE_OPM_COMMON
void extractSolutionLevelGrids(const Dune::CpGrid& grid,
                               const std::vector<std::vector<int>>& toOutput_refinedLevels,
                               const Opm::data::Solution& leafSolution,
                               std::vector<Opm::data::Solution>&);

template <typename Grid>
void extractRestartValueLevelGrids(const Grid& grid,
                                   const Opm::RestartValue& leafRestartValue,
                                   std::vector<Opm::RestartValue>& restartValue_levels);
#endif

} // namespace Lgr
} // namespace Opm

template <typename Container>
Container Opm::Lgr::reorderForOutput(const Container& simulatorContainer,
                                     const std::vector<int>& toOutput)
{
    // Use toOutput to reorder simulatorContainer
    Container outputContainer;
    outputContainer.resize(toOutput.size());
    for (std::size_t i = 0; i < toOutput.size(); ++i) {
        outputContainer[i] = simulatorContainer[toOutput[i]];
    }
    return outputContainer;
}

template <typename ScalarType>
ScalarType Opm::Lgr::processChildrenData(const std::vector<std::vector<ScalarType>>& levelVectors,
                                         const Dune::cpgrid::Entity<0>& element,
                                         const Dune::CpGrid& grid)
{
    const auto& [level, level_indices] = grid.currentData()[element.level()]->getChildrenLevelAndIndexList(element.index());

    // For now, we compute field output properties as follows:
    // - int fields: use the maximum value among the children data
    // - double fields: use the average of the children data
    //
    // In the future, this could be extended to use the std::string field name
    // to apply field-specific computation logic.
    auto childrenDataFunc = [](){
        if constexpr (std::is_same_v<ScalarType, double>)
            return AverageChildrenData<ScalarType>{};
        else
            return MaxChildrenData<ScalarType>{};
    }();

    for (const auto& levelIdx : level_indices) {
        childrenDataFunc(level, levelIdx, levelVectors);
    }
    return childrenDataFunc.getValue();
}

template <typename ScalarType>
void Opm::Lgr::populateDataVectorLevelGrids(const Dune::CpGrid& grid,
                                            int maxLevel,
                                            const std::vector<ScalarType>& leafVector,
                                            const std::vector<std::vector<int>>& toOutput_refinedLevels,
                                            std::vector<std::vector<ScalarType>>& levelVectors)
{
    for (int level = 0; level <= maxLevel; ++level) {
        levelVectors[level].resize(grid.levelGridView(level).size(0));
    }
    // For level cells that appear in the leaf, extract the data value from leafVector
    // and assign it to the equivalent level cell.
    for (const auto& element : Dune::elements(grid.leafGridView())) {
        levelVectors[element.level()][element.getLevelElem().index()] = leafVector[element.index()];
    }
    // Note that all cells from maxLevel have assigned values at this point.
    // Now, assign values for parent cells (for now, average of children values).
    if (maxLevel)  {
        for (int level = maxLevel-1; level >= 0; --level) {
            for (const auto& element : Dune::elements(grid.levelGridView(level))) {
                if (!element.isLeaf()) {
                    levelVectors[level][element.index()] = Opm::Lgr::processChildrenData(levelVectors,
                                                                                         element,
                                                                                         grid);
                }
            }
        }
    }
    // Use toOutput_levels to reorder in ascending level cartesian indices
    for (int level = 1; level<=maxLevel; ++level) { // exclude level zero (does not need reordering)
        levelVectors[level] = Opm::Lgr::reorderForOutput(levelVectors[level], toOutput_refinedLevels[level-1]);
    }
}

#if HAVE_OPM_COMMON
template <typename Grid>
void Opm::Lgr::extractRestartValueLevelGrids(const Grid& grid,
                                             const Opm::RestartValue& leafRestartValue,
                                             std::vector<Opm::RestartValue>& restartValue_levels)
{
    if constexpr (std::is_same_v<Grid, Dune::CpGrid>) {

        int maxLevel = grid.maxLevel();
        restartValue_levels.resize(maxLevel+1); // level 0, 1, ..., max level

        // To store leafRestartValue.extra data in the order expected
        // by outout files (increasing level Cartesian indices)
        std::vector<std::vector<int>> toOutput_refinedLevels{};
        toOutput_refinedLevels.resize(maxLevel); // exclude level zero (does not need reordering)

        const Opm::LevelCartesianIndexMapper<Dune::CpGrid> levelCartMapp(grid);
        for (int level = 1; level <= maxLevel; ++level) { // exclude level zero (does not need reordering)
            toOutput_refinedLevels[level-1] = mapLevelIndicesToCartesianOutputOrder(grid, levelCartMapp, level);
        }

        std::vector<Opm::data::Solution> dataSolutionLevels{};
        extractSolutionLevelGrids(grid,
                                  toOutput_refinedLevels,
                                  leafRestartValue.solution,
                                  dataSolutionLevels);



        for (int level = 0; level <= maxLevel; ++level) {
            restartValue_levels[level] = Opm::RestartValue(std::move(dataSolutionLevels[level]),
                                                           leafRestartValue.wells,
                                                           leafRestartValue.grp_nwrk,
                                                           leafRestartValue.aquifer,
                                                           level);
        }

        for (const auto& [rst_key, leafVector] : leafRestartValue.extra) {

            std::vector<std::vector<double>> levelVectors{};
            levelVectors.resize(maxLevel+1);

            if (rst_key.key == "OPMEXTRA") {
                // For OPMEXTRA, leafVector has size 1 instead of
                // grid.leafGridView().size(0)
                continue; // skip it
            }
            Opm::Lgr::populateDataVectorLevelGrids<double>(grid,
                                                           maxLevel,
                                                           leafVector,
                                                           toOutput_refinedLevels,
                                                           levelVectors);
            for (int level = 0; level <= maxLevel; ++level) {
                restartValue_levels[level].addExtra(rst_key.key, rst_key.dim, std::move(levelVectors[level]));
            }
        }
    }
}
#endif

#endif // OPM_GRID_CPGRID_LGROUTPUTHELPERS_HEADER_INCLUDED