Hybrid Newton Python tests and C++ refactoring

This commit is contained in:
lechevaa committed 2025-09-29 10:59:55 +02:00
1 parent a57d7b613e
commit 5796d57604
10 files changed
+861 -292

No files matched your search

+72 -104
View File
@@ -30,7 +30,6 @@
#include <opm/ml/ml_model.hpp>
#include <string>
#include <fstream>
#include <vector>
namespace Opm {
@@ -88,19 +87,21 @@ public:
validateFluidSystem();
// Load configs if not yet loaded
if (!configsLoaded_) {
configs_ = HybridNewtonConfig::parseHybridNewtonConfigs(Parameters::Get<Parameters::HyNeConfigFile>()); // note: plural
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;
validateAllConfigs();
}
Scalar current_time = simulator_.time();
// Find and apply all models that should run at this time
auto applicable_models = getApplicableModels(current_time);
if (!applicable_models.empty()) {
for (const auto& config : applicable_models) {
for (const auto& config : configs_) {
if (shouldApplyHybridNewton(current_time, config)) {
runHybridNewton(config);
}
}
@@ -123,14 +124,9 @@ private:
*/
void runHybridNewton(const HybridNewtonConfig& config)
{
// Load cell indices from file at runtime
std::vector<int> cell_indices = config.loadCellIndicesFromFile();
const std::size_t n_cells = cell_indices.size();
// Pass cells to each stage
auto input = constructInputTensor(config, cell_indices);
auto output = constructOutputTensor(input, config, n_cells);
updateInitialGuess(output, config, cell_indices);
auto input = constructInputTensor(config);
auto output = constructOutputTensor(input, config);
updateInitialGuess(output, config);
}
protected:
@@ -152,19 +148,6 @@ protected:
}
}
std::vector<HybridNewtonConfig> getApplicableModels(Scalar current_time)
{
std::vector<HybridNewtonConfig> applicable;
for (const auto& config : configs_) {
if (shouldApplyHybridNewton(current_time, config)) {
applicable.push_back(config);
}
}
return applicable;
}
/*!
* \brief Check whether the Hybrid Newton method should be applied at the given time.
*
@@ -205,30 +188,6 @@ protected:
return false;
}
void validateAllConfigs() const
{
for (const auto& config : configs_) {
// Check if RS or RV features appear in this config
bool hasRsFeature =
config.hasInputFeature("RS") ||
config.hasOutputFeature("RS") ||
config.hasOutputFeature("DELTA_RS");
bool hasRvFeature =
config.hasInputFeature("RV") ||
config.hasOutputFeature("RV") ||
config.hasOutputFeature("DELTA_RV");
// If RS or RV appear, ensure compositionSwitchEnabled is true
if ((hasRsFeature || hasRvFeature) && !compositionSwitchEnabled) {
OPM_THROW(std::runtime_error,
"HybridNewton: RS or RV features detected but compositionSwitchEnabled is false. "
"Composition support required for RS/RV."
);
}
}
}
/*!
* \brief Construct the input feature tensor for the Hybrid Newton model.
@@ -239,54 +198,66 @@ protected:
* a row-major 2D tensor.
*
* \param config The HybridNewtonConfig containing feature definitions.
* \param cell_indices The list of cell indices to process.
* \param n_cells The number of cells (typically cell_indices.size()).
* \return A tensor of shape [n_cells x n_input_features], where rows
* correspond to cells and columns correspond to input features.
*/
ML::Tensor<Evaluation> constructInputTensor(
const HybridNewtonConfig& config,
const std::vector<int>& cell_indices
const HybridNewtonConfig& config
)
{
const std::size_t n_cells = cell_indices.size();
{
const auto& features = config.input_features;
std::size_t num_scalar_features = 0, num_per_cell_features = 0;
std::size_t num_scalar_features = 0;
std::size_t num_per_cell_features = 0;
for (const auto& [name, spec] : features)
if (name == "TIMESTEP") ++num_scalar_features;
else ++num_per_cell_features;
for (const auto& feature : features) {
const FeatureSpec& spec = feature.second;
if (spec.actual_name == "TIMESTEP") {
++num_scalar_features;
} else {
++num_per_cell_features;
}
}
ML::Tensor<Evaluation> input(num_scalar_features + num_per_cell_features * n_cells);
ML::Tensor<Evaluation> input(num_scalar_features + num_per_cell_features * config.n_cells);
std::size_t offset = 0;
// Scalar features
for (const auto& [name, spec] : features)
if (name == "TIMESTEP")
input(offset++) = static_cast<Evaluation>(getScalarFeatureValue(name, spec));
// Per-cell features
for (const auto& [name, spec] : features) {
if (name == "TIMESTEP") continue;
for (std::size_t i = 0; i < n_cells; ++i)
input(offset + i) = static_cast<Evaluation>(getPerCellFeatureValue(name, spec, cell_indices[i]));
offset += n_cells;
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;
}
}
// Per-cell features
for (const auto& feature: features) {
const FeatureSpec& spec = feature.second;
if (spec.actual_name == "TIMESTEP") {
continue;
}
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 std::string& name, const FeatureSpec& spec)
Scalar getScalarFeatureValue(const FeatureSpec& spec)
{
Scalar value = 0.0;
if (name == "TIMESTEP") {
if (spec.actual_name == "TIMESTEP") {
value = simulator_.timeStepSize();
} else {
OPM_THROW(std::runtime_error, "Unknown scalar feature: " + name);
OPM_THROW(std::runtime_error, "Unknown scalar feature: " + spec.actual_name);
}
value = spec.transform.apply(value);
@@ -294,7 +265,7 @@ protected:
}
Scalar getPerCellFeatureValue(const std::string& name, const FeatureSpec& spec, int cell_index)
Scalar getPerCellFeatureValue(const FeatureSpec& spec, int cell_index)
{
const auto& intQuants = simulator_.model().intensiveQuantities(cell_index, 0);
const auto& fs = intQuants.fluidState();
@@ -302,33 +273,34 @@ protected:
Scalar value = 0.0;
if (name == "PRESSURE") {
if (spec.actual_name == "PRESSURE") {
value = getValue(fs.pressure(oilPhaseIdx));
value = unitSyst.from_si(UnitSystem::measure::pressure, value);
} else if (name == "SWAT") {
value = getValue(fs.saturation(waterPhaseIdx));
} else if (name == "SGAS") {
} else if (spec.actual_name == "SWAT") {
value = getValue(fs.saturation(waterPhaseIdx));
} else if (spec.actual_name == "SGAS") {
value = getValue(fs.saturation(gasPhaseIdx));
} else if (name == "SOIL") {
} else if (spec.actual_name == "SOIL") {
value = getValue(fs.saturation(oilPhaseIdx));
} else if (name == "RS") {
} else if (spec.actual_name == "RS") {
value = getValue(fs.Rs());
value = unitSyst.from_si(UnitSystem::measure::gas_oil_ratio, value);
} else if (name == "RV") {
} else if (spec.actual_name == "RV") {
value = getValue(fs.Rv());
value = unitSyst.from_si(UnitSystem::measure::oil_gas_ratio, value);
} else if (name == "PERMX") {
} 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: " + name);
OPM_THROW(std::runtime_error, "Unknown per-cell feature: " + spec.actual_name);
}
value = spec.transform.apply(value);
return spec.scaler.scale(value);
Scalar transformed = spec.transform.apply(value);
Scalar scaled = spec.scaler.scale(transformed);
return scaled;
}
@@ -350,7 +322,7 @@ protected:
* output tensor does not match the expected feature layout.
*/
ML::Tensor<Evaluation>
constructOutputTensor(const ML::Tensor<Evaluation>& input, const HybridNewtonConfig& config, std::size_t n_cells)
constructOutputTensor(const ML::Tensor<Evaluation>& input, const HybridNewtonConfig& config)
{
const auto& features = config.output_features;
const int n_features = features.size();
@@ -358,7 +330,7 @@ protected:
ML::NNModel<Evaluation> model;
model.loadModel(config.model_path);
ML::Tensor<Evaluation> output(1, n_cells * n_features);
ML::Tensor<Evaluation> output(1, config.n_cells * n_features);
model.apply(input, output);
return output;
@@ -378,27 +350,23 @@ protected:
* and columns correspond to output features.
* \param config The HybridNewtonConfig specifying model path and output
* feature definitions.
* \param cell_indices The list of cells whose initial guesses are updated.
* \param n_cells The number of cells (should match cell_indices.size()).
*
* \throws std::runtime_error if an expected output feature is missing
* or if state update fails for any cell.
*/
void updateInitialGuess(
ML::Tensor<Evaluation>& output,
const HybridNewtonConfig& config,
const std::vector<int>& cell_indices
const HybridNewtonConfig& config
)
{
const std::size_t n_cells = cell_indices.size();
const auto& features = config.output_features;
FeatureFlags flags = flagFeatures(features);
const auto& unitSyst = simulator_.vanguard().schedule().getUnits();
for (std::size_t i = 0; i < n_cells; ++i) {
const int cell_idx = cell_indices[i];
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();
@@ -412,7 +380,7 @@ protected:
for (const auto& [name, spec] : features) {
auto scaled_value = getValue(output(feature_idx * n_cells + i));
auto scaled_value = getValue(output(feature_idx * config.n_cells + i));
// Inverse scaling
Scalar raw_value = spec.scaler.unscale(scaled_value);
@@ -478,9 +446,9 @@ protected:
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));
sw = std::max(0.0, std::min(sw, 1.0));
so = std::max(0.0, std::min(so, 1.0));
sg = std::max(0.0, std::min(sg, 1.0));
Scalar sum = sw + so + sg;
if (sum <= 1e-12) {
+67 -61
View File
@@ -109,12 +109,46 @@ 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;
// Construct config from a PropertyTree (single model)
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; });
@@ -125,12 +159,31 @@ public:
[&](const auto& p) { return p.first == name; });
}
std::vector<int> loadCellIndicesFromFile() const {
std::vector<int> indices;
std::ifstream cellFile(cell_indices_file);
if (!cellFile.is_open()) {
throw std::runtime_error("Cannot open cell indices file: " + cell_indices_file);
// In HybridNewtonConfig
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:
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;
@@ -140,17 +193,16 @@ public:
try { indices.push_back(std::stoi(line)); }
catch (...) {
throw std::runtime_error("Invalid cell index at line " + std::to_string(lineNumber) +
" in file " + cell_indices_file + ": " + line);
" in file " + filename + ": " + line);
}
}
if (indices.empty()) {
throw std::runtime_error("No valid cell indices found in file: " + cell_indices_file);
}
if (indices.empty())
throw std::runtime_error("No valid cell indices found in file: " + filename);
return indices;
}
private:
void parseFeatures(const PropertyTree& pt, const std::string& path,
std::vector<std::pair<std::string, FeatureSpec>>& features) {
auto subtreeOpt = pt.get_child_optional(path);
@@ -160,6 +212,7 @@ private:
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")) {
@@ -173,63 +226,16 @@ private:
} else {
spec.scaler.type = Scaler::Type::None;
}
} else {
spec.scaler.type = Scaler::Type::None;
}
spec.is_delta = (name.size() >= 6 && std::memcmp(name.data(), "DELTA_", 6) == 0);
spec.actual_name = spec.is_delta ? name.substr(6) : name;
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));
}
}
public:
/*!
* \brief Parse a JSON config file containing one or multiple HybridNewton configs.
* \param path Path to the JSON configuration file.
* \return A vector of HybridNewtonConfig objects, one per model.
*/
static std::vector<HybridNewtonConfig> parseHybridNewtonConfigs(const std::string& path) {
std::vector<HybridNewtonConfig> configs;
PropertyTree pt(path);
for (const auto& model_key : pt.get_child_keys()) {
const PropertyTree mt = pt.get_child(model_key);
HybridNewtonConfig cfg;
cfg.model_path = mt.get<std::string>("model_path", "");
if (cfg.model_path.empty()) {
throw std::runtime_error("Missing 'model_path' for model " + model_key + " in file: " + path);
}
cfg.cell_indices_file = mt.get<std::string>("cell_indices_file", "");
if (cfg.cell_indices_file.empty()) {
throw std::runtime_error("Missing 'cell_indices_file' for model " + model_key + " in file: " + path);
}
auto applyTimesOpt = mt.get_child_optional("apply_times");
if (!applyTimesOpt) {
throw std::runtime_error("Missing 'apply_times' for model " + model_key + " in file: " + path);
}
for (const auto& key : applyTimesOpt->get_child_keys()) {
cfg.apply_times.push_back(applyTimesOpt->get_child(key).get<double>(""));
}
if (cfg.apply_times.empty()) {
throw std::runtime_error("'apply_times' must contain at least one value for model " + model_key);
}
cfg.parseFeatures(mt, "features.inputs", cfg.input_features);
cfg.parseFeatures(mt, "features.outputs", cfg.output_features);
configs.push_back(std::move(cfg));
}
if (configs.empty()) {
throw std::runtime_error("No models found in HybridNewton config file: " + path);
}
return configs;
}
};
} // namespace Opm
-4
View File
@@ -1,7 +1,3 @@
requirements (same as opm-common/python/opm/ml/ml_tools/requirements.txt)
# Python v.3.9.0 and above
View File
Whitespace-only changes.
+104
View File
@@ -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.
---
Whitespace-only changes.
@@ -1,2 +1,201 @@
from opm.io.ecl_state import EclipseState
from opm.simulators import BlackOilSimulator
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()
+414
View File
@@ -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},
],
},
]
@@ -1,121 +0,0 @@
import json
import os
def write_config(models, model_dir, multi_mode=False, case_label=None):
"""
Write one or multiple model configs to JSON.
models: list of dicts, each dict must contain:
- input_features
- output_features
- model_path
- active_cells
- apply_time
- feature_engineering_input (optional)
- scaling_input (optional)
- feature_engineering_output (optional)
- scaling_output (optional)
- input_scaling_params (optional)
- output_scaling_params (optional)
multi_mode: if True, write all configs into a single JSON.
case_label: used as filename for multi-mode JSON (required if multi_mode=True).
"""
if multi_mode and not case_label:
raise ValueError("case_label must be provided for multi_mode=True")
os.makedirs(model_dir, exist_ok=True)
all_configs = {}
for idx, m in enumerate(models):
# Normalize input/output features
input_features = m["input_features"]
output_features = m["output_features"]
if isinstance(input_features, str):
input_features = [input_features]
if isinstance(output_features, str):
output_features = [output_features]
n_input = len(input_features)
n_output = len(output_features)
# Optional fields (safe defaults)
fe_input = m.get("feature_engineering_input") or [None]*n_input
fe_output = m.get("feature_engineering_output") or [None]*n_output
scaling_input = m.get("scaling_input") or [None]*n_input
scaling_output = m.get("scaling_output") or [None]*n_output
input_scaling_params = m.get("input_scaling_params") or [None]*n_input
output_scaling_params = m.get("output_scaling_params") or [None]*n_output
# Build input block
input_block = {}
for i, fname in enumerate(input_features):
feat_eng = fe_input[i] if i < len(fe_input) else None
scale = scaling_input[i] if i < len(scaling_input) else None
feature_dict = {
"feature_engineering": feat_eng.lower() if feat_eng and feat_eng.lower() != "none" else None,
"scaling": scale.lower() if scale and scale.lower() != "none" else None,
}
feature_dict = {k: v for k, v in feature_dict.items() if v is not None}
if i < len(input_scaling_params) and input_scaling_params[i] is not None:
feature_dict["scaling_params"] = input_scaling_params[i]
input_block[fname] = feature_dict
# Build output block
output_block = {}
for i, fname in enumerate(output_features):
feat_eng = fe_output[i] if i < len(fe_output) else None
scale = scaling_output[i] if i < len(scaling_output) else None
feature_dict = {
"feature_engineering": feat_eng.lower() if feat_eng and feat_eng.lower() != "none" else None,
"scaling": scale.lower() if scale and scale.lower() != "none" else None,
}
feature_dict = {k: v for k, v in feature_dict.items() if v is not None}
if i < len(output_scaling_params) and output_scaling_params[i] is not None:
feature_dict["scaling_params"] = output_scaling_params[i]
output_block[fname] = feature_dict
# Save active cells
model_base = os.path.splitext(os.path.basename(m["model_path"]))[0]
cells_file = os.path.join(model_dir, model_base + "_active_cells.txt")
with open(cells_file, "w") as f:
for cell in m["active_cells"]:
f.write(f"{cell}\n")
# If m["apply_time"] is a single float, wrap as dict
if isinstance(m["apply_time"], float):
apply_times_dict = {"0": m["apply_time"]}
elif isinstance(m["apply_time"], list):
apply_times_dict = {str(i): t for i, t in enumerate(m["apply_time"])}
else:
raise ValueError("apply_time must be float or list of floats")
# Build config dict
cfg = {
"model_path": m["model_path"],
"cell_indices_file": cells_file,
"apply_times": apply_times_dict,
"features": {
"inputs": input_block,
"outputs": output_block
}
}
all_configs[str(idx)] = cfg
# Determine JSON filename
if multi_mode:
if not case_label:
raise ValueError("case_label is required for multi_mode")
json_name = case_label + ".json"
else:
json_name = os.path.splitext(os.path.basename(models[0]["model_path"]))[0] + ".json"
json_path = os.path.join(model_dir, json_name)
# Always write a list of configs
with open(json_path, "w") as f:
json.dump(all_configs, f, indent=2)
return json_path
@@ -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