Merge branch 'main' into cristina/psdred

CHANGES.rst

@@ -7,7 +7,7 @@ Version 0.0.7   (unreleased)
 ----------------------------
 
 • New module ``scico.trace`` for tracing function/method calls.
-• Support ``jaxlib`` and ``jax`` versions 0.4.13 to 0.4.37.
+• Support ``jaxlib`` and ``jax`` versions 0.4.13 to 0.5.0.
 • Support ``flax`` versions 0.8.0 to 0.10.2.
 
 

LICENSE

@@ -1,6 +1,6 @@
 BSD 3-Clause License
 
-Copyright (c) 2021-2024, Los Alamos National Laboratory
+Copyright (c) 2021-2025, Los Alamos National Laboratory
 All rights reserved.
 
 Redistribution and use in source and binary forms, with or without

README.md

@@ -61,7 +61,7 @@ License (see the `LICENSE` file for details).
 
 LANL open source approval reference C20091.
 
-\(c\) 2020-2024. Triad National Security, LLC. All rights reserved. This
+\(c\) 2020-2025. Triad National Security, LLC. All rights reserved. This
 program was produced under U.S. Government contract 89233218CNA000001
 for Los Alamos National Laboratory (LANL), which is operated by Triad
 National Security, LLC for the U.S. Department of Energy/National

docs/source/conf/10-project.py

@@ -2,7 +2,7 @@
 
 # General information about the project.
 project = "SCICO"
-copyright = "2020-2024, SCICO Developers"
+copyright = "2020-2025, SCICO Developers"
 
 # The version info for the project you're documenting, acts as replacement for
 # |version| and |release|, also used in various other places throughout the

requirements.txt

@@ -1,11 +1,11 @@
 typing_extensions
-numpy>=1.20.0
-scipy>=1.6.0
+numpy>=1.25.0
+scipy>=1.11.0
 imageio>=2.17
 tifffile
 matplotlib
-jaxlib>=0.4.13,<=0.4.37
-jax>=0.4.13,<=0.4.37
+jaxlib>=0.4.13,<=0.5.0
+jax>=0.4.13,<=0.5.0
 orbax-checkpoint>=0.5.0
 flax>=0.8.0,<=0.10.2
 pyabel>=0.9.0

scico/linop/xray/__init__.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-# Copyright (C) 2023-2024 by SCICO Developers
+# Copyright (C) 2023-2025 by SCICO Developers
 # All rights reserved. BSD 3-clause License.
 # This file is part of the SCICO package. Details of the copyright and
 # user license can be found in the 'LICENSE' file distributed with the
@@ -66,12 +66,10 @@
 
 import sys
 
+from ._util import rotate_volume
 from ._xray import XRayTransform2D, XRayTransform3D
 
-__all__ = [
-    "XRayTransform2D",
-    "XRayTransform3D",
-]
+__all__ = ["XRayTransform2D", "XRayTransform3D", "rotate_volume"]
 
 
 # Imported items in __all__ appear to originate in top-level xray module

scico/linop/xray/_util.py

