brkraw-sordino 0.1.3__py3-none-any.whl

brkraw_sordino/__init__.py

@@ -0,0 +1,7 @@
+"""BrkRaw SORDINO-ZTE converter hook package."""
+
+from .hook import get_dataobj
+
+__all__ = ["__version__", "get_dataobj"]
+
+__version__ = "0.1.3"

brkraw_sordino/brkraw_hook.yaml

@@ -0,0 +1,8 @@
+docs: docs.md
+specs:
+  - specs/info_spec.yaml
+  - specs/metadata_spec.yaml
+rules:
+  - rules/sordino.yaml
+transforms:
+  - specs/utils.py

brkraw_sordino/docs.md

@@ -0,0 +1,99 @@
+# brkraw-sordino
+
+SORDINO-ZTE reconstruction hook for BrkRaw.
+
+## Install
+
+```bash
+pip install -e .
+```
+
+## Hook install
+
+```bash
+brkraw hook install brkraw-sordino
+```
+
+This installs the hook rule from the package manifest (`brkraw_hook.yaml`).
+
+## Usage
+
+Once installed, `brkraw` applies the hook automatically when a dataset matches the rule.
+
+Basic conversion:
+
+```bash
+brkraw convert /path/to/study --scan-id 3 --reco-id 1
+```
+
+The hook behaves the same whether invoked via the CLI or via the Python API (the same hook entrypoint and arguments are used).
+
+To explicitly pass hook options (or override defaults), use `--hook-arg` / `--hook-args-yaml` below.
+
+## Hook options
+
+Hook arguments can be passed via `brkraw convert` using `--hook-arg` with the
+entrypoint name (`sordino`):
+
+```bash
+brkraw convert /path/to/study -s 3 -r 1 \
+  --hook-arg sordino:ext_factors=1.2 \
+  --hook-arg sordino:offset=2 \
+  --hook-arg sordino:split_ch=false
+```
+
+### Pass hook options via YAML (`--hook-args-yaml`)
+
+BrkRaw can also load hook arguments from YAML. Generate a template like this:
+
+```bash
+brkraw hook preset sordino -o hook_args.yaml
+```
+
+Edit the generated YAML, then pass it to `brkraw convert` (repeatable):
+
+```bash
+brkraw convert /path/to/study -s 3 -r 1 --hook-args-yaml hook_args.yaml
+```
+
+Example:
+
+```yaml
+hooks:
+  sordino:
+    ext_factors: 1.2
+    offset: 2
+    split_ch: false
+    # as_complex: true  # optional, return (real, imag)
+    # cache_dir: ~/.brkraw/cache/sordino  # optional (add manually if needed)
+```
+
+Notes:
+
+- CLI `--hook-arg` values override YAML.
+- YAML supports both `{hooks: {sordino: {...}}}` and `{sordino: {...}}` shapes.
+- You can also set `BRKRAW_CONVERT_HOOK_ARGS_YAML` (comma-separated paths).
+
+Supported keys:
+
+- `ext_factors`: scalar or 3-item sequence (default: 1.0)
+- `ignore_samples`: int (default: 1)
+- `offset`: int (default: 0)
+- `num_frames`: int or null (default: None)
+- `correct_spoketiming`: bool (default: false)
+- `correct_ramptime`: bool (default: true)
+- `offreso_ch`: int or null (default: None)
+- `offreso_freq`: float (default: 0.0)
+- `mem_limit`: float (default: 0.5)
+- `clear_cache`: bool (default: true)
+- `split_ch`: bool (default: false, merge channels)
+- `as_complex`: bool (default: false, return complex as (real, imag))
+- `cache_dir`: string path (default: ~/.brkraw/cache/sordino)
+
+## Notes
+
+- The hook reconstructs data using an adjoint NUFFT and returns magnitude images by default.
+- Multi-channel data defaults to merged channels; set `split_ch=true` to keep channels split.
+- When `split_ch=false`, magnitude uses RSS while complex uses coherent sum.
+- Orientation is normalized when the first 3D axes are spatial; see `notebooks/orientation.ipynb`.
+- Cache files live under `~/.brkraw/cache/sordino` (or `BRKRAW_CONFIG_HOME`).

brkraw_sordino/helper.py

@@ -0,0 +1,81 @@
+import logging
+import os
+import sys
+import time
+from typing import Any, Optional, TextIO, cast
+
+logger = logging.getLogger("brkraw.sordino")
+
+
+def progressbar(iterable: Any, *, desc: str = "", ncols: int = 100, disable: Optional[bool] = None):
+    """Lightweight progress iterator (tqdm-like) without external deps."""
+    if disable is None:
+        disable = logger.isEnabledFor(logging.DEBUG)
+    stream: TextIO = cast(
+        TextIO,
+        sys.__stderr__
+        or sys.stderr
+        or sys.__stdout__
+        or sys.stdout
+        or open(os.devnull, "w", encoding="utf-8"),
+    )
+    try:
+        is_tty = bool(getattr(stream, "isatty", lambda: False)())
+    except Exception:
+        is_tty = False
+    if disable or not is_tty or not logger.isEnabledFor(logging.INFO):
+        return iterable
+
+    try:
+        total = len(iterable)  # type: ignore[arg-type]
+    except Exception:
+        total = 0
+
+    if total <= 0:
+        return iterable
+
+    bar_width = max(10, min(40, ncols - max(0, len(desc)) - 20))
+    start = time.time()
+    last_emit = 0.0
+
+    def _emit(i: int) -> None:
+        nonlocal last_emit
+        now = time.time()
+        if now - last_emit < 0.1 and i < total:
+            return
+        last_emit = now
+        frac = min(1.0, max(0.0, i / total))
+        filled = int(bar_width * frac)
+        bar = "#" * filled + "-" * (bar_width - filled)
+        elapsed = max(0.001, now - start)
+        rate = i / elapsed if i > 0 else 0.0
+        remaining = max(0, total - i)
+        eta = int(remaining / rate) if rate > 0 else -1
+        eta_txt = f"{eta}s" if eta >= 0 else "?"
+        prefix = f"{desc} " if desc else ""
+        line = f"{prefix}[{bar}] {i}/{total} ETA {eta_txt}"
+        try:
+            stream.write("\r" + line)
+            stream.flush()
+        except Exception:
+            pass
+
+    def _done() -> None:
+        try:
+            stream.write("\r" + (" " * (ncols if ncols > 0 else 120)) + "\r\n")
+            stream.flush()
+        except Exception:
+            pass
+
+    def _iter():
+        for i, item in enumerate(iterable, start=1):
+            _emit(i)
+            yield item
+        _done()
+
+    return _iter()
+
+
+__all__ = [
+    "progressbar",
+]

brkraw_sordino/hook.py

