nabu 2023.2.1__py3-none-any.whl → 2024.1.0rc3__py3-none-any.whl

nabu/opencl/processing.py

@@ -1,16 +1,24 @@
-import numpy as np
-from ..misc.processing_base import ProcessingBase
+from ..processing.processing_base import ProcessingBase
+from ..utils import MissingComponentError
 from .utils import get_opencl_context, __has_pyopencl__
 
 if __has_pyopencl__:
+    from .kernel import OpenCLKernel
     import pyopencl as cl
     import pyopencl.array as parray
+    from pyopencl.tools import dtype_to_ctype
+
+    OpenCLArray = parray.Array
+else:
+    OpenCLArray = MissingComponentError("pyopencl")
+    dtype_to_ctype = MissingComponentError("pyopencl")
 
 
 class OpenCLProcessing(ProcessingBase):
-    _array_class = parray.Array
+    array_class = OpenCLArray
+    dtype_to_ctype = dtype_to_ctype
 
-    def __init__(self, ctx=None, device_type="all", queue=None, **kwargs):
+    def __init__(self, ctx=None, device_type="all", queue=None, profile=False, **kwargs):
         """
         Initialie a OpenCLProcessing instance.
 
@@ -31,10 +39,24 @@ class OpenCLProcessing(ProcessingBase):
         else:
             self.ctx = ctx
         if queue is None:
-            queue = cl.CommandQueue(self.ctx)
+            queue_init_kwargs = {}
+            if profile:
+                queue_init_kwargs = {"properties": cl.command_queue_properties.PROFILING_ENABLE}
+            queue = cl.CommandQueue(self.ctx, **queue_init_kwargs)
         self.queue = queue
 
     # TODO push_context, pop_context ?
 
     def _allocate_array_mem(self, shape, dtype):
         return parray.zeros(self.queue, shape, dtype)
+
+    def kernel(self, kernel_name, filename=None, src=None, automation_params=None, **build_kwargs):
+        return OpenCLKernel(
+            kernel_name,
+            None,
+            queue=self.queue,
+            filename=filename,
+            src=src,
+            automation_params=automation_params,
+            **build_kwargs,
+        )

nabu/opencl/src/backproj.cl

@@ -17,7 +17,7 @@ static inline float linear_interpolation(global float* arr, int Nx, float x, int
     if (x < 0 || x >= Nx) return 0.0f; // texture address mode CLAMP_TO_EDGE
     int xm = (int) floor(x);
     int xp = (int) ceil(x);
-    if (xm == xp) return arr[y*Nx+xm];
+    if ((xm == xp) || (xp >= Nx)) return arr[y*Nx+xm];
     else return (arr[y*Nx+xm] * (xp - x)) + (arr[y*Nx+xp] * (x - xm));
 }
 
@@ -35,8 +35,8 @@ kernel void backproj(
     float axis_position,
     int n_x,
     int n_y,
-    int offset_x,
-    int offset_y,
+    float offset_x,
+    float offset_y,
     global float* d_cos,
     global float* d_msin,
     #ifdef DO_AXIS_CORRECTION

nabu/opencl/src/fftshift.cl

@@ -1,24 +1,77 @@
 #include <pyopencl-complex.h>
-typedef cfloat_t complex;
 
+#ifndef DTYPE
+  #define DTYPE float
+#endif
 
-/* 
-  In-place one-dimensional fftshift, along fast (horizontal) dimension.
+
+static inline void swap(global DTYPE* arr, size_t idx, size_t idx2) {
+    DTYPE tmp = arr[idx];
+    arr[idx] = arr[idx2];
+    arr[idx2] = tmp;
+}
+
+
+/*
+  In-place one-dimensional fftshift, along horizontal dimension.
   The array can be 1D or 2D.
+
+  direction > 0 means fftshift, direction < 0 means ifftshift.
+
+  It works for even-sized arrays.
+  Odd-sized arrays need an additional step (see roll.cl: roll_forward_x)
+
 */
-__kernel void fftshift1(
-    __global DTYPE* array,
+
+__kernel void fftshift_x_inplace(
+    __global DTYPE* arr,
     int Nx,
-    int Ny
+    int Ny,
+    int direction
 ) {
     int x = get_global_id(0), y = get_global_id(1);
-    int N2 = Nx >> 1;
-    int N2b = Nx - N2;
-    if (x >= N2b || y >= Ny) return;
-    DTYPE tmp = array[y*Nx + x];
-    if (x < N2) array[y*Nx + x] = array[y*Nx + N2b + x]
-    array[y*Nx + N2 + x] = tmp
+
+    int shift = Nx / 2;
+    if (x >= shift) return;
+
+    // (i)fftshift on odd-sized arrays cannot be done in-place in one step - need another kernel after this one
+    if ((Nx & 1) && (direction > 0)) shift++;
+    size_t idx = y * Nx + x;
+    size_t idx_out = y * Nx + ((x + shift) % Nx);
+    swap(arr, idx, idx_out);
+}
+
+
+#ifdef DTYPE_OUT
+
+/*
+    Out-of-place fftshift, possibly with type casting - useful for eg. fft(ifftshift(array))
+*/
+__kernel void fftshift_x(global DTYPE* arr, global DTYPE_OUT* dst, int Nx, int Ny, int direction) {
+
+    int x = get_global_id(0), y = get_global_id(1);
+    if (x >= Nx || y >= Ny) return;
+
+    int shift = Nx / 2;
+    if ((Nx & 1) && (direction < 0)) shift++;
+
+    size_t idx = y * Nx + x;
+    size_t idx_out = y * Nx + ((x + shift) % Nx);
+
+    DTYPE_OUT out_item;
+    #ifdef CAST_TO_COMPLEX
+      out_item = cfloat_new(arr[idx], 0);
+    #else
+      #ifdef CAST_TO_REAL
+        out_item = cfloat_real(arr[idx]);
+      #else
+        out_item = (DTYPE_OUT) arr[idx];
+      #endif
+    #endif
+    dst[idx_out] = out_item;
+
 }
 
+#endif
 
 

nabu/opencl/src/padding.cl

@@ -13,8 +13,8 @@ __kernel void coordinate_transform(
     uint y = get_global_id(1);
     if ((x >= Nx_padded) || (y >= Ny_padded)) return;
     uint idx = y*Nx_padded  +  x;
-    int x2 = cols_inds[idx];
-    int y2 = rows_inds[idx];
+    int x2 = cols_inds[x];
+    int y2 = rows_inds[y];
     array_out[idx] = array_in[y2*Nx + x2];
 }
 

nabu/opencl/src/roll.cl

@@ -0,0 +1,96 @@
+#include <pyopencl-complex.h>
+
+
+static inline void swap(global DTYPE* arr, size_t idx, size_t idx2) {
+    DTYPE tmp = arr[idx];
+    arr[idx] = arr[idx2];
+    arr[idx2] = tmp;
+}
+
+
+/*
+    This code should work but it not used yet.
+    The first intent was to have an in-place fftshift for odd-sized arrays:
+      fftshift_odd = fftshift_even followed by roll(-1) on second half of the array
+      ifft_odd = fftshift_even followed by roll(1) on second half of the array
+
+    Roll elements (as in numpy.roll(arr, 1)) of an array, in-place.
+    Needs to be launched with a large horizontal work group.
+*/
+__kernel void roll_forward_x(
+    __global DTYPE* array,
+    int Nx,
+    int Ny,
+    int offset_x,
+    __local DTYPE* shmem
+) {
+
+    int Nx_tot = Nx;
+    if (offset_x > 0) {
+        Nx_tot = Nx;
+        Nx -= offset_x;
+    }
+
+    int x = get_global_id(0), y = get_global_id(1);
+    if ((x >= Nx / 2) || (y >= Ny)) return;
+
+    __global DTYPE* arr = array + y * Nx_tot + offset_x;
+
+    int lid = get_local_id(0);
+    int wg_size = get_local_size(0);
+
+    int n_steps = (int) ceil((Nx - (Nx & 1)) * 1.0f / (2*wg_size));
+
+    DTYPE previous, current, write_on_first;
+    int offset = 0;
+
+    for (int step = 0; step < n_steps; step++) {
+
+        int idx = 2*lid + 1;
+        if (offset + idx >= Nx) break;
+
+        previous = arr[offset + idx - 1];
+        current = arr[offset + idx];
+        arr[offset + idx] = previous;
+
+        if ((step == n_steps - 1) && (offset + idx + 1 >= Nx - 1)) {
+            if (Nx & 1) write_on_first = arr[offset + idx + 1];
+            else write_on_first = current;
+        }
+
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        if ((step > 0) && (lid == 0)) arr[offset + idx - 1] = shmem[0];
+        if ((lid == wg_size - 1) && (step < n_steps - 1)) shmem[0] = current;
+        else if (offset + idx + 1 <= Nx - 1) arr[offset + idx + 1] = current;
+
+        if ((step == n_steps - 1) && (offset + idx + 1 >= Nx - 1)) arr[0] = write_on_first;
+
+        barrier(CLK_LOCAL_MEM_FENCE);
+
+        offset += 2 * wg_size;
+    }
+
+}
+
+
+__kernel void revert_array_x(
+    __global DTYPE* array,
+    int Nx,
+    int Ny,
+    int offset_x
+) {
+    int x = get_global_id(0), y = get_global_id(1);
+
+    int Nx_tot = Nx;
+    if (offset_x > 0) {
+        Nx_tot = Nx;
+        Nx -= offset_x;
+    }
+
+    if ((x >= Nx / 2) || (y >= Ny)) return;
+
+    size_t idx = y * Nx_tot + offset_x + x;
+    size_t idx2 = y * Nx_tot + offset_x + (Nx - 1 - x); // Nx ?
+    swap(array, idx, idx2);
+}

nabu/opencl/src/transpose.cl

@@ -0,0 +1,16 @@
+#ifndef SRC_DTYPE
+  #define SRC_DTYPE float
+#endif
+#ifndef DST_DTYPE
+  #define DST_DTYPE float
+#endif
+
+#include <pyopencl-complex.h>
+
+__kernel void transpose(__global SRC_DTYPE* src, __global DST_DTYPE* dst, int src_width, int src_height) {
+    // coordinates for "dst"
+    uint x = get_global_id(0);
+    uint y = get_global_id(1);
+    if ((x >= src_height) || (y >= src_width)) return;
+    dst[y*src_height + x] = (DST_DTYPE) src[x*src_width + y];
+}

nabu/pipeline/config_validators.py

@@ -1,7 +1,7 @@
 import os
 
 path = os.path
-from ..utils import is_writeable
+from ..utils import check_supported, is_writeable
 from .params import *
 
 """