@@ -0,0 +1,58 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2024-2025 by SCICO Developers
+# All rights reserved. BSD 3-clause License.
+# This file is part of the SCICO package. Details of the copyright and
+# user license can be found in the 'LICENSE' file distributed with the
+# package.
+
+"""Utilities for CT data."""
+
+from typing import Optional
+
+import jax.numpy as jnp
+from jax.scipy.ndimage import map_coordinates
+from jax.scipy.spatial.transform import Rotation
+from jax.typing import ArrayLike
+
+
+def rotate_volume(
+    vol: ArrayLike,
+    rot: Rotation,
+    x: Optional[ArrayLike] = None,
+    y: Optional[ArrayLike] = None,
+    z: Optional[ArrayLike] = None,
+    center: Optional[ArrayLike] = None,
+):
+    """Rotate a 3D array.
+
+    Rotate a 3D array as specified by an instance of
+    :class:`~jax.scipy.spatial.transform.Rotation`. Any axis coordinates
+    that are not specified default to a range corresponding to the size
+    of the array on that axis, starting at zero.
+
+    Args:
+        vol: Array to be rotated.
+        rot: Rotation specification.
+        x: Coordinates for :code:`x` axis (axis 0).
+        y: Coordinates for :code:`y` axis (axis 1).
+        z: Coordinates for :code:`z` axis (axis 2).
+        center: A 3-vector specifying the center of rotation.
+           Defaults to the center of the array.
+
+    Returns:
+        Rotated array.
+    """
+    shape = vol.shape
+    if x is None:
+        x = jnp.arange(shape[0])
+    if y is None:
+        y = jnp.arange(shape[1])
+    if z is None:
+        z = jnp.arange(shape[2])
+    if center is None:
+        center = (jnp.array(shape, dtype=jnp.float32) - 1.0) / 2.0
+    gx, gy, gz = jnp.meshgrid(x - center[0], y - center[1], z - center[2], indexing="ij")
+    crd = jnp.stack((gx.ravel(), gy.ravel(), gz.ravel()))
+    rot_crd = rot.as_matrix() @ crd + center[:, jnp.newaxis]  # faster than rot.apply(crd.T)
+    rot_vol = map_coordinates(vol, rot_crd.reshape((3,) + shape), order=1)
+    return rot_vol

scico/test/linop/test_linop.py

@@ -109,8 +109,7 @@ def test_scalar_left(testobj, operator, scalar):
     assert isinstance(comp_op, linop.LinearOperator)  # Ensure we don't get a Map
     assert comp_op.input_dtype == testobj.A.dtype
     np.testing.assert_allclose(comp_mat @ testobj.x, comp_op @ testobj.x, rtol=5e-5)
-
-    np.testing.assert_allclose(comp_mat.conj().T @ testobj.y, comp_op.adj(testobj.y), rtol=5e-5)
+    np.testing.assert_allclose(comp_mat.conj().T @ testobj.y, comp_op.adj(testobj.y), rtol=1e-4)
 
 
 @pytest.mark.parametrize("operator", [op.mul, op.truediv])
@@ -218,7 +217,7 @@ def test_transpose_matvec(testobj):
 
     assert a.dtype == testobj.A.dtype
     assert b.dtype == testobj.A.dtype
-    np.testing.assert_allclose(a, comp_mat, rtol=5e-5)
+    np.testing.assert_allclose(a, comp_mat, rtol=1e-4)
     np.testing.assert_allclose(a, b, rtol=5e-5)
 
 
@@ -264,7 +263,7 @@ def test_adjoint_matvec(testobj):
     assert a.dtype == testobj.A.dtype
     assert b.dtype == testobj.A.dtype
     assert c.dtype == testobj.A.dtype
-    np.testing.assert_allclose(a, comp_mat, rtol=5e-5)
+    np.testing.assert_allclose(a, comp_mat, rtol=1e-4)
     np.testing.assert_allclose(a, b, rtol=5e-5)
     np.testing.assert_allclose(a, c, rtol=5e-5)
 

scico/test/linop/xray/test_astra.py

@@ -24,9 +24,9 @@
 
 
 N = 128
-RTOL_CPU = 5e-5
-RTOL_GPU = 7e-2
-RTOL_GPU_RANDOM_INPUT = 1.0
+RTOL_CPU = 1e-4
+RTOL_GPU = 1e-1
+RTOL_GPU_RANDOM_INPUT = 2.0
 
 
 def make_im(Nx, Ny, is_3d=True):

scico/test/linop/xray/test_xray_util.py

@@ -0,0 +1,12 @@
+import numpy as np
+
+from jax.scipy.spatial.transform import Rotation
+
+from scico.linop.xray import rotate_volume
+
+
+def test_rotate_volume():
+    vol = np.arange(27).reshape((3, 3, 3))
+    rot = Rotation.from_euler("XY", [90, 90], degrees=True)
+    vol_rot = rotate_volume(vol, rot)
+    np.testing.assert_allclose(vol.transpose((1, 2, 0)), vol_rot, rtol=1e-7)