@@ -0,0 +1,164 @@
+import os
+import logging
+import numpy as np
+import tempfile
+from pathlib import Path
+
+from typing import Any, Optional, Tuple, Dict, Union, cast, TYPE_CHECKING
+
+from brkraw.specs.remapper import load_spec, map_parameters
+from brkraw.resolver import fid as fid_resolver
+from brkraw.resolver import datatype as dtype_resolver
+from brkraw.core import config as config_core
+from brkraw.core.fs import DatasetFile
+from brkraw.core.zip import ZippedFile
+
+from numpy.typing import NDArray
+from .typing import Options
+from .traj import get_trajectory
+from .helper import progressbar
+from .recon import get_dataobj_shape, recon_dataobj
+from .spoketiming import prep_fid_segmentation, correct_spoketiming
+from .orientation import correct as correct_orientation
+
+
+FileIO = Union[DatasetFile, ZippedFile]
+logger = logging.getLogger("brkraw.sordino")
+
+
+def _normalize_ext_factors(value: Any) -> Tuple[float, float, float]:
+    if value is None:
+        return (1.0, 1.0, 1.0)
+    if isinstance(value, (int, float)):
+        val = float(value)
+        return (val, val, val)
+    if isinstance(value, (list, tuple, np.ndarray)):
+        items = list(value)
+        if len(items) == 1:
+            val = float(items[0])
+            return (val, val, val)
+        if len(items) == 3:
+            return (float(items[0]), float(items[1]), float(items[2]))
+    raise ValueError("ext_factors must be a scalar or a 3-item sequence")
+
+
+def _get_cache_dir(path: Optional[Union[str, Path]]) -> Path:
+    if path:
+        base = Path(path).expanduser()
+    else:
+        base = config_core.resolve_root(None) / "cache" / "sordino"
+    base.mkdir(parents=True, exist_ok=True)
+    return base
+
+
+def _build_options(kwargs: Dict[str, Any]) -> Options:
+    cache_dir = _get_cache_dir(kwargs.get("cache_dir"))
+    logger.debug("Cache dir: %s", cache_dir)
+
+    return Options(
+        ext_factors=_normalize_ext_factors(kwargs.get("ext_factors")),
+        ignore_samples=int(kwargs.get("ignore_samples", 1)),
+        offset=int(kwargs.get("offset", 0)),
+        num_frames=kwargs.get("num_frames"),
+        correct_spoketiming=bool(kwargs.get("correct_spoketiming", False)),
+        correct_ramptime=bool(kwargs.get("correct_ramptime", True)),
+        offreso_ch=kwargs.get("offreso_ch"),
+        offreso_freq=float(kwargs.get("offreso_freq", 0.0)),
+        mem_limit=float(kwargs.get("mem_limit", 0.5)),
+        clear_cache=bool(kwargs.get("clear_cache", True)),
+        split_ch=bool(kwargs.get("split_ch", False)),
+        cache_dir=cache_dir,
+        as_complex=bool(kwargs.get("as_complex", False)),
+    )
+
+
+def _parse_recon_info(scan):
+    spec_path = Path(__file__).parent / "specs" / "recon_spec.yaml"
+    spec, transforms = load_spec(spec_path, validate=True)
+    recon_info = map_parameters(scan, spec, transforms)
+    dtype_info = dtype_resolver.resolve(scan)
+    if not dtype_info or "dtype" not in dtype_info:
+        raise ValueError("Failed to resolve FID dtype from acqp.")
+    recon_info['FIDDataType'] = dtype_info["dtype"]
+    return recon_info
+
+
+def _get_fid_entry(scan: Any) -> FileIO:
+    fid_entry = fid_resolver.get_fid(scan)
+    if fid_entry is None:
+        logger.warning("No FID/rawdata entry found for scan %s.", 
+                       getattr(scan, "scan_id", "?"))
+    return cast(FileIO, fid_entry)
+
+
+def get_dataobj(
+        scan: Any, reco_id: Optional[int] = None, **kwargs: Any,
+    ) -> Optional[Union[np.ndarray, Tuple[np.ndarray, ...]]]:
+    
+    options = _build_options(kwargs)
+    recon_info = _parse_recon_info(scan)
+    fid_entry = _get_fid_entry(scan)
+
+    with fid_entry.open() as fid_fobj:
+        traj = get_trajectory(recon_info, options)
+
+        with tempfile.NamedTemporaryFile(   
+            mode="w+b", delete=False, dir=options.cache_dir
+        ) as img_fobj:  # where reconstructed dataobj will be stored.
+            
+            if options.correct_spoketiming:
+                with tempfile.NamedTemporaryFile(
+                    mode="w+b", delete=False, dir=options.cache_dir
+                ) as stc_fobj:  # where spoketiming corrected fid will be stored.
+                    
+                    segs = prep_fid_segmentation(fid_fobj, recon_info, options)
+                    logger.info("Spoketiming correction: %s segment(s).", segs.shape[0])
+
+                    stc_param = correct_spoketiming(segs, fid_fobj, stc_fobj, recon_info, options)
+                    dtype = recon_dataobj(stc_fobj, traj, recon_info, img_fobj, options, 
+                                          override_buffer_size=stc_param['buffer_size'],
+                                          override_dtype=stc_param['dtype'])
+            else:
+                dtype = recon_dataobj(fid_fobj, traj, recon_info, img_fobj, options)
+    
+    with open(img_fobj.name, "rb") as img_fobj:
+        dataobj_shape = get_dataobj_shape(recon_info, options)
+        dataobj = np.frombuffer(img_fobj.read(), dtype=dtype).reshape(dataobj_shape, order='F')
+    num_receivers = recon_info.get("EncNReceivers", 1)
+    if not options.as_complex:
+        dataobj = np.abs(dataobj)
+
+    is_multi = num_receivers > 1
+    if not options.split_ch and is_multi:
+        if options.as_complex:
+            dataobj = np.sum(dataobj, axis=0)
+            is_multi = False
+        else:
+            dataobj = np.sqrt(np.sum(dataobj ** 2, axis=0))
+            is_multi = False
+
+    if options.as_complex:
+        if is_multi:
+            dataobj_list = []
+            for receiver_data in dataobj:
+                receiver_arr = cast(NDArray[Any], receiver_data)
+                dataobj_list.extend([
+                    correct_orientation(np.real(receiver_arr), recon_info),
+                    correct_orientation(np.imag(receiver_arr), recon_info),
+                ])
+            return cast(Tuple[np.ndarray, ...], tuple(dataobj_list))
+        return (
+            correct_orientation(np.real(dataobj), recon_info),
+            correct_orientation(np.imag(dataobj), recon_info),
+        )
+
+    if is_multi:
+        return cast(
+            Tuple[np.ndarray, ...],
+            tuple(correct_orientation(ch, recon_info) for ch in dataobj),
+        )
+    return correct_orientation(cast(np.ndarray, dataobj), recon_info)
+
+HOOK = {"get_dataobj": get_dataobj}
+
+__all__ = ["HOOK", "get_dataobj"]

