Choise for target accel in transforms

direct/common/subsample_config.py

@@ -11,11 +11,11 @@
 @dataclass
 class MaskingConfig(BaseConfig):
     name: str = MISSING
-    accelerations: Tuple[int, ...] = (5,)  # Ideally Union[float, int].
+    accelerations: Tuple[float, ...] = (5.0,)  # Ideally Union[float, int].
     center_fractions: Optional[Tuple[float, ...]] = (0.1,)  # Ideally Optional[Tuple[float, ...]]
     uniform_range: bool = False
     image_center_crop: bool = False
     dynamic_mask: Optional[bool] = None
 
-    val_accelerations: Tuple[int, ...] = (5, 10)
+    val_accelerations: Tuple[float, ...] = (5.0, 10.0)
     val_center_fractions: Optional[Tuple[float, ...]] = (0.1, 0.05)

direct/data/datasets_config.py

@@ -71,6 +71,7 @@ class TransformsConfig(BaseConfig):
     image_recon_type: str = "rss"
     pad_coils: Optional[int] = None
     use_seed: bool = True
+    target_accelerations: Optional[tuple[float, ...]] = None
 
 
 @dataclass

direct/data/mri_transforms.py

@@ -339,7 +339,7 @@ def __call__(self, sample: Dict[str, Any]) -> Dict[str, Any]:
 
             for i in range(nz):
                 sampling_mask_z, acceleration_z, center_fraction_z = self.mask_func(
-                    shape=shape, seed=seed, return_acs=False, return_acceleration=True
+                    shape=shape, seed=dynamic_seeds[i], return_acs=False, return_acceleration=True
                 )
 
                 sampling_mask.append(sampling_mask_z.to(sample["kspace"].dtype))
@@ -1642,6 +1642,30 @@ def forward(self, sample: Dict[str, Any]) -> Dict[str, Any]:
         return sample
 
 
+class AddTargetAcceleration(DirectTransform):
+    """This will find the sample acceleration provided by the mask function.
+
+    Then it will add a tensor with the corresponding acceleration in the target accelerations."""
+    def __init__(self, mask_func: Callable, target_accelerations: tuple[int, ...]):
+        super().__init__()
+        if mask_func.uniform_range:
+            raise ValueError(f"Cannot apply this transform for `uniform_range`=True for the mask function.")
+        self.mask_func_accelerations = mask_func.accelerations
+        self.target_accelerations = target_accelerations
+
+    def __call__(self, sample: Dict[str, Any]):
+        # Convert tensor to Python scalar
+        sample_acceleration = sample["acceleration"].item()
+
+        # Find the index of the value in the list
+        ind = self.mask_func_accelerations.index(sample_acceleration)
+
+        target_acceleration = self.target_accelerations[ind]
+        sample["acceleration"] = torch.tensor([target_acceleration], dtype=sample["acceleration"].dtype)
+
+        return sample
+
+
 class ModuleWrapper:
     class SubWrapper:
         def __init__(self, transform, toggle_dims):
@@ -2055,6 +2079,7 @@ def build_mri_transforms(
     scaling_key: TransformKey = TransformKey.MASKED_KSPACE,
     scale_percentile: Optional[float] = 0.99,
     use_seed: bool = True,
+    target_accelerations: Optional[tuple[int,...]] = None,
 ) -> object:
     """Build transforms for MRI.
 
@@ -2228,6 +2253,7 @@ def build_mri_transforms(
                 return_acs=estimate_sensitivity_maps,
                 dynamic_mask=dynamic_mask,
             ),
+            AddTargetAcceleration(mask_func, target_accelerations),
         ]
     if compute_and_apply_padding:
         mri_transforms += [ApplyZeroPadding("sampling_mask", "padding")]

direct/nn/adaptive/config.py

@@ -20,6 +20,8 @@ class PolicyConfig(ModelConfig):
     use_softplus: bool = True
     slope: float = 10
     fix_sign_leakage: bool = True
+    acceleration: Optional[float] = None
+    center_fraction: Optional[float] = None
 
 
 @dataclass

direct/nn/adaptive/parameterized.py

@@ -43,6 +43,8 @@ def __init__(
         fix_sign_leakage: bool = True,
         st_slope: float = 10,
         st_clamp: bool = False,
+        acceleration: Optional[float] = None,
+        center_fraction: Optional[float] = None,
     ):
         """Inits :class:`ParameterizedPolicy`.
 
