diff --git a/.github/workflows/container.yml b/.github/workflows/container.yml new file mode 100644 index 000000000..177a435fc --- /dev/null +++ b/.github/workflows/container.yml @@ -0,0 +1,53 @@ +name: Build and Push Container Image + +on: + push: + tags: + - v* + +jobs: + build-and-push-image: + runs-on: ubuntu-latest + permissions: + contents: read + packages: write + strategy: + matrix: + platform: [full, cupy] + mpi: [openmpi, mpich] + steps: + - name: Checkout Code + uses: actions/checkout@v4 + + - name: Log in to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Extract Version from Tag + id: tags + run: echo version=$(echo "${{ github.ref_name }}" | cut -c 2-) >> $GITHUB_OUTPUT + + - name: Docker Metadata + id: meta + uses: docker/metadata-action@v5 + with: + images: ghcr.io/${{ github.repository }}_${{ matrix.platform }}_${{ matrix.mpi }} + tags: | + type=raw,value=${{ steps.tags.outputs.version}} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Build and Push Image + uses: docker/build-push-action@v6 + with: + push: true + context: . + tags: ${{ steps.meta.outputs.tags }} + target: runtime + build-args: | + PLATFORM=${{ matrix.platform}} + MPI=${{ matrix.mpi }} diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 1567443c1..50a890caa 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -51,7 +51,7 @@ jobs: miniforge-version: latest - name: Install ${{ matrix.conda-env }} dependencies run: | - sed -i 's/python/python=${{ matrix.python-version }}/' dependencies_${{ matrix.conda-env }}.yml + sed -i '/- python/d' dependencies_${{ matrix.conda-env }}.yml conda install -c conda-forge mpich conda env update --file dependencies_${{ matrix.conda-env }}.yml --name ptypy_env conda install -c conda-forge flake8 pytest pytest-cov diff --git a/Dockerfile b/Dockerfile index 880103111..dc75afe02 100644 --- a/Dockerfile +++ b/Dockerfile @@ -5,69 +5,92 @@ ARG MPI=openmpi ARG PLATFORM=cupy # Select CUDA version -ARG CUDAVERSION=12.4 +ARG CUDAVERSION=12.8 # Pull from mambaforge and install XML and ssh -FROM condaforge/mambaforge as base +FROM condaforge/mambaforge AS base ENV DEBIAN_FRONTEND=noninteractive RUN apt-get update && apt-get install -y libxml2 ssh # Pull from base image and install OpenMPI/MPICH -FROM base as mpi +FROM base AS mpi ARG MPI RUN mamba install -n base -c conda-forge ${MPI} -# Pull from MPI build install core dependencies -FROM base as core -COPY ./dependencies_core.yml ./dependencies.yml -RUN mamba env update -n base -f dependencies.yml +# Pull from MPI build and install core dependencies +FROM base AS core +RUN mamba install -n base -y -c conda-forge \ + python numpy scipy h5py pip # Pull from MPI build and install full dependencies -FROM mpi as full -COPY ./dependencies_full.yml ./dependencies.yml -RUN mamba env update -n base -f dependencies.yml +FROM mpi AS full +ARG MPI +RUN mamba install -n base -y -c conda-forge \ + python numpy scipy matplotlib h5py \ + pyzmq mpi4py[build=*${MPI}*] packaging \ + pillow pyfftw pyyaml pip # Pull from MPI build and install accelerate/pycuda dependencies -FROM mpi as pycuda -ARG CUDAVERSION -COPY ./ptypy/accelerate/cuda_pycuda/dependencies.yml ./dependencies.yml -COPY ./cufft/dependencies.yml ./dependencies_cufft.yml -RUN mamba install cuda-version=${CUDAVERSION} && \ - mamba env update -n base -f dependencies.yml && \ - mamba env update -n base -f dependencies_cufft.yml +FROM mpi AS pycuda +ARG CUDAVERSION MPI +RUN mamba install -n base -y -c conda-forge -c nvidia \ + python numpy scipy matplotlib h5py pyzmq mpi4py[build=*${MPI}*] \ + pillow pyfftw pyyaml compilers pip \ + reikna pycuda cuda-nvcc cuda-cudart-dev cuda-version=${CUDAVERSION} # Pull from MPI build and install accelerate/cupy dependencies -FROM mpi as cupy -ARG CUDAVERSION -COPY ./ptypy/accelerate/cuda_cupy/dependencies.yml ./dependencies.yml -COPY ./cufft/dependencies.yml ./dependencies_cufft.yml -RUN mamba install cuda-version=${CUDAVERSION} && \ - mamba env update -n base -f dependencies.yml && \ - mamba env update -n base -f dependencies_cufft.yml +FROM mpi AS cupy +ARG CUDAVERSION MPI +RUN mamba install -n base -y -c conda-forge \ + python numpy scipy matplotlib h5py pyzmq mpi4py[build=*${MPI}*] \ + pillow pyfftw pyyaml compilers pip \ + cupy cuda-version=${CUDAVERSION} +RUN mamba clean -y -a # Pull from platform specific image and install ptypy -FROM ${PLATFORM} as build +FROM ${PLATFORM} AS build COPY pyproject.toml ./ COPY ./templates ./templates COPY ./benchmark ./benchmark -COPY ./cufft ./cufft COPY ./ptypy ./ptypy RUN pip install . -# For core/full build, no post processing needed -FROM build as core-post -FROM build as full-post +# For core build, clean up conda env +FROM build AS core-post +RUN mamba clean -y -a + +# For full build, clean up conda env +FROM build AS full-post +RUN mamba clean -y -a # For pycuda build, install filtered cufft -FROM build as pycuda-post +FROM build AS pycuda-post +ARG CUDAVERSION +RUN mamba install -n base -y -c conda-forge -c nvidia \ + python cmake>=3.8.0 pybind11 compilers \ + cuda-nvcc cuda-cudart-dev libcufft-dev libcufft-static cuda-version=${CUDAVERSION} +COPY ./cufft ./cufft RUN pip install ./cufft +RUN mamba remove -n base -y \ + cmake pybind11 cuda-nvcc cuda-cudart-dev libcufft-dev libcufft-static +RUN mamba clean -y -a -# For pycuda build, install filtered cufft -FROM build as cupy-post +# For cupy build, install filtered cufft +FROM build AS cupy-post +ARG CUDAVERSION +RUN mamba install -n base -y -c conda-forge -c nvidia \ + python cmake>=3.8.0 pybind11 compilers \ + cuda-nvcc cuda-cudart-dev libcufft-dev libcufft-static cuda-version=${CUDAVERSION} +COPY ./cufft ./cufft RUN pip install ./cufft +RUN mamba remove -n base -y \ + cmake pybind11 libcufft-dev libcufft-static +RUN mamba clean -y -a # Platform specific runtime container -FROM ${PLATFORM}-post as runtime +FROM ${PLATFORM}-post AS runtime +RUN useradd --user-group ptypy-user +USER ptypy-user # Run PtyPy run script as entrypoint ENTRYPOINT ["ptypy.cli"] diff --git a/cufft/extensions.py b/cufft/extensions.py index 468d8d1cb..bf41f76cc 100644 --- a/cufft/extensions.py +++ b/cufft/extensions.py @@ -147,7 +147,6 @@ def build_extension(self, ext): if has_cu: old_compiler = self.compiler self.compiler = NvccCompiler(verbose=old_compiler.verbose, - dry_run=old_compiler.dry_run, force=old_compiler.force) # this is our bespoke compiler super(CustomBuildExt, self).build_extension(ext) self.compiler=old_compiler diff --git a/ptypy/accelerate/cuda_cupy/engines/ML_cupy.py b/ptypy/accelerate/cuda_cupy/engines/ML_cupy.py index 5f35fbd67..f386e170b 100644 --- a/ptypy/accelerate/cuda_cupy/engines/ML_cupy.py +++ b/ptypy/accelerate/cuda_cupy/engines/ML_cupy.py @@ -14,6 +14,7 @@ import numpy as np import cupy as cp import cupyx +import os from ptypy.engines import register from ptypy.accelerate.base.engines.ML_serial import ML_serial, BaseModelSerial @@ -35,6 +36,11 @@ MAX_BLOCKS = 99999 # MAX_BLOCKS = 3 # can be used to limit the number of blocks, simulating that they don't fit +# Estimate the device memory safety margin as fraction of total device memory +device_memory_fractional_safety_margin = 0.025 +if "PTYPY_DEVICE_MEM_SAFETY" in os.environ: + device_memory_fractional_safety_margin = float(os.environ["PTYPY_DEVICE_MEM_SAFETY"]) + @register() class ML_cupy(ML_serial): @@ -160,9 +166,11 @@ def _setup_kernels(self): # as both will be used for allocations mempool = cp.get_default_memory_pool() mem = cp.cuda.runtime.memGetInfo()[0] + mempool.total_bytes() - mempool.used_bytes() + tot = cp.cuda.runtime.memGetInfo()[1] + safety_margin = max(device_memory_fractional_safety_margin * tot, 400 * 1024 * 1024) - # leave 200MB room for safety - fit = int(mem - 200 * 1024 * 1024) // blk + # leave room for safety + fit = int(mem - safety_margin) // blk if not fit: log(1, "Cannot fit memory into device, if possible reduce frames per block. Exiting...") raise SystemExit("ptypy has been exited.") diff --git a/ptypy/accelerate/cuda_cupy/engines/projectional_cupy_stream.py b/ptypy/accelerate/cuda_cupy/engines/projectional_cupy_stream.py index 56d65d6de..b4a5b123a 100644 --- a/ptypy/accelerate/cuda_cupy/engines/projectional_cupy_stream.py +++ b/ptypy/accelerate/cuda_cupy/engines/projectional_cupy_stream.py @@ -13,6 +13,7 @@ :license: see LICENSE for details. """ +import os import numpy as np import cupy as cp import cupyx @@ -32,6 +33,11 @@ MAX_BLOCKS = 99999 # MAX_BLOCKS = 3 # can be used to limit the number of blocks, simulating that they don't fit +# Estimate the device memory safety margin as fraction of total device memory +device_memory_fractional_safety_margin = 0.025 +if "PTYPY_DEVICE_MEM_SAFETY" in os.environ: + device_memory_fractional_safety_margin = float(os.environ["PTYPY_DEVICE_MEM_SAFETY"]) + __all__ = ['DM_cupy_stream', 'RAAR_cupy_stream'] @@ -65,9 +71,11 @@ def _setup_kernels(self): # as both will be used for allocations mempool = cp.get_default_memory_pool() mem = cp.cuda.runtime.memGetInfo()[0] + mempool.total_bytes() - mempool.used_bytes() - - # leave 200MB room for safety - fit = int(mem - 200 * 1024 * 1024) // blk + tot = cp.cuda.runtime.memGetInfo()[1] + safety_margin = max(device_memory_fractional_safety_margin * tot, 400 * 1024 * 1024) + + # leave room for safety + fit = int(mem - safety_margin) // blk if not fit: log(1, "Cannot fit memory into device, if possible reduce frames per block. Exiting...") raise SystemExit("ptypy has been exited.") diff --git a/ptypy/cli/command_line_interface.py b/ptypy/cli/command_line_interface.py index 283223a36..2631a67a9 100644 --- a/ptypy/cli/command_line_interface.py +++ b/ptypy/cli/command_line_interface.py @@ -31,6 +31,8 @@ def parse(): help="Provide parameter configuration as a JSON string.") parser.add_argument('--output-folder', '-o', type=str, help="The path we want the outputs to exist in (will get created).") + parser.add_argument('--use-timestamp', action="store_true", + help="Turn on timestamping of output folder") parser.add_argument('--ptypy-level', '-l', default=5, type=str, help="The level we want to run to ptypy to.") parser.add_argument('--identifier', '-i', type=str, default=None, @@ -57,9 +59,9 @@ def run(args): # Load parameter tree from file or JSON string if args.file: - p = create_parameter_tree(load_config_as_dict_from_file(args.file)) + p, extra = create_parameter_tree(load_config_as_dict_from_file(args.file)) if args.json: - p = create_parameter_tree(json.loads(args.json)) + p, extra = create_parameter_tree(json.loads(args.json)) p.run = args.identifier # TODO @@ -74,7 +76,7 @@ def run(args): # if args.output_folder is not None: p.io.home = get_output_file_name(args) - p.io.rfile = "%s.ptyr" % get_output_file_name(args) + p.io.rfile = "%(run)s_%(engine)s_%(iterations)04d.ptyr" #parameters.io.autosave = u.Param(active=True) #log(3, "Autosave is on, with io going in {}, and the final reconstruction into {}".format(parameters.io.home, # parameters.io.rfile)) @@ -83,8 +85,8 @@ def run(args): p.io.autosave = u.Param(active=False) log("info", "Autosave is off. No output will be saved.") - # Substitute %(run) with in ptyscan - substitute_id_in_ptyscan(p) + # Substitute %(run) with in ptyscan + substitute_id_in_ptyscan(p, extra) # Run PtyPy to given level P = Ptycho(p, level=args.ptypy_level) @@ -110,9 +112,16 @@ def create_parameter_tree(params) -> u.Param: parameters_to_run.update(previous_parameters) if params['parameter_tree'] is not None: parameters_to_run.update(params['parameter_tree'], Convert=True) - return parameters_to_run -def substitute_id_in_ptyscan(params): + # Additional non-ptypy params + extra = {} + for k,v in params.items(): + if k != "parameter_tree": + extra[k] = v + + return parameters_to_run, extra + +def substitute_id_in_ptyscan(params, extra): def _substitute(d, p): for k, v in d.items(): if isinstance(v, MutableMapping): @@ -122,15 +131,16 @@ def _substitute(d, p): d[k] = v % p for scan_key, scan in params.scans.items(): data_entry = scan.data - _substitute(data_entry, params) + _substitute(data_entry, params | extra) def get_output_file_name(args): from datetime import datetime now = datetime.now() + output_path = "{}/scan".format(args.output_folder) if args.identifier is not None: - output_path = "{}/scan_{}_{}".format(args.output_folder, args.identifier, now.strftime("%Y%m%d%H%M%S")) - else: - output_path = "{}/scan_{}".format(args.output_folder, now.strftime("%Y%m%d%H%M%S")) + output_path += "_{}".format(args.identifier) + if args.use_timestamp: + output_path += "_{}".format(now.strftime("%Y%m%d%_H%M%S")) log("info", "Output is going in: {}".format(output_path)) return output_path diff --git a/ptypy/core/classes.py b/ptypy/core/classes.py index 622077a48..a755763b3 100644 --- a/ptypy/core/classes.py +++ b/ptypy/core/classes.py @@ -33,7 +33,7 @@ """ import numpy as np -import weakref +import weakref, os from collections import OrderedDict try: @@ -90,7 +90,7 @@ class Base(object): _PREFIX = BASE_PREFIX __slots__ = ['ID','numID','owner','_pool','_recs','_record'] - _fields = [('ID','= l: nl = l + 8192 if idx > 10000 else 2*l - recs = np.resize(recs,(nl,)) - self._recs[prefix] = recs - rec = recs[idx] - obj._record = rec - rec['ID'] = nID + self._recs[prefix].resize((nl,), refcheck=False) + # after .resize() all previous references to the records + # array are not guaranteed to point to correct memory + # therefore need to update all previous pointers + for v in self._pool[prefix].values(): + v._record = self._recs[prefix][v.numID] + obj._record = self._recs[prefix][idx] + self._recs[prefix][idx]['ID'] = nID return @@ -1141,15 +1143,15 @@ class View(Base): """ _fields = Base._fields + \ [('active', 'b1'), - ('dlayer', '=threshhold] = (I_low-dark_low)[I_high>=threshhold] + return frame + + def load_and_average_dark(fnames): + dark_frames = np.array([np.array(Image.open(self.info.folder_dark+x)) for x in fnames_dark_even]) + return np.mean(dark_frames, axis=0) + + raw, weights, positions = {}, {}, {} + + fnames_dark = np.array(os.listdir(self.info.folder_dark)) + fnames_dark_even = fnames_dark[0::2] + fnames_dark_odd = fnames_dark[1::2] + dark_even = load_and_average_dark(fnames_dark_even) + dark_odd = load_and_average_dark(fnames_dark_odd) + + fnames_diff = sorted(os.listdir(self.info.folder_diff), key=find_index) + fnames_diff_even = fnames_diff[0::2] + fnames_diff_odd = fnames_diff[1::2] + + for ind in indices: + frame_even = np.array(Image.open(self.info.folder_diff+fnames_diff_even[ind])) + frame_odd = np.array(Image.open(self.info.folder_diff+fnames_diff_odd[ind])) + + if self.info.high_first >=1: + raw[ind] = make_hdr_frame(I_low = frame_odd, + I_high = frame_even, + dark_low = dark_odd, + dark_high = dark_even, + threshhold = self.info.high_threshhold, + factor = self.info.factor_low_to_high) + else: + raw[ind] = make_hdr_frame(I_low = frame_even, + I_high = frame_odd, + dark_low = dark_even, + dark_high = dark_odd, + threshhold = self.info.high_threshhold, + factor = self.info.factor_low_to_high) + raw[ind][raw[ind]<0] = 0 + + return raw, positions, weights + + def load_weight(self): + """ + Provides the mask used for every diffraction pattern in the whole scan + This mask will have the shape of the first frame. + """ + + r, w, p = self.load(indices=(0,)) + data = r[0] + mask = np.ones_like(data) + + return mask + diff --git a/ptypy/experiment/hdf5_loader.py b/ptypy/experiment/hdf5_loader.py index 72846e734..21f2c75b0 100644 --- a/ptypy/experiment/hdf5_loader.py +++ b/ptypy/experiment/hdf5_loader.py @@ -320,6 +320,12 @@ class Hdf5Loader(PtyScan): type = list, ndarray help = This is the array or list with the re-ordered indices. + [nearfield_defocus] + default = None + type = float + help = Distance from sample to focus (for nearfield only) + doc = If set, magnification will be calculated automatically and applied to detector distance and pixelsize + """ def __init__(self, pars=None, swmr=False, **kwargs): @@ -378,6 +384,10 @@ def __init__(self, pars=None, swmr=False, **kwargs): if self.p.electron_data: self.meta.energy = u.m2keV(u.electron_wavelength(self.meta.energy)) + # For nearfield data, manipulate distance and psize + if self.p.nearfield_defocus: + self._prepare_nearfield() + # it's much better to have this logic here than in load! if (self._ismapped and (self._scantype == 'arb')): log(3, "This scan looks to be a mapped arbitrary trajectory scan.") @@ -590,6 +600,20 @@ def _prepare_meta_info(self): assert self.pad.size == 4, "self.p.padding needs to of size 4" log(3, "Padding the detector frames by {}".format(self.p.padding)) + def _prepare_nearfield(self): + """ + Calculate magnification and modify distance and psize + """ + defocus = self.p.nearfield_defocus + mag = self.meta.distance / defocus + dist_eff = (self.meta.distance - defocus) / mag + psize_eff = self.info.psize / mag + log(3, f"Nearfield: With defocus {defocus} m the magmification is {mag}") + log(3, f"Nearfield: The effective detector distance is {dist_eff}") + log(3, f"Nearfield: The effective pixel size is {psize_eff}") + self.meta.distance = dist_eff + self.info.psize = psize_eff + def _prepare_center(self): """ define how data should be loaded (center, cropping) diff --git a/ptypy/utils/__init__.py b/ptypy/utils/__init__.py index 37df5cf85..f2aaa2585 100644 --- a/ptypy/utils/__init__.py +++ b/ptypy/utils/__init__.py @@ -14,6 +14,7 @@ from .parameters import * from .verbose import * from .citations import * +from .metrics import * from . import descriptor from . import parallel from .. import __has_matplotlib__ as hmpl diff --git a/ptypy/utils/metrics.py b/ptypy/utils/metrics.py new file mode 100644 index 000000000..34215da8b --- /dev/null +++ b/ptypy/utils/metrics.py @@ -0,0 +1,348 @@ +# -*- coding: utf-8 -*- +""" +Metric utility functions: + +* radial_power_spectrum: returns the power spectrum of an image averaged azimuthally. +* frc: Fourier ring correlation between two images +* fsc: Foutier shell correlation between two volumes +* frc_thresdhold: Compute the threshold curve for fsc or frc +* compute_intersection: utility to find intersection of two curves. + +Typical usage for FSC between vol1 and vol2. Voxel size is vx + +``` +freq, fsc12, npts = fsc(vol1, vol2) +fsc_curve = frc_threshold(npts) +freq_intersect = compute_intersection(freq, fsc12, fsc_curve) +print(f"FSC resolution: {vx/freq_intersect}") +``` + +This file is part of the PTYPY package. + + :copyright: Copyright 2014 by the PTYPY team, see AUTHORS. + :license: see LICENSE for details. +""" +import numpy as np + +__all__ = ['radial_power_spectrum', 'frc', 'fsc', 'frc_threshold', 'compute_intersection'] + + +def fourier_ring_sum(f_input, ringthick=1, pixel_size=1.0): + """ + Ring (or shell) sum of the Fourier transform of an image or volume + + Parameters + ---------- + f_input : array-like + array containing the Fourier transform of the image or volume. + Note that the (0,0) frequency must be the first element of the array. + (not fftshifted) + + ringthick : int + thickness of the ring (or shell) in pixel units for the azimuthal averaging + + pixel_size : float + size of the pixels in the image or volume (default is 1.0) + + Returns + ------- + freq: array-like + 1D array containing the spatial frequencies + FRP : array-like + 1D array containing the sum in each ring (or shell) + npts : array-like + 1D array containing the number of points in each ring (or shell) + """ + sh = f_input.shape + ndim = f_input.ndim + size = f_input.size + + # Spatial frequencies + ulist = [ + np.fft.fftfreq(sh[i]).reshape([-1 if j==i else 1 for j in range(ndim)]) for i in range(ndim) + ] + + # Radial spatial frequency norm (flatten again) + unorm = np.sqrt(sum([u**2 for u in ulist])).reshape(size) + + # Ring/shell width in frequency space (largest frequency step is 1/min(sh)) + uw = ringthick/min(sh) + + # Spatial frequencies (center of rings/shells) + # and bins. A bit complicated to ensure that + # the first ring is centered at 0 + urange = np.arange(0, unorm.max() + 1.5*uw - 1e-12, uw) + frequencies = urange[:-1] + shell_bins = (urange - .5*uw).clip(0) + + # Number of bins + nbins = len(shell_bins) - 1 + + # create the rings/shells + shells = np.digitize(unorm, shell_bins) + + # Count number of point in each ring/shell + npts = np.bincount(shells)[1:] + + # Sum in each ring/shell + # Bincount does not work with complex numbers + f_flat = f_input.reshape(size) + if np.iscomplexobj(f_input): + FRP = np.bincount(shells, weights=f_flat.real)[1:] + 1j*np.bincount(shells, weights=f_flat.imag)[1:] + else: + FRP = np.bincount(shells, weights=f_flat)[1:] + + return frequencies, FRP, npts + + +def radial_power_spectrum(input, ringthick=1, pixel_size=1.0): + """ + Radial power spectrum of an image or volume + + Parameters + ---------- + input : array-like + array containing the image or volume + + ringthick : int + thickness of the ring (or shell) in pixel units for the azimuthal averaging + + pixel_size : float + size of the pixels in the image or volume (default is 1.0) + + Returns + ------- + freq: array-like + 1D array containing the spatial frequencies + RPS : array-like + 1D array containing the radial power spectrum + """ + # Fourier transform of the input + f_input = np.fft.fftn(input) + # Power spectrum + PS = np.abs(f_input)**2 + # Ring/shell sum + freq, RPS, npts = fourier_ring_sum(f_input=PS, ringthick=ringthick, pixel_size=pixel_size) + # Normalization by the number of points in each ring/shell + RPS /= npts + + return freq, RPS + + +def _fouriercorrelation(input1, input2, ringthick=1): + """ + Auxiliary function for the calculation of either FRC or FSC + """ + # Computation of the FFTs + F1 = np.fft.fftn(np.fft.ifftshift(input1)) # FFT of input1 + F2 = np.fft.fftn(np.fft.ifftshift(input2)) # FFT of input2 + + freq, C11, npts = fourier_ring_sum(np.abs(F1)**2, ringthick=ringthick) + _, C22, _ = fourier_ring_sum(np.abs(F2)**2, ringthick=ringthick) + _, C12, _ = fourier_ring_sum(F1 * np.conj(F2), ringthick=ringthick) + + FRC = np.abs(C12) / np.sqrt(C11 * C22) + + return freq, FRC, npts + + +def frc(input1, input2, ringthick=1, apod_width=0, align=False): + """ + Compute the Fourier Ring Correlation (FRC) between input1 and input2. + + Parameters + ---------- + input1 : array-like + array containing the first image, must be two-dimensional + + input2 : array-like + array containing the second image, must be two-dimensional + + ringthick : int + thickness of the ring for averaging the correlation + + apod_width: int + width of image apodization to limit cyclic discontinuities (default: 0 = no apodization) + NOT IMPLEMENTED + + align: bool + pre-align the images before comparison (default: False) + NOT IMPLEMENTED + + Returns + ------- + freq: array-like + 1D array containing the spatial frequencies + + FRC : array-like + 1D array containing the Fourier Ring Correlation values + + npts : array-like + 1D array containing the number of points in each ring + """ + + # Check if the arrays have 2 dimensions + if input1.ndim!=2 or input2.ndim!=2: + raise ValueError("The arrays must have 2 dimensions") + + # Check if the arrays have the same size + if input1.shape != input2.shape: + raise ValueError("The arrays must have the same size") + + # Image apodization + if apod_width > 0: + raise NotImplementedError('Apodization has not been implemented yet.') + + # Image alignment + if align: + raise NotImplementedError('Image pre-alignment has not been implemented yet.') + + # Run Fourier correlation + return _fouriercorrelation(input1, input2, ringthick) + + +def fsc(input1, input2, ringthick=1, apod_width=0, align=False): + """ + Compute the Fourier Shell Correlation (FSC) between input1 and input2. + + Parameters + ---------- + input1 : array-like + array containing the first image, must be three-dimensional + + input2 : array-like + array containing the second image, must be three-dimensional + + ringthick : int, optional + thickness of the ring for averaging the correlation + + apod_width: int + width of volume apodization to limit cyclic discontinuities (default: 0 = no apodization) + NOT IMPLEMENTED + + align: bool + pre-align the volumes before comparison (default: False) + NOT IMPLEMENTED + + Returns + ------- + freq: array-like + 1D array containing the spatial frequencies + + FRC : array-like + 1D array containing the FRC values + + npts : array-like + 1D array containing the number of points in each ring + """ + + # Check if the arrays have 3 dimensions + if input1.ndim!=3 or input2.ndim!=3: + raise ValueError("The arrays must have 3 dimensions") + + # Check if the arrays have the same size + if input1.shape != input2.shape: + raise ValueError("The arrays must have the same size") + + # Image apodization + if apod_width > 0: + raise NotImplementedError('Apodization has not been implemented yet.') + + # Image alignment + if align: + raise NotImplementedError('Volume pre-alignment has not been implemented yet.') + + # run Fourier correlation + return _fouriercorrelation(input1, input2, ringthick) + +def frc_threshold(npts, threshold = 'onebit'): + """ + Compute the FRC threshold curve + + Parameters + ---------- + npts : array-like + 1D array containing the number of points in each ring + + threshold : str, optional + The option `onebit` means 1 bit threshold with ``SNRt = 0.5``, which + should be used for two independent measurements. The option `halfbit` + means 1/2 bit threshold with ``SNRt = 0.2071``, which should be + use for split tomogram. The default option is ``onebit``. + + Returns + ------- + T : array-like + 1D array containing the FRC threshold values + + """ + # Setting the threshold + if threshold == 'halfbit' or threshold == 'half-bit': + print('Computing FRC using half-bit threshold') + snrt = 0.2071 + raise Warning('FRC should be used with a one-bit threshold. If you proceed with half-bit, you do it on your own responsibility.') + elif threshold == 'onebit' or threshold =='one-bit': + print('Computing FRC using 1-bit threshold') + snrt = 0.5 + else: + raise ValueError( + "You must choose between 'halfbit' or 'onebit' threshold" + ) + + Tnum = (snrt + (2*np.sqrt(snrt)/np.sqrt(npts)) + 1/np.sqrt(npts)) + Tden = (snrt + (2*np.sqrt(snrt)/np.sqrt(npts)) + 1) + + return Tnum / Tden + + +def compute_intersection(x, f1, f2, exclude_origin=True): + """ + Function used to extract the resolution given the FRC/FSC and the threshold curve + + Parameters + ---------- + f1 : ndarray + First curve (1-D set of values) + + f2 : ndarray + Second curve (1-D set of values) + + x : ndarray + x-axis (same size and shape as f1 and f2) + + exclude_origin : bool, optional + If True, the intersection point at the origin (if any) is ignored. + + Returns + ------- + + x : float + Associated with the intersection point + + """ + x = np.asarray(x) + f1 = np.asarray(f1) + f2 = np.asarray(f2) + + if x.ndim!=1 or f1.ndim!=1 or f2.ndim!=1: + raise ValueError('Parameters must be 1-D arrays') + + diff = f2 - f1 + + # Find indices where the sign of the difference changes + idx = np.where(np.diff(np.sign(diff)))[0] + + crossings = [] + for i in idx: + x1, x2 = x[i], x[i+1] + d1, d2 = diff[i], diff[i+1] + + # Linear interpolation to find the crossing point + x_cross = x1 - d1 * (x2 -x1) / (d2-d1) + crossings.append(x_cross) + + if exclude_origin and crossings[0]==x[0]: + crossings = crossings[1:] + + # Return the first crossing point, or None if there are no crossings + return crossings[0] if crossings else None diff --git a/ptypy/utils/plot_utils.py b/ptypy/utils/plot_utils.py index a4a81a596..fe569212a 100644 --- a/ptypy/utils/plot_utils.py +++ b/ptypy/utils/plot_utils.py @@ -772,7 +772,7 @@ def _update(self,renew_image=False): plt.setp(self.ax.get_yticklabels(), fontsize=self.fontsize) # determine number of points. v, h = self.shape - steps = [1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 1500, 2000, 3000] + steps = [1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 1500, 2000, 3000, 4000, 5000, 6000] Nindex = steps[max([v // s <= 4 for s in steps].index(True) - 1, 0)] self.ax.yaxis.set_major_locator(mpl.ticker.IndexLocator(Nindex, 0.5)) Nindex = steps[max([h // s <= 4 for s in steps].index(True) - 1, 0)] diff --git a/ptypy/utils/scripts.py b/ptypy/utils/scripts.py index f438bd226..1c6631725 100644 --- a/ptypy/utils/scripts.py +++ b/ptypy/utils/scripts.py @@ -621,7 +621,7 @@ def stxm_analysis(storage, probe=None): pod = list(v.pods.values())[0] if not pod.active: continue - t = pod.diff.sum() + t = (pod.diff * pod.mask).sum() if t > t2: t2=t ss = v.slice @@ -629,7 +629,7 @@ def stxm_analysis(storage, probe=None): # slice(v.roi[0,0], v.roi[1,0]), # slice(v.roi[0,1], v.roi[1,1])) # bufview = buf[ss] - m = mass_center(pod.diff) # + 1. + m = mass_center(pod.diff * pod.mask) # + 1. q = pod.di_view.storage._to_phys(m) dpc_row[ss] += q[0] * v.psize[0] * pr * 2 * np.pi / pod.geometry.lz dpc_col[ss] += q[1] * v.psize[1] * pr * 2 * np.pi / pod.geometry.lz @@ -723,7 +723,7 @@ def phase_from_dpc(dpc_row, dpc_col): """ py =- dpc_row - px =- dpc_col + px =+ dpc_col sh = px.shape sh = np.asarray(sh) fac = np.ones_like(sh) diff --git a/templates/misc/test_FRC.ipynb b/templates/misc/test_FRC.ipynb new file mode 100644 index 000000000..dc64a95b3 --- /dev/null +++ b/templates/misc/test_FRC.ipynb @@ -0,0 +1,240 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "2e2bc653-fb7f-4ea7-ab96-99200624e572", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from scipy.ndimage import fourier_shift, gaussian_filter\n", + "from skimage import data\n", + "\n", + "from ptypy.utils.metrics import frc, fsc, frc_threshold, compute_intersection" + ] + }, + { + "cell_type": "markdown", + "id": "4e673250-bd9c-473a-9ac5-8645cce47cf6", + "metadata": {}, + "source": [ + "### Test with 2D images of random values\n", + "Since the values are random, there must not be an important correlation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dd46805a-5c8c-4c2e-a732-3bba0840e7d8", + "metadata": {}, + "outputs": [], + "source": [ + "np.random.seed(1) \n", + "img1 = np.random.randn(301,301)\n", + "img2 = np.random.randn(301,301)\n", + "\n", + "X,FRC,N = frc(img1, img2, apod_width = 0, ringthick=2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "96151032-25ac-4804-aacd-fe199cb1ec97", + "metadata": {}, + "outputs": [], + "source": [ + "fig0 = plt.figure()\n", + "ax1 = fig0.add_subplot(121)\n", + "ax1.imshow(img1, cmap='bone')\n", + "ax1.set_axis_off()\n", + "ax2 = fig0.add_subplot(122)\n", + "ax2.imshow(img2, cmap='bone')\n", + "ax2.set_axis_off()\n", + "plt.tight_layout()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f068c940-2d75-4cd3-96bd-76e0b67c2fc6", + "metadata": {}, + "outputs": [], + "source": [ + "threshold = frc_threshold(N)\n", + "resolution = compute_intersection(X, FRC, threshold)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e766f128-10ac-4625-9660-540299cac35b", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure()\n", + "plt.plot(X, FRC, color=\"b\", label=\"FRC\")\n", + "plt.plot(X, threshold, color=\"r\", ls=\"--\", label=\"1 bit threshold\")\n", + "plt.xlim(0,0.5)\n", + "plt.legend()\n", + "if resolution is not None:\n", + " plt.axvline(resolution, color=\"k\", ls=\":\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "ede883b5-0140-416e-87f0-78a9f4879e3c", + "metadata": {}, + "source": [ + "### Test with 2D images: cameraman" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "febc679d-f2e6-45bd-8b2f-986d98d55a11", + "metadata": {}, + "outputs": [], + "source": [ + "testimg = data.camera()\n", + "img1 = testimg\n", + "img2 = gaussian_filter(testimg,sigma=2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a169a361-ca45-4a55-93a1-cfe83225547b", + "metadata": {}, + "outputs": [], + "source": [ + "fig1 = plt.figure(1)\n", + "ax1 = fig1.add_subplot(121)\n", + "ax1.imshow(img1, cmap='bone')\n", + "ax1.set_axis_off()\n", + "ax2 = fig1.add_subplot(122)\n", + "ax2.imshow(img2, cmap='bone')\n", + "ax2.set_axis_off()\n", + "plt.tight_layout()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3db935d2-8a5d-4c59-b86d-b11145e5464d", + "metadata": {}, + "outputs": [], + "source": [ + "X, FRC, N = frc(img1, img2, apod_width = 0, ringthick=2)\n", + "threshold = frc_threshold(N)\n", + "resolution = compute_intersection(X, FRC, threshold)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "059ce538-bed0-4b6b-bd77-b692d176768c", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure()\n", + "plt.plot(X, FRC, color=\"b\", label=\"FRC\")\n", + "plt.plot(X, threshold, color=\"r\", ls=\"--\", label=\"1 bit threshold\")\n", + "plt.xlim(0,0.5)\n", + "plt.legend()\n", + "if resolution is not None:\n", + " plt.axvline(resolution, color=\"k\", ls=\":\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "87105b1f-4a17-4e8f-b89d-a50f10662022", + "metadata": {}, + "source": [ + "### Another 2D image" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bf2e27ca-dee5-462d-af69-d60c223ea974", + "metadata": {}, + "outputs": [], + "source": [ + "testimg = data.chelsea().mean(axis=2)\n", + "img1 = testimg\n", + "img2 = testimg + 0.5*np.max(testimg)*np.random.rand(testimg.shape[0],testimg.shape[1])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f17b8dcf-3a99-447e-9265-460b6b3df99f", + "metadata": {}, + "outputs": [], + "source": [ + "fig2 = plt.figure(2)\n", + "ax1 = fig2.add_subplot(121)\n", + "ax1.imshow(img1, cmap='bone')\n", + "ax1.set_axis_off()\n", + "ax2 = fig2.add_subplot(122)\n", + "ax2.imshow(img2, cmap='bone')\n", + "ax2.set_axis_off()\n", + "plt.tight_layout()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "62f5857d-3a05-403b-9a7a-066b780cee02", + "metadata": {}, + "outputs": [], + "source": [ + "X, FRC, N = frc(img1, img2, apod_width = 0, ringthick=2)\n", + "threshold = frc_threshold(N)\n", + "resolution = compute_intersection(X, FRC, threshold)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d8f00e7d-f83c-4e55-b267-b9061faf44ab", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure()\n", + "plt.plot(X, FRC, color=\"b\", label=\"FRC\")\n", + "plt.plot(X, threshold, color=\"r\", ls=\"--\", label=\"1 bit threshold\")\n", + "plt.xlim(0,0.5)\n", + "plt.legend()\n", + "if resolution is not None:\n", + " plt.axvline(resolution, color=\"k\", ls=\":\")\n", + "plt.show()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/templates/misc/test_FSC.ipynb b/templates/misc/test_FSC.ipynb new file mode 100644 index 000000000..7aabeec10 --- /dev/null +++ b/templates/misc/test_FSC.ipynb @@ -0,0 +1,238 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "fd0b6ca2", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from ptypy.utils.metrics import fsc, frc_threshold, compute_intersection\n", + "from scipy.ndimage import gaussian_filter" + ] + }, + { + "cell_type": "markdown", + "id": "7b6d272e", + "metadata": {}, + "source": [ + "#### Generate the volumes" + ] + }, + { + "cell_type": "markdown", + "id": "ca003baa", + "metadata": {}, + "source": [ + "Here we generate:\n", + "- a first volume with random numbers\n", + "- a second volume by adding gaussian noise to the first" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "934fd932", + "metadata": {}, + "outputs": [], + "source": [ + "# Add gaussian noise to the image\n", + "def add_gaussian_noise(volume, mean=0.0, std=0.1):\n", + " noise = np.random.normal(mean, std, volume.shape)\n", + " return volume + noise\n", + "\n", + "original_vol = gaussian_filter(np.random.randn(110,120,100), sigma=3)\n", + "noisy_vol = add_gaussian_noise(original_vol, mean=0.0, std=0.03)" + ] + }, + { + "cell_type": "markdown", + "id": "53dcf675", + "metadata": {}, + "source": [ + "#### Visualise the volumes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "895f79e2", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_vol(noisy_vol, original, title=''):\n", + " \n", + " pshape = noisy_vol.shape[0]\n", + "\n", + " pos_limit = np.max(original) \n", + " neg_limit = np.min(original)\n", + " \n", + " fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(6,4), dpi=100)\n", + " for i in range(3):\n", + " for j in range(2):\n", + " ax = axes[j,i]\n", + " ax.set_xticks([])\n", + " ax.set_yticks([])\n", + " ax.set_xticklabels([])\n", + " ax.set_yticklabels([])\n", + " axes[0,0].set_title(\"slice(Z)\")\n", + " axes[0,1].set_title(\"slice(Y)\")\n", + " axes[0,2].set_title(\"slice(X)\")\n", + " axes[0,0].set_ylabel(\"Original vol\")\n", + " axes[0,0].imshow(original[pshape//2], vmin=neg_limit, vmax=pos_limit)\n", + " axes[0,1].imshow(original[:,pshape//2], vmin=neg_limit, vmax=pos_limit)\n", + " axes[0,2].imshow(original[:,:,pshape//2], vmin=neg_limit, vmax=pos_limit)\n", + " axes[1,0].set_ylabel(\"Noisy vol\")\n", + " axes[1,0].imshow(noisy_vol[pshape//2], vmin=neg_limit, vmax=pos_limit)\n", + " axes[1,1].imshow(noisy_vol[:,pshape//2], vmin=neg_limit, vmax=pos_limit)\n", + " im1 = axes[1,2].imshow(noisy_vol[:,:,pshape//2], vmin=neg_limit, vmax=pos_limit)\n", + " fig.suptitle('Volumes')\n", + " fig.colorbar(im1, ax=axes.ravel().tolist())\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4b7b22bc", + "metadata": {}, + "outputs": [], + "source": [ + "plot_vol(noisy_vol, original_vol)" + ] + }, + { + "cell_type": "markdown", + "id": "5e44da9b", + "metadata": {}, + "source": [ + "#### Compute the FSC and plot" + ] + }, + { + "cell_type": "markdown", + "id": "3b36c0ce", + "metadata": {}, + "source": [ + "The amount of noise added to generate volume 2 was relatively small, and this is reflected in the FSC curve." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "436d50c5", + "metadata": {}, + "outputs": [], + "source": [ + "X, FSC, N = fsc(\n", + " noisy_vol, \n", + " original_vol, \n", + " apod_width = 0, \n", + " ringthick=2, \n", + ")\n", + "threshold = frc_threshold(N)\n", + "resolution = compute_intersection(X, FSC, threshold)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cf1a07e7-ba8d-412d-855e-409baf51a9f2", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure()\n", + "plt.plot(X, FSC, color=\"b\", label=\"FSC\")\n", + "plt.plot(X, threshold, color=\"r\", ls=\"--\", label=\"1 bit threshold\")\n", + "plt.xlim(0,0.5)\n", + "plt.legend()\n", + "if resolution is not None:\n", + " plt.axvline(resolution, color=\"k\", ls=\":\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "5fa80415", + "metadata": {}, + "source": [ + "#### Increasing noise in the second volume" + ] + }, + { + "cell_type": "markdown", + "id": "5e18f35c", + "metadata": {}, + "source": [ + "Now we generate another volume, still by adding noise to the first volume, but a larger amount of noise. This change has a clear impact on the FSC curve." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6a3526a9", + "metadata": {}, + "outputs": [], + "source": [ + "noisy_vol2 = add_gaussian_noise(original_vol, mean=0.0, std=0.07)\n", + "plot_vol(noisy_vol2, original_vol)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d50c2b8c", + "metadata": {}, + "outputs": [], + "source": [ + "X, FSC, N = fsc(\n", + " noisy_vol2, \n", + " original_vol, \n", + " apod_width = 0, \n", + " ringthick=2, \n", + ")\n", + "threshold = frc_threshold(N)\n", + "resolution = compute_intersection(X, FSC, threshold)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7b43fbed-86d8-4b60-bdd5-eea4df67ce84", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure()\n", + "plt.plot(X, FSC, color=\"b\", label=\"FSC\")\n", + "plt.plot(X, threshold, color=\"r\", ls=\"--\", label=\"1 bit threshold\")\n", + "plt.xlim(0,0.5)\n", + "plt.legend()\n", + "if resolution is not None:\n", + " plt.axvline(resolution, color=\"k\", ls=\":\")\n", + "plt.show()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}