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()