File size: 2,122 Bytes
ae73c7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """
Synthetic Well-like spatiotemporal field generator.
MOVED HERE this session from env.py, where it was bundled alongside
MultiStepPoincareEnv/WellStreamDataset purely by file-proximity, not by
any real shared contract. That coupling had a concrete cost: env.py
imports gymnasium unconditionally at module level, so ANY call to
get_synthetic_dataset() in data_real.py (which needs only this plain
torch.utils.data.Dataset, no RL machinery at all) pulled in a hard
gymnasium dependency it never used. A reported test run without
gymnasium installed failed here, and the failure was correctly
diagnosed as a missing dependency in the wrong place, not a logic bug in
provenance.py's contracts. Fixed by giving this class its own module
with its own (much smaller) dependency footprint -- consistent with this
project's own "one contract per module" convention, applied to a case
that convention was previously violated for no real reason.
"""
from __future__ import annotations
import torch
from torch.utils.data import Dataset
class SyntheticWellLike(Dataset):
def __init__(
self,
n_samples: int = 512,
n_steps: int = 12,
height: int = 32,
width: int = 32,
n_channels: int = 2,
noise: float = 0.15,
):
self.n_samples = n_samples
self.n_steps = n_steps
self.H = height
self.W = width
self.C = n_channels
self.data = torch.randn(n_samples, n_steps, n_channels, height, width) * noise
for i in range(n_samples):
t = torch.linspace(0, 1, n_steps).view(-1, 1, 1)
x = torch.linspace(-1, 1, width).view(1, 1, -1)
y = torch.linspace(-1, 1, height).view(1, -1, 1)
self.data[i, :, 0] += 0.9 * torch.sin(3 * x + 2.5 * t) * torch.cos(2 * y - 1.2 * t)
self.data[i, :, 1] += 0.7 * torch.cos(2.2 * x - 0.8 * t) * torch.sin(3.1 * y + 0.6 * t)
self.data[i, :, 0] += 0.25 * torch.sin(4 * (x - 0.7 * t))
def __len__(self):
return self.n_samples
def __getitem__(self, idx):
return {"fields": self.data[idx], "idx": idx}
|