mbirtorch 0.0.1__tar.gz

mbirtorch-0.0.1/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2024, Charles A. Bouman and Gregery T. Buzzard
+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.

mbirtorch-0.0.1/PKG-INFO

@@ -0,0 +1,121 @@
+Metadata-Version: 2.4
+Name: mbirtorch
+Version: 0.0.1
+Summary: High-performance tomographic reconstruction using PyTorch
+Author-email: MBIRTORCH development team <buzzard@purdue.edu>
+License: BSD 3-Clause License
+        
+        Copyright (c) 2024, Charles A. Bouman and Gregery T. Buzzard
+        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.
+        
+Project-URL: source, https://github.com/cabouman/mbirtorch
+Keywords: tomography,tomographic reconstruction,computed tomography
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Information Analysis
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.13
+Requires-Dist: numpy
+Requires-Dist: matplotlib
+Requires-Dist: h5py
+Requires-Dist: scipy
+Requires-Dist: tifffile
+Requires-Dist: tqdm
+Requires-Dist: pywavelets
+Requires-Dist: osqp
+Requires-Dist: opencv-python
+Requires-Dist: olefile
+Requires-Dist: gdown
+Requires-Dist: scikit-learn
+Requires-Dist: psutil
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-xdist; extra == "test"
+Provides-Extra: docs
+Requires-Dist: sphinx; extra == "docs"
+Requires-Dist: sphinxcontrib-bibtex; extra == "docs"
+Requires-Dist: sphinx-book-theme>=1.0.1; extra == "docs"
+Requires-Dist: sphinx-design; extra == "docs"
+Requires-Dist: sphinx-copybutton; extra == "docs"
+Dynamic: license-file
+
+# mbirtorch
+
+[![CI](https://github.com/cabouman/mbirtorch/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/cabouman/mbirtorch/actions/workflows/ci.yml)
+[![Documentation Status](https://readthedocs.org/projects/mbirtorch/badge/?version=latest)](https://mbirtorch.readthedocs.io/en/latest/)
+
+A PyTorch port of [mbirjax](https://github.com/cabouman/mbirjax): high-performance
+tomographic reconstruction with multi-GPU support.
+
+This repo is the home for the port evaluated in
+`mbirjax_plans/plans/torch_port/port_plan.md` (checked out parallel to this
+repo).  The plan defines the motivation, the parity gates against mbirjax,
+the incremental phase plan, and a progress record; its findings pages and the
+supporting scripts live alongside it under `mbirjax_plans/plans/`.
+
+## Setup
+
+Create the conda environment and install the package in editable mode:
+
+```bash
+conda env create -f environment.yml
+conda activate mbirtorch
+pip install -e ".[test]"
+```
+
+## Tests
+
+```bash
+dev_scripts/run_tests.sh
+```
+
+## Caches
+
+mbirtorch keeps one on-disk cache: compiled `torch.compile` artifacts, under
+`~/.mbirtorch/torch_cache`.  It exists to make cold starts fast -- with it, a
+fresh process reuses prior compilations instead of recompiling (roughly 14 s
+down to 2 s for a first small reconstruction).  It grows with the number of
+distinct compiled shapes and typically stays in the tens of megabytes; it is
+never cleaned automatically.  To remove it:
+
+```python
+import mbirtorch
+mbirtorch.clear_cache()   # deletes ~/.mbirtorch entirely (recreated empty)
+```
+
+The location can be redirected by setting the `TORCHINDUCTOR_CACHE_DIR`
+environment variable before the first compile (e.g. to node-local or scratch
+storage on a cluster, where home quotas are tight); `clear_cache()` does not
+touch a redirected location.
+
+Everything else the package caches is in-memory only and is freed with the
+objects that hold it (e.g. the per-model pixel-index cache); nothing besides
+`~/.mbirtorch` is written to disk.

mbirtorch-0.0.1/README.md

@@ -0,0 +1,52 @@
+# mbirtorch
+
+[![CI](https://github.com/cabouman/mbirtorch/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/cabouman/mbirtorch/actions/workflows/ci.yml)
+[![Documentation Status](https://readthedocs.org/projects/mbirtorch/badge/?version=latest)](https://mbirtorch.readthedocs.io/en/latest/)
+
+A PyTorch port of [mbirjax](https://github.com/cabouman/mbirjax): high-performance
+tomographic reconstruction with multi-GPU support.
+
+This repo is the home for the port evaluated in
+`mbirjax_plans/plans/torch_port/port_plan.md` (checked out parallel to this
+repo).  The plan defines the motivation, the parity gates against mbirjax,
+the incremental phase plan, and a progress record; its findings pages and the
+supporting scripts live alongside it under `mbirjax_plans/plans/`.
+
+## Setup
+
+Create the conda environment and install the package in editable mode:
+
+```bash
+conda env create -f environment.yml
+conda activate mbirtorch
+pip install -e ".[test]"
+```
+
+## Tests
+
+```bash
+dev_scripts/run_tests.sh
+```
+
+## Caches
+
+mbirtorch keeps one on-disk cache: compiled `torch.compile` artifacts, under
+`~/.mbirtorch/torch_cache`.  It exists to make cold starts fast -- with it, a
+fresh process reuses prior compilations instead of recompiling (roughly 14 s
+down to 2 s for a first small reconstruction).  It grows with the number of
+distinct compiled shapes and typically stays in the tens of megabytes; it is
+never cleaned automatically.  To remove it:
+
+```python
+import mbirtorch
+mbirtorch.clear_cache()   # deletes ~/.mbirtorch entirely (recreated empty)
+```
+
+The location can be redirected by setting the `TORCHINDUCTOR_CACHE_DIR`
+environment variable before the first compile (e.g. to node-local or scratch
+storage on a cluster, where home quotas are tight); `clear_cache()` does not
+touch a redirected location.
+
+Everything else the package caches is in-memory only and is freed with the
+objects that hold it (e.g. the per-model pixel-index cache); nothing besides
+`~/.mbirtorch` is written to disk.

mbirtorch-0.0.1/mbirtorch/__init__.py

@@ -0,0 +1,130 @@
+"""mbirtorch: a PyTorch port of mbirjax (parallel-beam geometry).
+
+The public API mirrors mbirjax where implemented: numpy in, numpy out by
+default, with device tensors available via ``output_sharded=True`` (the name
+kept for API compatibility; here it means "return the device tensor").
+"""
+
+__version__ = "0.0.1"
+
+# ── persistent torch.compile cache ────────────────────────────────────────────
+# The inductor cache directory defaults to /tmp/torchinductor_<user>, which the
+# OS may clean; pin it to a stable per-user location so compiled artifacts
+# survive across processes and reboots (the mbirjax ~/.mbirjax/jax_cache
+# analog).  The FX-graph cache is what makes a FRESH PROCESS reuse prior
+# compilations; enable it explicitly for torch versions where it is not the
+# default.  setdefault keeps both overridable per-run via the environment, and
+# -- like mbirjax's TF_CPP_MIN_LOG_LEVEL -- this is effective only if mbirtorch
+# is imported before torch triggers its first compile, which any
+# import-mbirtorch-first program satisfies.  Dynamo TRACING still runs per
+# process (the cache skips inductor codegen, not tracing), so a cold process
+# keeps a small residual warmup.  ``mbirtorch.clear_cache()`` removes the
+# whole ~/.mbirtorch directory (see utilities.py).
+import os as _os
+
+_os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR",
+                       _os.path.expanduser("~/.mbirtorch/torch_cache"))
+_os.environ.setdefault("TORCHINDUCTOR_FX_GRAPH_CACHE", "1")
+
+from .parallel_beam import ParallelBeamModel
+from .cone_beam import ConeBeamModel
+from .translation_model import TranslationModel
+from .multiaxis_parallel import MultiAxisParallelModel, MultiAxisParallelBeamModel
+from .denoising import QGGMRFDenoiser
+from .tomography_model import TomographyModel
+from .autograd import (TorchProjector, forward_project_differentiable,
+                       back_project_differentiable)
+from .vcd_utils import (gen_weights, gen_weights_mar, gen_full_indices,
+                        gen_pixel_partition, gen_set_of_pixel_partitions,
+                        gen_partition_sequence, get_2d_ror_mask)
+from .denoising import median_filter3d
+from .qggmrf import (qggmrf_gradient_and_hessian_at_indices, get_b_from_nbr_wts,
+                     b_tilde_by_definition, qggmrf_loss)
+from .utilities import (generate_3d_shepp_logan_low_dynamic_range, clear_cache,
+                        makedirs, load_data_hdf5, save_data_hdf5,
+                        export_recon_hdf5, import_recon_hdf5,
+                        build_model, download_and_extract,
+                        copy_ct_model, stitch_arrays,
+                        get_ct_model, generate_demo_data,
+                        generate_3d_shepp_logan_reference, gen_cube_phantom,
+                        gen_translation_vectors, gen_translation_phantom,
+                        get_helical_half_rotation_slice_range,
+                        merge_log_files)
+from .memory_stats import get_memory_stats
+
+# __all__ is the DECLARED public surface, and autodoc honors it: every name here is
+# documented by ``automodule:: mbirtorch :members:``.  It is deliberately narrower than
+# the import list above -- the VCD and qGGMRF helpers stay importable as attributes
+# (mbirtorch.gen_full_indices still works, and the tests use that spelling) but are not
+# promised as public API, matching the surface mbirjax documents.  The same goes for
+# gen_cube_phantom and get_helical_half_rotation_slice_range: mbirjax's star imports put
+# them on the package, and its tests call them there, but its docs do not carry them.
+__all__ = [
+    "ParallelBeamModel", "ConeBeamModel", "TranslationModel",
+    "MultiAxisParallelModel", "TomographyModel", "QGGMRFDenoiser",
+    "TorchProjector", "forward_project_differentiable",
+    "back_project_differentiable", "gen_weights", "gen_weights_mar",
+    "median_filter3d", "download_and_extract", "build_model",
+    "save_data_hdf5", "load_data_hdf5", "export_recon_hdf5",
+    "import_recon_hdf5",
+    "generate_3d_shepp_logan_low_dynamic_range", "clear_cache",
+    "get_memory_stats", "SliceViewer", "VolumeStack", "slice_viewer",
+    "stitch_arrays", "get_ct_model", "copy_ct_model",
+    "generate_demo_data", "generate_3d_shepp_logan_reference",
+    # Documented hsnt and vcls names; these resolve lazily through __getattr__.
+    "hyper_denoise", "dehydrate", "rehydrate", "import_hsnt_data_hdf5",
+    "create_hsnt_metadata", "export_hsnt_data_hdf5", "generate_hyper_data",
+    "get_opt_views", "show_image_with_projection_rays",
+]
+
+# ── lazy exports (PEP 562) ───────────────────────────────────────────────────
+# The viewer names resolve on first attribute access so that a headless
+# `import mbirtorch` never imports matplotlib; most mbirtorch runs (batch
+# recons, tests) never open a viewer.  The preprocess, hsnt, and vcls
+# modules resolve the same way, so `import mbirtorch` never pays for their
+# dependency stacks (preprocess: osqp pulls scipy.sparse, plus cv2 and
+# tifffile; hsnt: scikit-learn, scipy, matplotlib; vcls: the model layer and
+# tqdm).  Both spellings keep working -- `mbirtorch.hsnt` resolves here, and
+# `import mbirtorch.hsnt` is an ordinary submodule import -- and the
+# star-exported FUNCTION names (mbirtorch.dehydrate, mbirtorch.get_opt_views,
+# ...) resolve through _LAZY_NAMES, so the public surface matches mbirjax's
+# eager star imports exactly; only WHEN each module loads changes.
+_VIEWER_EXPORTS = ("SliceViewer", "VolumeStack", "slice_viewer")
+
+_LAZY_MODULES = ("preprocess", "hsnt", "vcls")
+
+# The names mbirjax exposes at package level via `from .hsnt import *` and
+# `from .vcls import *`, mapped to their owning module (neither module
+# declares __all__, so this is their full public def list; a new public
+# function in either module gets a line here).
+_LAZY_NAMES = {
+    'hyper_denoise': 'hsnt', 'dehydrate': 'hsnt', 'rehydrate': 'hsnt',
+    'import_hsnt_data_hdf5': 'hsnt', 'create_hsnt_metadata': 'hsnt',
+    'export_hsnt_data_hdf5': 'hsnt', 'generate_hyper_data': 'hsnt',
+    'subsample_R_gamma': 'vcls', 'max_abs_neighbor_diff': 'vcls',
+    'get_opt_views': 'vcls', 'compute_view_basis_functions': 'vcls',
+    'compute_cov_matrix': 'vcls', 'compute_vcl': 'vcls',
+    'compute_opt_angle_subset': 'vcls', 'get_2d_subsampling_indices': 'vcls',
+    'show_image_with_projection_rays': 'vcls', 'reorder_by_priority': 'vcls',
+    # The blue-noise pattern (a 382 KB array literal), loaded on first use.
+    'bn256': 'bn256',
+}
+
+
+def __getattr__(name):
+    import importlib
+    if name in _VIEWER_EXPORTS:
+        from . import view_utils
+        value = getattr(view_utils, name)
+        globals()[name] = value  # cache: later accesses skip this hook
+        return value
+    if name in _LAZY_MODULES:
+        value = importlib.import_module('.' + name, __name__)
+        globals()[name] = value
+        return value
+    if name in _LAZY_NAMES:
+        module = importlib.import_module('.' + _LAZY_NAMES[name], __name__)
+        value = getattr(module, name)
+        globals()[name] = value
+        return value
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")