@@ -63,15 +63,32 @@ def convert_to_bool(val):
     if not error:
         res = val_int > 0
     else:
-        if val.lower() in ["yes", "true"]:
+        if val.lower() in ["yes", "true", "y"]:
             res = True
             error = None
-        if val.lower() in ["no", "false"]:
+        if val.lower() in ["no", "false", "n"]:
             res = False
             error = None
     return res, error
 
 
+def str2bool(val):
+    """This is an interface to convert_to_bool and it is meant
+    to work as a class: in argparse interface the type argument can be set to float, int .. in general to a class.
+    The argument value is then created, at parsing time, by typecasting the input string to the given class.
+    A possibly occuring exception then trigger, in case, the display explanation provided by the argparse library.
+    All what this methods does is simply trying to convert an argument into a bool, and return it,
+    or generate an exception if there is a problem
+    """
+    import argparse
+
+    res, error = convert_to_bool(val)
+    if error:
+        raise argparse.ArgumentTypeError(error)
+    else:
+        return res
+
+
 def convert_to_bool_noerr(val):
     res, err = convert_to_bool(val)
     if err is not None:
@@ -226,6 +243,20 @@ def binning_validator(val):
     return max(1, val_int)
 
 
+@validator
+def projections_subsampling_validator(val):
+    val = val.strip()
+    err_msg = "projections_subsampling: expected one positive integer or two integers in the format step:begin"
+    if ":" not in val:
+        val += ":0"
+    step, begin = val.split(":")
+    step_int, error1 = convert_to_int(step)
+    begin_int, error2 = convert_to_int(begin)
+    if error1 is not None or error2 is not None or step_int <= 0 or begin_int < 0:
+        raise ValueError(err_msg)
+    return step_int, begin_int
+
+
 @validator
 def optional_file_name_validator(val):
     if len(val) > 0:
