simulate_data.py

from __future__ import annotations

from pathlib import Path

import nibabel as nib
import numpy as np
import yaml

from src.models import InversionRecoveryModel


def create_brain_mask(shape: tuple[int, int, int], radius: float) -> np.ndarray:
    x_coords, y_coords, z_coords = np.indices(shape)
    center = (np.asarray(shape, dtype=np.float64) - 1.0) / 2.0
    squared_distance = (
        (x_coords - center[0]) ** 2 + (y_coords - center[1]) ** 2 + (z_coords - center[2]) ** 2
    )
    return squared_distance <= radius**2


def main() -> None:
    output_root = Path(__file__).resolve().parent
    data_shape = (10, 10, 10)
    inversion_times = np.asarray([100.0, 400.0, 1000.0, 2000.0], dtype=np.float64)
    model = InversionRecoveryModel()

    brain_mask = create_brain_mask(data_shape, radius=3.2)
    true_s0 = np.full(data_shape, 10.0, dtype=np.float64)
    true_t1 = np.full(data_shape, 1200.0, dtype=np.float64)

    true_s0[brain_mask] = 500.0
    true_t1[brain_mask] = 1200.0

    simulated_volume = np.empty(data_shape + (inversion_times.shape[0],), dtype=np.float64)
    for time_index, inversion_time in enumerate(inversion_times):
        simulated_volume[..., time_index] = model.compute_signal(inversion_time, true_s0, true_t1)

    rng = np.random.default_rng(seed=42)
    noisy_volume = simulated_volume + rng.normal(loc=0.0, scale=5.0, size=simulated_volume.shape)
    noisy_volume = np.clip(noisy_volume, a_min=0.0, a_max=None)

    nifti_path = output_root / "fake_data.nii.gz"
    config_path = output_root / "config.yaml"

    nib.save(nib.Nifti1Image(noisy_volume.astype(np.float32), np.eye(4)), str(nifti_path))

    config_data = {
        "input_nifti": nifti_path.name,
        "inversion_times": inversion_times.tolist(),
        "output_dir": "output",
        "mask_threshold": 50.0,
    }

    with config_path.open("w", encoding="utf-8") as handle:
        yaml.safe_dump(config_data, handle, sort_keys=False)

    print(f"Saved simulated NIfTI to: {nifti_path}")
    print(f"Saved config to: {config_path}")


if __name__ == "__main__":
    main()