augmentrum 0.0.2__tar.gz → 0.0.3__tar.gz

{augmentrum-0.0.2 → augmentrum-0.0.3}/MANIFEST.in

@@ -1,7 +1,6 @@
 include README.md
 include LICENSE
 recursive-include augmentrum *.py
-recursive-include augmentrum/resources/scans *
 recursive-include assets/figures *.png *.jpg *.svg
 recursive-exclude * __pycache__
 recursive-exclude * *.pyc

{augmentrum-0.0.2/augmentrum.egg-info → augmentrum-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: augmentrum
-Version: 0.0.2
+Version: 0.0.3
 Summary: A Data Augmentation Package for MR Spectroscopy
 Author-email: "John T. LaMaster" <jlamaste@gmail.com>, "Julian P. Merkofer" <j.p.merkofer@tue.nl>, "Kay C. Igwe" <kci2104@columbia.edu>
 License: Apache-2.0
@@ -23,14 +23,14 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy>=1.21.0
 Requires-Dist: scipy>=1.7.0
-Requires-Dist: torch>=1.10.0
-Requires-Dist: fsl-mrs>=2.1.0
 Requires-Dist: nifti-mrs>=0.1.0
+Requires-Dist: nifti-mrs-plus[ops]>=0.2.1
 Requires-Dist: spec2nii>=0.6.0
-Requires-Dist: torchkbnufft>=1.3.0
 Requires-Dist: tqdm>=4.62.0
 Requires-Dist: matplotlib>=3.4.0
 Requires-Dist: shutup>=0.2.0
+Provides-Extra: fsl
+Requires-Dist: fsl-mrs>=2.1.0; extra == "fsl"
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0.0; extra == "dev"
 Requires-Dist: pytest-cov>=3.0.0; extra == "dev"
@@ -38,20 +38,42 @@ Requires-Dist: pytest-mock>=3.10.0; extra == "dev"
 Requires-Dist: black>=22.0.0; extra == "dev"
 Requires-Dist: flake8>=4.0.0; extra == "dev"
 Requires-Dist: mypy>=0.950; extra == "dev"
+Provides-Extra: torch
+Requires-Dist: torch>=1.10.0; extra == "torch"
 Provides-Extra: jax
 Requires-Dist: jax>=0.4.0; extra == "jax"
 Requires-Dist: jaxlib>=0.4.0; extra == "jax"
+Provides-Extra: nufft
+Requires-Dist: torch>=1.10.0; extra == "nufft"
+Requires-Dist: torchkbnufft>=1.3.0; extra == "nufft"
+Provides-Extra: tensorflow
+Requires-Dist: tensorflow>=2.8.0; extra == "tensorflow"
+Provides-Extra: csm
+Requires-Dist: h5py>=3.0.0; extra == "csm"
 Provides-Extra: keras
 Requires-Dist: keras>=3.0.0; extra == "keras"
+Provides-Extra: seq
+Requires-Dist: pypulseq>=1.4.0; extra == "seq"
+Provides-Extra: docs
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: sphinx-autobuild; extra == "docs"
+Requires-Dist: sphinx-autodoc-typehints; extra == "docs"
+Requires-Dist: sphinx-copybutton; extra == "docs"
+Requires-Dist: myst-parser; extra == "docs"
+Requires-Dist: furo; extra == "docs"
 Provides-Extra: all
+Requires-Dist: torch>=1.10.0; extra == "all"
+Requires-Dist: torchkbnufft>=1.3.0; extra == "all"
+Requires-Dist: h5py>=3.0.0; extra == "all"
 Requires-Dist: tensorflow>=2.8.0; extra == "all"
 Requires-Dist: keras>=3.0.0; extra == "all"
 Requires-Dist: jax>=0.4.0; extra == "all"
 Requires-Dist: jaxlib>=0.4.0; extra == "all"
+Requires-Dist: pypulseq>=1.4.0; extra == "all"
 Dynamic: license-file
 
 <div align="center">
-  <img src="https://raw.githubusercontent.com/julianmer/Augmentrum/main/assets/figures/logo.png" alt="Augmentrum Logo" width="100" style="margin-bottom: -10px;"/>
+  <img src="https://raw.githubusercontent.com/julianmer/Augmentrum/main/assets/logo/logo.png" alt="Augmentrum Logo" width="100" style="margin-bottom: -10px;"/>
   <h1 style="margin-top: 5px; margin-bottom: 5px;">Augmentrum</h1>
   <p style="margin-top: 0px;"><em>A Data Augmentation Package for MR Spectroscopy</em></p>
   
@@ -131,6 +153,58 @@ augmenter = Augmentrum(
 train_data = augmenter.dataloader()
 ```
 
+### Supervised Pairs: Taps and `outputs`
+
+Any pipeline stage can be named with a `'tap:<name>'` marker and yielded by the
+dataloaders through the `outputs` spec — an arbitrarily nested tuple of stage
+tokens (`'data'` / `'water'` for the pipeline end, `'<tap>'` / `'<tap>.water'`
+for a tapped stage). The classic `(data, water)` pair stays the default.
+
+```python
+augmenter = Augmentrum(
+    data=data,
+    pipeline=['coil_sampling', 'macromolecules', 'tap:clean', 'undersampling'],
+    outputs=(('data', 'water'), ('clean', 'clean.water')),
+    backend='pytorch',
+)
+
+# supervised reconstruction pairs: input fully augmented, target frozen at the tap
+for (x, x_water), (y, y_water) in augmenter.dataloader():
+    train_step(x, y)
+```
+
+### Per-Step Parameters, Seeds and Subject-Wise Splits
+
+A keyword argument reaches every module whose constructor names it, so
+`lb_hz=(0, 5)` would move `LineBroadening` *and* `Apodization` (you get a
+warning). A pipeline entry can carry its own kwargs, which reach that step only
+and override the globals there; a kwarg no module accepts raises with the
+closest valid names instead of being ignored.
+
+```python
+augmenter = Augmentrum(
+    data=data,
+    groups=subject_ids,                     # one id per item: a subject never straddles splits
+    split_fractions={'val': 0.1, 'test': 0.1},
+    pipeline=[
+        'coil_sampling',
+        {'noise': {'sigma_frac': (0.0, 0.02)}},
+        {'line_broadening': {'lb_hz': (0.0, 5.0)}},
+        ('apodization', {'mode': 'exponential', 'lb_hz': 2.0}),
+        'baseline_bspline',                 # registry alias fixing mode='bspline'
+    ],
+    n_coils=(1, 8),                         # integer-typed ranges are inclusive
+    seed=1,                                 # fixes the split, subject draws, ranges and modules
+)
+augmenter.split_groups['val']               # the subject ids held out for validation
+```
+
+Without `groups`, a user-defined `SubjectID` header-extension field is used
+when every item carries one. `seed` makes the whole run replay — including
+`num_workers > 0` DataLoaders, where each worker is reseeded from it — and
+`augmenter.reseed(k)` restarts every stream; `'fixed'` mode draws its ranges
+once per split and keeps them across `dataloader()` calls.
+
 ---
 
 ## Module Reference & Backend Support
@@ -143,30 +217,36 @@ no manual conversion needed.
 | Module | Modes / Methods                                         | NIfTI | NumPy | PyTorch | TensorFlow | JAX | Keras |
 |:---|:--------------------------------------------------------|:-----:|:---:|:---:|:---:|:---:|:---:|
 | `AmplitudeScaling` | uniform, normal                                         |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `Apodization` | exponential, truncate                                   |   ✓   | ✓† | ✓† | ✓† | ✓† | ✓† |
+| `Apodization` | exponential, gaussian, hamming, truncate                |   ✓   | ✓† | ✓† | ✓† | ✓† | ✓† |
 | `ArtificialPeaks` | Lorentzian, Gaussian, Voigt                             |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `AverageSampler` | random, consecutive, strided, deterministic             |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
 | `BaselineAugmentation` | random_walk, bspline, polynomial                        |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `CoilAverageSampler` | random, deterministic                                   |   ✓   | ~ | ~ | ~ | ~ | ~ |
-| `EddyCurrent` | synthetic, water                                        |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `FrequencyShift` | —                                                       |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `GaussianNoise` | sigma, sigma_frac, snr, snr_db                          |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `KspaceReconstructor` | NUFFT + density compensation                            |   —   | — | ✓ | — | — | — |
+| `CoilSampler` | draw random, reweight, grow from CSMs |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `EddyCurrent` | synthetic, water (call-time reference or source library) |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `EditCombiner` | diff, sum along DIM_EDIT            |   ~   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `FrequencyShift` | shift_hz                                                |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `Noise` | sigma, sigma_frac, snr, snr_db                          |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `KspaceUndersampling` | cartesian, gridded, nufft, dual trajectory              |   ~   | ✓ | ✓ | ✓ | ✓ | ✓ |
 | `LineBroadening` | lorentzian, gaussian, voigt                             |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `NIfTI_RawProcessor` | coil combination, alignment, averaging, ECC, phase/freq |   ✓   | ~ | ~ | ~ | ~ | ~ |
+| `Macromolecules` | parametrized, measured, supplied     |   ~   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `RawProcessor` | coil combination, registration, ECC, ... |   ✓   | ✓‡ | ✓‡ | ✓‡ | ✓‡ | ✓‡ |
 | `PhaseShift` | zero_order, first_order                                 |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `ResidualWater` | —                                                       |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `SpatialAugmentations` | 2-D / 3-D affine, flip, zoom, shear                     |  ✓ ‡  | ✓ ‡ | ✓ | ~ | ~ | ~ |
-| `SpuriousEchoes` | replica, hybrid                                         |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `ResidualWater` | lobes, turco (7-Lorentzian WaterFit)                    |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `SpatialAugmentations` | 2-D / 3-D affine, flip, zoom, shear                     |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `SpuriousEchoes` | echo, replica, hybrid                                   |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `Tap` | marker: snapshot a stage for `outputs`         |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `TransientSynthesizer` | correlated drift, respiration, motion events            |   ~   | ✓ | ✓ | ✓ | ✓ | ✓ |
 | `ZeroFill` | pad FID to target length                                |   ✓   | ✓† | ✓† | ✓† | ✓† | ✓† |
 
 **✓** native — data tensor stays in the target framework throughout.  
 **~** automatic — the base class routes the call through the NIfTI-list path; the result is returned in the original backend format. Functional, but not natively differentiable.  
 **—** not supported / not applicable.
 
-> **†** `Apodization[truncate]` and `ZeroFill` both change `N_PTS`. Because `target_pts` / `n_pts` is a module-level scalar, the **same length is applied uniformly to every batch member** — the output is still a uniform tensor. On non-NIfTI backends the base class detects the shape change, rebuilds each NIfTI_MRS object at the new length, and emits a `RuntimeWarning`.
+Modes that differ in backend support get a row each, named `Module[mode]`.
 
-> **‡** `SpatialAugmentations` uses `F.grid_sample` (PyTorch) internally; NIfTI-list and NumPy inputs are converted to PyTorch for processing and converted back. TensorFlow, JAX, and Keras are handled by the base-class automatic routing (~).
+> **†** `Apodization[truncate]` and `ZeroFill` both change `N_PTS`. Because `target_pts` / `n_pts` is a module-level scalar, the **same length is applied uniformly to every batch member** — the output is still a uniform tensor.
 
+> **‡** `RawProcessor[remove_outliers]` on tensor backends **zeroes outlier transients instead of dropping them** — a batched tensor cannot go ragged. 
 
 ---
 

{augmentrum-0.0.2 → augmentrum-0.0.3}/README.md

@@ -1,5 +1,5 @@
 <div align="center">
-  <img src="https://raw.githubusercontent.com/julianmer/Augmentrum/main/assets/figures/logo.png" alt="Augmentrum Logo" width="100" style="margin-bottom: -10px;"/>
+  <img src="https://raw.githubusercontent.com/julianmer/Augmentrum/main/assets/logo/logo.png" alt="Augmentrum Logo" width="100" style="margin-bottom: -10px;"/>
   <h1 style="margin-top: 5px; margin-bottom: 5px;">Augmentrum</h1>
   <p style="margin-top: 0px;"><em>A Data Augmentation Package for MR Spectroscopy</em></p>
   
@@ -79,6 +79,58 @@ augmenter = Augmentrum(
 train_data = augmenter.dataloader()
 ```
 
+### Supervised Pairs: Taps and `outputs`
+
+Any pipeline stage can be named with a `'tap:<name>'` marker and yielded by the
+dataloaders through the `outputs` spec — an arbitrarily nested tuple of stage
+tokens (`'data'` / `'water'` for the pipeline end, `'<tap>'` / `'<tap>.water'`
+for a tapped stage). The classic `(data, water)` pair stays the default.
+
+```python
+augmenter = Augmentrum(
+    data=data,
+    pipeline=['coil_sampling', 'macromolecules', 'tap:clean', 'undersampling'],
+    outputs=(('data', 'water'), ('clean', 'clean.water')),
+    backend='pytorch',
+)
+
+# supervised reconstruction pairs: input fully augmented, target frozen at the tap
+for (x, x_water), (y, y_water) in augmenter.dataloader():
+    train_step(x, y)
+```
+
+### Per-Step Parameters, Seeds and Subject-Wise Splits
+
+A keyword argument reaches every module whose constructor names it, so
+`lb_hz=(0, 5)` would move `LineBroadening` *and* `Apodization` (you get a
+warning). A pipeline entry can carry its own kwargs, which reach that step only
+and override the globals there; a kwarg no module accepts raises with the
+closest valid names instead of being ignored.
+
+```python
+augmenter = Augmentrum(
+    data=data,
+    groups=subject_ids,                     # one id per item: a subject never straddles splits
+    split_fractions={'val': 0.1, 'test': 0.1},
+    pipeline=[
+        'coil_sampling',
+        {'noise': {'sigma_frac': (0.0, 0.02)}},
+        {'line_broadening': {'lb_hz': (0.0, 5.0)}},
+        ('apodization', {'mode': 'exponential', 'lb_hz': 2.0}),
+        'baseline_bspline',                 # registry alias fixing mode='bspline'
+    ],
+    n_coils=(1, 8),                         # integer-typed ranges are inclusive
+    seed=1,                                 # fixes the split, subject draws, ranges and modules
+)
+augmenter.split_groups['val']               # the subject ids held out for validation
+```
+
+Without `groups`, a user-defined `SubjectID` header-extension field is used
+when every item carries one. `seed` makes the whole run replay — including
+`num_workers > 0` DataLoaders, where each worker is reseeded from it — and
+`augmenter.reseed(k)` restarts every stream; `'fixed'` mode draws its ranges
+once per split and keeps them across `dataloader()` calls.
+
 ---
 
 ## Module Reference & Backend Support
@@ -91,30 +143,36 @@ no manual conversion needed.
 | Module | Modes / Methods                                         | NIfTI | NumPy | PyTorch | TensorFlow | JAX | Keras |
 |:---|:--------------------------------------------------------|:-----:|:---:|:---:|:---:|:---:|:---:|
 | `AmplitudeScaling` | uniform, normal                                         |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `Apodization` | exponential, truncate                                   |   ✓   | ✓† | ✓† | ✓† | ✓† | ✓† |
+| `Apodization` | exponential, gaussian, hamming, truncate                |   ✓   | ✓† | ✓† | ✓† | ✓† | ✓† |
 | `ArtificialPeaks` | Lorentzian, Gaussian, Voigt                             |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `AverageSampler` | random, consecutive, strided, deterministic             |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
 | `BaselineAugmentation` | random_walk, bspline, polynomial                        |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `CoilAverageSampler` | random, deterministic                                   |   ✓   | ~ | ~ | ~ | ~ | ~ |
-| `EddyCurrent` | synthetic, water                                        |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `FrequencyShift` | —                                                       |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `GaussianNoise` | sigma, sigma_frac, snr, snr_db                          |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `KspaceReconstructor` | NUFFT + density compensation                            |   —   | — | ✓ | — | — | — |
+| `CoilSampler` | draw random, reweight, grow from CSMs |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `EddyCurrent` | synthetic, water (call-time reference or source library) |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `EditCombiner` | diff, sum along DIM_EDIT            |   ~   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `FrequencyShift` | shift_hz                                                |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `Noise` | sigma, sigma_frac, snr, snr_db                          |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `KspaceUndersampling` | cartesian, gridded, nufft, dual trajectory              |   ~   | ✓ | ✓ | ✓ | ✓ | ✓ |
 | `LineBroadening` | lorentzian, gaussian, voigt                             |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `NIfTI_RawProcessor` | coil combination, alignment, averaging, ECC, phase/freq |   ✓   | ~ | ~ | ~ | ~ | ~ |
+| `Macromolecules` | parametrized, measured, supplied     |   ~   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `RawProcessor` | coil combination, registration, ECC, ... |   ✓   | ✓‡ | ✓‡ | ✓‡ | ✓‡ | ✓‡ |
 | `PhaseShift` | zero_order, first_order                                 |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `ResidualWater` | —                                                       |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
-| `SpatialAugmentations` | 2-D / 3-D affine, flip, zoom, shear                     |  ✓ ‡  | ✓ ‡ | ✓ | ~ | ~ | ~ |
-| `SpuriousEchoes` | replica, hybrid                                         |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `ResidualWater` | lobes, turco (7-Lorentzian WaterFit)                    |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `SpatialAugmentations` | 2-D / 3-D affine, flip, zoom, shear                     |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `SpuriousEchoes` | echo, replica, hybrid                                   |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `Tap` | marker: snapshot a stage for `outputs`         |   ✓   | ✓ | ✓ | ✓ | ✓ | ✓ |
+| `TransientSynthesizer` | correlated drift, respiration, motion events            |   ~   | ✓ | ✓ | ✓ | ✓ | ✓ |
 | `ZeroFill` | pad FID to target length                                |   ✓   | ✓† | ✓† | ✓† | ✓† | ✓† |
 
 **✓** native — data tensor stays in the target framework throughout.  
 **~** automatic — the base class routes the call through the NIfTI-list path; the result is returned in the original backend format. Functional, but not natively differentiable.  
 **—** not supported / not applicable.
 
-> **†** `Apodization[truncate]` and `ZeroFill` both change `N_PTS`. Because `target_pts` / `n_pts` is a module-level scalar, the **same length is applied uniformly to every batch member** — the output is still a uniform tensor. On non-NIfTI backends the base class detects the shape change, rebuilds each NIfTI_MRS object at the new length, and emits a `RuntimeWarning`.
+Modes that differ in backend support get a row each, named `Module[mode]`.
 
-> **‡** `SpatialAugmentations` uses `F.grid_sample` (PyTorch) internally; NIfTI-list and NumPy inputs are converted to PyTorch for processing and converted back. TensorFlow, JAX, and Keras are handled by the base-class automatic routing (~).
+> **†** `Apodization[truncate]` and `ZeroFill` both change `N_PTS`. Because `target_pts` / `n_pts` is a module-level scalar, the **same length is applied uniformly to every batch member** — the output is still a uniform tensor.
 
+> **‡** `RawProcessor[remove_outliers]` on tensor backends **zeroes outlier transients instead of dropping them** — a batched tensor cannot go ragged. 
 
 ---
 

augmentrum-0.0.3/augmentrum/__init__.py

@@ -0,0 +1,41 @@
+####################################################################################################
+#                                           __init__.py                                            #
+####################################################################################################
+#                                                                                                  #
+# Authors: J. T. LaMaster (john.t.lamaster@gmail.com)                                              #
+#          J. P. Merkofer (j.p.merkofer@tue.nl)                                                    #
+#          K. C. Igwe (kci2104@columbia.edu)                                                       #
+#                                                                                                  #
+# Created: 2025-10-25                                                                              #
+#                                                                                                  #
+# Purpose: Package entry point. Re-exports the public API: the Augmentrum container, the pipeline, #
+#          and every augmentation module.                                                          #
+#                                                                                                  #
+####################################################################################################
+
+__version__ = "0.0.3"
+__author__ = "John T. LaMaster, Julian P. Merkofer, Kay C. Igwe"
+__email__ = "j.p.merkofer@tue.nl"
+
+# Main API
+#*************#
+#   imports   #
+#*************#
+import nifti_mrs_plus
+
+# Claim the NIfTI-MRS provenance record before anything can write one. The
+# container is a shared transport, so left alone it would credit itself for
+# work Augmentrum did.
+nifti_mrs_plus.set_provenance("Augmentrum", __version__)
+
+from augmentrum.core import Augmentrum, NIfTI_MRS_Plus, Backend, BaseModule, AugmentationPipeline  # noqa: E402
+
+__all__ = [
+    "__version__",
+    "Augmentrum",
+    "NIfTI_MRS_Plus",
+    "Backend",
+    "BaseModule",
+    "AugmentationPipeline",
+]
+

augmentrum-0.0.3/augmentrum/augmentation/__init__.py

@@ -0,0 +1,56 @@
+####################################################################################################
+#                                     augmentation/__init__.py                                     #
+####################################################################################################
+#                                                                                                  #
+# Authors: K. C. Igwe (kci2104@columbia.edu)                                                       #
+#          J. P. Merkofer (j.p.merkofer@tue.nl)                                                    #
+#                                                                                                  #
+# Created: 2026-02-13                                                                              #
+#                                                                                                  #
+# Purpose: Re-exports the augmentation modules, each a BaseModule that perturbs MRS data in a      #
+#          physically motivated way.                                                               #
+#                                                                                                  #
+####################################################################################################
+
+#*************#
+#   imports   #
+#*************#
+from augmentrum.augmentation.noise import Noise
+from augmentrum.augmentation.line_broadening import LineBroadening
+from augmentrum.augmentation.baseline_augmentation import BaselineAugmentation
+from augmentrum.augmentation.residual_water import ResidualWater
+from augmentrum.augmentation.spurious_echoes import SpuriousEchoes
+from augmentrum.augmentation.artificial_peaks import ArtificialPeaks
+from augmentrum.augmentation.eddy_current import EddyCurrent
+from augmentrum.augmentation.apodization import Apodization
+from augmentrum.augmentation.phase_frequency import PhaseShift, FrequencyShift
+from augmentrum.augmentation.amplitude_scaling import AmplitudeScaling
+from augmentrum.augmentation.spatial_augmentations import SpatialAugmentations
+from augmentrum.augmentation.transient_synthesis import TransientSynthesizer
+from augmentrum.augmentation.zero_fill import ZeroFill
+from augmentrum.augmentation.macromolecules import (
+    Macromolecules, MMSource, Parametrized, SemiParametrized, Measured, Supplied,
+)
+
+__all__ = [
+    'Noise',
+    'LineBroadening',
+    'BaselineAugmentation',
+    'ResidualWater',
+    'Macromolecules',
+    'MMSource',
+    'Parametrized',
+    'SemiParametrized',
+    'Measured',
+    'Supplied',
+    'SpuriousEchoes',
+    'ArtificialPeaks',
+    'EddyCurrent',
+    'Apodization',
+    'PhaseShift',
+    'FrequencyShift',
+    'AmplitudeScaling',
+    'SpatialAugmentations',
+    'TransientSynthesizer',
+    'ZeroFill',
+]

{augmentrum-0.0.2 → augmentrum-0.0.3}/augmentrum/augmentation/amplitude_scaling.py

@@ -16,18 +16,23 @@
 #*************#
 #   imports   #
 #*************#
+from __future__ import annotations
+
 from typing import List, Optional, Union, Tuple
 import numpy as np
-from fsl_mrs.core import NIFTI_MRS
 
 from augmentrum.core.base_module import BaseModule
-from augmentrum.core.nifti_mrs_plus import Backend
-from augmentrum.utils.tensor_ops import match_backend
+from augmentrum.processing.utils import to_backend
+from nifti_mrs_plus import Backend
 
 
-####################################################################################################
-#                                      AmplitudeScaling Class                                      #
-####################################################################################################
+#**************************************************************************************************#
+#                                      Class AmplitudeScaling                                      #
+#**************************************************************************************************#
+#                                                                                                  #
+# Apply amplitude scaling to MRS data.                                                             #
+#                                                                                                  #
+#**************************************************************************************************#
 class AmplitudeScaling(BaseModule):
     """
     Apply amplitude scaling to MRS data.
@@ -65,6 +70,12 @@ class AmplitudeScaling(BaseModule):
     SUPPORTED_BACKENDS = [Backend.NIFTI_LIST, Backend.NUMPY, Backend.PYTORCH,
                          Backend.TENSORFLOW, Backend.JAX, Backend.KERAS]
 
+    # Acts on every coil and transient alike, so per-sample masks pass through.
+    MASKS = 'pass'
+
+    # A scale broadcasts, so a batch can carry one factor per sample.
+    PER_SAMPLE_PARAMS = ('scale_factor',)
+
     def __init__(
         self,
         scale_factor: Union[float, Tuple[float, float]] = 1.0,
@@ -73,29 +84,30 @@ class AmplitudeScaling(BaseModule):
         **kwargs
     ):
         """Initialize AmplitudeScaling module."""
-        super().__init__(**kwargs)
+        super().__init__()
 
         self.scale_factor = scale_factor
         self.distribution = distribution
         self.seed = seed
-        self.rng = np.random.default_rng(seed)
 
         # Validate distribution
         if distribution not in ['uniform', 'normal']:
             raise ValueError(f"distribution must be 'uniform' or 'normal', got {distribution}")
 
-    def _get_scale_factor(self) -> float:
-        """Get scale factor (sample if range is provided)."""
+    def _get_scale_factor(self, index: int = 0) -> float:
+        """Get scale factor (sample if a range, index if a per-sample vector)."""
+        if isinstance(self.scale_factor, np.ndarray) and self.scale_factor.ndim:
+            return float(self.sample_of(self.scale_factor, index))
         if isinstance(self.scale_factor, (tuple, list)):
             min_scale, max_scale = self.scale_factor
 
             if self.distribution == 'uniform':
-                return self.rng.uniform(min_scale, max_scale)
+                return self.rng.numpy_rng().uniform(min_scale, max_scale)
             elif self.distribution == 'normal':
                 # Use range as mean ± std
                 mean = (min_scale + max_scale) / 2.0
                 std = (max_scale - min_scale) / 4.0  # ~95% within range
-                return self.rng.normal(mean, std)
+                return self.rng.numpy_rng().normal(mean, std)
         else:
             return float(self.scale_factor)
 
@@ -117,9 +129,9 @@ class AmplitudeScaling(BaseModule):
         """
         scaled_data = []
 
-        for nifti in data_list:
-            # Get scale factor (sample if random)
-            scale = self._get_scale_factor()
+        for i, nifti in enumerate(data_list):
+            # Get scale factor (sample if random, index if per-sample)
+            scale = self._get_scale_factor(i)
 
             # Get data
             fid = nifti[:].copy()
@@ -154,14 +166,14 @@ class AmplitudeScaling(BaseModule):
         """
         # Get scale factor for each sample in batch
         batch_size = data.shape[0]
-        scales = np.array([self._get_scale_factor() for _ in range(batch_size)])
+        scales = np.array([self._get_scale_factor(i) for i in range(batch_size)])
 
         # Reshape scales for broadcasting
         scale_shape = [batch_size] + [1] * (data.ndim - 1)
         scales = scales.reshape(scale_shape)
 
         # Apply scaling
-        data_scaled = data * match_backend(scales, data)
+        data_scaled = data * to_backend(scales, data)
 
         return data_scaled, water