Merge pull request #29 from mrphys/develop

Release 0.21.0

.devcontainer/Dockerfile

@@ -1,4 +1,4 @@
-FROM ghcr.io/mrphys/tensorflow-manylinux:1.11.0
+FROM ghcr.io/mrphys/tensorflow-manylinux:1.12.0
 
 # To enable plotting.
 RUN apt-get update && \

.github/workflows/build-package.yml

@@ -18,7 +18,7 @@ jobs:
     runs-on: ubuntu-latest
 
     container:
-      image: ghcr.io/mrphys/tensorflow-manylinux:1.11.0
+      image: ghcr.io/mrphys/tensorflow-manylinux:1.12.0
 
     strategy:
       matrix:

Makefile

@@ -38,12 +38,18 @@ docs: $(TARGET)
 	$(PYTHON) tools/docs/create_documents.py
 	$(MAKE) -C tools/docs dirhtml PY_VERSION=$(PY_VERSION)
 
-test: $(wildcard tensorflow_mri/python/ops/*.py)
+test: $(wildcard tensorflow_mri/python/*.py)
 	$(PYTHON) -m unittest discover -v -p *_test.py
 
-lint: $(wildcard tensorflow_mri/python/ops/*.py)
+doctest: $(wildcard tensorflow_mri/python/*.py)
+	$(PYTHON) tools/docs/test_docs.py
+
+lint: $(wildcard tensorflow_mri/python/*.py)
 	pylint --rcfile=pylintrc tensorflow_mri/python
 
+api: $(wildcard tensorflow_mri/python/*.py)
+	$(PYTHON) tools/build/create_api.py
+
 clean:
 	rm -rf artifacts/
 	rm -rf $(TARGET)

README.rst

@@ -111,6 +111,8 @@ versions of TensorFlow and TensorFlow MRI according to the table below.
 ======================  ========================  ============
 TensorFlow MRI Version  TensorFlow Compatibility  Release Date
 ======================  ========================  ============
+v0.21.0                 v2.9.x                    Jul 24, 2022
+v0.20.0                 v2.9.x                    Jun 18, 2022
 v0.19.0                 v2.9.x                    Jun 1, 2022
 v0.18.0                 v2.8.x                    May 6, 2022
 v0.17.0                 v2.8.x                    Apr 22, 2022

RELEASE.rst

@@ -1,34 +1,61 @@
-Release 0.20.0
+Release 0.21.0
 ==============
 
+This release contains new functionality for wavelet decomposition and
+reconstruction and optimized Gram matrices for some linear operators. It also
+redesigns the convex optimization module and contains some improvements to the
+documentation.
+
+
+Breaking Changes
+----------------
+
+* ``tfmri.convex``:
+
+  * Argument ``ndim`` has been removed from all functions.
+  * All functions will now require the domain dimension to be
+    specified. Therefore, `domain_dimension` is now the first positional
+    argument in several functions including ``ConvexFunctionIndicatorBall``,
+    ``ConvexFunctionNorm`` and ``ConvexFunctionTotalVariation``. However, while
+    this parameter is no longer optional, it is now possible to pass dynamic
+    or static information as opposed to static only (at least in the general
+    case, but specific operators may have additional restrictions).
+  * For consistency and accuracy, argument ``axis`` of
+    ``ConvexFunctionTotalVariation`` has been renamed to ``axes``.
+
+
 Major Features and Improvements
 -------------------------------
 
-* ``tfmri.layers``:
+* ``tfmri.convex``:
 
-  * Added new layers ``MaxPooling1D``, ``MaxPooling2D``, ``MaxPooling3D``,
-    ``AveragePooling1D``, ``AveragePooling2D`` and ``AveragePooling3D``.
-    These are drop-in replacements for the core Keras layers of the same name,
-    but they also support complex values.
-  * Added new layers ``DWT1D``, ``DWT2D``, ``DWT3D``, ``IDWT1D``, ``IDWT2D``,
-    and ``IDWT3D`` to compute 1D, 2D and 3D forward and inverse discrete wavelet
-    transforms.
-  * Layer ``ConvBlock`` is now deprecated in favor of the new endpoints in
-    the ``tfmri.models`` submodule.
-  * Layer ``UNet`` is now deprecated in favor of the new endpoints in
-    the ``tfmri.models`` submodule.
+  * Added new class ``ConvexFunctionL1Wavelet``, which enables image/signal
+    reconstruction with L1-wavelet regularization.
+  * Added new argument ``gram_operator`` to ``ConvexFunctionLeastSquares``,
+    which allows the user to specify a custom, potentially more efficient Gram
+    matrix.
 
-* ``tfmri.models``:
+* ``tfmri.linalg``:
 
-  * Added new models ``ConvBlock1D``, ``ConvBlock2D`` and ``ConvBlock3D``. These
-    replace the previous ``ConvBlock`` layer, which is now deprecated.
-  * Added new models ``UNet1D``, ``UNet2D`` and ``UNet3D``. These replace
-    the previous ``UNet`` layer, which is now deprecated.
+  * Added new classes ``LinearOperatorNUFFT`` and ``LinearOperatorGramNUFFT``
+    to enable the use of NUFFT as a linear operator.
+  * Added new class ``LinearOperatorWavelet`` to enable the use of wavelets
+    as a linear operator.
+
+* ``tfmri.sampling``:
+
+  * Added new ordering type ``sorted_half`` to ``radial_trajectory``.
+
+* ``tfmri.signal``:
+
+  * Added new functions ``wavedec`` and ``waverec`` for wavelet decomposition
+    and reconstruction, as well as utilities ``wavelet_coeffs_to_tensor``,
+    ``tensor_to_wavelet_coeffs``, and ``max_wavelet_level``.
 
 
 Bug Fixes and Other Changes
 ---------------------------
 
-* ``tfmri.signal``:
+* ``tfmri.recon``:
 
-  * Improved static shape inference for ``dwt`` op.
+  * Improved error reporting for ``least_squares``.

tensorflow_mri/__about__.py

@@ -29,7 +29,7 @@
 __summary__ = "A collection of TensorFlow add-ons for computational MRI."
 __uri__ = "https://github.com/mrphys/tensorflow-mri"
 
-__version__ = "0.20.0"
+__version__ = "0.21.0"
 
 __author__ = "Javier Montalt Tordera"
 __email__ = "javier.montalt@outlook.com"

tensorflow_mri/__init__.py

@@ -23,6 +23,7 @@
 from tensorflow_mri import python
 
 # Import submodules.
+from tensorflow_mri._api import array
 from tensorflow_mri._api import callbacks
 from tensorflow_mri._api import coils
 from tensorflow_mri._api import convex

tensorflow_mri/_api/array/__init__.py

@@ -0,0 +1,5 @@
+# This file was automatically generated by tools/build/create_api.py.
+# Do not edit.
+"""Array processing operations."""
+
+from tensorflow_mri.python.ops.array_ops import update_tensor as update_tensor

tensorflow_mri/_api/convex/__init__.py

@@ -14,6 +14,7 @@
 from tensorflow_mri.python.ops.convex_ops import ConvexFunctionL2NormSquared as ConvexFunctionL2NormSquared
 from tensorflow_mri.python.ops.convex_ops import ConvexFunctionTikhonov as ConvexFunctionTikhonov
 from tensorflow_mri.python.ops.convex_ops import ConvexFunctionTotalVariation as ConvexFunctionTotalVariation
+from tensorflow_mri.python.ops.convex_ops import ConvexFunctionL1Wavelet as ConvexFunctionL1Wavelet
 from tensorflow_mri.python.ops.convex_ops import ConvexFunctionQuadratic as ConvexFunctionQuadratic
 from tensorflow_mri.python.ops.convex_ops import ConvexFunctionLeastSquares as ConvexFunctionLeastSquares
 from tensorflow_mri.python.ops.optimizer_ops import admm_minimize as admm_minimize

tensorflow_mri/_api/linalg/__init__.py

@@ -9,6 +9,10 @@
 from tensorflow_mri.python.util.linalg_imaging import LinearOperatorScaledIdentity as LinearOperatorScaledIdentity
 from tensorflow_mri.python.util.linalg_imaging import LinearOperatorDiag as LinearOperatorDiag
 from tensorflow_mri.python.util.linalg_imaging import LinearOperatorGramMatrix as LinearOperatorGramMatrix
-from tensorflow_mri.python.util.linalg_imaging import LinearOperatorFiniteDifference as LinearOperatorFiniteDifference
+from tensorflow_mri.python.ops.linalg_ops import LinearOperatorNUFFT as LinearOperatorNUFFT
+from tensorflow_mri.python.ops.linalg_ops import LinearOperatorGramNUFFT as LinearOperatorGramNUFFT
+from tensorflow_mri.python.ops.linalg_ops import LinearOperatorFiniteDifference as LinearOperatorFiniteDifference
+from tensorflow_mri.python.ops.linalg_ops import LinearOperatorWavelet as LinearOperatorWavelet
 from tensorflow_mri.python.ops.linalg_ops import LinearOperatorMRI as LinearOperatorMRI
+from tensorflow_mri.python.ops.linalg_ops import LinearOperatorGramMRI as LinearOperatorGramMRI
 from tensorflow_mri.python.ops.linalg_ops import conjugate_gradient as conjugate_gradient

tensorflow_mri/_api/signal/__init__.py

@@ -12,3 +12,8 @@
 from tensorflow_nufft.python.ops.nufft_ops import nufft as nufft
 from tensorflow_mri.python.ops.wavelet_ops import dwt as dwt
 from tensorflow_mri.python.ops.wavelet_ops import idwt as idwt
+from tensorflow_mri.python.ops.wavelet_ops import wavedec as wavedec
+from tensorflow_mri.python.ops.wavelet_ops import waverec as waverec
+from tensorflow_mri.python.ops.wavelet_ops import dwt_max_level as max_wavelet_level
+from tensorflow_mri.python.ops.wavelet_ops import coeffs_to_tensor as wavelet_coeffs_to_tensor
+from tensorflow_mri.python.ops.wavelet_ops import tensor_to_coeffs as tensor_to_wavelet_coeffs

tensorflow_mri/python/ops/array_ops.py

@@ -14,7 +14,12 @@
 # ==============================================================================
 """Array manipulation operations."""
 
+import numpy as np
 import tensorflow as tf
+import tensorflow.experimental.numpy as tnp
+from tensorflow.python.ops.numpy_ops import np_array_ops
+
+from tensorflow_mri.python.util import api_util
 
 
 def broadcast_static_shapes(*shapes):
@@ -274,3 +279,107 @@ def _compute_static_output_shape(input_shape, target_shape):
           output_shape.as_list(), input_shape.as_list())])
 
   return output_shape
+
+
+@api_util.export("array.update_tensor")
+def update_tensor(tensor, slices, value):
+  """Updates the values of a tensor at the specified slices.
+
+  This operator performs slice assignment.
+
+  .. note::
+    Equivalent to `tensor[slices] = value`.
+
+  .. warning::
+    TensorFlow does not support slice assignment because tensors are immutable.
+    This operator works around this limitation by creating a new tensor, which
+    may have performance implications.
+
+  Args:
+    tensor: A `tf.Tensor`.
+    slices: The indices or slices.
+    value: A `tf.Tensor`.
+
+  Returns:
+    An updated `tf.Tensor` with the same shape and type as `tensor`.
+  """
+  # Using a private implementation in the TensorFlow NumPy API.
+  # pylint: disable=protected-access
+  return _with_index_update_helper(np_array_ops._UpdateMethod.UPDATE,
+                                   tensor, slices, value)
+
+
+def _with_index_update_helper(update_method, a, slice_spec, updates):  # pylint: disable=missing-param-doc
+  """Implementation of ndarray._with_index_*."""
+  # Adapted from tensorflow/python/ops/numpy_ops/np_array_ops.py.
+  # pylint: disable=protected-access
+  if (isinstance(slice_spec, bool) or (isinstance(slice_spec, tf.Tensor) and
+                                       slice_spec.dtype == tf.dtypes.bool) or
+      (isinstance(slice_spec, (np.ndarray, tnp.ndarray)) and
+       slice_spec.dtype == np.bool_)):
+    slice_spec = tnp.nonzero(slice_spec)
+
+  if not isinstance(slice_spec, tuple):
+    slice_spec = np_array_ops._as_spec_tuple(slice_spec)
+
+  return np_array_ops._slice_helper(a, slice_spec, update_method, updates)
+
+
+def map_fn(fn, elems, batch_dims=1, **kwargs):
+  """Transforms `elems` by applying `fn` to each element.
+
+  .. note::
+    Similar to `tf.map_fn`, but it supports unstacking along multiple batch
+    dimensions.
+
+  For the parameters, see `tf.map_fn`. The only difference is that there is an
+  additional `batch_dims` keyword argument which allows specifying the number
+  of batch dimensions. The default is 1, in which case this function is equal
+  to `tf.map_fn`.
+  """
+  # This function works by reshaping any number of batch dimensions into a
+  # single batch dimension, calling the original `tf.map_fn`, and then
+  # restoring the original batch dimensions.
+  static_batch_dims = tf.get_static_value(batch_dims)
+
+  # Get batch shapes.
+  if static_batch_dims is None:
+    # We don't know how many batch dimensions there are statically, so we can't
+    # get the batch shape statically.
+    static_batch_shapes = tf.nest.map_structure(
+        lambda _: tf.TensorShape(None), elems)
+  else:
+    static_batch_shapes = tf.nest.map_structure(
+        lambda x: x.shape[:static_batch_dims], elems)
+  dynamic_batch_shapes = tf.nest.map_structure(
+      lambda x: tf.shape(x)[:batch_dims], elems)
+
+  # Flatten the batch dimensions.
+  elems = tf.nest.map_structure(
+      lambda x: tf.reshape(
+          x, tf.concat([[-1], tf.shape(x)[batch_dims:]], 0)), elems)
+
+  # Process each batch.
+  output = tf.map_fn(fn, elems, **kwargs)
+
+  # Unflatten the batch dimensions.
+  output = tf.nest.map_structure(
+      lambda x, dynamic_batch_shape: tf.reshape(
+          x, tf.concat([dynamic_batch_shape, tf.shape(x)[1:]], 0)),
+      output, dynamic_batch_shapes)
+
+  # Set the static batch shapes on the output, if known.
+  if static_batch_dims is not None:
+    output = tf.nest.map_structure(
+        lambda x, static_batch_shape: tf.ensure_shape(
+            x, static_batch_shape.concatenate(x.shape[static_batch_dims:])),
+        output, static_batch_shapes)
+
+  return output
+
+
+def slice_along_axis(tensor, axis, start, length):
+  """Slices a tensor along the specified axis."""
+  begin = tf.scatter_nd([[axis]], [start], [tensor.shape.rank])
+  size = tf.tensor_scatter_nd_update(tf.shape(tensor), [[axis]], [length])
+  return tf.slice(tensor, begin, size)

tensorflow_mri/python/ops/array_ops_test.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for module `array_ops`."""
+# pylint: disable=missing-class-docstring,missing-function-docstring
 
+from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
@@ -239,5 +241,24 @@ def graph_fn(x):
                         expected_output_shape)
 
 
+class UpdateTensorTest(test_util.TestCase):
+
+  @test_util.run_all_execution_modes
+  @parameterized.named_parameters(
+      # name, tensor, slices, value
+      ("test0", [0, 0, 0], (slice(1, 2),), [2]),
+      ("test1", [0, 0, 0], (slice(0, None, 2)), [2, 1]),
+      ("test2", [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
+       (slice(0, 2), slice(1, 3)), [[1, 2], [3, 4]])
+  )
+  def test_update_tensor(self, tensor, slices, value):
+    with tf.device('/cpu:0'):
+      expected = np.asarray(tensor)
+      expected[slices] = value
+      tensor = tf.constant(tensor)
+      result = array_ops.update_tensor(tensor, slices, value)
+      self.assertAllClose(expected, result)
+
+
 if __name__ == '__main__':
   tf.test.main()