MIRTorch 0.1.0__py3-none-any.whl

MIRTorch-0.1.0.dist-info/LICENSE

@@ -0,0 +1,80 @@
+Copyright (c) 2021, Guanhua Wang, Neel Shah, Keyue Zhu and Jeffrey Fessler
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+
+-------------------- LICENSE FOR fastMRI ------------------
+
+MIT License
+
+Copyright (c) Facebook, Inc. and its affiliates.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+
+-------------------- LICENSE FOR SigPy ------------------
+
+Copyright (c) 2016, Frank Ong.
+Copyright (c) 2016, The Regents of the University of California.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+----------------------------- LICENSE FOR numpy-ml --------------------------------
+GNU GENERAL PUBLIC LICENSE
+Version 3, 29 June 2007
+
+Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+Everyone is permitted to copy and distribute verbatim copies
+of this license document, but changing it is not allowed.
+
+----------------------------- LICENSE for pytorch_wavelets----------------------------
+This licence applies to any parts of this library which are novel in comparison
+to the original DTCWT MATLAB toolbox written by Nick Kingsbury and Cian
+Shaffrey. See the Provenance section of README.rst file for details on any further
+restrictions of use. If you wish to use the DTCWT, you should read that license as well.
+The DWT sections come under this license.
+
+MIT License
+
+Copyright (c) 2020 Fergal Cotter
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

MIRTorch-0.1.0.dist-info/METADATA