brkraw_sordino/orientation.py

@@ -0,0 +1,45 @@
+import numpy as np
+from numpy.typing import NDArray
+from typing import Dict, Any, Tuple, Optional, cast
+import logging
+
+logger = logging.getLogger("brkraw.sordino")
+
+
+PLANE_PERM = {
+    "axial": (0, 1, 2),
+    "coronal": (0, 2, 1),
+    "sagittal": (1, 2, 0),
+}
+
+
+def apply_axis_perm_flip(dataobj: NDArray, R: NDArray) -> NDArray[Any]:
+    perm = np.argmax(np.abs(R), axis=1)
+    data_t = np.transpose(dataobj, axes=tuple(perm) + tuple(range(3, dataobj.ndim)))
+    return data_t
+
+
+def apply_plane_fix(
+    dataobj: NDArray,
+    plane: str,
+) -> NDArray[Any]:
+    perm = PLANE_PERM[plane]
+    perm_full = tuple(perm) + tuple(range(3, dataobj.ndim))
+    dataobj = np.transpose(dataobj, perm_full)
+    return dataobj
+
+
+def correct(dataobj: NDArray, recon_info: Dict[str, Any]) -> NDArray[Any]:
+    R = cast(NDArray[Any], recon_info.get('GradientOrientation'))
+    plane = cast(str, recon_info.get('SliceOrientation'))
+    dataobj = apply_axis_perm_flip(dataobj, R.T)
+    dataobj = apply_plane_fix(
+        dataobj,
+        plane,
+    )
+    return dataobj
+
+
+__all__ = [
+    "correct",
+]

brkraw_sordino/recon.py

