utils.py

from __future__ import annotations

import os
import random
from dataclasses import dataclass
from typing import Dict, Tuple

import logging
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Sampler
from tqdm.auto import tqdm
import argparse
from functools import partial
import sys
from dataset import *


def none_or_str(value):
    if value.lower() == "none":
        return None
    return value


def add_bool_arg(group, name, default, help_text):
    group.add_argument(name, type=int, choices=(
        0, 1), default=default, help=help_text)


def positive_int(value):
    value = int(value)
    if value <= 0:
        raise argparse.ArgumentTypeError("must be a positive integer")
    return value


def nonnegative_int(value):
    value = int(value)
    if value < 0:
        raise argparse.ArgumentTypeError("must be a non-negative integer")
    return value


def positive_float(value):
    value = float(value)
    if value <= 0:
        raise argparse.ArgumentTypeError("must be positive")
    return value


def nonnegative_float(value):
    value = float(value)
    if value < 0:
        raise argparse.ArgumentTypeError("must be non-negative")
    return value


def probability(value):
    value = float(value)
    if not 0.0 <= value <= 1.0:
        raise argparse.ArgumentTypeError("must be between 0 and 1")
    return value


def arg_parser(argv=None):
    parser = argparse.ArgumentParser(
        description="Train SHORE coefficient prediction models.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    ############################
    # Paths
    ############################
    paths = parser.add_argument_group("paths")
    paths.add_argument("--checkpoint_path", type=str,
                       default="./checkpoints", help="Root path for experiment outputs.")
    paths.add_argument("--load_checkpoint_path", type=none_or_str,
                       default=None, help="Checkpoint to resume from, or None.")
    paths.add_argument("--data_path", type=str, default="./data",
                       help="Root directory of preprocessed subject data.")

    ############################
    # Training
    ############################
    training = parser.add_argument_group("training")
    training.add_argument("--n_epochs", type=positive_int, default=1,
                          help="Number of epochs for epoch-based training.")
    training.add_argument("--n_steps", type=positive_int, default=1000,
                          help="Number of steps for step-based training.")
    add_bool_arg(training, "--step_training", 0,
                 "Use step-based training instead of epoch-based training.")
    training.add_argument("--epoch_step_interval", type=positive_int, default=250,
                          help="Number of steps treated as one logical epoch.")
    training.add_argument("--train_batch_size", type=positive_int,
                          default=1, help="Training batch size.")
    training.add_argument("--val_batch_size", type=positive_int,
                          default=1, help="Validation/test batch size.")
    training.add_argument("--train_num_workers", type=nonnegative_int, default=max(
        (os.cpu_count() or 1) - 1, 0), help="Training dataloader workers.")
    training.add_argument("--train_prefetch_factor", type=positive_int, default=2,
                          help="Batches prefetched per training worker.")
    training.add_argument("--seed", type=int, default=12345,
                          help="Random seed. Use a negative value to disable seeding.")
    training.add_argument("--lr", type=positive_float,
                          default=1e-3, help="Learning rate.")
    training.add_argument("--optimizer", choices=("adam",
                          "adamw", "sgd"), default="adam", help="Optimizer.")
    training.add_argument("--weight_decay", type=nonnegative_float,
                          default=1e-4, help="Optimizer weight decay.")
    training.add_argument("--scheduler", type=none_or_str, choices=(None, "cosine", "step",
                          "plateau", "decay"), default=None, help="Learning rate scheduler, or None.")
    training.add_argument("--step_size", type=positive_int, default=10,
                          help="Logical epochs between StepLR learning-rate updates.")
    training.add_argument("--gamma", type=positive_float, default=0.1,
                          help="Learning-rate multiplier for step and plateau schedulers.")
    training.add_argument("--patience", type=nonnegative_int, default=5,
                          help="Logical epochs without improvement before plateau updates.")
    add_bool_arg(training, "--amp", 0, "Use automatic mixed precision.")
    add_bool_arg(training, "--compile", 1,
                 "Compile model and loss with torch.compile.")
    training.add_argument("--save_interval", type=positive_int, default=5,
                          help="Checkpoint interval in logical epochs.")
    training.add_argument("--grad_clip_norm", type=positive_float, default=12.0,
                          help="Maximum gradient norm used for clipping.")

    ############################
    # Model
    ############################
    model = parser.add_argument_group("model")
    model.add_argument(
        "--network",
        choices=("spatial_invariant_att_set_lite_strong",),
        default="spatial_invariant_att_set_lite_strong",
        help="Network architecture. The spatial masked-set model of the paper.",
    )
    model.add_argument(
        "--trainer",
        choices=(
            "BaseTrainerFinal",
            "BaseTrainerFinalMask",
        ),
        default="BaseTrainerFinalMask",
        help=(
            "Trainer class. BaseTrainerFinalMask performs sparse-input, "
            "dense-output training with shell-wise gradient dropping and is "
            "the setting used for the paper results. BaseTrainerFinal is the "
            "full-observation (no gradient dropping) variant."
        ),
    )
    model.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available()
                       else "cpu", help="Training device.")
    add_bool_arg(model, "--use_bval_as_input", 1,
                 "Append normalized b-value to each measurement feature.")
    add_bool_arg(model, "--use_spatial_context", 1,
                 "Enable the 3x3x3 spatial neighbor attention context.")

    ############################
    # Dataset
    ############################
    dataset = parser.add_argument_group("dataset")
    dataset.add_argument(
        "--dataset_type",
        choices=(
            "base_final_mmap_neighbor",
            "base_final_mmap_compact_neighbor",
            "base_final_ram_compact_neighbor",
        ),
        default="base_final_mmap_neighbor",
        help="Dataset implementation.",
    )
    dataset.add_argument("--patient_list", type=str, default="whole",
                         help="Subject split to use: 'whole' for the checked-in "
                              "HCP100 split, or a path to your own split JSON "
                              "with 'train', 'val' and 'test' keys.")
    dataset.add_argument("--rand_drop", type=probability, default=0.0,
                         help="Probability for random b-vector rotation when rotation is enabled.")
    add_bool_arg(dataset, "--shuffle", 0,
                 "Shuffle gradient order during training.")
    add_bool_arg(dataset, "--rotate_bvecs", 0,
                 "Randomly rotate b-vectors during training.")
    add_bool_arg(dataset, "--signal_drop", 0,
                 "Use masked-gradient signal dropping.")
    add_bool_arg(dataset, "--center_only", 0,
                 "Use only the center voxel signal. Must be 0 for the spatial "
                 "masked-set network, which requires the neighborhood.")
    add_bool_arg(dataset, "--whole_brain", 0, "Use whole-brain subject split.")
    add_bool_arg(
        dataset,
        "--eval_data_in_ram",
        0,
        "Load validation and test compact data into RAM when using the RAM dataset.",
    )

    ############################
    # Loss
    ############################
    loss = parser.add_argument_group("loss")
    loss.add_argument("--loss_function", choices=("shore_mse", "shore_huber", "signal_mse",
                      "combined", "combined_coeff_recon"), default="combined", help="Training loss.")
    loss.add_argument("--shore_coeff_loss", choices=("mse", "huber"), default="mse",
                      help="SHORE coefficient loss used inside combined loss.")
    loss.add_argument("--shore_huber_delta", type=positive_float, default=1.0,
                      help="Delta parameter for SHORE Huber loss.")
    loss.add_argument("--reconstructed_signal_loss", choices=("mse", "huber"), default="mse",
                      help="Signal-space loss used inside combined_coeff_recon.")
    loss.add_argument("--signal_huber_delta", type=positive_float, default=1.0,
                      help="Delta parameter for reconstructed-signal Huber loss.")
    loss.add_argument("--shore_loss_weight", type=nonnegative_float, default=0.5,
                      help="SHORE coefficient MSE weight for combined loss.")
    loss.add_argument("--signal_loss_weight", type=nonnegative_float, default=0.5,
                      help="Signal reconstruction MSE weight for combined loss.")
    add_bool_arg(loss, "--query_based_loss", 0,
                 "For masked training, compute signal loss on dropped/query gradients.")
    add_bool_arg(loss, "--all_signals_masked", 0,
                 "For masked training, compute signal loss over all gradients.")

    ############################
    # Evaluation
    ############################
    evaluation = parser.add_argument_group("evaluation")
    evaluation.add_argument(
        "--selection_metric",
        choices=("MSE_signal", "MSE_signal_patient_macro",
                 "robust_MSE_signal_patient_macro"),
        default="MSE_signal_patient_macro",
        help="Validation metric used for best-checkpoint selection.",
    )
    evaluation.add_argument(
        "--deterministic_rotated_val_count",
        type=nonnegative_int,
        default=0,
        help="Number of fixed rotations included in full-gradient validation.",
    )
    evaluation.add_argument(
        "--eval_num_workers",
        type=nonnegative_int,
        default=8,
        help="Validation/test dataloader workers.",
    )

    ############################
    # Misc
    ############################
    misc = parser.add_argument_group("misc")
    misc.add_argument("--name", type=none_or_str, default=None,
                      help="Optional suffix for the experiment name.")
    add_bool_arg(misc, "--show_progress", 1, "Show tqdm progress bars.")
    misc.add_argument("--progress_loss_interval", type=positive_int, default=50,
                      help="Steps between tqdm loss-value refreshes.")

    args = parser.parse_args(argv)
    _validate_args(parser, args)
    return args