@@ -524,6 +555,35 @@ def logging_validator(val):
     return name_range_checker(val, set(log_levels.values()), "logging level", replacements=log_levels)
 
 
+@validator
+def exclude_projections_validator(val):
+    val = val.strip()
+    if val == "":
+        return None
+    if path.isfile(val):
+        # previous/default behavior
+        return {"type": "indices", "file": val}
+    if "=" not in val:
+        raise ValueError(
+            "exclude_projections: expected either 'angles=angles_file.txt' or 'indices=indices_file.txt' or 'angular_range=[a,b]'"
+        )
+    excl_type, excl_val = val.split("=")
+    excl_type = excl_type.strip()
+    excl_val = excl_val.strip()
+    check_supported(excl_type, exclude_projections_type.keys(), "exclude_projections type")
+    if excl_type == "angular_range":
+
+        def _get_range(range_val):
+            for c in ["(", ")", "[", "]"]:
+                range_val = range_val.replace(c, "")
+            r_min, r_max = range_val.split(",")
+            return (float(r_min), float(r_max))
+
+        return {"type": "angular_range", "range": _get_range(excl_val)}
+    else:
+        return {"type": excl_type, "file": excl_val}
+
+
 @validator
 def no_validator(val):
     return val

nabu/pipeline/dataset_validator.py

@@ -157,7 +157,7 @@ class DatasetValidatorBase:
         """
         dataset_cfg = self.nabu_config["dataset"]
         self.binning = (dataset_cfg["binning"], dataset_cfg["binning_z"])
-        subsampling_factor = dataset_cfg["projections_subsampling"]
+        subsampling_factor, subsampling_begin = dataset_cfg["projections_subsampling"]
         self.subsampling_factor = subsampling_factor or 1
 
         if self.binning != (1, 1):
@@ -207,5 +207,5 @@ class DatasetValidatorBase:
             # TODO automatically compute binning/subsampling factors as a function of lowest memory (GPU)
             self.nabu_config["dataset"]["binning"] = 2
             self.nabu_config["dataset"]["binning_z"] = 2
-            self.nabu_config["dataset"]["projections_subsampling"] = 2
+            self.nabu_config["dataset"]["projections_subsampling"] = 2, 0
         # TODO handle other modes