@@ -0,0 +1,125 @@
+import numpy as np
+from typing import Any, Dict, Tuple
+from numpy.typing import NDArray
+import logging
+from .helper import progressbar
+from .typing import Options
+
+logger = logging.getLogger("brkraw.sordino")
+
+
+def parse_fid_info(recon_info: Dict[str, Any]) -> Tuple[np.ndarray, np.dtype]:
+    n_receivers = int(recon_info.get("EncNReceivers") or 0)
+    n_points = int(recon_info.get("NPoints") or 0)
+    n_pro = int(recon_info.get("NPro") or 0)
+    dtype = recon_info['FIDDataType']
+    if not all((n_receivers, n_points, n_pro)):
+        raise ValueError("Missing reconstruction dimensions in recon_spec output.")
+    return np.array([2, n_points, n_receivers, n_pro]), dtype
+
+
+def get_num_frames(recon_info: Dict[str, Any], options: Options):
+    """ Return number of data frames need to be reconstructed
+    """
+    total_frames = recon_info['NRepetitions']
+    offset = getattr(options, 'offset') or 0
+    avail_frames = total_frames - offset
+    set_frames = getattr(options, 'num_frames') or total_frames
+    
+    if set_frames > avail_frames:
+        diff = set_frames - avail_frames
+        set_frames -= diff
+    return set_frames
+
+
+def parse_volume_shape(recon_info: Dict[str, Any], 
+                       options: Options) -> NDArray[np.int_]:
+    matrix = recon_info.get("Matrix")
+    if matrix is None:
+        matrix = [int(recon_info.get("NPoints") or 0)] * 3
+        logger.warning("Matrix size missing; defaulting to %s.", matrix)
+    ext_factors = getattr(options, 'ext_factors', None)
+    if ext_factors is None: 
+        ext_factors = [1.0, 1.0, 1.0]
+    return np.asarray(matrix * np.asarray(ext_factors)).astype(int).tolist()
+
+
+def get_dataobj_shape(recon_info: Dict[str, Any], 
+                      options: Options):
+    num_receivers = parse_fid_info(recon_info)[0][2]
+    vol_shape = parse_volume_shape(recon_info, options)
+    num_frame = get_num_frames(recon_info, options)
+
+    if num_receivers > 1:
+        return [num_receivers] + vol_shape + [num_frame]
+    else:
+        return vol_shape + [num_frame]
+
+
+def nufft_adjoint(kspace, traj, volume_shape, operator='finufft'):
+    """Run nufft and return the reconstucted image"""
+    from mrinufft import get_operator
+    
+    dcf = np.sqrt(np.square(traj).sum(-1)).flatten() ** 2
+    dcf /= dcf.max()
+    traj = traj.copy() / 0.5 * np.pi
+    
+    nufft_op = get_operator(operator)(traj, shape=volume_shape, density=dcf)
+    complex_img = nufft_op.adj_op(kspace.flatten())
+    return complex_img
+
+
+def recon_dataobj(fid_fobj, 
+                  traj, 
+                  recon_info: Dict[str, Any],
+                  img_fobj,
+                  options: Options,
+                  override_buffer_size=None, 
+                  override_dtype=None):
+    img_fobj.seek(0)
+    fid_shape, fid_dtype = parse_fid_info(recon_info)
+    volume_shape = parse_volume_shape(recon_info, options)
+    
+    offset = getattr(options, 'offset') or 0
+    num_frames = get_num_frames(recon_info, options)
+    ignore_samples = getattr(options, 'ignore_samples') or 1
+
+    if all(arg != None for arg in [override_buffer_size, override_buffer_size]):
+        fid_fobj.seek(0)
+        buffer_size = override_buffer_size
+        fid_dtype = override_dtype
+    else:
+        buffer_size = int(np.prod(fid_shape) * fid_dtype.itemsize)
+        buf_offset = offset * buffer_size
+        fid_fobj.seek(buf_offset)
+    
+    trimmed_traj = traj[:, ignore_samples:, ...]
+    logger.debug("Reconstruction traj shape: %s", trimmed_traj.shape)
+    
+    dtype = None
+    for n in progressbar(range(num_frames), desc='frames', ncols=100):
+        buffer = fid_fobj.read(buffer_size)
+        vol = np.frombuffer(buffer, dtype=fid_dtype).reshape(fid_shape, order='F')
+        vol = (vol[0] + 1j * vol[1])[np.newaxis, ...]
+        k_space = vol.squeeze().T[..., ignore_samples:]
+        logger.debug("Reconstruction k-space shape: %s", k_space.shape)
+        n_receivers = fid_shape[2]
+        
+        if n_receivers > 1:
+            recon_vol = []
+            for ch_id in range(n_receivers):
+                _k_space = k_space[:, ch_id, :]
+                _vol = nufft_adjoint(_k_space, trimmed_traj, volume_shape)
+                recon_vol.append(_vol)
+            recon_vol = np.stack(recon_vol, axis=0)
+        else:
+            recon_vol = nufft_adjoint(k_space, trimmed_traj, volume_shape)
+        if n == 0:
+            dtype = recon_vol.dtype
+        img_fobj.write(recon_vol.T.flatten(order="C").tobytes())
+    return dtype
+
+__all__ = [
+    'recon_dataobj',
+    'get_dataobj_shape',
+]

brkraw_sordino/rules/sordino.yaml

@@ -0,0 +1,35 @@
+converter_hook:
+  - name: "sordino-hook"
+    description: "SORDINO-ZTE reconstruction hook."
+    when:
+      method:
+        sources:
+          - file: method
+            key: Method
+    if:
+      regex: ["$method", "sordino"]
+    use: "sordino"
+
+info_spec:
+  - name: "sordino-info"
+    description: "SORDINO-ZTE scan info."
+    when:
+      method:
+        sources:
+          - file: method
+            key: Method
+    if:
+      regex: ["$method", "sordino"]
+    use: "sordino_info"
+
+metadata_spec:
+  - name: "sordino-metadata"
+    description: "SORDINO-ZTE metadata (draft BEP alignment)."
+    when:
+      method:
+        sources:
+          - file: method
+            key: Method
+    if:
+      regex: ["$method", "sordino"]
+    use: "sordino_metadata"