def _validate_args(parser, args):
    if args.trainer == "BaseTrainerFinalMask" and not args.signal_drop:
        parser.error(f"{args.trainer} requires --signal_drop=1")
    if args.center_only:
        parser.error(f"{args.network} requires --center_only=0")
    if args.query_based_loss and args.all_signals_masked:
        parser.error(
            "--query_based_loss and --all_signals_masked are mutually exclusive")
    if args.eval_data_in_ram and args.dataset_type != "base_final_ram_compact_neighbor":
        parser.error(
            "--eval_data_in_ram=1 requires "
            "--dataset_type=base_final_ram_compact_neighbor"
        )


def print_args(args):
    # Log the output to the file
    logging.info("*" * 10)
    logging.info("Arguments:")

    for arg in vars(args):
        logging.info(f"{arg}: {getattr(args, arg)}")

    logging.info("*" * 10)


def setup_logging(log_file_path="app_log.txt"):
    """
    Configures the logging module to log both to the console and to a specified file path.

    :param log_file_path: The path where the log file will be saved. Defaults to 'app_log.txt'.
    """

    # Check if logging has already been configured
    if hasattr(setup_logging, "done"):
        return

    # Set the log formatter
    formatter = logging.Formatter(
        "[%(levelname)s] %(asctime)s - %(message)s", datefmt="%d-%b-%y %H:%M:%S"
    )

    # Create a stream handler for console output
    stream_handler = logging.StreamHandler(sys.stdout)
    stream_handler.setFormatter(formatter)

    # Ensure the directory exists for the log file
    log_dir = os.path.dirname(log_file_path)
    if log_dir and not os.path.exists(log_dir):
        os.makedirs(log_dir)  # Create the directory if it does not exist

    # Create a file handler to save logs to the provided file path
    file_handler = logging.FileHandler(
        log_file_path)  # Log file path from parameter
    file_handler.setFormatter(formatter)

    # Get the logger and remove any existing handlers
    logger = logging.getLogger("root")
    if logger.handlers:
        for h in logger.handlers:
            logger.removeHandler(h)

    # Set the log level and add both handlers
    logger.setLevel(os.getenv("LOG_LEVEL", "INFO"))
    logger.addHandler(stream_handler)
    logger.addHandler(file_handler)

    # Mark the logging setup as done
    setattr(setup_logging, "done", True)


