From cb97047ade626ae559ed301a57202939f5067f70 Mon Sep 17 00:00:00 2001 From: Sam Welborn Date: Wed, 11 Sep 2024 14:58:24 -0400 Subject: [PATCH] add electronCount for single frame add individual frame counting - no support for row-dark currently - had to change ctor for vectorToPyArray, see here: https://github.com/pybind/pybind11/issues/1042#issuecomment-2262727578 . This may be a numpy 2 thing - outputs will be a SparseArray with one frame/scan shape = (1, 1) and no metadata. This is to take advantage of methods in SparseArray --- python/image.cpp | 47 +++++++++++++++++++++- python/stempy/image/__init__.py | 68 ++++++++++++++++++++++++++++++++ stempy/electron.cpp | 26 +++++++++++++ stempy/electron.h | 5 +++ tests/test_image.py | 69 ++++++++++++++++++++++++++++++++- 5 files changed, 213 insertions(+), 2 deletions(-) diff --git a/python/image.cpp b/python/image.cpp index 894caadf..52484505 100644 --- a/python/image.cpp +++ b/python/image.cpp @@ -25,7 +25,8 @@ py::array_t vectorToPyArray(std::vector&& v) auto deleter = [](void* v) { delete reinterpret_cast*>(v); }; auto* ptr = new std::vector(std::move(v)); auto capsule = py::capsule(ptr, deleter); - return py::array(ptr->size(), ptr->data(), capsule); + py::array_t arr({ ptr->size() }, { sizeof(T) }, ptr->data(), capsule); + return arr; } struct ElectronCountedDataPyArray @@ -214,6 +215,44 @@ ElectronCountedDataPyArray electronCount(Reader* reader, return electronCount(reader, options.toCpp()); } +// Function to process individual frames +template +py::array_t electronCount( + py::array_t& frame, Dimensions2D frameDimensions, + const ElectronCountOptionsClassicPy& options) +{ + py::buffer_info frameBufferInfo = frame.request(); + + if (frameBufferInfo.ndim != 2 || + frameBufferInfo.format != py::format_descriptor::format()) { + throw std::runtime_error( + "Input frame must be a 2D array of the correct type."); + } + + const ElectronCountOptionsClassic cppOptions = options.toCpp(); + + // Convert the buffer to a std::vector + // TODO: is there a way to avoid this copy? For e.g span impl: + // https://github.com/pybind/pybind11/issues/1042#issuecomment-663154709 + std::vector frameVec(static_cast(frameBufferInfo.ptr), + static_cast(frameBufferInfo.ptr) + + frameBufferInfo.size); + + // Call the electronCount function with the std::vector + if (!cppOptions.darkReference) { + return vectorToPyArray( + electronCount(frameVec, frameDimensions, cppOptions)); + } else { + return vectorToPyArray( + electronCount(frameVec, frameDimensions, cppOptions)); + } + + std::vector result = + electronCount(frameVec, frameDimensions, cppOptions); + + return vectorToPyArray(std::move(result)); +} + // Explicitly instantiate version for py::array_t template std::vector createSTEMImages( const std::vector>>& sparseData, @@ -490,6 +529,12 @@ PYBIND11_MODULE(_image, m) electronCount, py::call_guard()); + // Count individual frame + m.def("electron_count_frame", + (py::array_t(*)(py::array_t&, Dimensions2D, + const ElectronCountOptionsClassicPy&)) & + electronCount); + // Calculate thresholds, with gain m.def( "calculate_thresholds", diff --git a/python/stempy/image/__init__.py b/python/stempy/image/__init__.py index 1909c89f..dcbd3556 100644 --- a/python/stempy/image/__init__.py +++ b/python/stempy/image/__init__.py @@ -348,6 +348,74 @@ def electron_count(reader, darkreference=None, number_of_samples=40, return array + +def electron_count_frame( + frame: np.ndarray, + options=None, + darkreference=None, + background_threshold=4.0, + xray_threshold=2000.0, + gain=None, +): + """Generate a list of coordinates of electron hits for a single 2D numpy array. + + :param frame: the frame. + :type frame: numpy.ndarray + :param options: the options to use for electron counting. If set, all other + parameters are ignored. + :type options: stempy.image.ElectronCountOptionsClassic + :param darkreference: the dark reference to subtract, potentially generated + via stempy.image.calculate_average(). + :type darkreference: stempy.image.ImageArray or numpy.ndarray + :param background_threshold: the threshold for background + :type background_threshold: float + :param xray_threshold: the threshold for x-rays + :type xray_threshold: float + :param gain: the gain mask to apply. Must match the frame dimensions. + :type gain: numpy.ndarray (2D) + + :return: the coordinates of the electron hits for the frame. + :rtype: numpy.ndarray + """ + + if gain is not None: + # Invert, as we will multiply in C++ + # It also must be a float32 + gain = np.power(gain, -1) + gain = _safe_cast(gain, np.float32, "gain") + + if options is None: + if isinstance(darkreference, np.ndarray): + # Must be float32 for correct conversions + darkreference = _safe_cast(darkreference, np.float32, "dark reference") + + options = _image.ElectronCountOptionsClassic() + + options.dark_reference = darkreference + options.background_threshold = background_threshold + options.x_ray_threshold = xray_threshold + options.gain = gain + options.apply_row_dark_subtraction = False + options.optimized_mean = 0.0 + options.apply_row_dark_use_mean = False + else: + if options.apply_row_dark_subtraction: + print("Warning: apply_row_dark_subtraction is not supported " + "for single frame electron counting. Ignoring this option.") + options.apply_row_dark_subtraction = False + options.gain = gain + + electron_counts = _image.electron_count_frame(frame, frame.shape, options) + np_data = np.empty((1, 1), dtype=object) + np_data[0, 0] = np.array(electron_counts, copy=False) + kwargs = { + "data": np_data, + "scan_shape": (1, 1), + "frame_shape": frame.shape, + } + return SparseArray(**kwargs) + + def radial_sum(reader, center=(-1, -1), scan_dimensions=(0, 0)): """Generate a radial sum from which STEM images can be generated. diff --git a/stempy/electron.cpp b/stempy/electron.cpp index f9e73df7..ee8c0315 100644 --- a/stempy/electron.cpp +++ b/stempy/electron.cpp @@ -560,6 +560,24 @@ std::vector electronCount( return maximalPoints(frame, frameDimensions); } +template +std::vector electronCount(std::vector& frame, + Dimensions2D frameDimensions, + const ElectronCountOptionsClassic& options) +{ + auto* darkReference = options.darkReference; + auto backgroundThreshold = options.backgroundThreshold; + auto xRayThreshold = options.xRayThreshold; + auto* gain = options.gain; + auto applyRowDarkSubtraction = options.applyRowDarkSubtraction; + auto applyRowDarkUseMean = options.applyRowDarkUseMean; + auto optimizedMean = options.optimizedMean; + + return electronCount( + frame, frameDimensions, darkReference, backgroundThreshold, xRayThreshold, + gain, applyRowDarkSubtraction, optimizedMean, applyRowDarkUseMean); +} + template ElectronCountedData electronCount(Reader* reader, const ElectronCountOptions& options) @@ -874,4 +892,12 @@ template ElectronCountedData electronCount( SectorStreamMultiPassThreadedReader* reader, const ElectronCountOptions& options); +template std::vector electronCount( + std::vector& frame, Dimensions2D frameDimensions, + const ElectronCountOptionsClassic& options); + +template std::vector electronCount( + std::vector& frame, Dimensions2D frameDimensions, + const ElectronCountOptionsClassic& options); + } // end namespace stempy diff --git a/stempy/electron.h b/stempy/electron.h index 4b9fd5f3..3b100e14 100644 --- a/stempy/electron.h +++ b/stempy/electron.h @@ -68,6 +68,11 @@ template ElectronCountedData electronCount(InputIt first, InputIt last, const ElectronCountOptionsClassic& options); +template +std::vector electronCount(std::vector& frame, + Dimensions2D frameDimensions, + const ElectronCountOptionsClassic& options); + template ElectronCountedData electronCount(Reader* reader, const ElectronCountOptions& options); diff --git a/tests/test_image.py b/tests/test_image.py index 4c18a907..491a9ab3 100644 --- a/tests/test_image.py +++ b/tests/test_image.py @@ -2,7 +2,7 @@ import numpy as np -from stempy.image import com_dense, com_sparse, radial_sum_sparse +from stempy.image import com_dense, com_sparse, electron_count_frame, radial_sum_sparse from stempy.io.sparse_array import SparseArray @@ -61,3 +61,70 @@ def test_com_sparse_parameters(simulate_sparse_array): # No counts will be in the center so all positions will be np.nan com2 = com_sparse(sp, crop_to=(10,10), init_center=(1,1)) assert np.isnan(com2[0,0,0]) + + +def test_electron_count_frame(): + # Create a synthetic 2D numpy array (frame) + frame = np.array( + [ + [2000, 0, 1000, 0, 0], + [0, 0, 0, 200, 0], + [0, 0, 1000, 0, 0], + [0, 200, 0, 200, 0], + [0, 0, 1000, 0, 0], + ], + dtype=np.uint16, + ) + + dark = np.ones_like(frame) * 100 + + # Define expected electron hits (coordinates) + expected_hits = np.array( + [ + [ + [ + [1, 0, 1, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 1, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 1, 0, 0], + ] + ] + ] + ) + + # Electron + electron_hits = electron_count_frame( + frame, xray_threshold=10000, background_threshold=1, darkreference=dark + ) + assert np.array_equal( + electron_hits.to_dense(), expected_hits + ), f"Expected {expected_hits}, but got {electron_hits}" + + # Test with no dark reference + electron_hits = electron_count_frame( + frame, xray_threshold=10000, background_threshold=1 + ) + + # Test where dark reference removes some points + dark = np.ones_like(frame) * 1000 + electron_hits = electron_count_frame( + frame, xray_threshold=10000, background_threshold=1, darkreference=dark + ) + expected_hits = np.array( + [ + [ + [ + [1, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + ] + ] + ] + ) + + assert np.array_equal( + electron_hits.to_dense(), expected_hits + ), f"Expected {expected_hits}, but got {electron_hits}"