Merge pull request #11 from psavery/separate-pyxrf-out

Allow pyxrf to be ran in its own environment

Author: celiafish

.github/workflows/build_and_test.yml

@@ -65,6 +65,10 @@ jobs:
         conda install -y --override-channels -c conda-forge --file tomviz/tests/python/requirements-dev.txt
         pip install --no-build-isolation --no-deps -U tomviz/tomviz/python
         pip install --no-build-isolation --no-deps -U tomviz/acquisition
+        # These do not have to be in the same environment as Tomviz, but the dependencies
+        # are so far compatible and we'll just let it go for the tests.
+        conda install -y --override-channels -c conda-forge --file tomviz/.github/workflows/pyxrf_requirements.txt
+        pip install --no-build-isolation --no-deps -U tomviz/tomviz/python/tomviz/pyxrf/pyxrf-utils
 
     - name: Run Tests
       env:

.github/workflows/build_requirements.txt

@@ -6,3 +6,5 @@ vtk
 libxml2-devel=2.14*
 # This is needed for building the tests
 gtest
+# This is needed for installing stuff on Python3.13
+setuptools

.github/workflows/pyxrf_requirements.txt

@@ -0,0 +1,9 @@
+# These are PyXRF workflow dependencies
+hxntools
+pyxrf
+xrf-tomo
+xraylib
+scikit-beam
+tomopy
+# These are required for installing the pyxrf-utils environment
+hatchling

.github/workflows/runtime_requirements.txt

@@ -7,11 +7,4 @@ jsonpointer
 pyfftw
 pygments
 pystackreg
-# These are PyXRF workflow dependencies
 h5py
-hxntools
-pyxrf
-xrf-tomo
-xraylib
-scikit-beam
-tomopy

pixi.toml

@@ -9,17 +9,8 @@ platforms = ["linux-64"]
 python = "==3.12"
 tomviz = ">=2.1.3"
 h5py = "*"
-xraylib = "*"
 tomopy = "*"
 
-[pypi-dependencies]
-bluesky-tiled-plugins = ">=2.0.0rc1"
-pyxrf = "*"
-hxntools = "*"
-xrf-tomo = "*"
-scikit-beam = "*"
-tiled = { version = ">=0.2.0", extras = ["client"] }
-
 [tasks]
 # Launch Tomviz from the pixi environment
 tomviz = "tomviz"

tests/cxx/PyXRFWorkflowTest.cxx

@@ -130,9 +130,12 @@ private slots:
     // After accepting, a modal dialog will appear. Start posting
     // events to check it and accept it when it appears.
     bool found = false;