def seed_everything(seed: int = 0) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # torch.backends.cudnn.deterministic = False
    # torch.backends.cudnn.benchmark = True


def worker_init_fn(worker_id: int, num_workers: int, seed: int, rank: int = 0) -> None:
    worker_seed = seed + num_workers * rank + worker_id
    random.seed(worker_seed)
    np.random.seed(worker_seed)
    torch.manual_seed(worker_seed)


def create_seeded_dataloader(
    args, dataset, non_verbose=False, **dataloader_args
) -> DataLoader:
    """
    Creates a dataloader object from a dataset, setting the seeds for the workers (if `--seed` is set).

    Args:
        args: the arguments of the program
        dataset: the dataset to be loaded
        verbose: whether to print the number of workers
        dataloader_args: external arguments of the dataloader

    Returns:
        the dataloader object
    """

    n_cpus = 4 if not hasattr(os, "sched_getaffinity") else len(
        os.sched_getaffinity(0))
    default_workers = min(8, n_cpus) if getattr(
        args, "train_num_workers", None) is None else args.train_num_workers
    dataloader_args.setdefault("num_workers", default_workers)

    if not non_verbose:
        logging.info(
            f'Using {dataloader_args["num_workers"]} workers for the dataloader.')

    if args.seed is not None and args.seed >= 0 and dataloader_args["num_workers"] > 0:
        g = torch.Generator()
        g.manual_seed(int(args.seed))
        dataloader_args.setdefault("generator", g)
        dataloader_args.setdefault("worker_init_fn", partial(worker_init_fn, seed=int(
            args.seed), num_workers=dataloader_args["num_workers"]))

    return DataLoader(dataset, **dataloader_args)


