You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
332 lines
13 KiB
Python
332 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
import h5py
|
|
import joblib
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from src.common.config import Config
|
|
from src.data.curve_processing import resample_curve_to_features_with_time_and_mask
|
|
from src.data.preprocess import MaskedStandardScaler
|
|
from src.data.preprocess import preprocess_dataset
|
|
from src.evaluation.autofit_objective import (
|
|
prediction_curve_time_from_meta,
|
|
timed_dual_log_objective,
|
|
)
|
|
from src.training.train_forward import (
|
|
CurveStats,
|
|
LossConfig,
|
|
LossContext,
|
|
SampleReweightConfig,
|
|
compute_weighted_loss,
|
|
load_processed_dataset,
|
|
smooth_l1_per_sample,
|
|
)
|
|
def _write_minimal_raw_h5(
|
|
path: Path,
|
|
curve_time: np.ndarray,
|
|
*,
|
|
include_mask: bool,
|
|
) -> None:
|
|
n_samples, n_time_points = curve_time.shape
|
|
with h5py.File(path, "w") as f:
|
|
f.create_dataset("params", data=np.ones((n_samples, 7), dtype=np.float32))
|
|
f.create_dataset("schedule", data=np.zeros((n_samples, 4), dtype=np.float32))
|
|
f.create_dataset(
|
|
"curve",
|
|
data=np.ones((n_samples, 2 * n_time_points), dtype=np.float32),
|
|
)
|
|
f.create_dataset("curve_time", data=curve_time.astype(np.float32))
|
|
if include_mask:
|
|
f.create_dataset(
|
|
"curve_valid_mask",
|
|
data=np.ones((n_samples, n_time_points), dtype=np.uint8),
|
|
)
|
|
|
|
|
|
class FixedTimePipelineTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.cfg = Config(ROOT / "configs" / "data_gen_family_random_v2_q.yaml")
|
|
|
|
def test_all_curves_use_one_physical_time_grid(self) -> None:
|
|
grids = []
|
|
masks = []
|
|
for end_time in (12.0, 30.0, 72.0, 120.0, 140.0):
|
|
time = np.geomspace(1.6e-3, end_time, 240)
|
|
pressure = 10.0 + np.log1p(time)
|
|
derivative = 0.1 + np.sqrt(time)
|
|
features, grid, valid_mask = resample_curve_to_features_with_time_and_mask(
|
|
self.cfg,
|
|
time,
|
|
pressure,
|
|
derivative,
|
|
)
|
|
self.assertEqual(features.shape, (320,))
|
|
grids.append(grid)
|
|
masks.append(valid_mask)
|
|
|
|
for grid in grids[1:]:
|
|
np.testing.assert_array_equal(grid, grids[0])
|
|
self.assertAlmostEqual(float(grids[0][0]), 1.6e-3, places=8)
|
|
self.assertAlmostEqual(float(grids[0][-1]), 139.0, places=5)
|
|
self.assertLess(int(masks[0].sum()), int(masks[-1].sum()))
|
|
|
|
def test_fixed_grid_point_before_curve_start_is_invalid(self) -> None:
|
|
time = np.geomspace(1.61e-3, 139.0, 240)
|
|
_, grid, valid_mask = resample_curve_to_features_with_time_and_mask(
|
|
self.cfg,
|
|
time,
|
|
np.ones_like(time),
|
|
np.ones_like(time),
|
|
)
|
|
|
|
self.assertAlmostEqual(float(grid[0]), 1.6e-3, places=8)
|
|
self.assertEqual(int(valid_mask[0]), 0)
|
|
self.assertEqual(int(valid_mask[1]), 1)
|
|
|
|
def test_masked_scaler_ignores_placeholder_values(self) -> None:
|
|
rng = np.random.RandomState(7)
|
|
values = rng.normal(size=(12, 8))
|
|
mask = np.ones_like(values, dtype=bool)
|
|
mask[1:, 5:] = False
|
|
|
|
changed = values.copy()
|
|
changed[~mask] = 1.0e9
|
|
scaler_a = MaskedStandardScaler().fit(values, mask)
|
|
scaler_b = MaskedStandardScaler().fit(changed, mask)
|
|
|
|
np.testing.assert_allclose(scaler_a.mean_, scaler_b.mean_, atol=0.0, rtol=0.0)
|
|
np.testing.assert_allclose(scaler_a.scale_, scaler_b.scale_, atol=0.0, rtol=0.0)
|
|
|
|
def test_masked_loss_has_no_gradient_outside_coverage(self) -> None:
|
|
pred = torch.tensor([[0.2, 0.4, 3.0, 4.0]], dtype=torch.float32, requires_grad=True)
|
|
target = torch.zeros_like(pred)
|
|
mask = torch.tensor([[1.0, 1.0, 0.0, 0.0]], dtype=torch.float32)
|
|
loss = smooth_l1_per_sample(pred, target, beta=0.05, valid_mask=mask).mean()
|
|
loss.backward()
|
|
|
|
self.assertGreater(float(pred.grad[0, 0]), 0.0)
|
|
self.assertGreater(float(pred.grad[0, 1]), 0.0)
|
|
self.assertEqual(float(pred.grad[0, 2]), 0.0)
|
|
self.assertEqual(float(pred.grad[0, 3]), 0.0)
|
|
|
|
def test_composite_training_loss_respects_time_mask(self) -> None:
|
|
pred = torch.zeros((2, 320), dtype=torch.float32, requires_grad=True)
|
|
target = torch.ones_like(pred)
|
|
valid_mask = torch.zeros((2, 160), dtype=torch.float32)
|
|
valid_mask[:, :80] = 1.0
|
|
context = LossContext(
|
|
slices={"log_pressure": slice(0, 160), "log_derivative": slice(160, 320)},
|
|
curve_stats=CurveStats(
|
|
mean_raw=torch.zeros(320, dtype=torch.float32),
|
|
scale_raw=torch.ones(320, dtype=torch.float32),
|
|
),
|
|
loss_cfg=LossConfig(),
|
|
reweight_cfg=SampleReweightConfig(enabled=False),
|
|
)
|
|
losses = compute_weighted_loss(pred, target, valid_mask, context)
|
|
losses["loss"].backward()
|
|
|
|
self.assertTrue(torch.isfinite(losses["loss"]))
|
|
self.assertGreater(float(pred.grad[:, :80].abs().sum()), 0.0)
|
|
self.assertGreater(float(pred.grad[:, 160:240].abs().sum()), 0.0)
|
|
self.assertEqual(float(pred.grad[:, 80:160].abs().sum()), 0.0)
|
|
self.assertEqual(float(pred.grad[:, 240:].abs().sum()), 0.0)
|
|
|
|
def test_perfect_fit_composite_loss_has_finite_gradient(self) -> None:
|
|
pred = torch.zeros((1, 4), dtype=torch.float32, requires_grad=True)
|
|
target = torch.zeros_like(pred)
|
|
valid_mask = torch.ones((1, 2), dtype=torch.float32)
|
|
context = LossContext(
|
|
slices={"log_pressure": slice(0, 2), "log_derivative": slice(2, 4)},
|
|
curve_stats=CurveStats(
|
|
mean_raw=torch.zeros(4, dtype=torch.float32),
|
|
scale_raw=torch.ones(4, dtype=torch.float32),
|
|
),
|
|
loss_cfg=LossConfig(),
|
|
reweight_cfg=SampleReweightConfig(enabled=False),
|
|
)
|
|
|
|
losses = compute_weighted_loss(pred, target, valid_mask, context)
|
|
losses["loss"].backward()
|
|
|
|
self.assertTrue(torch.isfinite(losses["loss"]))
|
|
self.assertTrue(torch.isfinite(pred.grad).all())
|
|
self.assertEqual(float(pred.grad.abs().sum()), 0.0)
|
|
|
|
def test_surrogate_metadata_requires_fixed_time(self) -> None:
|
|
layout = {
|
|
"parts": [
|
|
{"name": "log_pressure", "start": 0, "end": 3},
|
|
{"name": "log_derivative", "start": 3, "end": 6},
|
|
]
|
|
}
|
|
with self.assertRaises(ValueError):
|
|
prediction_curve_time_from_meta({"curve_time_mode": "missing"}, layout)
|
|
with self.assertRaises(ValueError):
|
|
prediction_curve_time_from_meta({"curve_time_mode": "per_sample"}, layout)
|
|
|
|
result = prediction_curve_time_from_meta(
|
|
{
|
|
"curve_time_mode": "fixed",
|
|
"prediction_curve_time": [0.1, 1.0, 10.0],
|
|
},
|
|
layout,
|
|
)
|
|
np.testing.assert_array_equal(result, np.asarray([0.1, 1.0, 10.0]))
|
|
|
|
def test_fixed_grid_objective_aligns_on_physical_time(self) -> None:
|
|
prediction_time = np.asarray([0.1, 1.0, 2.0, 4.0], dtype=np.float64)
|
|
prediction_pressure = 2.0 + 3.0 * prediction_time
|
|
prediction_derivative = 1.0 + 0.5 * prediction_time
|
|
prediction_curve = np.concatenate(
|
|
[np.log(prediction_pressure), np.log(prediction_derivative)]
|
|
)
|
|
target_time = np.asarray([0.3, 0.8, 1.5, 2.3, 3.0], dtype=np.float64)
|
|
layout = {
|
|
"parts": [
|
|
{"name": "log_pressure", "start": 0, "end": 4},
|
|
{"name": "log_derivative", "start": 4, "end": 8},
|
|
]
|
|
}
|
|
|
|
objective = timed_dual_log_objective(
|
|
target_time=target_time,
|
|
target_pressure=2.0 + 3.0 * target_time,
|
|
target_derivative=1.0 + 0.5 * target_time,
|
|
pred_time=prediction_time,
|
|
pred_curve=prediction_curve,
|
|
curve_layout=layout,
|
|
)
|
|
|
|
self.assertLess(objective["dual_log_objective"], 1.0e-12)
|
|
|
|
def test_preprocess_preserves_fixed_grid_and_masks(self) -> None:
|
|
rng = np.random.RandomState(11)
|
|
n_groups = 10
|
|
n = n_groups * 5
|
|
grid = np.geomspace(1.6e-3, 139.0, 160).astype(np.float32)
|
|
params = np.column_stack(
|
|
[
|
|
rng.uniform(1.0e-3, 1.0, n),
|
|
rng.uniform(-2.0, 2.0, n),
|
|
rng.uniform(1.0e-4, 0.1, n),
|
|
rng.uniform(0.01, 0.2, n),
|
|
rng.uniform(2.0, 20.0, n),
|
|
rng.uniform(1.0e-4, 2.0e-3, n),
|
|
np.tile(np.arange(1, 6, dtype=np.float32), n_groups),
|
|
]
|
|
).astype(np.float32)
|
|
schedule = rng.normal(size=(n, 12)).astype(np.float32)
|
|
curve = rng.normal(size=(n, 320)).astype(np.float32)
|
|
mask = np.zeros((n, 160), dtype=np.uint8)
|
|
for row in range(n):
|
|
mask[row, : 80 + 20 * (row % 5)] = 1
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
h5_path = Path(temp_dir) / "fixed.h5"
|
|
pkl_path = Path(temp_dir) / "fixed.pkl"
|
|
with h5py.File(h5_path, "w") as f:
|
|
f.create_dataset("params", data=params)
|
|
f.create_dataset("schedule", data=schedule)
|
|
f.create_dataset("curve", data=curve)
|
|
f.create_dataset("curve_time", data=np.tile(grid, (n, 1)))
|
|
f.create_dataset("curve_valid_mask", data=mask)
|
|
f.create_dataset("group_id", data=np.repeat(np.arange(n_groups), 5))
|
|
f.attrs["param_names"] = np.asarray(
|
|
["k", "skin", "wellboreC", "phi", "h", "Cf", "solverType"],
|
|
dtype="S",
|
|
)
|
|
|
|
preprocess_dataset(h5_path, pkl_path, test_size=0.2, val_size=0.2, random_seed=3)
|
|
processed = joblib.load(pkl_path)
|
|
|
|
self.assertEqual(processed["meta"]["curve_time_mode"], "fixed")
|
|
self.assertEqual(processed["meta"]["curve_valid_mask_mode"], "explicit")
|
|
np.testing.assert_array_equal(
|
|
np.asarray(processed["meta"]["prediction_curve_time"], dtype=np.float32),
|
|
grid,
|
|
)
|
|
train_mask = processed["curve_valid_mask_train"].astype(bool)
|
|
full_train_mask = np.concatenate([train_mask, train_mask], axis=1)
|
|
self.assertTrue(np.all(processed["Y_curve_train"][~full_train_mask] == 0.0))
|
|
self.assertTrue(np.all(np.asarray(processed["scaler_curve"].n_samples_seen_) > 0))
|
|
|
|
def test_preprocess_keeps_240_dim_two_channel_curve(self) -> None:
|
|
fixed_grid = np.geomspace(1.6e-3, 139.0, 120)
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
h5_path = Path(temp_dir) / "two_channel_240.h5"
|
|
pkl_path = Path(temp_dir) / "two_channel_240.pkl"
|
|
_write_minimal_raw_h5(
|
|
h5_path,
|
|
np.tile(fixed_grid, (20, 1)),
|
|
include_mask=True,
|
|
)
|
|
|
|
preprocess_dataset(h5_path, pkl_path, random_seed=5)
|
|
processed = joblib.load(pkl_path)
|
|
|
|
self.assertEqual(processed["meta"]["curve_dim"], 240)
|
|
self.assertEqual(processed["meta"]["curve_layout"]["n_time_points"], 120)
|
|
self.assertFalse(processed["meta"]["dropped_legacy_slope"])
|
|
self.assertEqual(processed["Y_curve_train"].shape[1], 240)
|
|
self.assertEqual(processed["curve_valid_mask_train"].shape[1], 120)
|
|
np.testing.assert_allclose(
|
|
processed["meta"]["prediction_curve_time"],
|
|
fixed_grid,
|
|
rtol=1.0e-6,
|
|
atol=1.0e-10,
|
|
)
|
|
|
|
def test_preprocess_rejects_per_sample_time_grid(self) -> None:
|
|
fixed_grid = np.geomspace(1.6e-3, 139.0, 160)
|
|
other_grid = np.geomspace(1.6e-3, 120.0, 160)
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
h5_path = Path(temp_dir) / "per_sample.h5"
|
|
_write_minimal_raw_h5(
|
|
h5_path,
|
|
np.stack([fixed_grid, other_grid], axis=0),
|
|
include_mask=True,
|
|
)
|
|
|
|
with self.assertRaisesRegex(ValueError, "one fixed curve_time grid"):
|
|
preprocess_dataset(h5_path, Path(temp_dir) / "unused.pkl")
|
|
|
|
def test_preprocess_rejects_missing_curve_mask(self) -> None:
|
|
fixed_grid = np.geomspace(1.6e-3, 139.0, 160)
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
h5_path = Path(temp_dir) / "missing_mask.h5"
|
|
_write_minimal_raw_h5(
|
|
h5_path,
|
|
np.tile(fixed_grid, (2, 1)),
|
|
include_mask=False,
|
|
)
|
|
|
|
with self.assertRaisesRegex(ValueError, "requires curve_valid_mask"):
|
|
preprocess_dataset(h5_path, Path(temp_dir) / "unused.pkl")
|
|
|
|
def test_training_rejects_legacy_processed_dataset(self) -> None:
|
|
legacy = {}
|
|
for split in ("train", "val", "test"):
|
|
legacy[f"X_params_{split}"] = np.zeros((1, 1), dtype=np.float32)
|
|
legacy[f"X_schedule_{split}"] = np.zeros((1, 1), dtype=np.float32)
|
|
legacy[f"Y_curve_{split}"] = np.zeros((1, 4), dtype=np.float32)
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
path = Path(temp_dir) / "legacy.pkl"
|
|
joblib.dump(legacy, path)
|
|
with self.assertRaisesRegex(ValueError, "curve_time_mode='fixed'"):
|
|
load_processed_dataset(path)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|