diff --git a/src/cditools/analysis_scripts/cosmics.py b/src/cditools/analysis_scripts/cosmics.py new file mode 100644 index 0000000..d27fb0b --- /dev/null +++ b/src/cditools/analysis_scripts/cosmics.py @@ -0,0 +1,153 @@ +""" +Cosmics for a counting detector will trigger values of 1 or 2, +which are hard to see in default color settings, so we mask out +high and zero values so that cosmics stand out more. + +Data is accessed like data[image][row][column], +rather than data[image][x][y] +""" + +from __future__ import annotations + +from typing import Any + +import h5py +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np +from numpy.typing import NDArray +from scipy.stats import poisson +from skimage.measure import label + + +def setup( + det_name: str, date_dir: str, file_name: str, proposal_id: str +) -> NDArray[np.floating[Any]]: + """Sets up the data to use later for analysis. + + Args: + det_name (str): name of directory for the detector, e.g. "merlines-1" + date_dir (str): data in format "YYYY/MM/DD" + file_name (str): file name without the .h5 extension, e.g. "scan_0000" + proposal_id (str): current PASS proposal ID + + Returns: + data: dataset to be analyzed, in the form of a 3D array with dimensions (num_images, rows, cols) + """ + file_path = f"/nsls2/data/cdi/proposals/commissioning/pass-{proposal_id}/assets/{det_name}/{date_dir}/{file_name}.h5" + with h5py.File(file_path, "r") as f: + return np.asarray(f["entry"]["data"]["data"]) + + +def movie(data: NDArray[np.floating[Any]], vmin: int = 1, vmax: int = 5) -> None: + """Creates a movie of the images in the dataset, with a colormap that highlights values between vmin and vmax. + + Args: + data (NDArray[np.floating[Any]]): data set with entries in the form of (num_images, rows, cols) + vmin (int, optional): minimum value for cosmics. Defaults to 1. + vmax (int, optional): maximum value for cosmics. Defaults to 5. + """ + cmap = mpl.colormaps["gray"].with_extremes(under="aliceblue", over="r") + norm = mpl.colors.Normalize(vmin, vmax) + for n in range(data.shape[0]): + plt.imshow(data[n], cmap=cmap, norm=norm) + plt.title(f"Image {n}") + plt.draw() + plt.pause(0.1) + + +def label_cosmics(data: NDArray[np.floating[Any]]) -> tuple[list[int], float]: + """Labels the cosmics in each image and counts the number of cosmics per image, + then calculates the average number of cosmics across all images. + + Args: + data (NDArray[np.floating[Any]]): data set with entries in the form of (num_images, rows, cols) + + Returns: + tuple[list[int], float]: A tuple containing a list of counts of cosmics per image and the average number of cosmics across all images. + """ + counts = [] + cosmic_sum = 0 + for n in range(data.shape[0]): + image = data[n].copy() + # find indices of top 3 pixels + flat_idx = np.argpartition(image, -3, axis=None)[-3:] + coords = np.unravel_index(flat_idx, image.shape) + # set them to 0 + image[coords] = 0 + _, num = label(image, return_num=True) + counts.append(num) + cosmic_sum += num + + average = cosmic_sum / data.shape[0] + return counts, average + + +def plot_counts(counts: NDArray[np.floating[Any]]) -> None: + """Plot the histogram of the cosmic counts and fit a poisson distribution + + Args: + counts (NDArray[np.floating[Any]]): array of counts of cosmics per image + """ + counts = np.array(counts) + k_vals = np.arange(counts.max() + 1) + fit_vals = poisson.pmf(k_vals, mu=np.mean(counts)) + + plt.hist(counts, bins=np.arange(counts.max() + 2) - 0.5, density=True, alpha=0.6) + plt.plot(k_vals, fit_vals, "o-", label="Poisson fit") + plt.xlabel("Counts per image") + plt.ylabel("Probability") + plt.title("Histogram of counts") + plt.legend() + plt.show() + + +def check_empty(data: NDArray[np.floating[Any]]) -> list[int]: + """Checks for images that have values of all zero + + Args: + data (NDArray[np.floating[Any]]): data set with entries in the form of (num_images, rows, cols) + Returns: + list[int]: list of indices of images that have all zero values + """ + zero_images = [] + for n in range(data.shape[0]): + if not np.any(data[n]): + zero_images.append(n) + return zero_images + + +def find_cosmics( + data: NDArray[np.floating[Any]], vmin: int = 1, vmax: int = 5 +) -> set[int]: + """Masks images to find the cosmics in + + Args: + data (NDArray[np.floating[Any]]): data set with entries in the form of (num_images, rows, cols) + vmin (int, optional): minimum value for cosmics. Defaults to 1. + vmax (int, optional): maximum value for cosmics. Defaults to 5. + + Returns: + set[int]: set of points that are identified as cosmics, based on the values in the data set being between vmin and vmax + """ + cosmic_mask = (data >= vmin) & (data <= vmax) + return set(np.flatnonzero(np.any(cosmic_mask, axis=(1, 2)))) + + +def show_image( + data: NDArray[np.floating[Any]], image_index: int, vmin: int = 1, vmax: int = 5 +): + """Display the desired image by index + + Args: + data (NDArray[np.floating[Any]]): data set with entries in the form of (num_images, rows, cols) + image_index (int): index of the image to display + vmin (int, optional): minimum value for cosmics. Defaults to 1. + vmax (int, optional): maximum value for cosmics. Defaults to 5. + """ + cmap = mpl.colormaps["gray"].with_extremes(under="aliceblue", over="r") + norm = mpl.colors.Normalize(vmin, vmax) + plt.imshow(data[image_index], cmap=cmap, norm=norm) + plt.title(f"Image {image_index}") + plt.colorbar() + plt.show() diff --git a/tests/test_cosmics.py b/tests/test_cosmics.py new file mode 100644 index 0000000..0e9d346 --- /dev/null +++ b/tests/test_cosmics.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +import numpy as np + +from cditools.analysis_scripts.cosmics import check_empty, find_cosmics, label_cosmics + + +def test_check_empty(): + data = np.array([[[0, 0], [0, 0]], [[1, 0], [0, 0]], [[0, 0], [0, 0]]]) + empty_indices = check_empty(data) + assert empty_indices == [0, 2] + + +def test_find_cosmics(): + data = np.array([[[0, 0], [0, 0]], [[1, 0], [0, 0]], [[0, 0], [5, 0]]]) + cosmic_points = find_cosmics(data, vmin=1, vmax=5) + assert cosmic_points == {1, 2} + + +def test_label_cosmics(): + data = np.array([[[0, 0], [0, 0]], [[1, 0], [0, 0]], [[0, 0], [5, 0]]]) + labeled_data = label_cosmics(data) + expected_labeled_data = ([0, 0, 0], 0.0) + assert labeled_data == expected_labeled_data