-    auto checkFunc = [&found](){
+
+    std::function<void()> checkFunc;
+    checkFunc = [&found, &checkFunc](){
       auto* dialog = findWidget<SelectItemsDialog>();
       if (!dialog) {
+        QTimer::singleShot(1000, checkFunc);
         return;
       }
 
@@ -150,7 +153,6 @@ private slots:
     int maxTime = 30;
     while (!found && timeElapsed < maxTime) {
       QThread::sleep(1);
-      QTimer::singleShot(0, checkFunc);
       QApplication::processEvents();
       timeElapsed += 1;
     }

tests/python/CMakeLists.txt

@@ -3,3 +3,7 @@ include(PythonTests.cmake)
 add_python_test(operator)
 add_python_test(external)
 add_python_test(pystackreg)
+add_python_test(multi_array)
+add_python_test(xcorr)
+add_python_test(tilt_axis_shift)
+add_python_test(constraint_dft)

tests/python/conftest.py

@@ -9,6 +9,9 @@
 
 import numpy as np
 
+from tomviz.executor import load_dataset
+from tomviz.external_dataset import Dataset
+
 from utils import download_file, download_and_unzip_file
 
 DATA_URL = 'https://data.kitware.com/api/v1/file'
@@ -30,6 +33,25 @@ def hxn_xrf_example_output_dir(data_dir: Path) -> Path:
     return output_dir
 
 
+@pytest.fixture(scope='function')
+def hxn_xrf_example_dataset(hxn_xrf_example_output_dir: Path) -> Dataset:
+    example_files = [
+        'Pt_L.h5',
+        'Zn_K.h5',
+    ]
+    example_files = [
+        hxn_xrf_example_output_dir / 'extracted_elements' / name
+        for name in example_files
+    ]
+    dataset = load_dataset(example_files[0])
+    for new_file in example_files[1:]:
+        new_dataset = load_dataset(new_file)
+        new_name = new_dataset.active_name
+        dataset.arrays[new_name] = new_dataset.active_scalars
+
+    return dataset
+
+
 @pytest.fixture
 def pystackreg_reference_output(data_dir: Path) -> dict[str, np.ndarray]:
     filepath = data_dir / 'test_pystackreg_reference_output.npz'
@@ -41,6 +63,39 @@ def pystackreg_reference_output(data_dir: Path) -> dict[str, np.ndarray]:
     return np.load(filepath)
 
 
+@pytest.fixture
+def xcorr_reference_output(data_dir: Path) -> dict[str, np.ndarray]:
+    filepath = data_dir / 'test_xcorr_reference_output.npz'
+    if not filepath.exists():
+        # Download it
+        url = DATA_URL + '/6920b212cdb169d7a021d769/download'
+        download_file(url, filepath)
+
+    return np.load(filepath)
+
+
+@pytest.fixture
+def tilt_axis_shift_reference_output(data_dir: Path) -> dict[str, np.ndarray]:
+    filepath = data_dir / 'test_tilt_axis_shift_reference_output.npz'
+    if not filepath.exists():
+        # Download it
+        url = DATA_URL + '/6920b333cdb169d7a021d76c/download'
+        download_file(url, filepath)
+
+    return np.load(filepath)
+
+
+@pytest.fixture
+def constraint_dft_reference_output(data_dir: Path) -> dict[str, np.ndarray]:
+    filepath = data_dir / 'test_constraint_dft_reference_output.npz'
+    if not filepath.exists():
+        # Download it
+        url = DATA_URL + '/6920b750cdb169d7a021d773/download'
+        download_file(url, filepath)
+
+    return np.load(filepath)
+
+
 @pytest.fixture(scope="module")
 def test_state_file(tmpdir_factory):
     tmpdir = tmpdir_factory.mktemp('state')

tests/python/constraint_dft_test.py

@@ -0,0 +1,48 @@
+import numpy as np
+
+from utils import load_operator_class, load_operator_module
+
+from tomviz.external_dataset import Dataset
+
+
+def test_constraint_dft(hxn_xrf_example_dataset: Dataset,
+                        constraint_dft_reference_output: dict[str, np.ndarray]):
+    dataset = hxn_xrf_example_dataset
+    reference_output = constraint_dft_reference_output
+
+    # Load the operator module
+    module = load_operator_module('Recon_DFT_constraint')
+    operator = load_operator_class(module)
+
+    # Verify the keys match
+    assert list(sorted(reference_output)) == list(sorted(dataset.arrays))
+
+    # But the arrays should not match. In fact, the shapes should not match.
+    assert not all(
+        np.array_equal(reference_output[key].shape, dataset.arrays[key].shape)
+        for key in reference_output
+    )
+
+    # Apply the transformation
+    results = operator.transform(
+        dataset,
+        Niter=10,
+        Niter_update_support=50,
+        supportSigma=0.1,
+        supportThreshold=10,
+        seed=0,
+    )
+
+    recon_dataset = results['reconstruction']
+
+    # Verify the keys match
+    assert list(sorted(reference_output)) == list(sorted(recon_dataset.arrays))
+
+    # Arrays should now match
+    assert all(
+        np.allclose(reference_output[key], recon_dataset.arrays[key], atol=1e-2)
+        for key in reference_output
+    )
+
+    # Save a new reference output:
+    # np.savez_compressed('test_constraint_dft_reference_output.npz', **recon_dataset.arrays)

tests/python/multi_array_test.py

@@ -0,0 +1,68 @@
+import numpy as np
+
+from utils import load_operator_class, load_operator_module
+
+from tomviz.external_dataset import Dataset
+
+
+def test_multi_arrays(hxn_xrf_example_dataset: Dataset):
+    dataset = hxn_xrf_example_dataset
+
+    # Test a few operators that use `@apply_to_each_array`, and verify
+    # that they do indeed apply to each array.
+    # Load the operator module
+    bin_module = load_operator_module('BinTiltSeriesByTwo')
+
+    # Verify the tilt angles won't change, but image shape decreases by 2x
+    orig_shapes = {k: v.shape for k, v in dataset.arrays.items()}
+    orig_tilt_angles = dataset.tilt_angles
+
+    # Apply the transformation
+    bin_module.transform(dataset)
+
+    # New shapes should be binned by 2 in x and y, but not z
+    # Tilt angles should be unaffected
+    binned_shapes = {k: v.shape for k, v in dataset.arrays.items()}
+    expected_shapes = {k: (shape[0] // 2, shape[1] // 2, shape[2])
+                       for k, shape in orig_shapes.items()}
+
+    assert binned_shapes == expected_shapes
+    assert np.allclose(orig_tilt_angles, dataset.tilt_angles)
+
+    # Delete 3 slices. This will delete tilt angles too.
+    delete_slices_module = load_operator_module('DeleteSlices')
+
+    first_slice = 3
+    last_slice = 5
+    num_deleted = last_slice - first_slice + 1
+    delete_slices_module.transform(dataset,
+                                   firstSlice=first_slice,
+                                   lastSlice=last_slice,
+                                   axis=2)
+
+    # New shapes should be the same, other than 3 less in Z
+    expected_shapes = {k: (shape[0], shape[1], shape[2] - num_deleted)
+                       for k, shape in binned_shapes.items()}
+    deleted_shapes = {k: v.shape for k, v in dataset.arrays.items()}
+
+    expected_angles = np.hstack((orig_tilt_angles[:first_slice],
+                                 orig_tilt_angles[last_slice + 1:]))
+
+    assert deleted_shapes == expected_shapes
+    assert np.allclose(dataset.tilt_angles, expected_angles)
+
+    # Now run a reconstruction and verify we get all results
+    recon_sirt_module = load_operator_module('Recon_SIRT')
+    operator = load_operator_class(recon_sirt_module)
+
+    results = operator.transform(dataset, Niter=2)
+
+    recon_dataset = results['reconstruction']
+
+    # We should have the same scalar names
+    assert sorted(recon_dataset.scalars_names) == sorted(dataset.scalars_names)
+
+    # Verify the output appears valid. We are not testing the actual
+    # reconstruction operator here, just that it ran on multiple arrays.
+    means = [np.mean(array) for array in recon_dataset.arrays.values()]
+    assert np.all(np.array(means) > 0)