@@ -0,0 +1,158 @@
+Metadata-Version: 2.1
+Name: MIRTorch
+Version: 0.1.0
+Summary: a PyTorch-based image reconstruction toolbox
+Author-email: Guanhua Wang <guanhuaw@umich.edu>
+License: BSD-3-Clause
+Project-URL: repository, https://github.com/guanhuaw/MIRTorch
+Keywords: signal processing,inverse problems
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Operating System :: OS Independent
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Topic :: Scientific/Engineering :: Image Processing
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch >=1.13
+Requires-Dist: torchvision
+Requires-Dist: torchkbnufft >=1.4.0
+Requires-Dist: numpy >=1.24
+Requires-Dist: dominate
+Requires-Dist: scipy
+Requires-Dist: PyWavelets
+Requires-Dist: requests
+Requires-Dist: einops
+Requires-Dist: matplotlib
+Requires-Dist: importlib-metadata ; python_version < "3.8"
+Provides-Extra: dev
+Requires-Dist: check-manifest ; extra == 'dev'
+Provides-Extra: test
+Requires-Dist: coverage ; extra == 'test'
+Requires-Dist: pytest ; extra == 'test'
+
+# MIRTorch
+
+![GitHub release (latest by date including pre-releases)](https://img.shields.io/github/v/release/guanhuaw/mirtorch?include_prereleases)
+![Read the Docs](https://img.shields.io/readthedocs/mirtorch)
+
+A Py***Torch***-based differentiable ***I***mage ***R***econstruction ***T***oolbox, developed at the University of ***M***ichigan.
+
+The work is inspired by [MIRT](https://github.com/JeffFessler/mirt), a well-acclaimed toolbox for medical imaging reconstruction.
+
+The main objective is to facilitate rapid, data-driven medical image reconstruction using CPUs and GPUs, for fast prototyping. Researchers can conveniently develop new model-based and learning-based methods (e.g., unrolled neural networks) with abstraction layers. The availability of auto-differentiation enables optimization of imaging protocols and reconstruction parameters using gradient methods.
+
+Documentation: https://mirtorch.readthedocs.io/en/latest/
+
+------
+
+### Installation
+
+We recommend to [pre-install `PyTorch` first](https://pytorch.org/).
+Use `pip install mirtorch` to install.
+To install the `MIRTorch` locally, after cloning the repo, please try `pip install -e .`(one may modify the package locally with this option.)
+
+------
+
+### Features
+
+#### Linear maps
+
+The `LinearMap` class overloads common matrix operations, such as `+, - , *`.
+
+Instances include basic linear operations (like convolution), classical imaging processing, and MRI system matrix (Cartesian and Non-Cartesian, sensitivity- and B0-informed system models). ***NEW!*** MIRTorch recently adds the support for SPECT and CT.
+
+Since the Jacobian matrix of a linear operator is itself, the toolbox can actively calculate such Jacobians during backpropagation, avoiding the large cache cost required by auto-differentiation.
+
+When defining linear operators, please make sure that all torch tensors are on the same device and compatible. For example, `torch.cfloat` are compatible with `torch.float` but not `torch.double`. Similarly, `torch.chalf` is compatible with `torch.half`.
+When the data is image, there are 2 empirical formats: `[num_batch, num_channel, nx, ny, (nz)]` and `[nx, ny, (nz)]`.
+For some LinearMaps, there is a boolean `batchmode` to control the shape.
+
+#### Proximal operators
+
+The toolbox contains common proximal operators such as soft thresholding. These operators also support the regularizers that involve multiplication with diagonal or unitary matrices, such as orthogonal wavelets.
+
+#### Iterative reconstruction (MBIR) algorithms
+
+Currently, the package includes the conjugate gradient (CG), fast iterative thresholding (FISTA), optimized gradient method (POGM), forward-backward primal-dual (FBPD) algorithms for image reconstruction.
+
+#### Dictionary learning
+
+For dictionary learning-based reconstruction, we implemented an efficient dictionary learning algorithm ([SOUP-DIL](https://arxiv.org/abs/1511.06333)) and orthogonal matching pursuit ([OMP](https://ieeexplore.ieee.org/abstract/document/342465/?casa_token=aTDkQVCM9WEAAAAA:5rXu9YikP822bCBvkhYxKWlBTJ6Fn6baTQJ9kuNrU7K-64EmGOAczYvF2dTW3al3PfPdwJAiYw)). Due to PyTorch’s limited support of sparse matrices, we use SciPy as the backend.
+
+#### Multi-GPU support
+
+Currently, MIRTorch uses `torch.DataParallel` to support multiple GPUs. One may re-package the `LinearMap`, `Prox` or `Alg` inside a `torch.nn.Module` to enable data parallel. See [this tutorial](https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html) for detail.
+
+------
+
+### Usage and examples
+
+Generally, MIRTorch solves the image reconstruction problems that have the cost function $\textit{argmin}_{x} \|Ax-y\|_2^2 + \lambda \textit{R}(x)$. $A$ stands for the system matrix. When it is linear, one may use `LinearMap` to efficiently compute it. `y` usually denotes measurements. $\textit{R}(\cdot)$ denotes regularizers, which determines which `Alg` to be used. One may refer to [1](https://web.eecs.umich.edu/~fessler/book/), [2](https://web.stanford.edu/~boyd/cvxbook/bv_cvxbook.pdf) and [3](https://www.youtube.com/watch?v=J6_5rPYnr_s) for more tutorials on optimization.
+
+Here we provide several notebook tutorials focused on MRI, where $A$ is FFT or NUFFT.
+
+- `/example/demo_mnist.ipynb` shows the LASSO on MNIST with FISTA and POGM.
+- `/example/demo_mri.ipynb` contains the SENSE (CG-SENSE) and **B0**-informed reconstruction with penalized weighted least squares (*PWLS*).
+- `/example/demo_3d.ipynb` contains the 3d non-Cartesian MR reconstruction. *New!* Try the Toeplitz-embedding version of B0-informed reconstruction, which reduce hour-long recon to 5 secs.
+- `/example/demo_cs.ipynb` shows the compressed sensing reconstruction of under-determined MRI signals.
+- `/example/demo_dl.ipynb` exhibits the dictionary learning results.
+- `/example/demo_mlem` showcase SPECT recon algorithms, including EM and CNN.
+
+Since MIRTorch is differentiable, one may use AD to update many parameters. For example, updating the reconstruction neural network's weights. More importantly, one may update the imaging system itself via gradient-based and data-driven methods. As a user case, [Bjork repo](https://github.com/guanhuaw/Bjork) contains MRI sampling pattern optimization examples. One may use the reconstruction loss as the objective function to jointly optimize reconstruction algorithms and the sampling pattern. See [this video](https://www.youtube.com/watch?v=sLFOf5EvVAs) on how to jointly optimize reconstruction and acquisition.
+
+------
+
+### Acknowledgments
+
+This work is inspired by (but not limited to):
+
+* SigPy: https://github.com/mikgroup/sigpy
+
+* MIRT: https://github.com/JeffFessler/mirt
+
+* MIRT.jl: https://github.com/JeffFessler/MIRT.jl
+
+* PyLops: https://github.com/PyLops/pylops
+
+If the code is useful to your research, please consider citing:
+
+```bibtex
+@article{wang:22:bjork,
+  author={Wang, Guanhua and Luo, Tianrui and Nielsen, Jon-Fredrik and Noll, Douglas C. and Fessler, Jeffrey A.},
+  journal={IEEE Transactions on Medical Imaging},
+  title={B-spline Parameterized Joint Optimization of Reconstruction and K-space Trajectories ({BJORK}) for Accelerated {2D} {MRI}},
+  year={2022},
+  pages={1-1},
+  doi={10.1109/TMI.2022.3161875}}
+```
+
+```bibtex
+@inproceedings{wang:22:mirtorch,
+  title={{MIRTorch}: A {PyTorch}-powered Differentiable Toolbox for Fast Image Reconstruction and Scan Protocol Optimization},
+  author={Wang, Guanhua and Shah, Neel and Zhu, Keyue and Noll, Douglas C. and Fessler, Jeffrey A.},
+  booktitle={Proc. Intl. Soc. Magn. Resonance. Med. (ISMRM)},
+  pages={4982},
+  year={2022}
+}
+```
+If you use the SPECT code, please consider citing:
+
+```bibtex
+@ARTICLE{li:23:tet,
+  author={Li, Zongyu and Dewaraja, Yuni K. and Fessler, Jeffrey A.},
+  journal={IEEE Transactions on Radiation and Plasma Medical Sciences},
+  title={Training End-to-End Unrolled Iterative Neural Networks for SPECT Image Reconstruction},
+  year={2023},
+  volume={7},
+  number={4},
+  pages={410-420},
+  doi={10.1109/TRPMS.2023.3240934}}
+```
+
+
+------
+
+### License
+
+This package uses the BSD3 license.

MIRTorch-0.1.0.dist-info/RECORD

@@ -0,0 +1,43 @@
+mirtorch/__init__.py,sha256=M7cAwxRzIFuSAR126GlDF8wBsUbMMm7AtUi_unf1BXQ,148
+mirtorch/alg/__init__.py,sha256=COQtui_W-ewqJgkdPack-qGwUvKvHBSB_j8dkeNpwE8,274
+mirtorch/alg/cg.py,sha256=HJf1ar4Kik2csEAB9XmpKVzQzG_0JRfR5SPx2bgR9-U,4812
+mirtorch/alg/fbpd.py,sha256=RSL7b1PRkDHegWVDB9IL5fv0-_hAcn02zWRo_22eJ9Y,3022
+mirtorch/alg/fista.py,sha256=aQfhpvZ_-aQeU4NNtSOTvTYngzSUE9AGCEIXRkb_0B8,2462
+mirtorch/alg/pogm.py,sha256=2gPvoWnCq6qzpntd7Rc5SmFxGIVWdqh3sRZyZGbbKDo,2836
+mirtorch/alg/spectral.py,sha256=pxXDNXq6hbZ9e6PKtgVyX-MD3_-oAsyK7lfv9p79FB4,1138
+mirtorch/alg/util.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mirtorch/dic/__init__.py,sha256=z1wRn1Uu-UtcDq_1vmg3p3C6Aw7PFr7DKhyPWhGodzg,170
+mirtorch/dic/soup.py,sha256=BGPnoGmfU8lmdvA5ZlTAFlrQ3NzQxock7fXnzodlKMU,5005
+mirtorch/dic/util.py,sha256=ENoServ5AVIWkJKiiCVbN-r-namcU7k1_hmG576rGRc,1631
+mirtorch/linear/__init__.py,sha256=W4KMwJpVBcvJDIoTcZSD1LPuJ2V_xlUTVtGjM_kaa6Y,902
+mirtorch/linear/basics.py,sha256=FVo724snDeAA8HqoqHsuS29nudpW-ImqUDYDgrBktqI,17406
+mirtorch/linear/linearmaps.py,sha256=EKJiJ_vZ6dsIiRKA-0za-yA1CSZ4rB4PWv4iQzQMv_U,12743
+mirtorch/linear/mri.py,sha256=WjKocD7cFs3JPgtto1-vgdbW1gHUbH75-5ySdMkU618,25988
+mirtorch/linear/spect.py,sha256=XGoBvRy9AwE29BhwlgR5UTcnQl61UbR2vk_RTiO7ego,4750
+mirtorch/linear/util.py,sha256=ycEoLKeSSzgkWeA3tBLReQkj2PyHaUDJJ0GrbbIVsIg,8283
+mirtorch/linear/wavelets.py,sha256=KajzFw87_UznsApjQvtwT4Pubb8FBTYNOPTiVJ8WeU8,7494
+mirtorch/prox/__init__.py,sha256=UWEScui-abN47vly_Reg-7TOYf5-FhOZQis1t5kyeso,319
+mirtorch/prox/prox.py,sha256=QnFioBgOKnESvQ9bKZsHA1OGixHEkEoORU1iw-fdC78,9681
+mirtorch/vendors/pytorch_wavelets/__init__.py,sha256=ukA3rF0vBjLlpd9CaxxW4wFPXTm5Q2eh-rlcRH2EBs4,667
+mirtorch/vendors/pytorch_wavelets/_version.py,sha256=T9mHhV6eqKSKtCKDrltAYNNplh9SlVM1qRWp6oBp9CM,95
+mirtorch/vendors/pytorch_wavelets/utils.py,sha256=fGRjZwndXEdQqp8xBnf3rYL5UeLj5Qb3LVTZisvbe5E,7762
+mirtorch/vendors/pytorch_wavelets/dtcwt/__init__.py,sha256=6j3Ds3Bk_3EMQSAY-S96U8itjCKh3bIjay0iFy0JC5g,173
+mirtorch/vendors/pytorch_wavelets/dtcwt/coeffs.py,sha256=-wy3-p0YxGzEOkcHY-35qGj1TwMnT7zYVFLszb4Q5-I,4174
+mirtorch/vendors/pytorch_wavelets/dtcwt/lowlevel.py,sha256=uH8JR7DF8oC2rtVGN7dG06u_V-C1AK7h7V77_hCWCX0,10511
+mirtorch/vendors/pytorch_wavelets/dtcwt/lowlevel2.py,sha256=xBgYeJlN-2615L34Bk2CMP0ZO7e8kS9dUKzjTCFj3E4,24984
+mirtorch/vendors/pytorch_wavelets/dtcwt/transform2d.py,sha256=gPeLeiR58tNWHdtCfEHgZOfJYMMpROWpqIbKvYMYuDs,11185
+mirtorch/vendors/pytorch_wavelets/dtcwt/transform_funcs.py,sha256=-4QL4Ri8K5wufsFV8CFB2iTla_OldNQ2vd0oBLkxzkc,17008
+mirtorch/vendors/pytorch_wavelets/dwt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mirtorch/vendors/pytorch_wavelets/dwt/lowlevel.py,sha256=NlIIE7COdbYUUIF3zE6hctRz_VodlJDO0mtdu7ynxME,36173
+mirtorch/vendors/pytorch_wavelets/dwt/swt_inverse.py,sha256=wnPc2Qaj6skSOMS1MJP5J89jee9NjicX885DwrtPviw,8187
+mirtorch/vendors/pytorch_wavelets/dwt/transform1d.py,sha256=-hlErWgapPh073AXvDYua5waOUY5dWdf-iEB6LmRGog,3784
+mirtorch/vendors/pytorch_wavelets/dwt/transform2d.py,sha256=0BvE8suR9PDCihMXMxbgSS_BieC8xv1csLtQmpsrkKM,8070
+mirtorch/vendors/pytorch_wavelets/scatternet/__init__.py,sha256=XfquxUhTyrDV_bvL6FV4eUsqK6C4eA0ib-kZUEfk5Fk,83
+mirtorch/vendors/pytorch_wavelets/scatternet/layers.py,sha256=Tu1eDSoxypnaDJMcpjxpIzZjrXHW-7yaxMZkKyz0Dug,7851
+mirtorch/vendors/pytorch_wavelets/scatternet/lowlevel.py,sha256=kKzagY0QOM5y2mRQtqWfxHQzW3iqaTvluZyObgrelNQ,25205
+MIRTorch-0.1.0.dist-info/LICENSE,sha256=vlr6JcRY9pCrasK99KDlTAMdtJPiiy3-LtQahoWrTGU,6112
+MIRTorch-0.1.0.dist-info/METADATA,sha256=oOHvZurBaKEGsbmx9JcjPRm-13Ly65WLxTDXnahOtLo,8168
+MIRTorch-0.1.0.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+MIRTorch-0.1.0.dist-info/entry_points.txt,sha256=iq9Cu1AYMLWEh9jhFJ5t5hC3ma-xEfRs9gSPRWeyRzI,57
+MIRTorch-0.1.0.dist-info/top_level.txt,sha256=flms4ycMQuWT7P7BxFBIeoTRBbpzLoa-R9zpYypDiQU,9
+MIRTorch-0.1.0.dist-info/RECORD,,

MIRTorch-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (72.1.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

MIRTorch-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+my-script = my_package.module:function

MIRTorch-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+mirtorch

mirtorch/__init__.py

@@ -0,0 +1,6 @@
+from mirtorch import prox
+from mirtorch import linear
+from mirtorch import alg
+from mirtorch import dic
+
+__all__ = ["linear", "prox", "alg", "dic"]

mirtorch/alg/__init__.py

@@ -0,0 +1,12 @@
+import logging
+
+from .cg import CG
+from .spectral import power_iter
+from .pogm import POGM
+from .fista import FISTA
+from .fbpd import FBPD
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+__all__ = ["CG", "power_iter", "POGM", "FISTA", "FBPD"]

mirtorch/alg/cg.py

@@ -0,0 +1,153 @@
+import torch
+import logging
+from torch import Tensor
+
+logger = logging.getLogger(__name__)
+
+
+class CG_func(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, b: Tensor, A, max_iter, tol, alert, x0, eval_func, P):
+        ctx.save_for_backward(b)
+        ctx.A = A
+        ctx.max_iter = max_iter
+        ctx.tol = tol
+        ctx.alert = alert
+        ctx.eval_func = eval_func
+        ctx.P = P
+        return cg_block(x0, b, A, tol, max_iter, alert, eval_func, P)
+
+    @staticmethod
+    def backward(ctx, dx):
+        b = ctx.saved_tensors[0]
+        # a better initialization?
+        return (
+            cg_block(
+                b, dx, ctx.A, ctx.tol, ctx.max_iter, ctx.alert, ctx.eval_func, ctx.P
+            ),
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+def cg_block(x0, b, A, tol, max_iter, alert, eval_func, P):
+    # solver for PSD Ax = b
+    if P is None:
+        r0 = b - A * x0
+        rk = r0
+        p0 = r0.detach().clone()
+        pk = p0
+        xk = x0.detach().clone()
+        rktrk = torch.square(torch.norm(rk))
+        num_loop = 0
+        if eval_func is not None:
+            saved = []
+        while rktrk.item() > tol and num_loop < max_iter:
+            pktapk = torch.sum(pk.conj() * (A * pk)).abs()
+            alpha = rktrk / pktapk
+            xk1 = xk.add_(alpha * pk)
+            rk1 = rk.sub_(alpha * A * pk)
+            rk1trk1 = torch.square(torch.norm(rk1))
+            beta = rk1trk1 / rktrk
+            pk1 = (pk.mul_(beta)).add_(rk1)
+            xk = xk1
+            rk = rk1
+            pk = pk1
+            rktrk = rk1trk1
+            num_loop = num_loop + 1
+            if eval_func is not None:
+                saved.append(eval_func(rk))
+            if alert:
+                logger.info(
+                    "Residual at %dth iter in forward CG: %10.3e." % (num_loop, rktrk)
+                )
+    else:
+        r0 = b - A * x0
+        rk = r0
+        zk = P * rk
+        pk = zk.clone()
+        xk = x0.detach().clone()
+        rktzk = (rk.conj() * zk).sum().abs()
+        num_loop = 0
+        if eval_func is not None:
+            saved = []
+        while torch.square(torch.norm(rk)).item() > tol and num_loop < max_iter:
+            pktapk = torch.sum(pk.conj() * (A * pk)).abs()
+            alpha = rktzk / pktapk
+            xk1 = xk.add_(alpha * pk)
+            rk1 = rk.sub_(alpha * A * pk)
+            zk1 = P * rk1
+            rk1tzk1 = (rk1.conj() * zk1).sum().abs()
+            beta = rk1tzk1 / rktzk
+            pk1 = (pk.mul_(beta)).add_(zk1)
+            xk = xk1
+            rk = rk1
+            pk = pk1
+            rktzk = rk1tzk1
+            num_loop = num_loop + 1
+            if eval_func is not None:
+                saved.append(eval_func(rk))
+            if alert:
+                logger.info(
+                    "Residual at %dth iter in CG backpropagation: %10.3e."
+                    % (num_loop, rktzk)
+                )
+                if torch.cuda.is_available():
+                    logging.info(
+                        "GPU memory usage at %dth iter in CG backpropagation: %10.3e."
+                        % (num_loop, torch.cuda.memory_allocated() / 1024 / 1024 / 1024)
+                    )
+
+    if eval_func is not None:
+        return xk, saved
+    else:
+        return xk
+
+
+class CG:
+    r"""
+    Solve the equation :math:`Ax = b` with the conjugdate gradient (CG) method, where A is a positive semi-definite operator.
+    The backpropagation still calls the CG to calculate the Jacobian to save the memory.
+
+    Attributes:
+        A: LinearMap of a PSD matrix
+        tol: float, exiting tolerance
+        max_iter: int, max number of iterations
+        alert: bool, print the norm of residuals at the end
+        eval_func: user-defined function to calculate the loss at each iteration.
+        P: LinearMap of a Preconditioner
+
+    Methods:
+        run: run the CG algorithm
+    """
+
+    def __init__(self, A, max_iter=20, tol=1e-2, P=None, alert=False, eval_func=None):
+        self.solver = CG_func.apply
+        self.A = A
+        self.max_iter = max_iter
+        self.tol = tol
+        self.solver = CG_func.apply
+        self.alert = alert
+        self.eval_func = eval_func
+        self.P = P
+
+    def run(self, x0, b):
+        r"""Run the CG iterations.
+        Args:
+            x0: Initialization
+            b: RHS
+
+        Returns:
+            xk: results
+            saved: (optional) a list of intermediate results, calculated by the eval_func.
+        """
+        if list(self.A.size_out) != list(b.shape):
+            raise ValueError("The size of A and b do not match.")
+        return self.solver(
+            b, self.A, self.max_iter, self.tol, self.alert, x0, self.eval_func, self.P
+        )

mirtorch/alg/fbpd.py

@@ -0,0 +1,95 @@
+from typing import Callable
+import logging
+
+import torch
+from mirtorch.prox import Prox, Conj
+from mirtorch.linear import LinearMap
+
+logger = logging.getLogger(__name__)
+
+
+class FBPD:
+    r"""Forward-backward primal dual (FBPD) algorithm.
+
+    Ref:
+    L. Condat, A primal dual splitting method for convex optimization involving
+    Lipschitzian, proximable and linear composite terms. Journal of Optimization Theory
+    and Applications, 158(2):460-479, 2013.
+
+    The cost function is:
+
+    .. math::
+
+        arg \min_x f(x) + g(x) + h(Gx)
+
+    where f and h are proper convex functions, and g is a convex function with a L-Lipschitz continuous gradient.
+
+    Attributes:
+        g_grad: Callable to calculate the gradient of g
+        f_prox: Prox: proximal operator of f
+        h_prox: Prox: proximal operator of h
+        g_L: float, Lipschitz value of g_grad
+        G_norm: float of the norm of G'G, can be solved by power_iter()
+        tau: float, step size
+        max_iter: int, number of iterations to run
+        eval_func: user-defined function to calculate the loss at each iteration.
+    """
+
+    def __init__(
+        self,
+        g_grad: Callable,
+        f_prox: Prox,
+        h_prox: Prox,
+        g_L: float,
+        G_norm: float,
+        G: LinearMap | None = None,
+        tau: float | None = None,
+        max_iter: int = 10,
+        eval_func: Callable | None = None,
+        p: int = 1,
+    ):
+        self.max_iter = max_iter
+        self.g_grad = g_grad
+        self.f_prox = f_prox
+        self.h_prox = h_prox
+        self.h_conj_prox = Conj(self.h_prox)
+        self.g_L = g_L
+        self.G = G
+        self.G_norm = G_norm
+        self.p = p
+        if tau is None:
+            self.tau = 2.0 / (g_L + 2.0)
+        else:
+            self.tau = tau
+        self.sigma = (1.0 / self.tau - self.g_L / 2.0) / self.G_norm
+        self.eval_func = eval_func
+
+    def run(self, x0: torch.Tensor):
+        r"""
+        Run the algorithm
+        Args:
+            x0: tensor, initialization
+        Returns:
+            xk: tensor, results
+            saved: (optional) a list of intermediate results, calcuated by the eval_func.
+        """
+        uold = self.G * x0
+        xold = x0
+        if self.eval_func is not None:
+            saved = []
+        for i in range(1, self.max_iter + 1):
+            xold_bar = self.g_grad(xold) + self.G.H * uold
+            xnew = self.f_prox(xold - self.tau * xold_bar, self.tau)
+            uold_bar = self.G * (2 * xnew - xold)
+            unew = self.h_conj_prox(uold + self.sigma * uold_bar, self.sigma)
+            xold = self.p * xnew + (1 - self.p) * xold
+            uold = self.p * unew + (1 - self.p) * uold
+            if self.eval_func is not None:
+                saved.append(self.eval_func(xold))
+                logger.info(
+                    "The cost function at %dth iter in FBPD: %10.3e.", i, saved[-1]
+                )
+        if self.eval_func is not None:
+            return xold, saved
+        else:
+            return xold

mirtorch/alg/fista.py

@@ -0,0 +1,87 @@
+from mirtorch.prox import Prox
+import numpy as np
+import torch
+from typing import Callable
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class FISTA:
+    r"""
+    Fast Iterative Soft Thresholding Algorithm (FISTA) / Fast Proximal Gradient Method (FPGM)
+
+    .. math::
+
+        arg \min_x f(x) + g(x)
+
+    where grad(f(x)) is L-Lipschitz continuous and g is proximal-friendly function.
+
+    Attributes:
+        max_iter (int): number of iterations to run
+        f_grad (Callable): gradient of f
+        f_L (float): L-Lipschitz value of f_grad
+        g_prox (Prox): proximal operator g
+        restart (Union[...]): restart strategy, not yet implemented
+        eval_func: user-defined function to calculate the loss at each iteration.
+    """
+
+    def __init__(
+        self,
+        f_grad: Callable,
+        f_L: float,
+        g_prox: Prox,
+        max_iter: int = 10,
+        restart=False,
+        eval_func: Callable = None,
+    ):
+        self.max_iter = max_iter
+        self.f_grad = f_grad
+        self.f_L = f_L
+        self.prox = g_prox
+        self._alpha = 1 / self.f_L  # value for 1/L
+        self.eval_func = eval_func
+        if restart:
+            raise NotImplementedError
+        self.restart = restart
+
+    def run(self, x0: torch.Tensor):
+        r"""
+        Run the algorithm
+
+        Args:
+            x0: initialization
+
+        Returns:
+            xk: results
+            saved: (optional) a list of intermediate results, calcuated by the eval_func.
+        """
+
+        def _update_momentum():
+            nonlocal told, beta
+            tnew = 0.5 * (1 + np.sqrt(1 + 4 * told**2))
+            beta = (told - 1) / tnew
+            told = tnew
+
+        # initialize parameters
+        xold = x0
+        yold = x0
+        told = 1.0
+        beta = 0.0
+        if self.eval_func is not None:
+            saved = []
+        for i in range(1, self.max_iter + 1):
+            fgrad = self.f_grad(xold)
+            ynew = self.prox(xold - self._alpha * fgrad, self._alpha)
+            _update_momentum()
+            xnew = ynew + beta * (ynew - yold)
+            xold = xnew
+            yold = ynew
+            # log the cost function
+            if self.eval_func is not None:
+                saved.append(self.eval_func(xold))
+                logger.info(f"Cost function at {i}th iteration: {self.eval_func(xold)}")
+        if self.eval_func is not None:
+            return xold, saved
+        else:
+            return xold