Commit 4dc200d8 authored by Pierre Paleo's avatar Pierre Paleo
Browse files

WIP memory estimation

parent aaa5c367
......@@ -4,16 +4,14 @@
computations.py: determine computational needs, chunking method to be used, etc.
"""
from silx.image.tomography import get_next_power
def get_chunk_size(nabu_config, dataset, available_mem):
"""
Return the maximum chunk size that can be used with a given amount of memory.
"""
pass
def estimate_required_memory(nabu_config, dataset, chunk_size=None):
......@@ -33,21 +31,60 @@ def estimate_required_memory(nabu_config, dataset, chunk_size=None):
Na = dataset.n_angles
if chunk_size is not None:
Ny = chunk_size
if nabu_config["dataset"]["binning"]:
Nx = Nx // nabu_config["dataset"]["binning"]
if nabu_config["dataset"]["binning_z"]:
Ny = Ny // nabu_config["dataset"]["binning_z"]
if nabu_config["dataset"]["projections_subsampling"]:
Na = Na // nabu_config["dataset"]["projections_subsampling"]
data_volume_size = Nx * Ny * Na * 2 # assuming uint16
total_memory_needed = 0
# Read data
# ----------
binning = nabu_config["dataset"]["binning"]
binning_z = nabu_config["dataset"]["binning_z"]
projections_subsampling = nabu_config["dataset"]["projections_subsampling"]
Nx = Nx // binning
Ny = Ny // binning_z
Na = Na // projections_subsampling
data_volume_size = Nx * Ny * Na * 4
total_memory_needed += data_volume_size
# CCD processing
# ---------------
if nabu_config["preproc"]["flatfield_enabled"]:
# Flat-field is done in-place, but still need to load darks/flats
n_darks = len(dataset.darks)
n_flats = len(dataset.flats)
darks_size = n_darks * Nx * Ny * 2 # uint16
flats_size = n_flats * Nx * Ny * 4 # f32
total_memory_needed += darks_size + flats_size
if nabu_config["preproc"]["ccd_filter_enabled"]:
# CCD filter is "batched 2D"
total_memory_needed += data_volume_size # TODO multiple batches of size < chunk_size
# Phase retrieval
# ---------------
if nabu_config["phase"]["method"] not in ["none", ""]:
# Phase retrieval is done image-wise, so near in-place, but needs to
# allocate some images, fft plans, and so on
Nx_p = get_next_power(2 * Nx)
Ny_p = get_next_power(2 * Ny)
img_size_real = 2 * 4 * Nx_p * Ny_p
img_size_cplx = 2 * 8 * ((Nx_p * Ny_p)//2 + 1)
total_memory_needed += 2 * img_size_real + 3 * img_size_cplx
if abs(float(nabu_config["phase"]["unsharp_coeff"])) > 1e-7:
# Unsharp mask is done image-wise, near-in-place
pass
# Reconstruction
# ---------------
rec_config = nabu_config["reconstruction"]
reconstructed_volume_size = 0
if rec_config["method"] != "none":
Nx_rec = rec_config["end_x"] - rec_config["start_x"] + 1
Ny_rec = rec_config["end_y"] - rec_config["start_y"] + 1
Nz_rec = rec_config["end_z"] - rec_config["start_z"] + 1
if rec_config["method"] not in ["none", ""]:
Nx_rec = (rec_config["end_x"] - rec_config["start_x"] + 1)//binning
Ny_rec = (rec_config["end_y"] - rec_config["start_y"] + 1)//binning
Nz_rec = (rec_config["end_z"] - rec_config["start_z"] + 1)//binning_z
if chunk_size:
Nz_rec = chunk_size//binning_z
reconstructed_volume_size = Nx_rec * Ny_rec * Nz_rec * 4 # float32
total_memory_needed += reconstructed_volume_size
return total_memory_needed
......@@ -134,16 +134,16 @@ nabu_config = {
"hst_key": "paganin_marge",
"validator": integer_validator,
},
"paganin_unsharp_sigma": {
"unsharp_coeff": {
"default": "0",
"help": "Standard deviation of the Gaussian filter when applying an unsharp mask\nafter the Paganin filtering. Disabled if set to 0.",
"hst_key": "pus",
"help": "Unsharp mask strength. The unsharped image is equal to\n UnsharpedImage = (1 + coeff)*originalPaganinImage - coeff * ConvolutedImage. Setting this coefficient to zero means that no unsharp mask will be applied.",
"hst_key": "puc",
"validator": float_validator,
},
"paganin_unsharp_coeff": {
"unsharp_sigma": {
"default": "0",
"help": "Unsharp mask strength. The unsharped image is equal to\n UnsharpedImage = (1 + coeff)*originalPaganinImage - coeff * ConvolutedImage",
"hst_key": "puc",
"help": "Standard deviation of the Gaussian filter when applying an unsharp mask\nafter the Paganin filtering. Disabled if set to 0.",
"hst_key": "pus",
"validator": float_validator,
},
"paganin_padding_type": {
......
......@@ -203,7 +203,7 @@ def padding_mode_validator(val):
@validator
def reconstruction_method_validator(val):
avail_rec_methods = ["fbp", "none"] # dfi, iterative...
avail_rec_methods = ["fbp", "none", ""] # dfi, iterative...
return name_range_checker(val, avail_rec_methods, "reconstruction method")
@validator
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment