Skip to content

Commit

Permalink
Adding sampling masks (MagicMask, VariableDensityPoissonMask) (#216, C…
Browse files Browse the repository at this point in the history
…loses #185, #213)

Adding new sampling patterbs:

* Variable density Poisson masking function (cython implementation)
    * setup.py changes
* Fastmri's Magic masking function

Bug fix/update:
* `torch.where` output needed to be made contiguous to be inputted to fft, due to new torch version
  • Loading branch information
georgeyiasemis committed Jul 31, 2022
1 parent d107efe commit c4b1ca7
Show file tree
Hide file tree
Showing 10 changed files with 718 additions and 124 deletions.
1 change: 1 addition & 0 deletions .github/workflows/tox.yml
Original file line number Diff line number Diff line change
Expand Up @@ -21,5 +21,6 @@ jobs:
run: |
python -m pip install --upgrade pip setuptools wheel
pip install tox tox-gh-actions
python -m pip install -e ".[dev]"
- name: Test with tox
run: tox
9 changes: 7 additions & 2 deletions Makefile
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
.PHONY: clean clean-test clean-pyc clean-build docs help
.PHONY: clean clean-test clean-pyc clean-cpy clean-build docs help
.DEFAULT_GOAL := help

define BROWSER_PYSCRIPT
Expand Down Expand Up @@ -26,7 +26,7 @@ BROWSER := python -c "$$BROWSER_PYSCRIPT"
help:
@python -c "$$PRINT_HELP_PYSCRIPT" < $(MAKEFILE_LIST)

clean: clean-build clean-pyc clean-test clean-docs ## remove all build, test, coverage, docs and Python artifacts
clean: clean-build clean-pyc clean-cpy clean-test clean-docs ## remove all build, test, coverage, docs and Python and cython artifacts

clean-build: ## remove build artifacts
rm -fr build/
Expand All @@ -41,6 +41,11 @@ clean-pyc: ## remove Python file artifacts
find . -name '*~' -exec rm -f {} +
find . -name '__pycache__' -exec rm -fr {} +

clean-cpy: ## remove cython file artifacts
find . -name '*.c' -exec rm -f {} +
find . -name '*.cpp' -exec rm -f {} +
find . -name '*.so' -exec rm -f {} +

clean-test: ## remove test and coverage artifacts
rm -fr .tox/
rm -f .coverage
Expand Down
123 changes: 123 additions & 0 deletions direct/common/_poisson.pyx
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
#cython: cdivision=True
#cython: boundscheck=False
#cython: nonecheck=False
#cython: wraparound=False
#cython: overflowcheck=False
#cython: unraisable_tracebacks=False

import numpy as np
cimport numpy as cnp
from libc.math cimport cos, pi, sin
from libc.stdlib cimport RAND_MAX, rand, srand

cnp.import_array()


cdef double random_uniform() nogil:
"""Produces a random number in (0, 1)."""
cdef double r = rand()
return r / RAND_MAX


cdef int randint(int upper) nogil:
"""Produces a random integer in {0, 1, ..., upper-1}."""
return int(random_uniform() * (upper))


cdef inline Py_ssize_t fmax(Py_ssize_t one, Py_ssize_t two) nogil:
"""Max(a, b)."""
return one if one > two else two


cdef inline Py_ssize_t fmin(Py_ssize_t one, Py_ssize_t two) nogil:
"""Min(a, b)."""
return one if one < two else two


def poisson(
int nx,
int ny,
int max_attempts,
cnp.ndarray[cnp.int_t, ndim=2, mode='c'] mask,
cnp.ndarray[cnp.float64_t, ndim=2, mode='c'] radius_x,
cnp.ndarray[cnp.float64_t, ndim=2, mode='c'] radius_y,
int seed
):
"""
Notes
-----
* Code inspired and modified from [1]_ with BSD-3 licence, Copyright (c) 2016, Frank Ong, Copyright (c) 2016,
The Regents of the University of California [2]_.
References
----------
.. [1] https://github.com/mikgroup/sigpy/blob/1817ff849d34d7cbbbcb503a1b310e7d8f95c242/sigpy/mri/samp.py#L158
.. [2] https://github.com/mikgroup/sigpy/blob/master/LICENSE
"""

cdef int x, y, num_actives, i, k
cdef float rx, ry, v, t, qx, qy, distance
cdef Py_ssize_t startx, endx, starty, endy, px, py

# initialize active list
cdef cnp.ndarray[cnp.int_t, ndim=1, mode='c'] pxs = np.empty(nx * ny, dtype=int)
cdef cnp.ndarray[cnp.int_t, ndim=1, mode='c'] pys = np.empty(nx * ny, dtype=int)

srand(seed)

with nogil:

pxs[0] = randint(nx)
pys[0] = randint(ny)

num_actives = 1

while num_actives > 0:
# Select a sample from active list
i = randint(num_actives)
px = pxs[i]
py = pys[i]
rx = radius_x[px, py]
ry = radius_y[px, py]

# Attempt to generate point
done = False
k = 0

while not done and k < max_attempts:

# Generate point randomly from r and 2 * r
v = random_uniform() + 1
t = 2 * pi * random_uniform()
qx = px + v * rx * cos(t)
qy = py + v * ry * sin(t)

# Reject if outside grid or close to other points
if qx >= 0 and qx < nx and qy >= 0 and qy < ny:
startx = fmax(int(qx - rx), 0)
endx = fmin(int(qx + rx + 1), nx)
starty = fmax(int(qy - ry), 0)
endy = fmin(int(qy + ry + 1), ny)

done = True
for x in range(startx, endx):
for y in range(starty, endy):
distance = ((qx - x) / radius_x[x, y]) ** 2 + ((qy - y) / (radius_y[x, y])) ** 2
if (mask[x, y] == 1) and (distance < 1):
done = False
break

k += 1

# Add point if done else remove from active list
if done:
pxs[num_actives] = int(qx)
pys[num_actives] = int(qy)
mask[pxs[num_actives], pys[num_actives]] = 1
num_actives += 1
else:
num_actives -= 1
pxs[i] = pxs[num_actives]
pys[i] = pys[num_actives]
Loading

0 comments on commit c4b1ca7

Please sign in to comment.