Skip to content

Commit

Permalink
Triangular distribution added to favor higher accels
Browse files Browse the repository at this point in the history
  • Loading branch information
@georgeyiasemis
georgeyiasemis committed Feb 15, 2024
1 parent e4f4a57 commit f469f2a
Show file tree
Hide file tree
Showing 2 changed files with 121 additions and 7 deletions.
61 changes: 54 additions & 7 deletions direct/common/subsample.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
from direct.common._poisson import poisson as _poisson # pylint: disable=no-name-in-module
from direct.environment import DIRECT_CACHE_DIR
from direct.types import Number
from direct.utils import str_to_class
from direct.utils import str_to_class, triangular_distribution
from direct.utils.io import download_url

# pylint: disable=arguments-differ
Expand Down Expand Up @@ -63,6 +63,7 @@ def __init__(
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[float], Tuple[float, ...]]] = None,
uniform_range: bool = True,
linear_range: bool = False,
):
"""
Parameters
Expand All @@ -76,6 +77,9 @@ def __init__(
is True, then two values should be given. Default: None.
uniform_range: bool
If True then an acceleration will be uniformly sampled between the two values. Default: True.
linear_range : bool
If True, an acceleration will be sampled from a triangular distribution between the two values. Ignored if
`uniform_range` is True. Default: False.
"""
if center_fractions is not None:
if len([center_fractions]) != len([accelerations]):
Expand All @@ -88,10 +92,12 @@ def __init__(
self.accelerations = accelerations

self.uniform_range = uniform_range
if uniform_range and (len(center_fractions) != 2 or len(accelerations) != 2):
self.linear_range = linear_range
if (uniform_range or linear_range) and (len(center_fractions) != 2 or len(accelerations) != 2):
raise ValueError(
f"When `uniform_range` is True, both `center_fractions` and `accelerations` should have "
f"a length of two. Received center_fractions={center_fractions} and accelerations={accelerations}."
f"When any of `uniform_range` or `linear_range` is True, both `center_fractions` and `accelerations` "
f"should have a length of two. Received center_fractions={center_fractions} "
f"and accelerations={accelerations}."
)

self.rng = np.random.RandomState()
Expand All @@ -107,6 +113,11 @@ def choose_acceleration(self):
center_fraction = self.rng.uniform(
low=min(self.center_fractions), high=max(self.center_fractions), size=1
)[0]
elif self.linear_range:
acceleration = triangular_distribution(min(self.accelerations), max(self.accelerations), 1, self.rng)[0]
center_fraction = self.rng.uniform(
low=min(self.center_fractions), high=max(self.center_fractions), size=1
)[0]
else:
choice = self.rng.randint(0, len(self.accelerations))
acceleration = self.accelerations[choice]
Expand Down Expand Up @@ -144,11 +155,13 @@ def __init__(
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[float], Tuple[float, ...]]] = None,
uniform_range: bool = False,
linear_range: bool = False,
):
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
linear_range=linear_range,
)

@staticmethod
Expand Down Expand Up @@ -204,11 +217,13 @@ def __init__(
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[float], Tuple[float, ...]]] = None,
uniform_range: bool = False,
linear_range: bool = False,
):
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
linear_range=linear_range,
)

def mask_func(
Expand Down Expand Up @@ -273,6 +288,7 @@ def __init__(
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[int], Tuple[int, ...]]] = None,
uniform_range: bool = False,
linear_range: bool = False,
):
"""Inits :class:`CartesianRandomMaskFunc`.
Expand All @@ -286,11 +302,15 @@ def __init__(
If multiple values are provided, then one of these numbers is chosen uniformly each time.
uniform_range: bool
If True then an acceleration will be uniformly sampled between the two values. Default: True.
linear_range : bool
If True, an acceleration will be sampled from a triangular distribution between the two values. Ignored if
`uniform_range` is True. Default: False.
"""
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
linear_range=linear_range,
)

def mask_func(
Expand Down Expand Up @@ -367,11 +387,13 @@ def __init__(
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[float], Tuple[float, ...]]] = None,
uniform_range: bool = False,
linear_range: bool = False,
):
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
linear_range=linear_range,
)

def mask_func(
Expand Down Expand Up @@ -440,6 +462,7 @@ def __init__(
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[int], Tuple[int, ...]]] = None,
uniform_range: bool = False,
linear_range: bool = False,
):
"""Inits :class:`CartesianEquispacedMaskFunc`.
Expand All @@ -453,11 +476,15 @@ def __init__(
If multiple values are provided, then one of these numbers is chosen uniformly each time.
uniform_range: bool
If True then an acceleration will be uniformly sampled between the two values. Default: True.
linear_range : bool
If True, an acceleration will be sampled from a triangular distribution between the two values. Ignored if
`uniform_range` is True. Default: False.
"""
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
linear_range=linear_range,
)

def mask_func(
Expand Down Expand Up @@ -534,11 +561,13 @@ def __init__(
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[float], Tuple[float, ...]]] = None,
uniform_range: bool = False,
linear_range: bool = False,
):
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
linear_range=linear_range,
)

