diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..fccad86
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,33 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+.eggs/
+dist/
+build/
+*.so
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Training outputs (large files)
+outputs/
+*.pth
+
+# Temporary data files
+/tmp/
+*.log
+
+# OS
+.DS_Store
+Thumbs.db
+
+# IDE
+.idea/
+.vscode/
+*.swp
diff --git a/README.md b/README.md
index 5c14159..7b9075e 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,126 @@
-# Test
-a project for me to learn how to using github
+# U-Net for Seismic Data Interpolation & Reconstruction
+
+A PyTorch implementation of a **U-Net** network that reconstructs missing
+traces in 2-D seismic sections.  Missing traces are simulated by applying a
+random (or regular) subsampling mask during training; the network learns to
+fill in the gaps from the surrounding context.
+
+---
+
+## Overview
+
+| File | Purpose |
+|---|---|
+| `unet_model.py` | U-Net architecture (encoder → bottleneck → decoder with skip connections) |
+| `dataset.py` | `SeismicDataset` + mask generators (`random_mask`, `regular_mask`) |
+| `utils.py` | Loss function, SNR / PSNR metrics, visualisation, checkpoint helpers |
+| `train.py` | End-to-end training script (synthetic data built-in for quick demos) |
+| `predict.py` | Inference script – reconstruct and evaluate on new data |
+| `tests/` | Pytest unit tests for all components |
+| `requirements.txt` | Python dependencies |
+
+---
+
+## Quick Start
+
+### 1 – Install dependencies
+
+```bash
+pip install -r requirements.txt
+```
+
+### 2 – Train (synthetic data, no files needed)
+
+```bash
+python train.py --epochs 50 --batch_size 8
+```
+
+Checkpoints and a CSV training log are saved to `./outputs/`.
+
+### 3 – Train on your own data
+
+Your data should be a NumPy `.npy` file of shape `(N, T, X)` (N sections,
+T time samples, X traces) or `(T, X)` for a single section.
+SEG-Y support is available when `segyio` is installed.
+
+```bash
+python train.py \
+    --train_data data/train.npy \
+    --val_data   data/val.npy   \
+    --missing_ratio 0.5        \
+    --mask_type random         \
+    --epochs 100
+```
+
+### 4 – Reconstruct missing traces
+
+```bash
+python predict.py \
+    --checkpoint outputs/best_model.pth \
+    --input      data/incomplete.npy    \
+    --output     data/reconstructed.npy \
+    --plot
+```
+
+---
+
+## Model Architecture
+
+```
+Input  (B, 2, T, X)      ← seismic data  +  binary mask
+  │
+  ▼  Encoder
+  DoubleConv  64  ──────────────────────────────────────┐ skip 4
+  MaxPool → DoubleConv 128  ────────────────────────────┤ skip 3
+  MaxPool → DoubleConv 256  ────────────────────────────┤ skip 2
+  MaxPool → DoubleConv 512  ────────────────────────────┤ skip 1
+  MaxPool → DoubleConv 512 (bottleneck)
+  │
+  ▼  Decoder
+  Up → DoubleConv 256  ←─ skip 1
+  Up → DoubleConv 128  ←─ skip 2
+  Up → DoubleConv  64  ←─ skip 3
+  Up → DoubleConv  64  ←─ skip 4
+  │
+  OutConv 1×1
+  │
+Output (B, 1, T, X)      ← reconstructed seismic section
+```
+
+The input is a **2-channel** tensor: channel 0 is the incomplete seismic
+section (missing traces zeroed out) and channel 1 is the binary acquisition
+mask (1 = present, 0 = missing).
+
+---
+
+## Training Details
+
+| Hyperparameter | Default |
+|---|---|
+| Optimizer | AdamW |
+| Learning rate | 1 × 10⁻³ |
+| LR schedule | Cosine annealing |
+| Loss | 0.5 × L1 + 0.5 × MSE (missing traces up-weighted ×2) |
+| Missing ratio | 50 % (random mask) |
+
+---
+
+## Metrics
+
+* **SNR** (Signal-to-Noise Ratio) in dB – higher is better
+* **PSNR** (Peak SNR) in dB – higher is better
+
+---
+
+## Running Tests
+
+```bash
+python -m pytest tests/ -v
+```
+
+---
+
+## References
+
+* Ronneberger, O., Fischer, P., & Brox, T. (2015). *U-Net: Convolutional Networks for Biomedical Image Segmentation*. MICCAI 2015.
+* Liu, D., Wang, J., et al. (2022). *Seismic Data Reconstruction Using Deep Learning*. Geophysics.
diff --git a/dataset.py b/dataset.py
new file mode 100644
index 0000000..50593f6
--- /dev/null
+++ b/dataset.py
@@ -0,0 +1,292 @@
+"""
+Dataset utilities for seismic data interpolation / reconstruction.
+
+Supports loading seismic data from:
+  - NumPy .npy files  (shape: [time, trace] or [n_samples, time, trace])
+  - SEG-Y files via the `segyio` library (optional dependency)
+
+A random subsampling mask is applied during training to simulate missing
+traces (irregular or regular undersampling).
+"""
+
+from __future__ import annotations
+
+import os
+import random
+from typing import Callable, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+# ---------------------------------------------------------------------------
+# Mask generators
+# ---------------------------------------------------------------------------
+
+def random_mask(n_traces: int, missing_ratio: float, seed: int = None) -> np.ndarray:
+    """
+    Return a 1-D binary mask of shape (n_traces,).
+
+    1 = trace present, 0 = trace missing.
+
+    Parameters
+    ----------
+    n_traces : int
+        Total number of traces.
+    missing_ratio : float
+        Fraction of traces to remove (0 < missing_ratio < 1).
+    seed : int, optional
+        Random seed for reproducibility.
+    """
+    rng = np.random.default_rng(seed)
+    mask = np.ones(n_traces, dtype=np.float32)
+    n_missing = int(n_traces * missing_ratio)
+    idx = rng.choice(n_traces, size=n_missing, replace=False)
+    mask[idx] = 0.0
+    return mask
+
+
+def regular_mask(n_traces: int, keep_every: int = 2) -> np.ndarray:
+    """
+    Return a 1-D binary mask that keeps every *keep_every*-th trace.
+
+    Parameters
+    ----------
+    n_traces : int
+        Total number of traces.
+    keep_every : int
+        Decimation factor (2 = keep 50 %, 4 = keep 25 %, …).
+    """
+    mask = np.zeros(n_traces, dtype=np.float32)
+    mask[::keep_every] = 1.0
+    return mask
+
+
+# ---------------------------------------------------------------------------
+# Core dataset
+# ---------------------------------------------------------------------------
+
+class SeismicDataset(Dataset):
+    """
+    PyTorch Dataset for seismic interpolation / reconstruction.
+
+    Each sample is a 2-D seismic section (time × traces).  The dataset
+    applies a random subsampling mask at sample time, returning:
+
+        ``data``   – incomplete section (masked traces zeroed out)
+        ``mask``   – binary mask  (1 = present, 0 = missing)
+        ``target`` – full (ground-truth) section
+
+    Parameters
+    ----------
+    data_path : str or list of str
+        Path to a .npy file that contains a 2-D array of shape
+        ``(n_time, n_traces)`` or a 3-D array of shape
+        ``(n_samples, n_time, n_traces)``.  Alternatively, pass a list
+        of paths so that multiple files are concatenated along the
+        sample axis.
+    missing_ratio : float
+        Fraction of traces to randomly mask during training (0–1).
+    mask_type : {"random", "regular"}
+        Strategy used to generate the subsampling mask.
+    keep_every : int
+        Used only when ``mask_type == "regular"``.
+    patch_size : tuple of int, optional
+        If provided, each section is split into non-overlapping patches of
+        shape ``(patch_time, patch_traces)``.
+    normalize : bool
+        If True, each section is normalised to zero mean / unit variance.
+    transform : callable, optional
+        Additional transform applied to the sample dict after masking.
+    """
+
+    def __init__(
+        self,
+        data_path: Union[str, list],
+        missing_ratio: float = 0.5,
+        mask_type: str = "random",
+        keep_every: int = 2,
+        patch_size: Optional[Tuple[int, int]] = None,
+        normalize: bool = True,
+        transform: Optional[Callable] = None,
+    ):
+        super().__init__()
+        self.missing_ratio = missing_ratio
+        self.mask_type = mask_type
+        self.keep_every = keep_every
+        self.patch_size = patch_size
+        self.normalize = normalize
+        self.transform = transform
+
+        # Load data
+        paths = [data_path] if isinstance(data_path, str) else data_path
+        arrays = [self._load(p) for p in paths]
+        data = np.concatenate(arrays, axis=0)   # (N, T, X)
+        self.data = data.astype(np.float32)
+
+        # Optionally split into patches
+        if patch_size is not None:
+            self.data = self._patchify(self.data, patch_size)
+
+        self.n_samples, self.n_time, self.n_traces = self.data.shape
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _load(path: str) -> np.ndarray:
+        ext = os.path.splitext(path)[-1].lower()
+        if ext == ".npy":
+            arr = np.load(path)
+            if arr.ndim == 2:
+                arr = arr[np.newaxis]      # (1, T, X)
+            elif arr.ndim != 3:
+                raise ValueError(f"Expected 2-D or 3-D array, got shape {arr.shape}")
+            return arr
+        elif ext in (".segy", ".sgy"):
+            return SeismicDataset._load_segy(path)
+        else:
+            raise ValueError(f"Unsupported file format: {ext}")
+
+    @staticmethod
+    def _load_segy(path: str) -> np.ndarray:
+        try:
+            import segyio
+        except ImportError as exc:
+            raise ImportError(
+                "segyio is required to read SEG-Y files.  "
+                "Install it with:  pip install segyio"
+            ) from exc
+        with segyio.open(path, ignore_geometry=True) as f:
+            data = segyio.tools.collect(f.trace[:])   # (n_traces, n_time)
+        return data.T[np.newaxis]   # (1, n_time, n_traces)
+
+    @staticmethod
+    def _patchify(data: np.ndarray, patch_size: Tuple[int, int]) -> np.ndarray:
+        """Split each section into non-overlapping patches."""
+        n, t, x = data.shape
+        pt, px = patch_size
+        t_patches = t // pt
+        x_patches = x // px
+        patches = []
+        for i in range(n):
+            for ti in range(t_patches):
+                for xi in range(x_patches):
+                    patch = data[i,
+                                 ti * pt:(ti + 1) * pt,
+                                 xi * px:(xi + 1) * px]
+                    patches.append(patch)
+        return np.stack(patches, axis=0)   # (N_patches, pt, px)
+
+    # ------------------------------------------------------------------
+    # Dataset interface
+    # ------------------------------------------------------------------
+
+    def __len__(self) -> int:
+        return self.n_samples
+
+    def __getitem__(self, idx: int) -> dict:
+        section = self.data[idx].copy()   # (T, X)
+
+        # Normalise
+        if self.normalize:
+            std = section.std()
+            if std > 0:
+                section = (section - section.mean()) / std
+
+        # Build mask  (1-D over trace axis → broadcast to (T, X))
+        if self.mask_type == "random":
+            mask_1d = random_mask(self.n_traces, self.missing_ratio)
+        else:
+            mask_1d = regular_mask(self.n_traces, self.keep_every)
+
+        mask_2d = np.broadcast_to(mask_1d[np.newaxis, :], section.shape).copy()
+
+        # Apply mask
+        masked = section * mask_2d
+
+        # Convert to tensors and add channel dimension → (1, T, X)
+        target = torch.from_numpy(section[np.newaxis])
+        data_tensor = torch.from_numpy(masked[np.newaxis])
+        mask_tensor = torch.from_numpy(mask_2d[np.newaxis])
+
+        sample = {
+            "data": data_tensor,      # (1, T, X)  incomplete input
+            "mask": mask_tensor,      # (1, T, X)  binary mask
+            "target": target,         # (1, T, X)  ground truth
+        }
+
+        if self.transform is not None:
+            sample = self.transform(sample)
+
+        return sample
+
+
+# ---------------------------------------------------------------------------
+# Convenience factory
+# ---------------------------------------------------------------------------
+
+def build_loaders(
+    train_path: Union[str, list],
+    val_path: Union[str, list],
+    batch_size: int = 8,
+    missing_ratio: float = 0.5,
+    mask_type: str = "random",
+    patch_size: Optional[Tuple[int, int]] = None,
+    num_workers: int = 4,
+    **dataset_kwargs,
+):
+    """
+    Return (train_loader, val_loader) DataLoader pair.
+
+    Parameters
+    ----------
+    train_path, val_path : str or list of str
+        Paths to training / validation .npy or .segy files.
+    batch_size : int
+        Mini-batch size.
+    missing_ratio : float
+        Fraction of traces to mask.
+    mask_type : str
+        "random" or "regular".
+    patch_size : tuple, optional
+        Patch dimensions ``(n_time, n_traces)``.
+    num_workers : int
+        DataLoader worker processes.
+    **dataset_kwargs
+        Extra keyword arguments forwarded to :class:`SeismicDataset`.
+    """
+    from torch.utils.data import DataLoader
+
+    train_ds = SeismicDataset(
+        train_path,
+        missing_ratio=missing_ratio,
+        mask_type=mask_type,
+        patch_size=patch_size,
+        **dataset_kwargs,
+    )
+    val_ds = SeismicDataset(
+        val_path,
+        missing_ratio=missing_ratio,
+        mask_type=mask_type,
+        patch_size=patch_size,
+        **dataset_kwargs,
+    )
+
+    train_loader = DataLoader(
+        train_ds,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=True,
+    )
+    val_loader = DataLoader(
+        val_ds,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=True,
+    )
+    return train_loader, val_loader
diff --git a/predict.py b/predict.py
new file mode 100644
index 0000000..ddfc346
--- /dev/null
+++ b/predict.py
@@ -0,0 +1,229 @@
+"""
+Inference script: apply a trained U-Net to reconstruct missing seismic traces.
+
+Usage
+-----
+    python predict.py --checkpoint outputs/best_model.pth \
+                      --input path/to/incomplete.npy \
+                      --output path/to/reconstructed.npy \
+                      --mask_path path/to/mask.npy        # optional
+
+If ``--mask_path`` is not supplied, a random mask is generated using
+``--missing_ratio`` (default 0.5).
+
+The script can also visualise results with ``--plot``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+import numpy as np
+import torch
+
+from unet_model import UNet
+from utils import load_checkpoint, compute_metrics, plot_comparison
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def pad_to_divisible(arr: np.ndarray, divisor: int = 16) -> tuple[np.ndarray, tuple]:
+    """
+    Pad a 2-D array (T, X) so that both dimensions are divisible by *divisor*.
+
+    Returns the padded array and the original shape for later cropping.
+    """
+    t, x = arr.shape
+    pt = (divisor - t % divisor) % divisor
+    px = (divisor - x % divisor) % divisor
+    padded = np.pad(arr, ((0, pt), (0, px)), mode="reflect")
+    return padded, (t, x)
+
+
+def crop_to_original(arr: np.ndarray, original_shape: tuple) -> np.ndarray:
+    """Remove zero-padding added by :func:`pad_to_divisible`."""
+    t, x = original_shape
+    return arr[:t, :x]
+
+
+# ---------------------------------------------------------------------------
+# Reconstruct a single 2-D section
+# ---------------------------------------------------------------------------
+
+@torch.no_grad()
+def reconstruct(
+    model: torch.nn.Module,
+    data: np.ndarray,
+    mask: np.ndarray,
+    device: torch.device,
+    normalize: bool = True,
+) -> np.ndarray:
+    """
+    Reconstruct a single seismic section.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+        Trained U-Net.
+    data : ndarray  (T, X)
+        Incomplete seismic section (missing traces already zeroed).
+    mask : ndarray  (T, X)
+        Binary mask – 1 = present, 0 = missing.
+    device : torch.device
+    normalize : bool
+        Normalise the section before inference (mirrors training behaviour).
+
+    Returns
+    -------
+    ndarray  (T, X)
+        Reconstructed seismic section (in the original amplitude scale
+        when ``normalize=True``).
+    """
+    # Remember original statistics for de-normalisation
+    mean_ = data.mean()
+    std_ = data.std()
+
+    if normalize:
+        if std_ > 0:
+            data_n = (data - mean_) / std_
+        else:
+            data_n = data.copy()
+    else:
+        data_n = data.copy()
+
+    # Pad so spatial dims are divisible by 16 (4 pooling layers × 2)
+    data_padded, orig_shape = pad_to_divisible(data_n)
+    mask_padded, _ = pad_to_divisible(mask.astype(np.float32))
+
+    # Build input tensor: (1, 2, T, X)
+    data_t = torch.from_numpy(data_padded[np.newaxis, np.newaxis]).float().to(device)
+    mask_t = torch.from_numpy(mask_padded[np.newaxis, np.newaxis]).float().to(device)
+    net_in = torch.cat([data_t, mask_t], dim=1)
+
+    model.eval()
+    pred = model(net_in)   # (1, 1, T_pad, X_pad)
+
+    pred_np = pred.squeeze().cpu().numpy()   # (T_pad, X_pad)
+    pred_np = crop_to_original(pred_np, orig_shape)
+
+    # De-normalise
+    if normalize and std_ > 0:
+        pred_np = pred_np * std_ + mean_
+
+    return pred_np
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="Reconstruct missing seismic traces with a trained U-Net"
+    )
+    p.add_argument("--checkpoint", required=True,
+                   help="Path to model checkpoint (.pth)")
+    p.add_argument("--input", required=True,
+                   help="Path to input .npy file  (shape: (T,X) or (N,T,X))")
+    p.add_argument("--output", default="reconstructed.npy",
+                   help="Path to save reconstructed data (.npy)")
+    p.add_argument("--mask_path", default=None,
+                   help="Path to binary mask .npy file (same shape as input). "
+                        "If omitted, a random mask is generated.")
+    p.add_argument("--missing_ratio", type=float, default=0.5,
+                   help="Missing-trace ratio when --mask_path is not provided")
+    p.add_argument("--mask_type", choices=["random", "regular"], default="random")
+    p.add_argument("--keep_every", type=int, default=2,
+                   help="Trace decimation factor for regular mask")
+    p.add_argument("--base_features", type=int, default=32)
+    p.add_argument("--no_bilinear", action="store_true", default=False,
+                   help="Use transposed convolutions instead of bilinear upsampling")
+    p.add_argument("--no_normalize", action="store_true", default=False,
+                   help="Skip per-section amplitude normalisation")
+    p.add_argument("--plot", action="store_true", default=False,
+                   help="Save comparison figures to --output_dir/figures/")
+    p.add_argument("--output_dir", default="./outputs")
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    # ---- Load model ---------------------------------------------------------
+    model = UNet(
+        in_channels=2,
+        out_channels=1,
+        base_features=args.base_features,
+        bilinear=not args.no_bilinear,
+    ).to(device)
+
+    ckpt = load_checkpoint(model, optimizer=None, checkpoint_path=args.checkpoint, device=device)
+    print(f"Loaded checkpoint from epoch {ckpt.get('epoch', '?')}")
+
+    # ---- Load input data ----------------------------------------------------
+    raw = np.load(args.input).astype(np.float32)
+    if raw.ndim == 2:
+        raw = raw[np.newaxis]          # (1, T, X)
+    n_samples, n_time, n_traces = raw.shape
+    print(f"Input shape: {raw.shape}")
+
+    # ---- Load or generate masks ---------------------------------------------
+    if args.mask_path is not None:
+        masks = np.load(args.mask_path).astype(np.float32)
+        if masks.ndim == 2:
+            masks = np.broadcast_to(masks[np.newaxis], raw.shape).copy()
+    else:
+        from dataset import random_mask, regular_mask
+        masks = np.ones_like(raw)
+        for i in range(n_samples):
+            if args.mask_type == "random":
+                m1d = random_mask(n_traces, args.missing_ratio, seed=i)
+            else:
+                m1d = regular_mask(n_traces, args.keep_every)
+            masks[i] = np.broadcast_to(m1d[np.newaxis, :], (n_time, n_traces))
+
+    # ---- Reconstruct --------------------------------------------------------
+    reconstructed = np.zeros_like(raw)
+    for i in range(n_samples):
+        section = raw[i] * masks[i]   # apply mask
+        pred = reconstruct(model, section, masks[i], device, normalize=not args.no_normalize)
+        reconstructed[i] = pred
+        if (i + 1) % max(1, n_samples // 10) == 0:
+            print(f"  Reconstructed {i + 1}/{n_samples} …")
+
+    # ---- Save results -------------------------------------------------------
+    np.save(args.output, reconstructed)
+    print(f"Saved reconstructed data → {args.output}")
+
+    # ---- Optional metrics ---------------------------------------------------
+    from utils import signal_to_noise_ratio, peak_signal_to_noise_ratio
+    snrs = [signal_to_noise_ratio(reconstructed[i], raw[i]) for i in range(n_samples)]
+    psnrs = [peak_signal_to_noise_ratio(reconstructed[i], raw[i]) for i in range(n_samples)]
+    print(f"Mean SNR  : {np.mean(snrs):.2f} dB")
+    print(f"Mean PSNR : {np.mean(psnrs):.2f} dB")
+
+    # ---- Optional visualisation --------------------------------------------
+    if args.plot:
+        fig_dir = os.path.join(args.output_dir, "figures")
+        os.makedirs(fig_dir, exist_ok=True)
+        n_plot = min(5, n_samples)
+        for i in range(n_plot):
+            incomplete = raw[i] * masks[i]
+            plot_comparison(
+                incomplete,
+                reconstructed[i],
+                raw[i],
+                title=f"Sample {i}  |  SNR={snrs[i]:.1f} dB",
+                save_path=os.path.join(fig_dir, f"sample_{i:04d}.png"),
+            )
+        print(f"Figures saved to {fig_dir}/")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..d1b73ef
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,6 @@
+torch>=1.13.0
+numpy>=1.23.0
+matplotlib>=3.5.0
+tqdm>=4.64.0
+pytest>=7.0.0
+segyio>=1.9.0; platform_system != "Windows"   # optional: only needed for SEG-Y support
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_seismic_unet.py b/tests/test_seismic_unet.py
new file mode 100644
index 0000000..da7f709
--- /dev/null
+++ b/tests/test_seismic_unet.py
@@ -0,0 +1,289 @@
+"""
+Tests for the U-Net seismic interpolation project.
+
+Run with:
+    python -m pytest tests/ -v
+
+These tests are designed to be fast (CPU-only, small tensor sizes) and cover:
+  - U-Net forward pass and output shape
+  - Mask generators
+  - Dataset loading
+  - Loss functions
+  - Metric calculations
+  - Predict-time padding helpers
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from unet_model import UNet, DoubleConv, Down, Up, OutConv
+from dataset import SeismicDataset, random_mask, regular_mask
+from utils import (
+    ReconstructionLoss,
+    signal_to_noise_ratio,
+    peak_signal_to_noise_ratio,
+    compute_metrics,
+)
+from predict import pad_to_divisible, crop_to_original
+
+
+# ---------------------------------------------------------------------------
+# U-Net model tests
+# ---------------------------------------------------------------------------
+
+class TestUNet:
+    """Tests for the U-Net architecture."""
+
+    @pytest.mark.parametrize("in_ch,out_ch", [(1, 1), (2, 1)])
+    def test_output_shape(self, in_ch, out_ch):
+        model = UNet(in_channels=in_ch, out_channels=out_ch, base_features=8, bilinear=True)
+        x = torch.randn(2, in_ch, 64, 64)
+        y = model(x)
+        assert y.shape == (2, out_ch, 64, 64), f"Expected (2,{out_ch},64,64), got {y.shape}"
+
+    def test_non_square_input(self):
+        """Model should handle non-square spatial dimensions."""
+        model = UNet(in_channels=2, out_channels=1, base_features=8)
+        x = torch.randn(1, 2, 96, 64)
+        y = model(x)
+        assert y.shape == (1, 1, 96, 64)
+
+    def test_bilinear_false(self):
+        """Transposed-convolution upsampling path."""
+        model = UNet(in_channels=2, out_channels=1, base_features=8, bilinear=False)
+        x = torch.randn(1, 2, 64, 64)
+        y = model(x)
+        assert y.shape == (1, 1, 64, 64)
+
+    def test_no_nan_in_output(self):
+        model = UNet(in_channels=2, out_channels=1, base_features=8)
+        x = torch.randn(2, 2, 64, 64)
+        y = model(x)
+        assert not torch.isnan(y).any(), "Output contains NaN values"
+
+    def test_gradients_flow(self):
+        """Ensure backward pass does not raise and gradients are non-zero."""
+        model = UNet(in_channels=2, out_channels=1, base_features=8)
+        x = torch.randn(1, 2, 64, 64, requires_grad=False)
+        y = model(x)
+        loss = y.mean()
+        loss.backward()
+        grad_sum = sum(p.grad.abs().sum().item()
+                       for p in model.parameters() if p.grad is not None)
+        assert grad_sum > 0, "Gradients are all zero"
+
+    def test_double_conv_shape(self):
+        layer = DoubleConv(4, 8)
+        x = torch.randn(2, 4, 32, 32)
+        assert layer(x).shape == (2, 8, 32, 32)
+
+    def test_down_shape(self):
+        layer = Down(8, 16)
+        x = torch.randn(2, 8, 32, 32)
+        assert layer(x).shape == (2, 16, 16, 16)
+
+    def test_out_conv_shape(self):
+        layer = OutConv(16, 1)
+        x = torch.randn(2, 16, 32, 32)
+        assert layer(x).shape == (2, 1, 32, 32)
+
+
+# ---------------------------------------------------------------------------
+# Mask generator tests
+# ---------------------------------------------------------------------------
+
+class TestMaskGenerators:
+    def test_random_mask_shape(self):
+        m = random_mask(100, 0.5)
+        assert m.shape == (100,)
+
+    def test_random_mask_missing_ratio(self):
+        m = random_mask(100, 0.3, seed=0)
+        assert abs(m.sum() - 70) <= 1   # ~70 present
+
+    def test_random_mask_all_ones_zero_ratio(self):
+        m = random_mask(50, 0.0)
+        assert m.sum() == 50
+
+    def test_random_mask_reproducibility(self):
+        m1 = random_mask(100, 0.4, seed=42)
+        m2 = random_mask(100, 0.4, seed=42)
+        np.testing.assert_array_equal(m1, m2)
+
+    def test_regular_mask_pattern(self):
+        m = regular_mask(10, keep_every=2)
+        expected = np.array([1, 0, 1, 0, 1, 0, 1, 0, 1, 0], dtype=np.float32)
+        np.testing.assert_array_equal(m, expected)
+
+    def test_regular_mask_keep_every_1(self):
+        m = regular_mask(8, keep_every=1)
+        assert m.sum() == 8
+
+
+# ---------------------------------------------------------------------------
+# Dataset tests
+# ---------------------------------------------------------------------------
+
+@pytest.fixture
+def synthetic_npy(tmp_path):
+    """Write a small synthetic seismic .npy to tmp_path and return its path."""
+    rng = np.random.default_rng(42)
+    data = rng.standard_normal((20, 64, 64)).astype(np.float32)
+    p = str(tmp_path / "data.npy")
+    np.save(p, data)
+    return p
+
+
+class TestSeismicDataset:
+    def test_length(self, synthetic_npy):
+        ds = SeismicDataset(synthetic_npy, missing_ratio=0.5)
+        assert len(ds) == 20
+
+    def test_sample_keys(self, synthetic_npy):
+        ds = SeismicDataset(synthetic_npy, missing_ratio=0.5)
+        sample = ds[0]
+        assert set(sample.keys()) == {"data", "mask", "target"}
+
+    def test_sample_shapes(self, synthetic_npy):
+        ds = SeismicDataset(synthetic_npy, missing_ratio=0.5)
+        sample = ds[0]
+        assert sample["data"].shape == (1, 64, 64)
+        assert sample["mask"].shape == (1, 64, 64)
+        assert sample["target"].shape == (1, 64, 64)
+
+    def test_mask_values(self, synthetic_npy):
+        ds = SeismicDataset(synthetic_npy, missing_ratio=0.5)
+        sample = ds[0]
+        unique = torch.unique(sample["mask"])
+        assert set(unique.tolist()).issubset({0.0, 1.0}), "Mask contains values other than 0 and 1"
+
+    def test_masked_traces_are_zero(self, synthetic_npy):
+        ds = SeismicDataset(synthetic_npy, missing_ratio=0.5)
+        sample = ds[0]
+        # Where mask is 0 the data should be exactly 0
+        missing = sample["mask"] == 0
+        assert (sample["data"][missing] == 0).all()
+
+    def test_regular_mask(self, synthetic_npy):
+        ds = SeismicDataset(synthetic_npy, mask_type="regular", keep_every=2)
+        sample = ds[0]
+        unique = torch.unique(sample["mask"])
+        assert set(unique.tolist()).issubset({0.0, 1.0})
+
+    def test_2d_npy_input(self, tmp_path):
+        """A 2-D .npy file should be treated as a single sample."""
+        data = np.random.randn(64, 64).astype(np.float32)
+        p = str(tmp_path / "single.npy")
+        np.save(p, data)
+        ds = SeismicDataset(p, missing_ratio=0.3)
+        assert len(ds) == 1
+
+    def test_patch_size(self, tmp_path):
+        data = np.random.randn(4, 128, 128).astype(np.float32)
+        p = str(tmp_path / "large.npy")
+        np.save(p, data)
+        ds = SeismicDataset(p, patch_size=(64, 64))
+        # 4 sections × (128/64)^2 = 4 × 4 = 16 patches
+        assert len(ds) == 16
+        assert ds[0]["data"].shape == (1, 64, 64)
+
+
+# ---------------------------------------------------------------------------
+# Loss function tests
+# ---------------------------------------------------------------------------
+
+class TestReconstructionLoss:
+    def test_zero_loss_for_perfect_prediction(self):
+        crit = ReconstructionLoss()
+        t = torch.randn(2, 1, 32, 32)
+        loss = crit(t, t)
+        assert loss.item() < 1e-6
+
+    def test_loss_decreases_toward_target(self):
+        crit = ReconstructionLoss()
+        t = torch.ones(2, 1, 32, 32)
+        far = torch.zeros(2, 1, 32, 32)
+        close = torch.full((2, 1, 32, 32), 0.9)
+        assert crit(close, t).item() < crit(far, t).item()
+
+    def test_loss_with_mask(self):
+        crit = ReconstructionLoss(missing_weight=2.0)
+        t = torch.ones(2, 1, 32, 32)
+        p = torch.zeros(2, 1, 32, 32)
+        mask = torch.zeros(2, 1, 32, 32)          # all missing
+        mask_ones = torch.ones(2, 1, 32, 32)      # all present
+        # Higher weight on missing traces → larger loss
+        loss_missing = crit(p, t, mask)
+        loss_present = crit(p, t, mask_ones)
+        assert loss_missing.item() > loss_present.item()
+
+    def test_loss_is_tensor(self):
+        crit = ReconstructionLoss()
+        loss = crit(torch.randn(1, 1, 16, 16), torch.randn(1, 1, 16, 16))
+        assert isinstance(loss, torch.Tensor)
+        assert loss.ndim == 0   # scalar
+
+
+# ---------------------------------------------------------------------------
+# Metrics tests
+# ---------------------------------------------------------------------------
+
+class TestMetrics:
+    def test_snr_perfect(self):
+        a = np.ones((32, 32))
+        assert signal_to_noise_ratio(a, a) == float("inf")
+
+    def test_snr_positive_for_good_reconstruction(self):
+        t = np.random.randn(32, 32)
+        p = t + 0.01 * np.random.randn(32, 32)
+        assert signal_to_noise_ratio(p, t) > 10
+
+    def test_psnr_perfect(self):
+        a = np.random.randn(32, 32)
+        assert peak_signal_to_noise_ratio(a, a) == float("inf")
+
+    def test_compute_metrics_keys(self):
+        pred = torch.randn(2, 1, 32, 32)
+        target = torch.randn(2, 1, 32, 32)
+        m = compute_metrics(pred, target)
+        assert "snr" in m and "psnr" in m
+
+    def test_compute_metrics_returns_floats(self):
+        pred = torch.randn(2, 1, 32, 32)
+        target = pred + 0.01 * torch.randn_like(pred)
+        m = compute_metrics(pred, target)
+        assert isinstance(m["snr"], float)
+        assert isinstance(m["psnr"], float)
+
+
+# ---------------------------------------------------------------------------
+# Padding helpers tests (predict.py)
+# ---------------------------------------------------------------------------
+
+class TestPaddingHelpers:
+    @pytest.mark.parametrize("t,x,div", [
+        (100, 100, 16),
+        (128, 128, 16),
+        (63, 77, 16),
+        (1, 1, 16),
+    ])
+    def test_padded_dims_divisible(self, t, x, div):
+        arr = np.zeros((t, x))
+        padded, _ = pad_to_divisible(arr, divisor=div)
+        assert padded.shape[0] % div == 0
+        assert padded.shape[1] % div == 0
+
+    def test_round_trip(self):
+        arr = np.random.randn(100, 77).astype(np.float32)
+        padded, orig_shape = pad_to_divisible(arr, divisor=16)
+        recovered = crop_to_original(padded, orig_shape)
+        np.testing.assert_array_equal(recovered, arr)
+
+    def test_already_divisible(self):
+        arr = np.zeros((128, 64))
+        padded, orig = pad_to_divisible(arr, divisor=16)
+        assert padded.shape == (128, 64)
+        assert orig == (128, 64)
diff --git a/train.py b/train.py
new file mode 100644
index 0000000..4f33e0a
--- /dev/null
+++ b/train.py
@@ -0,0 +1,310 @@
+"""
+Training script for the U-Net seismic interpolation / reconstruction model.
+
+Usage
+-----
+Train with default options (uses synthetic data when no paths are given):
+
+    python train.py
+
+Train on real data:
+
+    python train.py --train_data path/to/train.npy --val_data path/to/val.npy
+
+See ``python train.py --help`` for all options.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+
+import numpy as np
+import torch
+import torch.optim as optim
+from torch.utils.data import DataLoader
+
+from dataset import SeismicDataset, build_loaders
+from unet_model import UNet
+from utils import ReconstructionLoss, compute_metrics, save_checkpoint, set_seed
+
+
+# ---------------------------------------------------------------------------
+# Synthetic data generator (for quick demos / smoke tests)
+# ---------------------------------------------------------------------------
+
+def _make_synthetic_data(n_samples: int = 200, n_time: int = 128, n_traces: int = 128) -> str:
+    """
+    Generate a small synthetic seismic dataset and save it to a temp .npy file.
+
+    Each section is a superposition of dipping plane waves with random slopes
+    and amplitudes – a simple but representative seismic model.
+
+    Returns the path to the saved .npy file.
+    """
+    rng = np.random.default_rng(0)
+    t = np.arange(n_time)
+    x = np.arange(n_traces)
+    data = np.zeros((n_samples, n_time, n_traces), dtype=np.float32)
+
+    for i in range(n_samples):
+        n_events = rng.integers(2, 6)
+        for _ in range(n_events):
+            slope = rng.uniform(-0.5, 0.5)   # samples / trace
+            t0 = rng.integers(10, n_time - 10)
+            amp = rng.uniform(0.5, 1.5) * rng.choice([-1, 1])
+            freq = rng.uniform(0.05, 0.2)    # cycles / sample
+            tt = t0 + slope * x              # (n_traces,)  broadcast below
+            envelope = np.exp(-0.5 * ((t[:, None] - tt[None, :]) / 8) ** 2)
+            wavelet = np.sin(2 * np.pi * freq * (t[:, None] - tt[None, :]))
+            data[i] += amp * (envelope * wavelet).astype(np.float32)
+
+    path = "/tmp/seismic_synthetic.npy"
+    np.save(path, data)
+    return path
+
+
+# ---------------------------------------------------------------------------
+# Training loop
+# ---------------------------------------------------------------------------
+
+def train_one_epoch(
+    model: torch.nn.Module,
+    loader: DataLoader,
+    optimizer: torch.optim.Optimizer,
+    criterion: ReconstructionLoss,
+    device: torch.device,
+    scaler,
+) -> dict:
+    model.train()
+    total_loss = 0.0
+    total_snr = 0.0
+    total_psnr = 0.0
+
+    for batch in loader:
+        data = batch["data"].to(device)    # (B, 1, T, X)  incomplete
+        mask = batch["mask"].to(device)    # (B, 1, T, X)
+        target = batch["target"].to(device)
+
+        # Network input: concatenate data + mask along channel axis → (B, 2, T, X)
+        net_input = torch.cat([data, mask], dim=1)
+
+        optimizer.zero_grad()
+
+        if scaler is not None:
+            with torch.cuda.amp.autocast():
+                pred = model(net_input)
+                loss = criterion(pred, target, mask)
+            scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            pred = model(net_input)
+            loss = criterion(pred, target, mask)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+
+        metrics = compute_metrics(pred, target)
+        total_loss += loss.item()
+        total_snr += metrics["snr"]
+        total_psnr += metrics["psnr"]
+
+    n = len(loader)
+    return {
+        "loss": total_loss / n,
+        "snr": total_snr / n,
+        "psnr": total_psnr / n,
+    }
+
+
+@torch.no_grad()
+def validate(
+    model: torch.nn.Module,
+    loader: DataLoader,
+    criterion: ReconstructionLoss,
+    device: torch.device,
+) -> dict:
+    model.eval()
+    total_loss = 0.0
+    total_snr = 0.0
+    total_psnr = 0.0
+
+    for batch in loader:
+        data = batch["data"].to(device)
+        mask = batch["mask"].to(device)
+        target = batch["target"].to(device)
+
+        net_input = torch.cat([data, mask], dim=1)
+        pred = model(net_input)
+        loss = criterion(pred, target, mask)
+
+        metrics = compute_metrics(pred, target)
+        total_loss += loss.item()
+        total_snr += metrics["snr"]
+        total_psnr += metrics["psnr"]
+
+    n = len(loader)
+    return {
+        "loss": total_loss / n,
+        "snr": total_snr / n,
+        "psnr": total_psnr / n,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="Train a U-Net for seismic data interpolation / reconstruction"
+    )
+    # Data
+    p.add_argument("--train_data", default=None,
+                   help="Path to training .npy or .segy file (or comma-separated list)")
+    p.add_argument("--val_data", default=None,
+                   help="Path to validation .npy or .segy file (or comma-separated list)")
+    p.add_argument("--missing_ratio", type=float, default=0.5,
+                   help="Fraction of traces to remove (0–1)")
+    p.add_argument("--mask_type", choices=["random", "regular"], default="random")
+    p.add_argument("--patch_size", type=int, nargs=2, default=None, metavar=("T", "X"),
+                   help="Patch height and width, e.g. --patch_size 128 128")
+
+    # Model
+    p.add_argument("--base_features", type=int, default=32,
+                   help="Base feature maps in the first U-Net encoder stage")
+    p.add_argument("--no_bilinear", action="store_true", default=False,
+                   help="Use transposed convolutions instead of bilinear upsampling")
+
+    # Training
+    p.add_argument("--epochs", type=int, default=50)
+    p.add_argument("--batch_size", type=int, default=8)
+    p.add_argument("--lr", type=float, default=1e-3)
+    p.add_argument("--weight_decay", type=float, default=1e-4)
+    p.add_argument("--amp", action="store_true", default=False,
+                   help="Use automatic mixed precision (AMP) training")
+    p.add_argument("--num_workers", type=int, default=4)
+
+    # Output
+    p.add_argument("--output_dir", default="./outputs",
+                   help="Directory for checkpoints and logs")
+    p.add_argument("--save_every", type=int, default=10,
+                   help="Save a checkpoint every N epochs")
+    p.add_argument("--seed", type=int, default=42)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    set_seed(args.seed)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # ---- Data ---------------------------------------------------------------
+    if args.train_data is None or args.val_data is None:
+        print("No data paths provided – generating synthetic seismic data …")
+        syn_path = _make_synthetic_data()
+        # Use 80 / 20 split of the synthetic data
+        data = np.load(syn_path)
+        n = len(data)
+        split = int(0.8 * n)
+        np.save("/tmp/seismic_train.npy", data[:split])
+        np.save("/tmp/seismic_val.npy", data[split:])
+        train_path = "/tmp/seismic_train.npy"
+        val_path = "/tmp/seismic_val.npy"
+    else:
+        train_path = args.train_data.split(",") if "," in args.train_data else args.train_data
+        val_path = args.val_data.split(",") if "," in args.val_data else args.val_data
+
+    patch_size = tuple(args.patch_size) if args.patch_size is not None else None
+
+    train_loader, val_loader = build_loaders(
+        train_path,
+        val_path,
+        batch_size=args.batch_size,
+        missing_ratio=args.missing_ratio,
+        mask_type=args.mask_type,
+        patch_size=patch_size,
+        num_workers=args.num_workers,
+    )
+    print(f"Training samples : {len(train_loader.dataset)}")
+    print(f"Validation samples: {len(val_loader.dataset)}")
+
+    # ---- Model --------------------------------------------------------------
+    model = UNet(
+        in_channels=2,
+        out_channels=1,
+        base_features=args.base_features,
+        bilinear=not args.no_bilinear,
+    ).to(device)
+
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"Trainable parameters: {n_params:,}")
+
+    # ---- Optimiser & loss ---------------------------------------------------
+    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
+    criterion = ReconstructionLoss()
+    scaler = torch.cuda.amp.GradScaler() if args.amp and torch.cuda.is_available() else None
+
+    # ---- Training loop ------------------------------------------------------
+    best_val_loss = float("inf")
+    log_rows = []
+
+    for epoch in range(1, args.epochs + 1):
+        t0 = time.time()
+        train_stats = train_one_epoch(model, train_loader, optimizer, criterion, device, scaler)
+        val_stats = validate(model, val_loader, criterion, device)
+        scheduler.step()
+
+        elapsed = time.time() - t0
+        print(
+            f"Epoch {epoch:3d}/{args.epochs} | "
+            f"train loss={train_stats['loss']:.4f}  SNR={train_stats['snr']:.2f} dB | "
+            f"val loss={val_stats['loss']:.4f}  SNR={val_stats['snr']:.2f} dB | "
+            f"{elapsed:.1f}s"
+        )
+        log_rows.append({
+            "epoch": epoch,
+            "train_loss": train_stats["loss"],
+            "train_snr": train_stats["snr"],
+            "train_psnr": train_stats["psnr"],
+            "val_loss": val_stats["loss"],
+            "val_snr": val_stats["snr"],
+            "val_psnr": val_stats["psnr"],
+        })
+
+        # Save best model
+        if val_stats["loss"] < best_val_loss:
+            best_val_loss = val_stats["loss"]
+            save_checkpoint(
+                model, optimizer, epoch, val_stats["loss"],
+                os.path.join(args.output_dir, "best_model.pth"),
+            )
+
+        # Periodic checkpoint
+        if epoch % args.save_every == 0:
+            save_checkpoint(
+                model, optimizer, epoch, val_stats["loss"],
+                os.path.join(args.output_dir, f"checkpoint_epoch{epoch:04d}.pth"),
+            )
+
+    # ---- Save training log --------------------------------------------------
+    log_path = os.path.join(args.output_dir, "training_log.csv")
+    import csv
+    with open(log_path, "w", newline="") as fh:
+        writer = csv.DictWriter(fh, fieldnames=log_rows[0].keys())
+        writer.writeheader()
+        writer.writerows(log_rows)
+    print(f"\nTraining complete.  Log saved to {log_path}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/unet_model.py b/unet_model.py
new file mode 100644
index 0000000..4215875
--- /dev/null
+++ b/unet_model.py
@@ -0,0 +1,152 @@
+"""
+U-Net model for seismic data interpolation and reconstruction.
+
+Architecture based on:
+  Ronneberger et al., "U-Net: Convolutional Networks for Biomedical Image
+  Segmentation", MICCAI 2015.
+
+Adapted for seismic data: the network takes an incomplete seismic section
+(traces set to zero where data is missing) concatenated with a binary mask
+that marks the missing locations, and outputs a fully reconstructed section.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Building blocks
+# ---------------------------------------------------------------------------
+
+class DoubleConv(nn.Module):
+    """Two consecutive (Conv -> BatchNorm -> ReLU) blocks."""
+
+    def __init__(self, in_channels: int, out_channels: int, mid_channels: int = None):
+        super().__init__()
+        if mid_channels is None:
+            mid_channels = out_channels
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.block(x)
+
+
+class Down(nn.Module):
+    """Max-pool downsampling followed by DoubleConv."""
+
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.block(x)
+
+
+class Up(nn.Module):
+    """Bilinear upsampling followed by DoubleConv (with skip connection)."""
+
+    def __init__(self, in_channels: int, out_channels: int, bilinear: bool = True):
+        super().__init__()
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels)
+
+    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        x1 = self.up(x1)
+        # Pad x1 to match x2 spatial dimensions if needed
+        diff_h = x2.size(2) - x1.size(2)
+        diff_w = x2.size(3) - x1.size(3)
+        x1 = F.pad(x1, [diff_w // 2, diff_w - diff_w // 2,
+                        diff_h // 2, diff_h - diff_h // 2])
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+
+class OutConv(nn.Module):
+    """1×1 convolution to map feature maps to the output channel count."""
+
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+# ---------------------------------------------------------------------------
+# U-Net
+# ---------------------------------------------------------------------------
+
+class UNet(nn.Module):
+    """
+    U-Net for seismic data interpolation / reconstruction.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.  Use 2 when the input is a (data, mask)
+        concatenation, or 1 when only data is passed.
+    out_channels : int
+        Number of output channels (typically 1 for a single seismic section).
+    base_features : int
+        Number of feature maps in the first encoder stage.  Subsequent stages
+        double this number (up to a factor of 8 at the bottleneck by default).
+    bilinear : bool
+        When True, use bilinear upsampling; otherwise use transposed convolutions.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 2,
+        out_channels: int = 1,
+        base_features: int = 64,
+        bilinear: bool = True,
+    ):
+        super().__init__()
+        f = base_features
+        factor = 2 if bilinear else 1
+
+        # Encoder
+        self.inc = DoubleConv(in_channels, f)
+        self.down1 = Down(f, f * 2)
+        self.down2 = Down(f * 2, f * 4)
+        self.down3 = Down(f * 4, f * 8)
+        self.down4 = Down(f * 8, f * 16 // factor)
+
+        # Decoder
+        self.up1 = Up(f * 16, f * 8 // factor, bilinear)
+        self.up2 = Up(f * 8, f * 4 // factor, bilinear)
+        self.up3 = Up(f * 4, f * 2 // factor, bilinear)
+        self.up4 = Up(f * 2, f, bilinear)
+
+        self.outc = OutConv(f, out_channels)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Encoder path
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+
+        # Decoder path with skip connections
+        x = self.up1(x5, x4)
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+
+        return self.outc(x)
diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000..34be21a
--- /dev/null
+++ b/utils.py
@@ -0,0 +1,264 @@
+"""
+Utility functions for seismic data interpolation / reconstruction.
+
+Covers:
+  - Loss functions (MSE, L1, perceptual, combined)
+  - Metrics      (SNR, SSIM, PSNR)
+  - Visualisation helpers
+  - Checkpoint save / load
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# Loss functions
+# ---------------------------------------------------------------------------
+
+class ReconstructionLoss(nn.Module):
+    """
+    Weighted combination of L1 and MSE losses.
+
+    The loss is evaluated on *all* output samples but can optionally
+    up-weight the contribution from the missing (masked-out) traces.
+
+    Parameters
+    ----------
+    l1_weight : float
+        Weight of the L1 component.
+    mse_weight : float
+        Weight of the MSE component.
+    missing_weight : float
+        Extra multiplier applied to the loss at missing-trace locations.
+        Set to 1.0 to treat all locations equally.
+    """
+
+    def __init__(
+        self,
+        l1_weight: float = 0.5,
+        mse_weight: float = 0.5,
+        missing_weight: float = 2.0,
+    ):
+        super().__init__()
+        self.l1_weight = l1_weight
+        self.mse_weight = mse_weight
+        self.missing_weight = missing_weight
+
+    def forward(
+        self,
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        pred : Tensor  (B, 1, T, X)
+        target : Tensor  (B, 1, T, X)
+        mask : Tensor  (B, 1, T, X), optional
+            Binary mask – 1 for observed traces, 0 for missing.
+        """
+        diff = pred - target
+
+        if mask is not None:
+            # Up-weight missing trace locations
+            weight = torch.ones_like(mask)
+            weight[mask == 0] = self.missing_weight
+        else:
+            weight = torch.ones_like(diff)
+
+        l1 = (weight * diff.abs()).mean()
+        mse = (weight * diff.pow(2)).mean()
+
+        return self.l1_weight * l1 + self.mse_weight * mse
+
+
+# ---------------------------------------------------------------------------
+# Metrics
+# ---------------------------------------------------------------------------
+
+def signal_to_noise_ratio(pred: np.ndarray, target: np.ndarray) -> float:
+    """
+    Signal-to-Noise Ratio (SNR) in dB.
+
+    SNR = 10 * log10( ||target||^2 / ||pred - target||^2 )
+    """
+    noise_power = np.mean((pred - target) ** 2)
+    signal_power = np.mean(target ** 2)
+    if noise_power == 0:
+        return float("inf")
+    return 10.0 * np.log10(signal_power / noise_power)
+
+
+def peak_signal_to_noise_ratio(pred: np.ndarray, target: np.ndarray) -> float:
+    """
+    Peak Signal-to-Noise Ratio (PSNR) in dB.
+
+    Uses the actual dynamic range of the target as peak value.
+    """
+    data_range = target.max() - target.min()
+    if data_range == 0:
+        return float("inf")
+    mse = np.mean((pred - target) ** 2)
+    if mse == 0:
+        return float("inf")
+    return 20.0 * np.log10(data_range) - 10.0 * np.log10(mse)
+
+
+def compute_metrics(pred: torch.Tensor, target: torch.Tensor) -> dict:
+    """
+    Compute SNR and PSNR for a batch.
+
+    Parameters
+    ----------
+    pred, target : Tensor  (B, 1, T, X)
+
+    Returns
+    -------
+    dict with keys "snr" and "psnr" (mean over batch, in dB).
+    """
+    pred_np = pred.detach().cpu().numpy()
+    target_np = target.detach().cpu().numpy()
+    snrs, psnrs = [], []
+    for p, t in zip(pred_np, target_np):
+        snrs.append(signal_to_noise_ratio(p, t))
+        psnrs.append(peak_signal_to_noise_ratio(p, t))
+    finite_snrs = [s for s in snrs if np.isfinite(s)]
+    finite_psnrs = [s for s in psnrs if np.isfinite(s)]
+    return {
+        "snr": float(np.mean(finite_snrs)) if finite_snrs else float("inf"),
+        "psnr": float(np.mean(finite_psnrs)) if finite_psnrs else float("inf"),
+    }
+
+
+# ---------------------------------------------------------------------------
+# Visualisation
+# ---------------------------------------------------------------------------
+
+def plot_comparison(
+    data: np.ndarray,
+    pred: np.ndarray,
+    target: np.ndarray,
+    title: str = "",
+    save_path: Optional[str] = None,
+    clip_percentile: float = 98.0,
+):
+    """
+    Side-by-side wiggle / image plot of input, prediction, and target.
+
+    Parameters
+    ----------
+    data, pred, target : ndarray  (T, X)
+        2-D seismic sections.
+    title : str
+        Figure title.
+    save_path : str, optional
+        If provided, save the figure to this path.
+    clip_percentile : float
+        Amplitude clip percentile for display.
+    """
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError as exc:
+        raise ImportError("matplotlib is required for visualisation.") from exc
+
+    vmax = np.percentile(np.abs(target), clip_percentile)
+    vmin = -vmax
+
+    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+    panels = [
+        (data, "Incomplete Input"),
+        (pred, "U-Net Reconstruction"),
+        (target, "Ground Truth"),
+    ]
+    for ax, (arr, label) in zip(axes, panels):
+        im = ax.imshow(
+            arr,
+            aspect="auto",
+            cmap="seismic",
+            vmin=vmin,
+            vmax=vmax,
+            interpolation="nearest",
+        )
+        ax.set_title(label, fontsize=12)
+        ax.set_xlabel("Trace index")
+        ax.set_ylabel("Time sample")
+        plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+
+    if title:
+        fig.suptitle(title, fontsize=14)
+    plt.tight_layout()
+
+    if save_path is not None:
+        os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
+        plt.savefig(save_path, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint helpers
+# ---------------------------------------------------------------------------
+
+def save_checkpoint(
+    model: nn.Module,
+    optimizer: torch.optim.Optimizer,
+    epoch: int,
+    loss: float,
+    save_path: str,
+    **extra,
+):
+    """Save a training checkpoint to *save_path*."""
+    os.makedirs(os.path.dirname(os.path.abspath(save_path)), exist_ok=True)
+    state = {
+        "epoch": epoch,
+        "model_state_dict": model.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict(),
+        "loss": loss,
+        **extra,
+    }
+    torch.save(state, save_path)
+
+
+def load_checkpoint(
+    model: nn.Module,
+    optimizer: Optional[torch.optim.Optimizer],
+    checkpoint_path: str,
+    device: torch.device = None,
+) -> dict:
+    """
+    Load a checkpoint saved by :func:`save_checkpoint`.
+
+    Returns
+    -------
+    dict
+        The full checkpoint dict (useful for resuming training state).
+    """
+    if device is None:
+        device = torch.device("cpu")
+    ckpt = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(ckpt["model_state_dict"])
+    if optimizer is not None and "optimizer_state_dict" in ckpt:
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+    return ckpt
+
+
+# ---------------------------------------------------------------------------
+# Seed / reproducibility
+# ---------------------------------------------------------------------------
+
+def set_seed(seed: int = 42):
+    """Set random seeds for reproducibility."""
+    import random
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)