@@ -129,6 +131,9 @@ def __init__(
 
         self.sampling_type = sampling_type
 
+        self.acceleration = acceleration
+        self.center_fraction = center_fraction
+
 
 class ParameterizedStaticPolicy(ParameterizedPolicy):
     """Base Parameterized policy model for non dynamic 2D or 3D data."""
@@ -142,6 +147,8 @@ def __init__(
         fix_sign_leakage: bool = True,
         st_slope: float = 10,
         st_clamp: bool = False,
+        acceleration: Optional[float] = None,
+        center_fraction: Optional[float] = None,
     ):
         """Inits :class:`ParameterizedStaticPolicy`.
 
@@ -173,6 +180,8 @@ def __init__(
             fix_sign_leakage=fix_sign_leakage,
             st_slope=st_slope,
             st_clamp=st_clamp,
+            acceleration=acceleration,
+            center_fraction=center_fraction,
         )
 
     @abstractmethod
@@ -244,6 +253,16 @@ def forward(
         nonzero_idcs = (mask == 0).nonzero(as_tuple=True)
         probs_to_norm = masked_prob_mask[nonzero_idcs].reshape(batch_size, -1)
 
+        if (self.acceleration is not None) and (self.center_fraction is not None):
+            acceleration = self.acceleration
+            center_fraction = center_fraction
+        else:
+            if (acceleration is None) or (center_fraction is None):
+                raise ValueError(f"One of `acceleration` or `center_fraction` received None for a value. "
+                                 f"This should not be None when `StraightThroughPolicy` is initialized "
+                                 f"with `acceleration` or `center_fraction` with None value."
+                                 )
+
         # Rescale probabilities to desired sparsity.
         budget = self.num_actions / acceleration - self.num_actions * center_fraction
         if isinstance(budget, float):
@@ -289,6 +308,8 @@ def __init__(
         fix_sign_leakage: bool = True,
         st_slope: float = 10,
         st_clamp: bool = False,
+        acceleration: Optional[float] = None,
+        center_fraction: Optional[float] = None,
     ):
         super().__init__(
             kspace_shape=kspace_shape,
@@ -298,6 +319,8 @@ def __init__(
             fix_sign_leakage=fix_sign_leakage,
             st_slope=st_slope,
             st_clamp=st_clamp,
+            acceleration=acceleration,
+            center_fraction=center_fraction,
         )
 
     def dim_check(self, kspace: torch.Tensor) -> None:
@@ -317,6 +340,8 @@ def __init__(
         fix_sign_leakage: bool = True,
         st_slope: float = 10,
         st_clamp: bool = False,
+        acceleration: Optional[float] = None,
+        center_fraction: Optional[float] = None,
     ):
         super().__init__(
             kspace_shape=kspace_shape,
@@ -326,6 +351,8 @@ def __init__(
             fix_sign_leakage=fix_sign_leakage,
             st_slope=st_slope,
             st_clamp=st_clamp,
+            acceleration=acceleration,
+            center_fraction=center_fraction,
         )
 
     def dim_check(self, kspace: torch.Tensor) -> None:
@@ -348,6 +375,8 @@ def __init__(
         fix_sign_leakage: bool = True,
         st_slope: float = 10,
         st_clamp: bool = False,
+        acceleration: Optional[float] = None,
+        center_fraction: Optional[float] = None,
     ):
         """Inits :class:`ParameterizedDynamicOrMultislice2dPolicy`.
 
@@ -385,6 +414,8 @@ def __init__(
             fix_sign_leakage=fix_sign_leakage,
             st_slope=st_slope,
             st_clamp=st_clamp,
+            acceleration=acceleration,
+            center_fraction=center_fraction,
         )
 
     def forward(
@@ -423,6 +454,16 @@ def forward(
         batch_size, _, slices, height, width, _ = kspace.shape  # batch, coils, time, height, width, complex
         masks = [mask.expand(batch_size, 1, slices, height, width, 1)]
 
+        if (self.acceleration is not None) and (self.center_fraction is not None):
+            acceleration = self.acceleration
+            center_fraction = center_fraction
+        else:
+            if (acceleration is None) or (center_fraction is None):
+                raise ValueError(f"One of `acceleration` or `center_fraction` received None for a value. "
+                                 f"This should not be None when `StraightThroughPolicy` is initialized "
+                                 f"with `acceleration` or `center_fraction` with None value."
+                                 )
+
         budget = self.num_actions / acceleration - self.num_actions * center_fraction
         if isinstance(budget, float):
             budget = int(budget)
@@ -574,6 +615,8 @@ def __init__(
         st_slope: float = 10,
         st_clamp: bool = False,
         non_uniform: bool = False,
+        acceleration: Optional[float] = None,
+        center_fraction: Optional[float] = None,
     ):
         """Inits :class:`ParameterizedDynamic2dPolicy`.
 
@@ -608,6 +651,8 @@ def __init__(
             fix_sign_leakage=fix_sign_leakage,
             st_slope=st_slope,
             st_clamp=st_clamp,
+            acceleration=acceleration,
+            center_fraction=center_fraction,
         )
         self.non_uniform = non_uniform
 
@@ -626,6 +671,8 @@ def __init__(
         st_slope: float = 10,
         st_clamp: bool = False,
         non_uniform: bool = False,
+        acceleration: Optional[float] = None,
+        center_fraction: Optional[float] = None,
     ):
         """Inits :class:`ParameterizedMultislice2dPolicy`.
 
@@ -660,5 +707,7 @@ def __init__(
             fix_sign_leakage=fix_sign_leakage,
             st_slope=st_slope,
             st_clamp=st_clamp,
+            acceleration=acceleration,
+            center_fraction=center_fraction,
         )
         self.non_uniform = non_uniform

direct/nn/adaptive/policy.py

@@ -505,6 +505,8 @@ def __init__(
         sampler_cwn_conv: bool = False,
         num_time_steps: Optional[int] = None,
         num_slices: Optional[int] = None,
+        acceleration: Optional[float] = None,
+        center_fraction: Optional[float] = None,
     ):
         super().__init__()
 
@@ -572,14 +574,17 @@ def __init__(
         for _ in range(num_layers):
             self.layers.append(st_policy_block(**st_policy_block_kwargs))
 
+        self.acceleration = acceleration
+        self.center_fraction = center_fraction
+
     def forward(
         self,
         masked_kspace: torch.Tensor,
         mask: torch.Tensor,
         sensitivity_map: torch.Tensor,
         kspace: torch.Tensor,
-        acceleration: float | torch.Tensor,
-        center_fraction: float | torch.Tensor,
+        acceleration: Optional[float | torch.Tensor] = None,
+        center_fraction: Optional[float | torch.Tensor] = None,
         padding: Optional[torch.Tensor] = None,
     ):
         if self.sampling_type in [
@@ -594,6 +599,16 @@ def forward(
         masks = [mask]
         prob_masks = []
 
+        if (self.acceleration is not None) and (self.center_fraction is not None):
+            acceleration = self.acceleration
+            center_fraction = center_fraction
+        else:
+            if (acceleration is None) or (center_fraction is None):
+                raise ValueError(f"One of `acceleration` or `center_fraction` received None for a value. "
+                                 f"This should not be None when `StraightThroughPolicy` is initialized "
+                                 f"with `acceleration` or `center_fraction` with None value."
+                                 )
+
         budget = self.num_actions / acceleration - self.num_actions * center_fraction
         if isinstance(budget, float):
             budget = int(budget)