def get_datasets(args):
    dataset_classes = {
        "base_final_mmap_neighbor": BaseDatasetFinalMMapNeighbour,
        "base_final_mmap_compact_neighbor": BaseDatasetFinalMMapCompactNeighbour,
    }
    if args.dataset_type == "base_final_ram_compact_neighbor":
        train_dataset_class = BaseDatasetFinalRAMCompactNeighbour
        eval_dataset_class = (
            BaseDatasetFinalRAMCompactNeighbour
            if getattr(args, "eval_data_in_ram", False)
            else BaseDatasetFinalMMapCompactNeighbour
        )
    elif args.dataset_type in dataset_classes:
        train_dataset_class = dataset_classes[args.dataset_type]
        eval_dataset_class = dataset_classes[args.dataset_type]
    else:
        raise ValueError(f"Unknown dataset type '{args.dataset_type}'")

    train_extra_kwargs = {}
    eval_extra_kwargs = {}
    if args.dataset_type == "base_final_ram_compact_neighbor":
        train_extra_kwargs["show_progress"] = args.show_progress
        if getattr(args, "eval_data_in_ram", False):
            eval_extra_kwargs["show_progress"] = args.show_progress
    train_dataset = train_dataset_class(root_dir=args.data_path, mode="train", shuffle=args.shuffle, rand_bvec_rotate=args.rand_drop, rotate_bvecs=args.rotate_bvecs,
                                        drop_signals=args.signal_drop, center_only=args.center_only, whole_brain=args.whole_brain, patient_list=args.patient_list,
                                        **train_extra_kwargs)
    val_dataset = eval_dataset_class(root_dir=args.data_path, mode="val", shuffle=False, rand_bvec_rotate=0, rotate_bvecs=args.rotate_bvecs,
                                     drop_signals=args.signal_drop, center_only=args.center_only, whole_brain=args.whole_brain, patient_list=args.patient_list,
                                     **eval_extra_kwargs)
    test_dataset = eval_dataset_class(root_dir=args.data_path, mode="test", shuffle=False, rand_bvec_rotate=0, rotate_bvecs=args.rotate_bvecs,
                                      drop_signals=args.signal_drop, center_only=args.center_only, whole_brain=args.whole_brain, patient_list=args.patient_list,
                                      **eval_extra_kwargs)

    return train_dataset, val_dataset, test_dataset


class InfiniteUniformBatchSampler(Sampler[list[int]]):
    """Yield uniformly sampled batches forever.

    Each batch is sampled uniformly without replacement, so a batch never
    contains duplicate indices. Batches are independent, so an index can appear
    again in later batches.
    """

    def __init__(self, data_source, batch_size: int, seed: int = 0):
        self.data_source = data_source
        self.batch_size = int(batch_size)
        self.seed = None if seed is None or int(seed) < 0 else int(seed)

        if self.batch_size <= 0:
            raise ValueError("batch_size must be positive.")

    def __iter__(self):
        n = len(self.data_source)
        if n <= 0:
            raise ValueError("Dataset is empty.")
        if self.batch_size > n:
            raise ValueError(
                f"batch_size={self.batch_size} is larger than dataset size={n}; "
                "unique indices within each batch are impossible."
            )

        generator = torch.Generator()
        if self.seed is None:
            generator.seed()
        else:
            generator.manual_seed(self.seed)

        while True:
            yield self._sample_unique_batch(n, generator)

    def _sample_unique_batch(self, n: int, generator: torch.Generator) -> list[int]:
        if self.batch_size > n // 2:
            return torch.randperm(n, generator=generator)[:self.batch_size].tolist()

        selected = []
        seen = set()
        draw_size = min(n, max(self.batch_size * 2, 32))

        while len(selected) < self.batch_size:
            needed = self.batch_size - len(selected)
            candidates = torch.randint(
                n, (max(draw_size, needed),), generator=generator).tolist()
            for idx in candidates:
                if idx not in seen:
                    seen.add(idx)
                    selected.append(idx)
                    if len(selected) == self.batch_size:
                        break

        return selected

    def __len__(self):
        n = len(self.data_source)
        if n <= 0:
            return 0
        return (n + self.batch_size - 1) // self.batch_size