def mask_func(
Expand Down Expand Up @@ -636,6 +665,7 @@ def __init__(
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[int], Tuple[int, ...]]] = None,
uniform_range: bool = False,
linear_range: bool = False,
):
"""Inits :class:`CartesianMagicMaskFunc`.
Expand All @@ -649,11 +679,15 @@ def __init__(
If multiple values are provided, then one of these numbers is chosen uniformly each time.
uniform_range: bool
If True then an acceleration will be uniformly sampled between the two values. Default: True.
linear_range : bool
If True, an acceleration will be sampled from a triangular distribution between the two values. Ignored if
`uniform_range` is True. Default: False.
"""
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
linear_range=linear_range,
)

def mask_func(
Expand Down Expand Up @@ -861,12 +895,15 @@ def __init__(
self,
accelerations: Union[List[Number], Tuple[Number, ...]],
subsampling_scheme: CIRCUSSamplingMode,
uniform_range: bool = False,
linear_range: bool = False,
**kwargs,
):
super().__init__(
accelerations=accelerations,
center_fractions=tuple(0 for _ in range(len(accelerations))),
uniform_range=False,
uniform_range=uniform_range,
linear_range=linear_range,
)
if subsampling_scheme not in ["circus-spiral", "circus-radial"]:
raise NotImplementedError(
Expand Down Expand Up @@ -1113,6 +1150,8 @@ def __init__(
max_attempts: Optional[int] = 10,
tol: Optional[float] = 0.2,
slopes: Optional[Union[List[float], Tuple[float, ...]]] = None,
uniform_range: bool = False,
linear_range: bool = False,
**kwargs,
):
"""Inits :class:`VariableDensityPoissonMaskFunc`.
Expand All @@ -1139,7 +1178,8 @@ def __init__(
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=False,
uniform_range=uniform_range,
linear_range=linear_range,
)
self.crop_corner = crop_corner
self.max_attempts = max_attempts
Expand Down Expand Up @@ -1275,11 +1315,13 @@ def __init__(
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[float], Tuple[float, ...]]] = None,
uniform_range: bool = False,
linear_range: bool = False,
):
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
linear_range=linear_range,
)

def mask_func(
Expand Down Expand Up @@ -1342,11 +1384,13 @@ def __init__(
accelerations: Union[List[Number], Tuple[Number, ...]],
center_fractions: Optional[Union[List[float], Tuple[float, ...]]] = None,
uniform_range: bool = False,
linear_range: bool = False,
):
super().__init__(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
linear_range=linear_range,
)

def mask_func(
Expand Down Expand Up @@ -1461,12 +1505,15 @@ def mask_func(self, data, return_acs=False):
return self.data_dictionary[data]


def build_masking_function(name, accelerations, center_fractions=None, uniform_range=False, **kwargs):
def build_masking_function(
name, accelerations, center_fractions=None, uniform_range=False, linear_range=False, **kwargs
):
MaskFunc: BaseMaskFunc = str_to_class("direct.common.subsample", name + "MaskFunc") # noqa
mask_func = MaskFunc(
accelerations=accelerations,
center_fractions=center_fractions,
uniform_range=uniform_range,
linear_range=linear_range,
)

return mask_func
67 changes: 67 additions & 0 deletions direct/utils/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

"""direct.utils module."""

from __future__ import annotations

import abc
import ast
Expand Down Expand Up @@ -550,3 +551,69 @@ def dict_flatten(in_dict: DictOrDictConfig, dict_out: Optional[DictOrDictConfig]
continue
dict_out[k] = v
return dict_out


def triangular_distribution(a: float, b: float, n: int, rng: np.random.RandomState | None = None) -> np.ndarray:
"""Samples from a triangular distribution with endpoints `a` and `b`.
The triangular distribution is defined such that the probability density function (PDF)
is proportional to `x` between `a` and `b`, such that:
.. math::
p(b) = \frac{b}{a} \cdot p(a)
The PDF of the triangular distribution is given by:
.. math::
p(x) = \frac{2x}{b^2 - a^2}
where `x` ranges from `a` to `b`.
This function uses inverse transform sampling to generate samples from the
triangular distribution based on its cumulative distribution function (CDF).
Parameters
----------
a : float
Left endpoint of the triangular distribution.
b : float
Right endpoint of the triangular distribution.
n : int
Number of samples to generate.
rng: np.random.RandomState, optional
Random generator object. Default: None.
Returns
-------
samples : ndarray
Array of `n` samples from the triangular distribution.
"""

def inverse_cdf(u: np.ndarray) -> np.ndarray:
"""Computes the inverse of the cumulative distribution function (CDF) of the distribution.
Parameters
----------
u : ndarray
Value at which to compute the inverse CDF.
Returns
-------
ndarray
The value of the inverse CDF at `u`.
"""
return np.sqrt(u * (b**2 - a**2) + a**2)

if rng is None:
rng = np.random.RandomState()

# Generate uniform samples
uniform_samples = rng.uniform(0, 1, n)

# Use inverse transform sampling
samples = inverse_cdf(uniform_samples)

return samples

0 comments on commit f469f2a

Please sign in to comment.