src/synthetic_fields.py

"""
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}