def make_infinite_uniform_dataloader(
    dataset,
    batch_size: int,
    num_workers: int = 0,
    seed: int = 0,
    pin_memory: bool = True,
    persistent_workers: bool = True,
    prefetch_factor: int = 2,
    drop_last: bool = True,
):
    if drop_last is False:
        logging.warning(
            "drop_last=False has no effect for the infinite uniform batch sampler.")
    if prefetch_factor <= 0:
        raise ValueError("prefetch_factor must be positive.")

    batch_sampler = InfiniteUniformBatchSampler(
        dataset,
        batch_size=batch_size,
        seed=seed,
    )

    generator = None
    worker_fn = None
    if seed is not None and seed >= 0:
        generator = torch.Generator()
        generator.manual_seed(int(seed))
        worker_fn = partial(worker_init_fn, seed=int(seed),
                            num_workers=num_workers)

    kwargs = dict(
        dataset=dataset,
        batch_sampler=batch_sampler,
        num_workers=num_workers,
        pin_memory=pin_memory,
    )

    if generator is not None:
        kwargs["generator"] = generator
        kwargs["worker_init_fn"] = worker_fn

    if num_workers > 0:
        kwargs["persistent_workers"] = persistent_workers
        kwargs["prefetch_factor"] = prefetch_factor

    return DataLoader(**kwargs)


def _format_count(value):
    if value >= 1_000_000 and value % 1_000_000 == 0:
        return f"{value // 1_000_000}m"
    if value >= 1_000 and value % 1_000 == 0:
        return f"{value // 1_000}k"
    return str(value)


def _format_schedule_name(args):
    schedule = (
        f"st{_format_count(args.n_steps)}"
        if args.step_training
        else f"ep{_format_count(args.n_epochs)}"
    )
    if args.scheduler is not None:
        schedule += f"-{args.scheduler}"
    return schedule


def _format_loss_name(args):
    subtype_names = {"mse": "MSE", "huber": "Hub"}
    loss_names = {
        "shore_mse": "shMSE",
        "shore_huber": "shHub",
        "signal_mse": "sigMSE",
        "combined": "cmb",
        "combined_coeff_recon": "cmbRec",
    }
    loss = loss_names[args.loss_function]

    if args.loss_function in {"combined", "combined_coeff_recon"}:
        loss += f"-shw{args.shore_loss_weight:g}-sgw{args.signal_loss_weight:g}"
        if args.loss_function == "combined_coeff_recon":
            loss += f"-rec{subtype_names[args.reconstructed_signal_loss]}"
            if args.reconstructed_signal_loss == "huber":
                loss += f"{args.signal_huber_delta:g}"
        if args.shore_coeff_loss != "mse":
            loss += f"-sh{subtype_names[args.shore_coeff_loss]}{args.shore_huber_delta:g}"
    elif args.loss_function == "shore_huber":
        loss += f"{args.shore_huber_delta:g}"

    if args.all_signals_masked:
        loss += "-asm"
    elif args.query_based_loss:
        loss += "-qry"

    return loss


def _format_dataset_name(args):
    if args.dataset_type not in {
        "base_final_mmap_neighbor",
        "base_final_mmap_compact_neighbor",
        "base_final_ram_compact_neighbor",
    }:
        raise ValueError(f"Unknown dataset_type: {args.dataset_type}")

    parts = [f"ds{args.patient_list}"]
    if args.dataset_type in {
        "base_final_mmap_compact_neighbor",
        "base_final_ram_compact_neighbor",
    }:
        parts.append("compact")
    if args.dataset_type == "base_final_ram_compact_neighbor":
        parts.append("ram")
    if args.center_only:
        parts.append("ctr")
    if args.shuffle:
        parts.append("shuf")
    if args.signal_drop:
        parts.append("drop")
    if args.rotate_bvecs:
        parts.append(f"rot{args.rand_drop:g}")
    return "-".join(parts)


def get_experiment_name(args):
    experiment_parts = [
        f"lr{args.lr:g}",
        f"bs{args.train_batch_size}",
        _format_schedule_name(args),
        _format_loss_name(args),
        _format_dataset_name(args),
    ]

    if args.use_bval_as_input:
        experiment_parts.append("bval")
    experiment_parts.append(f"ctx{args.use_spatial_context}")

    if args.name:
        experiment_parts.append(str(args.name))

    experiment_name = "_".join(experiment_parts)
    network_names = {
        "spatial_invariant_att_set_lite_strong": "spatialAttLiteStrong",
    }
    trainer_names = {
        "BaseTrainerFinal": "base",
        "BaseTrainerFinalMask": "mask",
    }
    checkpoint_dir = os.path.join(
        args.checkpoint_path,
        f"{network_names[args.network]}-{trainer_names[args.trainer]}"
        f"{'-wb' if args.whole_brain else ''}",
    )

    return experiment_name, checkpoint_dir


def get_trainer_class(args, datasets, writer, experiment_name):
    from trainers import BaseTrainerFinal, BaseTrainerFinalMask

    trainers = {
        "BaseTrainerFinal": BaseTrainerFinal,
        "BaseTrainerFinalMask": BaseTrainerFinalMask,
    }

    if args.trainer not in trainers:
        raise ValueError(
            f"Unknown trainer '{args.trainer}'. Available trainers: {list(trainers.keys())}"
        )

    if args.trainer == "BaseTrainerFinalMask" and not args.signal_drop:
        raise ValueError(
            "BaseTrainerFinalMask requires --signal_drop=1 for training.")

    if args.center_only:
        raise ValueError(
            "The spatial masked-set network requires neighborhood signals. "
            "Use --center_only=0."
        )

    trainer_class = trainers[args.trainer]
    return trainer_class(args, datasets, writer, experiment_name)


def batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
    out = {}
    for k, v in batch.items():
        if torch.is_tensor(v):
            out[k] = v.to(device, non_blocking=True)
        else:
            out[k] = v
    return out


def build_inputs(batch: Dict[str, torch.Tensor], use_bval_as_input, dim: int | None = None) -> torch.Tensor:
    q_direction = batch["q_direction"]  # (B, G, 3)
    norm_sig = batch["norm_sig"]  # (B, *S, G, 1) or center-only (B, G, 1)

    spatial_dims = 0 if dim is None else dim
    if spatial_dims not in {0, 3}:
        raise ValueError(f"dim must be None or 3; got {dim}.")

    batch_size, n_gradients = q_direction.shape[:2]
    spatial_shape = norm_sig.shape[1:-2]
    if len(spatial_shape) != spatial_dims:
        raise ValueError(
            f"norm_sig has spatial shape {tuple(spatial_shape)}, expected "
            f"{spatial_dims} spatial dimensions."
        )

    # Add singleton spatial axes, then broadcast q features across the signal
    # neighborhood without allocating. torch.cat/output assignment performs the
    # single required materialization of the final dense model input.
    q_view = q_direction.view(
        batch_size, *([1] * spatial_dims), n_gradients, 3
    )
    q_spatial = q_view.expand(batch_size, *spatial_shape, n_gradients, 3)

    if use_bval_as_input:
        bval = batch["bval_normalized"]  # (B, G, 1)
        bval_view = bval.view(
            batch_size, *([1] * spatial_dims), n_gradients, 1
        )
        bval_spatial = bval_view.expand_as(norm_sig)
        # Keep signal last: [q_x, q_y, q_z, normalized_bval, signal].
        # in older version of the build we used [q_x, q_y, q_z, signal, normalized_bval], but putting signal last allows for more convenient slicing of the input features if needed.
        return torch.cat((q_spatial, bval_spatial, norm_sig), dim=-1)

    # A direct output allocation is faster than concatenating an expanded
    # q-direction view with the signal when no b-value feature is requested.
    inputs = norm_sig.new_empty(*norm_sig.shape[:-1], 4)
    inputs[..., :3] = q_spatial
    inputs[..., 3:4] = norm_sig
    return inputs


def signal_recon_from_shore_coeffs(
    shore_coeff: torch.Tensor,
    shore_phi_matrix: torch.Tensor,
) -> torch.Tensor:
    # shore_phi_matrix: (B,G,C), shore_coeff: (B,C) => pred_sig: (B,G)
    # pred_sig = torch.einsum("bgc,bc->bg", shore_phi_matrix, shore_coeff)
    pred_sig = torch.bmm(
        shore_phi_matrix,
        shore_coeff.unsqueeze(2)
    ).squeeze(2)
    return pred_sig