diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..9016df2 Binary files /dev/null and b/.DS_Store differ diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..f2ed587 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,97 @@ +name: CI + +on: + pull_request: + branches: + - main + - develop + push: + branches: + - main + - develop + tags: '*' + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} + +jobs: + test: + name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }} + runs-on: ${{ matrix.os }} + + strategy: + fail-fast: false + matrix: + version: + - '1.10' + os: + - ubuntu-latest + arch: + - x64 + + steps: + - uses: actions/checkout@v6 + + - uses: julia-actions/setup-julia@v3 + with: + version: ${{ matrix.version }} + arch: ${{ matrix.arch }} + + - uses: julia-actions/cache@v2 + + - name: Add unregistered dependencies + run: | + julia --project=. -e ' + using Pkg + Pkg.develop(PackageSpec( + url="https://github.com/rainerheintzmann/EvalMultiPoly.jl" + )) + Pkg.instantiate() + ' + + - uses: julia-actions/julia-buildpkg@v1 + + - uses: julia-actions/julia-runtest@v1 + + - uses: julia-actions/julia-processcoverage@v1 + + - uses: codecov/codecov-action@v4 + env: + CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }} + with: + file: lcov.info + + docs: + name: Documentation + runs-on: ubuntu-latest + + permissions: + contents: write + + steps: + - uses: actions/checkout@v6 + + - uses: julia-actions/setup-julia@v3 + with: + version: '1.10' + + - uses: julia-actions/cache@v2 + + - name: Install documentation dependencies + run: | + julia --project=docs -e ' + using Pkg + Pkg.develop([ + PackageSpec( + url="https://github.com/rainerheintzmann/EvalMultiPoly.jl" + ), + PackageSpec(path=pwd()) + ]) + Pkg.instantiate() + ' + + - name: Build and deploy documentation + uses: julia-actions/julia-docdeploy@releases/v1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} \ No newline at end of file diff --git a/.gitignore b/.gitignore index 29126e4..32f88c7 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,19 @@ docs/site/ # committed for packages, but should be committed for applications that require a static # environment. Manifest.toml +examples/Manifest.toml +examples/*.tiff +examples/image_registration_2dpolim.jl + +*.mp4 +*.png +examples/figures/* +examples/*.tif +examples/*.ipynb +*.jpeg +*.gif +*.jpg +*.txt +*.txt +.*tiff +examples/6-4-Rigis.jl diff --git a/Project.toml b/Project.toml index 7c59800..c7f24bf 100644 --- a/Project.toml +++ b/Project.toml @@ -1,8 +1,38 @@ name = "DataToFunctions" uuid = "64cfdffa-4d02-49ee-ae8b-a805370874f5" -authors = ["RainerHeintzmann "] version = "0.1.0" +authors = ["RainerHeintzmann "] [deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +EvalMultiPoly = "c78649ec-f9ed-405e-be6d-0472f43586aa" FourierTools = "b18b359b-aebc-45ac-a139-9c0ccbb2871e" Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" + +[compat] +Adapt = "4.7.0" +Aqua = "0.8" +BenchmarkTools = "1" +CUDA = "6" +ChainRulesCore = "1.26.1" +EvalMultiPoly = "0.1" +ForwardDiff = "1" +FourierTools = "0.4, 0.5" +Interpolations = "0.15.1, 0.16" +StaticArrays = "1.9.6, 1.10" +Test = "1.9.3" +Zygote = "0.6, 0.7" +julia = "1" + +[extras] +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[targets] +test = ["Aqua", "Test", "Zygote", "BenchmarkTools", "ForwardDiff", "CUDA"] diff --git a/README.md b/README.md index 8f912d3..5c43c20 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,90 @@ # DataToFunctions.jl -Represents (measured) data as a function, which supports scaling and shifting. It is intended to be used as a tool for fitting data, where the fitting function is given by measured data. + +[![CI](https://github.com/RainerHeintzmann/DataToFunctions.jl/actions/workflows/ci.yml/badge.svg)](https://github.com/RainerHeintzmann/DataToFunctions.jl/actions/workflows/ci.yml) +[![Development Documentation](https://img.shields.io/badge/docs-dev-blue.svg)](https://rainerheintzmann.github.io/DataToFunctions.jl/dev/) + +`DataToFunctions.jl` represents data as continuously evaluated, parameterized functions. + +The package is intended for fitting, image registration, and inverse problems +where measured data itself forms part of the model. It provides efficient +interpolation-based coordinate transformations with support for affine, +polynomial, and user-defined transformations. + +## Features + +- Affine transformations in arbitrary dimensions using homogeneous coordinates; +- A convenient seven-parameter affine model for 2-D data; +- Polynomial coordinate transformations based on + [`EvalMultiPoly.jl`](https://github.com/RainerHeintzmann/EvalMultiPoly.jl); +- User-defined Cartesian and homogeneous-coordinate transformations; +- Allocation-free in-place CPU evaluation for performance-critical workflows; +- Automatic differentiation of affine and polynomial transformation parameters; +- Custom `ChainRulesCore` reverse rules for efficient Zygote gradients; +- Compatibility with gradient-based optimization such as `Optim.jl` with + `LBFGS`; +- Backend-generic execution, including CUDA arrays for supported interpolation + methods. + +## Quick example + +```julia +using DataToFunctions + +data = rand(Float32, 128, 128) + +f = get_function_affine(data) + +p = ( + 0.2f0, # x shift + -0.1f0, # y shift + 1.0f0, # x scale + 1.0f0, # y scale + 0.0f0, # xy shear + 0.0f0, # yx shear + 0.01f0, # rotation [rad] +) + +warped = f(p) +``` + +For repeated forward evaluations, an in-place version is available: + +```julia +f! = get_function_affine_inplace(data) + +out = similar(data) +f!(out, p) +``` + +On CPU, the in-place transformation path is designed to run without heap +allocations after compilation when using concrete tuple parameters and a +preallocated output array. + +## Automatic differentiation + +The affine and polynomial transformation interfaces provide custom reverse-mode +rules for differentiation with respect to transformation parameters. + +For example: + +```julia +using Zygote + +target = f(p) + +loss(q) = sum(abs2, f(q) .- target) + +gradient = Zygote.gradient(loss, p)[1] +``` + +These gradients can be used directly in optimization workflows. + +## Documentation + +- [Development documentation](https://rainerheintzmann.github.io/DataToFunctions.jl/dev/) + +The development documentation follows the `develop` branch. + + diff --git a/docs/Project.toml b/docs/Project.toml new file mode 100644 index 0000000..d2cdc05 --- /dev/null +++ b/docs/Project.toml @@ -0,0 +1,11 @@ +[deps] +DataToFunctions = "64cfdffa-4d02-49ee-ae8b-a805370874f5" +Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +DocumenterTools = "35a29f4d-8980-5a13-9543-d66fff28ecb8" +EvalMultiPoly = "c78649ec-f9ed-405e-be6d-0472f43586aa" +GR = "28b8d3ca-fb5f-59d9-8090-bfdbd6d07a71" +Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +SyntheticObjects = "e7028c27-0967-45e9-8fdb-dbc10ccb2b0a" diff --git a/docs/make.jl b/docs/make.jl new file mode 100644 index 0000000..0cd5922 --- /dev/null +++ b/docs/make.jl @@ -0,0 +1,23 @@ +import Pkg +Pkg.activate(@__DIR__) + +cd(@__DIR__) # go into `docs` folder + +using Documenter, Literate, DataToFunctions + +# convert tutorial/examples to markdown +# Literate.markdown("./src/tutorial.jl", "./src") +# Which markdown files to compile to HTML +# (which is also the sidebar and the table +# of contents for your documentation) + +pages = Any[ + "Introduction" => "index.md", + "Tutorial" => "tutorial.md", + "API" => "api.md", + ] + +# compile to HTML: +makedocs(; sitename="DataToFunctions.jl", pages, modules = [DataToFunctions], checkdocs = :exports) + +deploydocs(repo = "github.com/RainerHeintzmann/DataToFunctions.jl.git", devbranch = "develop") \ No newline at end of file diff --git a/docs/src/api.md b/docs/src/api.md new file mode 100644 index 0000000..49dfd5f --- /dev/null +++ b/docs/src/api.md @@ -0,0 +1,42 @@ +# API Reference + +This page lists the public transformation interface provided by +`DataToFunctions.jl`. + +## Affine transformations + +```@docs +DataToFunctions.get_function_affine +DataToFunctions.get_function_affine_inplace +``` + +## Polynomial transformations + +```@docs +DataToFunctions.get_function_poly +DataToFunctions.get_function_poly_inplace +``` + +## User-defined transformations + +```@docs +DataToFunctions.get_function_tuple +DataToFunctions.get_function_tuple_inplace +DataToFunctions.get_function_homogen +DataToFunctions.get_function_homogen_inplace +``` + +## Mode-based interface + +```@docs +DataToFunctions.get_interpolated_function +``` + +## Transformation modes + +If `AffineMode` and `PolynomialMode` have docstrings, they can be included here: + +```@docs +DataToFunctions.AffineMode +DataToFunctions.PolynomialMode +``` diff --git a/docs/src/index.md b/docs/src/index.md new file mode 100644 index 0000000..18c02ed --- /dev/null +++ b/docs/src/index.md @@ -0,0 +1,134 @@ +# DataToFunctions.jl + +`DataToFunctions.jl` turns sampled array data into continuously evaluated, +parameterized functions. + +It is intended for applications such as image registration, fitting, and inverse +problems where measured data itself forms part of the model. + +The package provides efficient coordinate transformations combined with +interpolation, including: + +- affine transformations, +- polynomial transformations, +- user-defined Cartesian-coordinate transformations, +- user-defined homogeneous-coordinate transformations, +- allocation-free in-place CPU evaluation, +- automatic differentiation support for affine and polynomial parameters, +- GPU-compatible execution for supported array and interpolation backends. + +## Quick start + +```julia +using DataToFunctions + +data = rand(Float32, 128, 128) + +f = get_function_affine(data) + +p = ( + 0.2f0, # x shift + -0.1f0, # y shift + 1.0f0, # x scale + 1.0f0, # y scale + 0.0f0, # xy shear + 0.0f0, # yx shear + 0.01f0, # rotation [rad] +) + +warped = f(p) +``` + +For repeated forward evaluations, use the in-place interface: + +```julia +f! = get_function_affine_inplace(data) + +out = similar(data) + +f!(out, p) +``` + +The in-place CPU implementation is designed to avoid heap allocations after +compilation when using concrete tuple parameters and a preallocated output. + +## Polynomial transformations + +Polynomial coordinate transformations are provided through `EvalMultiPoly.jl`: + +```julia +using EvalMultiPoly + +f = get_function_poly(data, Val(2)) + +coeffs = get_identity_multipoly_coeffs( + Val(2), + Val(2), +) + +warped = f(coeffs) +``` + +An allocation-minimized version is also available: + +```julia +f! = get_function_poly_inplace(data, Val(2)) + +out = similar(data) + +f!(out, coeffs) +``` + +## Automatic differentiation + +The allocating affine and polynomial interfaces support differentiation with +respect to their transformation parameters. + +For example: + +```julia +using Zygote + +target = f(p) + +loss(q) = sum(abs2, f(q) .- target) + +g = Zygote.gradient(loss, p)[1] +``` + +The affine and polynomial implementations provide custom reverse-mode rules +using `ChainRulesCore`, avoiding construction of the full image-to-parameter +Jacobian. + +This makes the resulting gradients suitable for optimization packages such as +`Optim.jl`. + +## GPU arrays + +The transformation implementation is backend-generic. + +CPU arrays use specialized CPU paths, while other `AbstractArray` backends use +broadcast and reduction operations. Supported interpolation objects are adapted +to the target array backend using `Adapt.jl`. + +For example, with CUDA.jl: + +```julia +using CUDA +using DataToFunctions + +CUDA.allowscalar(false) + +data_gpu = CuArray(data) + +f_gpu = get_function_affine(data_gpu) + +warped_gpu = f_gpu(p) +``` + +GPU compatibility depends on the selected interpolation method and array +backend. `BSpline(Linear())`, the default interpolation mode, is the primary +supported GPU path. + + +See the [API Reference](@ref) for the full list of exported functions. diff --git a/docs/src/tutorial.md b/docs/src/tutorial.md new file mode 100644 index 0000000..cc9aee2 --- /dev/null +++ b/docs/src/tutorial.md @@ -0,0 +1,172 @@ +# Tutorial + +This short tutorial shows the main workflow of `DataToFunctions.jl`: create a +parameterized transformation from sampled data, evaluate it efficiently, and +differentiate it for optimization. + +## Affine transformation + +Start from a 2-D array: + +```julia +using DataToFunctions + +data = rand(Float32, 128, 128) +``` + +Create an affine transformation function: + +```julia +f = get_function_affine(data) +``` + +For 2-D data, the transformation parameters are + +```julia +p = ( + 0.2f0, # shift x + -0.1f0, # shift y + 1.0f0, # scale x + 1.0f0, # scale y + 0.0f0, # shear xy + 0.0f0, # shear yx + 0.01f0, # rotation [rad] +) +``` + +Evaluate the transformed data with: + +```julia +warped = f(p) +``` + +For repeated forward evaluations, use the in-place version: + +```julia +f! = get_function_affine_inplace(data) + +out = similar(data) + +f!(out, p) +``` + +The in-place CPU path is intended for preallocated, allocation-free repeated +evaluation after compilation. + +## Polynomial transformation + +Polynomial coordinate transformations are created by specifying the polynomial +order: + +```julia +using EvalMultiPoly + +f_poly = get_function_poly(data, Val(2)) + +coeffs = get_identity_multipoly_coeffs( + Val(2), + Val(2), +) + +warped_poly = f_poly(coeffs) +``` + +For repeated evaluation: + +```julia +f_poly! = get_function_poly_inplace(data, Val(2)) + +out_poly = similar(data) + +f_poly!(out_poly, coeffs) +``` + +## Automatic differentiation + +The allocating affine and polynomial interfaces can be differentiated with +respect to their transformation parameters. + +For example: + +```julia +using Zygote + +target = f(( + 0.1f0, + -0.05f0, + 1.01f0, + 0.99f0, + 0.0f0, + 0.0f0, + 0.005f0, +)) + +loss(q) = sum(abs2, f(q) .- target) + +g = Zygote.gradient(loss, p)[1] +``` + +The package provides custom reverse-mode rules for the affine and polynomial +parameterizations, making these gradients suitable for iterative optimization. + +## Optimization with Optim.jl + +A gradient from Zygote can be used directly with `Optim.LBFGS`: + +```julia +using Optim + +p0 = Float32[ + 0.0, + 0.0, + 1.0, + 1.0, + 0.0, + 0.0, + 0.0, +] + +loss_vec(q) = sum(abs2, f(q) .- target) + +function g!(G, q) + G .= Zygote.gradient(loss_vec, q)[1] + return G +end + +result = optimize( + loss_vec, + g!, + p0, + LBFGS(), +) +``` + +The fitted parameters are available with: + +```julia +Optim.minimizer(result) +``` + +## GPU arrays + +The same high-level API can be used with GPU-backed arrays when the selected +interpolation mode is supported by the backend. + +For CUDA: + +```julia +using CUDA + +CUDA.allowscalar(false) + +data_gpu = CuArray(data) + +f_gpu = get_function_affine(data_gpu) + +warped_gpu = f_gpu(p) +``` + +The implementation keeps CUDA optional: `DataToFunctions.jl` does not require +CUDA.jl unless GPU execution is used. + +For the complete list of public functions, see the [API Reference](@ref). diff --git a/examples/PSF_fitting_general.jl b/examples/PSF_fitting_general.jl new file mode 100644 index 0000000..6bc5d8f --- /dev/null +++ b/examples/PSF_fitting_general.jl @@ -0,0 +1,412 @@ +using DataToFunctions +using Optim, StaticArrays, LinearAlgebra +using PointSpreadFunctions +using Zygote +using ForwardDiff, LineSearches, Plots, Printf +using View5D +using Distributions, Rotations +using Plots +using TestImages +using BenchmarkTools +#using InverseModeling +import Random +using Noise, Images, CSV, TiffImages +using ProgressBars + +""" + perform_fit_general(loss_function, fitting_data::AbstractArray) + +Performs a fit to the fitting data using a loss function defined by the user + +# Arguments +`loss_function`: User-defined loss function which is minimized +`fitting_data`: The data which is being fitted + +# Returns +a vector of 7 parameters: 2 for the shift, 2 for the scaling, 2 for shear, and 1 for rotation angle + +# Example +there is an example of this function in the `examples/star_fitting_genaral.jl` + +""" +function perform_fit_general(loss_function, fitting_data::AbstractArray, init_x::AbstractArray{T}) where T + # guess the shift parameters by taking the maximum values of the array and + # centering the positions + ##a, b = Tuple(argmax(fitting_data)) .- size(fitting_data) ./2.0 .- 1.0 + #print("INSIDE!!! hehe") + # assigning the initial parameter estimates + # init_x = vec([0.5, -1.5, 1.0, 1.0, 0.0, 0.0, pi/5]) #ndims(fitting_data)+1, ndims(fitting_data)+1)) + # reshape(Matrix(1.0*I, ndims(fitting_data)+1, ndims(fitting_data)+1), 1, 9)) #[a, b, 1.0, 1.0, 0.001, 0.001, 0.001] + + # setting the lower and upper boundary of the parameter values based on their limits + lower = T[-1*size(fitting_data)[1], -1*size(fitting_data)[2], 0.0, 0.0, -0.01, -0.01, 0.0] + upper = T[size(fitting_data)[1], size(fitting_data)[2], size(fitting_data)[1], size(fitting_data)[2], 0.01, 0.01, pi/2.0] + + # initializing the LBFGS optimizer + inner_optimizer = BFGS()#; m=3, linesearch=LineSearches.BackTracking(order=3)) + + # Computer, Optimize! :D + res = optimize( + loss_function, + + #LBFGS(), + lower, upper, + init_x, + Fminbox(inner_optimizer), + Optim.Options(store_trace = true, extended_trace = true, iterations=5000), + autodiff = :forward + ) + + # return the estimated parameters + return Optim.minimizer(res), res +end + +function perform_fit(loss_function, init_x::AbstractArray{T}) where T + + # Computer, Optimize! :D + res = optimize( + loss_function, + init_x, + #Newton(), + #BFGS(),#; linesearch=LineSearches.BackTracking(order=3)), + LBFGS(),#; linesearch=LineSearches.BackTracking(order=3)), + #lower, upper, + #init_x, + #Fminbox(inner_optimizer), + Optim.Options(store_trace = true, extended_trace = true, iterations=5000), + autodiff = :forward + ) + + # return the estimated parameters + return Optim.minimizer(res), res +end + +""" + apply_transform(;matrix=true, sz=64, dtype=Float32, noise_level=1/20.0) + +This function is designed for applying a transformation to a sample data using either matrix transformations or +parametric transformations based on the provided arguments. + +# Arguments +`matrix`: if true, the fitting is done using a matrix transformation, otherwise, the fitting is done using a parametric transformation +`sz`: the size of the sample data +`dtype`: the data type of the sample data +`noise_level`: the noise level to add to the sample data + +# Returns +a tuple of two arrays: the first array is the fitting data, and the second array is the estimated fitting data + +# Example +apply_transform(matrix=true, sz=64, dtype=Float32, noise_level=1/20.0) + +""" +function apply_transform(;matrix=false, sz=64, dtype=Float32, n_photons=1000, pure_rand=false, from_params=true, plotting=false) + Random.seed!(14) + + # defining the mean and the varixance of the test normal (Gaussian) distribution + μ = [0, 0] + Σ = [sz/10 0.0; + 0.0 sz/10] + + # initializing the multivariate normal distribution + p = MvNormal(μ, Σ) + + # to define the sample array based on a 2D normal distribution + X = -1*sz/2.0:1*sz/2.0 + Y = -1*sz/2.0:1*sz/2.0 + + z = [pdf(p, [x,y]) for y in Y, x in X] + + + # creating a PSF for the widefield microscope + sz_psf = (sz, sz, 100) + sampling = (0.040, 0.040, 0.050) + # simulate a confocal PSF + aberrations = Aberrations([Zernike_HorizontalComa],[0.8]); + pp = PSFParams(0.5,1.4,1.52, method=MethodPropagateIterative, aberrations=aberrations); + + #pp_ex = PSFParams(pp_em; λ=0.488);#, method=MethodPropagateIterative, aplanatic=aplanatic_illumination, aberrations=aberrations); + p_psf_3d = psf(sz_psf, pp, sampling=sampling); + p_psf = p_psf_3d[:, :, 50] + #sample_data = p_psf ./ maximum(p_psf) + + # normalizing the sample data + sample_data = p_psf ./ maximum(p_psf) + # sample_data = dtype.(z[1:sz, 1:sz]./maximum(z)) + # sample_data = dtype.(TestImages.shepp_logan(sz)) + # sample_data = rand(dtype, (sz, sz)) + + # sample_data .+= rand(dtype, (size(sample_data)...)).*noise_level; + p_img = n_photons .* (sample_data);# ./ maximum(sample_data)) + n_img = dtype.(poisson(Float64.(p_img))) + + x_cen, y_cen = (size(n_img) ./ 2.0) + t_to_origin = dtype[1.0 0.0 1*x_cen; 0.0 1.0 y_cen; 0.0 0.0 1.0]; + t_to_center = dtype[1.0 0.0 -1.0*x_cen; 0.0 1.0 -1.0*y_cen; 0.0 0.0 1.0]; + + true_vals = dtype[rand(-4.0:0.001:4.0), rand(-4.0:0.001:4.0), 1.0, 1.0, 0.0, 0.0, 0.0];#rand(0.9:0.001:1.1),rand(0.9:0.001:1.1), 0.0, 0.0, 0.0];#rand(0.001:0.001:pi/2.001)] + + if !pure_rand + + shear_mat = dtype[1.0 true_vals[5] 0.0; true_vals[6] 1.0 0.0; 0.0 0.0 1.0]; + scale_mat = dtype[1.0/true_vals[3] 0.0 0.0; 0.0 1/true_vals[4] 0.0; 0.0 0.0 1.0]; + + shift_mat = dtype[1.0 0.0 true_vals[1]; 0.0 1.0 true_vals[2]; 0.0 0.0 1.0]; + # converting the data to function (DataToFunctions.get_function) + + rot_mat = dtype[cos(true_vals[7]) -1.0*sin(true_vals[7]) 0.0; sin(true_vals[7]) cos(true_vals[7]) 0.0; 0.0 0.0 1.0]; + + matrix_c = (t_to_origin * scale_mat * shear_mat * rot_mat * shift_mat * t_to_center) + else + matrix_c = dtype.(t_to_origin * rand(0.1:0.001:1.0, (3, 3)) * t_to_center) + end + + f_affine_sim_img = get_function_affine(sample_data);#; super_sampling=1);#, extrapolation_bc=0.0); + if matrix + t_img = f_affine_sim_img(SMatrix{3, 3}(matrix_c))#, fitting_data); #.+ dtype.(rand(size(sample_data)...))./5.0; + else + t_img = f_affine_sim_img(true_vals)#, fitting_data); #.+ dtype.(rand(size(sample_data)...))./5.0; + end + fitting_data = dtype.(poisson(Float64.(t_img ./ maximum(t_img) .* n_photons))) #.+= rand(dtype, size(p_img)...).*noise_level; + + heatmap(fitting_data, aspect_ratio=1, size=(600, 600), title="fitting data", titlefont = font(20), legend=:none, axis=([], false)) + annotate!(vec(map(x -> Tuple((reverse(Tuple(x))..., text(@sprintf("%.0f", fitting_data[x]), :center, font(5), :white))), CartesianIndices(sample_data)))) + savefig("figures/fitting/sample_data_1.png") + + # plot(heatmap(sample_data, aspect_ratio=1), heatmap(fitting_data, aspect_ratio=1)) + + contour(sample_data, length=200, fill=false, title="Sample data", titlefont = font(20), legend=:none, aspect_ratio=1, size=(600, 600)) + savefig("figures/fitting/sample_data_1_contour.png") + return sample_data, fitting_data +end + + +function gauss_psf_comp() + x = -10.0:0.01:10.0 + p = Normal(0.0, 1.0) + y = pdf(p, x) + y_psf(x) = (sin(x) /x)^2 + plot(x, y_psf.(x), label="PSF of a circular aperture", title="Comparison of a Gaussian and a PSF", titlefont=20, size=(800, 400)) + plot!(x, y./maximum(y), label="Gaussian with μ=0.0 & σ=1.0") + savefig("figures/fitting/comp_psf_gaussian_1.png") + + plot(x, map(x -> (gradient(y_psf, x)[1]), x), label="Gradient of the PSF") + plot!(x, map(x -> (gradient(x -> pdf(p, x), x)[1]), x), label="Gradient of the Gaussian", title="Comparison of the Gradients", titlefont=20, size=(800, 400)) + savefig("figures/fitting/comp_psf_gaussian_1_gradients.png") +end + +""" + main_fitting(;matrix=true, sz=64, dtype=Float32, iterations=20, noise_level=1/20.0, pure_rand=false, from_params=true) + +This function is designed for performing fitting operations on sample data using either matrix transformations or +parametric transformations based on the provided arguments. + +# Arguments +`matrix`: if true, the fitting is done using a matrix transformation, otherwise, the fitting is done using a parametric transformation +`sz`: the size of the sample data +`dtype`: the data type of the sample data +`iterations`: the number of iterations to perform the fitting +`noise_level`: the noise level to add to the sample data +`pure_rand`: if true, the fitting is done using a random matrix transformation +`from_params`: if true, the fitting is done using the true values as the initial values + +# Returns +a tuple of two arrays: the first array is the fitting data, and the second array is the estimated fitting data + +# Example +x, y = main_fitting(matrix=true, sz=32, dtype=Float32, iterations=10, noise_level=1/20.0, pure_rand=false, from_params=true); + + +""" +function main_fitting(;matrix=true, sz=64, dtype=Float32, iterations=20, use_psf=true, n_photons=1000, pure_rand=false, from_params=true, plotting=false) + Random.seed!(14) + + # defining the mean and the varixance of the test normal (Gaussian) distribution + μ = [0, 0] + Σ = [sz/30 0.0; + 0.0 sz/30] + + # initializing the multivariate normal distribution + p = MvNormal(μ, Σ) + + # to define the sample array based on a 2D normal distribution + X = -1*sz/2.0:1*sz/2.0 + Y = -1*sz/2.0:1*sz/2.0 + + z = [pdf(p, [x,y]) for y in Y, x in X] + + + # creating a PSF for the widefield microscope + sz_psf = (sz, sz, 100) + sampling = (0.040, 0.040, 0.050) + # simulate a confocal PSF + aberrations = Aberrations([Zernike_HorizontalComa],[0.8]); + pp = PSFParams(0.5,1.4,1.52, method=MethodPropagateIterative, aberrations=aberrations); + + #pp_ex = PSFParams(pp_em; λ=0.488);#, method=MethodPropagateIterative, aplanatic=aplanatic_illumination, aberrations=aberrations); + p_psf_3d = psf(sz_psf, pp, sampling=sampling); + p_psf = p_psf_3d[:, :, 50] + #sample_data = p_psf ./ maximum(p_psf) + + # normalizing the sample data + sample_data = p_psf ./ maximum(p_psf) + sample_data_gaussian = dtype.(z[1:sz, 1:sz]./maximum(z)) + # sample_data = dtype.(TestImages.shepp_logan(sz)) + # sample_data = rand(dtype, (sz, sz)) + + # sample_data .+= rand(dtype, (size(sample_data)...)).*noise_level; + p_img = n_photons .* (sample_data);# ./ maximum(sample_data)) + n_img = dtype.(poisson(Float64.(p_img))) + + y = similar(n_img, (size(n_img)..., iterations)) + x = similar(n_img, (size(n_img)..., iterations)) + pos_res = zeros(Float32, iterations, 2) + pos_arr = zeros(Float32, iterations, 2) + + for i in ProgressBar(1:iterations) + # println("iteration: ", i) + + x_cen, y_cen = (size(n_img) ./ 2.0) + t_to_origin = dtype[1.0 0.0 1*x_cen; 0.0 1.0 y_cen; 0.0 0.0 1.0]; + t_to_center = dtype[1.0 0.0 -1.0*x_cen; 0.0 1.0 -1.0*y_cen; 0.0 0.0 1.0]; + + true_vals = dtype[rand(-4.0:0.001:4.0), rand(-4.0:0.001:4.0), 1.0, 1.0, 0.0, 0.0, 0.0];#rand(0.9:0.001:1.1),rand(0.9:0.001:1.1), 0.0, 0.0, 0.0];#rand(0.001:0.001:pi/2.001)] + + if !pure_rand + + shear_mat = dtype[1.0 true_vals[5] 0.0; true_vals[6] 1.0 0.0; 0.0 0.0 1.0]; + scale_mat = dtype[1.0/true_vals[3] 0.0 0.0; 0.0 1/true_vals[4] 0.0; 0.0 0.0 1.0]; + + shift_mat = dtype[1.0 0.0 true_vals[1]; 0.0 1.0 true_vals[2]; 0.0 0.0 1.0]; + # converting the data to function (DataToFunctions.get_function) + + rot_mat = dtype[cos(true_vals[7]) -1.0*sin(true_vals[7]) 0.0; sin(true_vals[7]) cos(true_vals[7]) 0.0; 0.0 0.0 1.0]; + + matrix_c = (t_to_origin * scale_mat * shear_mat * rot_mat * shift_mat * t_to_center) + else + matrix_c = dtype.(t_to_origin * rand(0.1:0.001:1.0, (3, 3)) * t_to_center) + end + + f_affine_sim_img = get_function_affine(sample_data);#; super_sampling=1);#, extrapolation_bc=0.0); + if matrix + t_img = f_affine_sim_img(SMatrix{3, 3}(matrix_c))#, fitting_data); #.+ dtype.(rand(size(sample_data)...))./5.0; + else + t_img = f_affine_sim_img(true_vals)#, fitting_data); #.+ dtype.(rand(size(sample_data)...))./5.0; + end + fitting_data = dtype.(poisson(Float64.(t_img ./ maximum(t_img) .* n_photons))) #.+= rand(dtype, size(p_img)...).*noise_level; + + + if use_psf + f_affine = get_function_affine(p_img);#; super_sampling=1);#, extrapolation_bc=0.0); + else + f_affine = get_function_affine(n_photons .* sample_data_gaussian);#; super_sampling=1);#, extrapolation_bc=0.0); + end + #f_affine = get_function_affine(n_img);#; super_sampling=1);#, extrapolation_bc=0.0); + # defining the loss function based on the gaussian noise + loss_m(p1::AbstractMatrix) = sum(abs2.(f_affine(SMatrix{size(p1)...}(p1)) .- fitting_data)) + # loss_m(p1::AbstractVector) = sum(abs2.(f_affine(p1::AbstractVector) .- fitting_data)) + loss_m(p1::AbstractVector) = sum(abs2.(f_affine([p1[1], p1[2], 0.0, 0.0, 0.0, 0.0, 0.0]) .- fitting_data)) + + + if matrix + if from_params + st_vals = dtype[1.0 0.0 -1.0*(argmax(fitting_data)[1]-size(fitting_data)[1]/2.0); 0.0 1.0 -1.0*(argmax(fitting_data)[1]-size(fitting_data)[2]/2.0); 0.0 0.0 1.0] + else + st_vals = dtype[1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0] + end + elseif from_params + st_vals = dtype[argmax(fitting_data)[1]-size(fitting_data)[1]/2.0, argmax(fitting_data)[2]-size(fitting_data)[2]/2.0, 1.0, 1.0, 0.0, 0.0, 0.0];#pi/8.0] + else + st_vals = dtype[0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0];#pi/8.0] + end + + # perform the main fit to the fitting data by minimizing the loss function + stats = @timed output, res = perform_fit(loss_m, st_vals) + + + if !matrix + pos_arr[i, :] = true_vals[1:2] + pos_res[i, :] = output[1:2] + else + pos_arr[i, :] = matrix_c[1:2, 3] + pos_res[i, :] = output[1:2, 3] + end + + y[:, :, i] = matrix ? f_affine(SMatrix{size(matrix_c)...}(output)) : f_affine(output) + x[:, :, i] = fitting_data + + # plotting the output of the fitting pocedure for further illustration + if plotting + begin + p00 = heatmap(n_img, aspect_ratio=1.0, title="Simulated sample PSF", colormap= :gist_gray); + p01 = heatmap(fitting_data, aspect_ratio=1.0, title="Simulated PSF", colormap= :gist_gray); + p02 = heatmap(matrix ? f_affine(SMatrix{size(matrix_c)...}(output)) : f_affine(output), aspect_ratio=1.0, title="Estimated fit", colormap= :gist_gray); + p03 = heatmap(fitting_data .- (matrix ? f_affine(SMatrix{size(matrix_c)...}(output)) : f_affine(output)), aspect_ratio=1.0, title="Residuals", colormap= :bwr, clim=(-maximum((abs.(fitting_data .- (matrix ? f_affine(SMatrix{size(matrix_c)...}(output)) : f_affine(output))))), maximum((abs.(fitting_data .- (matrix ? f_affine(SMatrix{size(matrix_c)...}(output)) : f_affine(output))))))); + + plot(p00, p01, p03, p02, layout=@layout([A B; C D]), + #framestyle=nothing, + #showaxis=false, + #xticks=false, yticks=false, + size=(1200, 1200), + plot_title=" $(if matrix "Matrix" else "Parametric" end) fitting +True vals: $(map(x -> @sprintf("%.3f",x), (matrix ? matrix_c : true_vals))) +fitted vals: $(map(x -> @sprintf("%.3f",x), output)) +n. of photons: $(@sprintf("%.0f", n_photons)) +time elapsed: $(@sprintf("%.1f", 1000.0*stats.time))ms, loss: $(@sprintf("%.2f", res.trace[1].value)) -> $(@sprintf("%.2f", res.trace[end].value))", + plot_titlevspan=0.14 + ) + savefig("figures/fitting/$(matrix ? "Matrix" : "Parametric")_fitting_$(i).png") + end + end + end + return x, y, pos_arr, pos_res +end + +x, y, pos_arr, pos_res = main_fitting(matrix=false, iterations=1000, sz=65, pure_rand=false, n_photons=10, from_params=true, plotting=false); println(mean(pos_res[:, 1] .- pos_arr[:, 1])); +println(std(pos_res[:, 1] .- pos_arr[:, 1])); + +#, title="Positional errors", markersize=2.0, xlabel="X error (pixels)", ylabel="Y error (pixels)", legend=:none, size=(600, 600), xlim=(-1.0, 1.0), ylim=(-1.0, 1.0), alpha=0.4) + +histogram2d(pos_res[:, 1] .- pos_arr[:, 1], pos_res[:, 2] .- pos_arr[:, 2], title="Positional errors histogram for 100 photons", xlabel="X error (pixels)", ylabel="Y error (pixels)", xlim=(-1.0, 1.0), ylim=(-1.0, 1.0), bins=20, aspect_ratio=1) +# + + +imgg = Gray{N0f16}.(x./maximum(x)); +ff = TiffImages.DenseTaggedImage(imgg); +TiffImages.save("test_4_aberrated.tif", ff); + +res_fiji = CSV.File(open(raw"C:\Users\ho82nat\Desktop\thunderstorm_res_100photons.csv")) +res_fiji_x = res_fiji["x [nm]"] ./ 80.0 .- 65.0 ./ 2.0; +res_fiji_y = res_fiji["y [nm]"] ./ 80.0 .- 65.0 ./ 2.0; + + +scatter(pos_res[:, 1] .- pos_arr[:, 1], pos_res[:, 2] .- pos_arr[:, 2], markershape= :circle, title="Positional errors", markersize=2.0, xlabel="X error (pixels)", ylabel="Y error (pixels)", legend=:none, size=(600, 600), label="DataToFunctions fitting")#, xlim=(-1.0, 1.0), ylim=(-1.0, 1.0), alpha=0.4) +#scatter!(res_fiji_y .- pos_arr[:, 1], res_fiji_x .- pos_arr[:, 2], markershape= :rect, markersize=2.0, alpha=0.2, label="ThunderSTORM fitting") + +println((std(pos_res[:, 1] .- pos_arr[:, 1]), std(pos_res[:, 2] .- pos_arr[:, 2])), (mean(pos_res[:, 1] .- pos_arr[:, 1]), mean(pos_res[:, 2] .- pos_arr[:, 2]))); +#println((std(res_fiji_y .- pos_arr[:, 1]), std(res_fiji_x .- pos_arr[:, 2])), (mean(res_fiji_y .- pos_arr[:, 1]), mean(res_fiji_x .- pos_arr[:, 2]))); + +#anim = @animate for i1 in 1:length(Optim.x_trace(res)) +# +# begin +# p00 = heatmap(sample_data, aspect_ratio=1.0, clim=(0.0, 1.0), title="Sample data", legend = :none); +# p01 = heatmap(fitting_data, aspect_ratio=1.0, clim=(0.0,1.0), title="Fitting data", legend = :none); +# p02 = heatmap(f_general(Optim.x_trace(res)[i1]), aspect_ratio=1.0, clim=(0.0,1.0), title="estimated fit", legend = :none); +# p03 = heatmap(fitting_data .- f_general(Optim.x_trace(res)[i1]), aspect_ratio=1.0, clim=(0.0, 0.3), title="discrepancy", legend = :none); +# +# plot(p00, p01, p02, p03, layout=@layout([A B C D]), +# framestyle=nothing, showaxis=false, +# xticks=false, yticks=false, +# size=(1200, 500), +# plot_title="iteration: $(Int(i1))/$(length(Optim.x_trace(res))), +# estimation: $(Optim.x_trace(res)[i1]) +# true vals : $(true_vals)", +# plot_titlevspan=0.25 +# ) +# end +# +# +#end; +# +#gif(anim, "DataToFunctions.jl/examples/anim_general_generalized.mp4", fps=2) +# diff --git a/examples/PSF_fitting_new.jl b/examples/PSF_fitting_new.jl new file mode 100644 index 0000000..dc9012f --- /dev/null +++ b/examples/PSF_fitting_new.jl @@ -0,0 +1,17 @@ +using PointSpreadFunctions +using Plots + +λ_em = 0.5; NA = 1.4; n = 1.52 +λ_ex = 0.488 # only needed for some PointSpreadFunctions, such as confocal, ISM or TwoPhoton +pp = PSFParams(λ_em, NA, n; pol=pol_x) + +sz = (256, 256, 256) +sampling = (0.020,0.020,0.020) + +aberr_sp = Aberrations([Zernike_VerticalAstigmatism],[1.0]); sz=(256,256,256) +pp_sp = PSFParams(λ_em, NA, n; method=MethodPropagateIterative, aberrations= aberr_sp) +p_sp = psf(sz, pp_sp; sampling=sampling); + +psf_example = sum(p_sp, dims=3)[:,:,1] + +heatmap(p_sp[:, :, 128], aspect_ratio=1) \ No newline at end of file diff --git a/examples/Project.toml b/examples/Project.toml new file mode 100644 index 0000000..a5b381c --- /dev/null +++ b/examples/Project.toml @@ -0,0 +1,46 @@ +[deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +AllocationViewer = "e68bee48-ccd2-4290-be16-570446adc5cc" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" +CoordinateTransformations = "150eb455-5306-5404-9cee-2592286d6298" +DataToFunctions = "64cfdffa-4d02-49ee-ae8b-a805370874f5" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +EvalMultiPoly = "c78649ec-f9ed-405e-be6d-0472f43586aa" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +FourierTools = "b18b359b-aebc-45ac-a139-9c0ccbb2871e" +ImageShow = "4e3cecfd-b093-5904-9786-8bbb286a6a31" +Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0" +IndexFunArrays = "613c443e-d742-454e-bfc6-1d7f8dd76566" +Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" +InverseModeling = "ce844058-9528-415d-a63d-06f3dd08b29f" +LBFGSB = "5be7bae1-8223-5378-bac3-9e7378a2f6e6" +LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +NDTools = "98581153-e998-4eef-8d0d-5ec2c052313d" +NLSolversBase = "d41bc354-129a-5804-8e4c-c37616107c6c" +Noise = "81d43f40-5267-43b7-ae1c-8b967f377efa" +OhMyREPL = "5fb14364-9ced-5910-84b2-373655c76a03" +Optim = "429524aa-4258-5aef-a3af-852621145aeb" +PProf = "e4faabce-9ead-11e9-39d9-4379958e3056" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +Pluto = "c3e4b0f8-55cb-11ea-2926-15256bba5781" +PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8" +PointSpreadFunctions = "e8810a93-244e-46c5-8da3-35c5dd956001" +Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +ProgressBars = "49802e3a-d2f1-5c88-81d8-b72133a6f568" +RandomExtensions = "fb686558-2515-59ef-acaa-46db3789a887" +Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc" +SeparableFunctions = "c8c7ead4-852c-491e-a42d-3d43bc74259e" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +SyntheticObjects = "e7028c27-0967-45e9-8fdb-dbc10ccb2b0a" +TaylorSeries = "6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea" +TestImages = "5e47fb64-e119-507b-a336-dd2b206d9990" +TiffImages = "731e570b-9d59-4bfa-96dc-6df516fadf69" +TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" +View5D = "90d841e0-6953-4e90-9f3a-43681da8e949" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/examples/benchmark_transformators.jl b/examples/benchmark_transformators.jl new file mode 100644 index 0000000..c9fa81d --- /dev/null +++ b/examples/benchmark_transformators.jl @@ -0,0 +1,147 @@ +using BenchmarkTools +using DataToFunctions +using EvalMultiPoly +using ForwardDiff +using Zygote + +function benchmark_image(::Type{T}=Float64, n::Int=128) where {T} + return [ + sin(T(0.071) * i) + cos(T(0.053) * j) + T(0.001) * i * j + for i in 1:n, j in 1:n + ] +end + +n = 128 +data = benchmark_image(Float64, n) +interior = 4:n-3 + +println("="^72) +println("DataToFunctions transform benchmark") +println("image size: $(n)×$(n), Float64") +println("="^72) + +# ----------------------------------------------------------------------------- +# Affine +# ----------------------------------------------------------------------------- + +println("\nAFFINE") +println("-"^72) + +f_aff = get_function_affine(data; extrapolation_bc=0.0) +f_aff! = get_function_affine_inplace(data; extrapolation_bc=0.0) +p_aff = (0.25, -0.15, 1.01, 0.99, 0.002, -0.003, 0.01) +p_aff_vec = collect(p_aff) +out_aff = similar(data) + +# warmup +f_aff(p_aff) +f_aff!(out_aff, p_aff) + +println("Allocating value:") +@btime $f_aff($p_aff) +# 371.700 μs (3 allocations: 128.08 KiB) + +println("In-place value:") +@btime $f_aff!($out_aff, $p_aff) +# 148.600 μs (0 allocations: 0 bytes) + +println("In-place allocated bytes:") +@show @allocated f_aff!(out_aff, p_aff) +# 0 + +affine_loss(p) = begin + y = f_aff(p) + sum(abs2, y[interior, interior] .- data[interior, interior]) +end + +# Zygote uses the package rrule. ForwardDiff is an independent comparison. +Zygote.gradient(affine_loss, p_aff_vec) +ForwardDiff.gradient(affine_loss, p_aff_vec) + +println("Zygote gradient (custom rrule):") +@btime Zygote.gradient($affine_loss, $p_aff_vec)[1] +# 637.400 μs (110 allocations: 845.17 KiB) + +println("ForwardDiff gradient:") +@btime ForwardDiff.gradient($affine_loss, $p_aff_vec) +# 1.502 ms (22 allocations: 2.93 MiB) + +println("Zygote gradient allocated bytes:") +@show @allocated Zygote.gradient(affine_loss, p_aff_vec) +# 865537 + +println("ForwardDiff gradient allocated bytes:") +@show @allocated ForwardDiff.gradient(affine_loss, p_aff_vec) +# 3074556 + +# ----------------------------------------------------------------------------- +# Polynomial +# ----------------------------------------------------------------------------- + +println("\nPOLYNOMIAL (2-D, order 2)") +println("-"^72) + +f_poly = get_function_poly(data, Val(2); extrapolation_bc=0.0) +f_poly! = get_function_poly_inplace(data, Val(2); extrapolation_bc=0.0) + +c0 = map(Float64, EvalMultiPoly.get_identity_multipoly_coeffs(Val(2), Val(2))) +M = length(c0) +c_poly = ntuple(k -> c0[k] + 1e-6 * k, M) +c_poly_vec = collect(c_poly) +out_poly = similar(data) + +# warmup +f_poly(c_poly) +f_poly!(out_poly, c_poly) + +println("Allocating value:") +@btime $f_poly($c_poly) +# 275.500 μs (3 allocations: 128.08 KiB) + +println("In-place value:") +@btime $f_poly!($out_poly, $c_poly) +# 275.200 μs (0 allocations: 0 bytes) + +println("In-place allocated bytes:") +@show @allocated f_poly!(out_poly, c_poly) +# 0 + +poly_loss(c) = begin + y = f_poly(c) + sum(abs2, y[interior, interior] .- data[interior, interior]) +end + +Zygote.gradient(poly_loss, c_poly_vec) +ForwardDiff.gradient(poly_loss, c_poly_vec) + +println("Zygote gradient (custom rrule):") +@btime Zygote.gradient($poly_loss, $c_poly_vec)[1] +# 3.091 ms (123 allocations: 848.67 KiB) + +println("ForwardDiff gradient:") +@btime ForwardDiff.gradient($poly_loss, $c_poly_vec) +# 2.159 ms (40 allocations: 4.70 MiB) + +println("Zygote gradient allocated bytes:") +@show @allocated Zygote.gradient(poly_loss, c_poly_vec) +# 869137 + +println("ForwardDiff gradient allocated bytes:") +@show @allocated ForwardDiff.gradient(poly_loss, c_poly_vec) +# 4927884 +# ----------------------------------------------------------------------------- +# Numerical gradient agreement +# ----------------------------------------------------------------------------- + +println("\nGRADIENT AGREEMENT") +println("-"^72) + +g_aff_zyg = Zygote.gradient(affine_loss, p_aff_vec)[1] +g_aff_fwd = ForwardDiff.gradient(affine_loss, p_aff_vec) +println("Affine max |Zygote - ForwardDiff|: ", maximum(abs, g_aff_zyg .- g_aff_fwd)) +# Affine max |Zygote - ForwardDiff|: 8.640199666842818e-12 + +g_poly_zyg = Zygote.gradient(poly_loss, c_poly_vec)[1] +g_poly_fwd = ForwardDiff.gradient(poly_loss, c_poly_vec) +println("Polynomial max |Zygote - ForwardDiff|: ", maximum(abs, g_poly_zyg .- g_poly_fwd)) +# Polynomial max |Zygote - ForwardDiff|: 1.7462298274040222e-9 \ No newline at end of file diff --git a/examples/deform_testimage.jl b/examples/deform_testimage.jl new file mode 100644 index 0000000..3e96f79 --- /dev/null +++ b/examples/deform_testimage.jl @@ -0,0 +1,160 @@ +### A Pluto.jl notebook ### +# v0.19.42 + +using Markdown +using InteractiveUtils + +# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error). +macro bind(def, element) + quote + local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end + local el = $(esc(element)) + global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el) + el + end +end + +# ╔═╡ 28975586-853e-4e19-b9eb-65c41fa61a43 +using Pkg + +# ╔═╡ 0ae2da4f-3f75-47bb-a899-9e89c5c3f17c +Pkg.activate(".") + +# ╔═╡ 4af0c13d-fc42-4fe7-97e6-2248e36b63e2 +Pkg.add("PlutoUI") + +# ╔═╡ a2d75cfb-feab-4130-8439-30c543618d04 +using DataToFunctions, ImageShow, TestImages, PlutoUI, Images + +# ╔═╡ 1c744bec-f085-4812-ab1a-40a32c2ac176 +import PlutoUI: combine + +# ╔═╡ 5ac1123d-5df3-4c9d-aff1-ffe91d931497 +data = Float32.(testimage("resolution_test_512")) + +# ╔═╡ b0c8d15e-1bd3-4e66-b2b9-57885636eb48 +simshow(data) + +# ╔═╡ 2fce0208-732f-4259-a47d-7f78921bfd87 +f = get_interpolated_function(data, AffineMode, super_sampling=1) + +# ╔═╡ dd535378-3f2a-4429-914c-8b7608b99706 +@bind shift_x Slider(-100:0.02:100, default=0) + +# ╔═╡ 3e26d4e8-b5d4-4414-8936-379ec63cb4d2 +@bind shift_y Slider(-100:0.02:100, default=0) + +# ╔═╡ bc3f3659-aa70-4e7f-bef6-4d05229a03c4 +@bind zoom_x Slider(0.2:0.02:4, default=1) + +# ╔═╡ 39f11dc5-351e-4c6d-8fc4-468222e99976 +@bind zoom_y Slider(0.2:0.02:4, default=1) + +# ╔═╡ 64b64d3b-092f-4553-916d-f7db1fdfa428 +simshow(f((shift_x, shift_y, 1/zoom_x, 1/zoom_y, 0.0, 0.0, 0.0))) + +# ╔═╡ c3fe6b25-f4c0-4a80-9d15-9ee30136d43b +typeof(f((shift_x, shift_y, 1/zoom_x, 1/zoom_y, 0.0, 0.0, 0.0))) + +# ╔═╡ ff143f0d-b070-4220-8fa6-0b5a93a56303 +typeof(data) + +# ╔═╡ 6676ba35-3efd-49f5-9819-411ad8f8a95c +md""" +# Polynomial transformations +""" + +# ╔═╡ 13461a95-95ea-4bad-8673-e94e06776254 +md""" +## First order polynomial + +First order polynomial transformation which is as follows: + +``x^{\prime} = c_{1} + c_{2}{x}^{1} + c_{3}{y}^{1}`` + +``y^{\prime} = c_{4} + c_{5}{x}^{1} + c_{6}{y}^{1}`` +""" + +# ╔═╡ 87be45c0-8b2e-4d49-abd1-a274b3c1815e +h = get_interpolated_function(data, PolynomialMode, 1); + +# ╔═╡ 68f771aa-2cde-41cd-990c-9ec7dc2146a4 +function coeffs_input(coeffs::Vector) + + return combine() do Child + + inputs = [ + md""" $(name): $( + Child(name, Slider(-2f0:0.05f0:2f0, default=0, show_value=true)) + )""" + + for name in coeffs + ] + + md""" + #### Transform coefficients + $(inputs) + """ + end +end; + +# ╔═╡ 69331d73-75a3-4727-acda-e79779a2bd03 +@bind c coeffs_input(["c1", "c2", "c3", "c4", "c5", "c6"]) + +# ╔═╡ 74228a9d-6cc2-4aaf-97da-67f32670341e +simshow(h((c.c1, c.c2, c.c3, c.c4, c.c5, c.c6)), cmap=:turbo)#,0f0,0f0,0f0,0f0,0f0,0f0,0f0,0f0,0f0))) + +# ╔═╡ 687a198b-9020-4959-8e27-fd0896d4b1fc +maximum(h((c.c1,c.c2,c.c3,c.c4,c.c5,c.c6))) + +# ╔═╡ 6584aacc-440e-457b-bf52-83f8db40c999 +h((c.c1,c.c2,c.c3,c.c4,c.c5,c.c6)) + +# ╔═╡ 74cf9f27-f559-4fe2-9a30-20cb7cf6fe81 +md""" +## Second order polynomial + +Second order polynomial transformation is as follows: + +``x^{\prime} = c_{1} + c_{2}{x}^{1} + c_{3}{x}^{2} + c_{4}{y}^{1} + c_{5}{x}{y} + c_{6}{y}^{2}`` + +``y^{\prime} = c_{7} + c_{8}{x}^{1} + c_{9}{x}^{2} + c_{10}{y}^{1} + c_{11}{x}{y} + c_{12}{y}^{2}`` +""" + +# ╔═╡ 1f440592-9bbb-429c-8ba3-d018f5b354b0 +h2 = get_interpolated_function(data, PolynomialMode, 2); + +# ╔═╡ 12f59a6d-4600-4379-b536-f3b4701bdbe4 +@bind c2 coeffs_input(["c1", "c2", "c3", "c4", "c5", "c6", "c7", "c8", "c9", "c10", "c11", "c12"]) + +# ╔═╡ fd60df55-2a80-48a9-95ec-eeeb1cfe7491 +simshow(h2((c2.c1, c2.c2, c2.c3, c2.c4, c2.c5, c2.c6, c2.c7, c2.c8, c2.c9, c2.c10, c2.c11, c2.c12)), cmap=:turbo)#,0f0,0f0,0f0,0f0,0f0,0f0,0f0,0f0,0f0))) + +# ╔═╡ Cell order: +# ╠═28975586-853e-4e19-b9eb-65c41fa61a43 +# ╠═0ae2da4f-3f75-47bb-a899-9e89c5c3f17c +# ╟─4af0c13d-fc42-4fe7-97e6-2248e36b63e2 +# ╠═1c744bec-f085-4812-ab1a-40a32c2ac176 +# ╠═a2d75cfb-feab-4130-8439-30c543618d04 +# ╠═5ac1123d-5df3-4c9d-aff1-ffe91d931497 +# ╠═b0c8d15e-1bd3-4e66-b2b9-57885636eb48 +# ╠═2fce0208-732f-4259-a47d-7f78921bfd87 +# ╠═dd535378-3f2a-4429-914c-8b7608b99706 +# ╠═3e26d4e8-b5d4-4414-8936-379ec63cb4d2 +# ╠═bc3f3659-aa70-4e7f-bef6-4d05229a03c4 +# ╠═39f11dc5-351e-4c6d-8fc4-468222e99976 +# ╠═64b64d3b-092f-4553-916d-f7db1fdfa428 +# ╠═c3fe6b25-f4c0-4a80-9d15-9ee30136d43b +# ╠═ff143f0d-b070-4220-8fa6-0b5a93a56303 +# ╟─6676ba35-3efd-49f5-9819-411ad8f8a95c +# ╟─13461a95-95ea-4bad-8673-e94e06776254 +# ╠═87be45c0-8b2e-4d49-abd1-a274b3c1815e +# ╠═69331d73-75a3-4727-acda-e79779a2bd03 +# ╠═68f771aa-2cde-41cd-990c-9ec7dc2146a4 +# ╠═74228a9d-6cc2-4aaf-97da-67f32670341e +# ╠═687a198b-9020-4959-8e27-fd0896d4b1fc +# ╠═6584aacc-440e-457b-bf52-83f8db40c999 +# ╟─74cf9f27-f559-4fe2-9a30-20cb7cf6fe81 +# ╠═1f440592-9bbb-429c-8ba3-d018f5b354b0 +# ╠═12f59a6d-4600-4379-b536-f3b4701bdbe4 +# ╠═fd60df55-2a80-48a9-95ec-eeeb1cfe7491 diff --git a/examples/deform_testimgage backup 1.jl b/examples/deform_testimgage backup 1.jl new file mode 100644 index 0000000..2a37851 --- /dev/null +++ b/examples/deform_testimgage backup 1.jl @@ -0,0 +1,50 @@ +### A Pluto.jl notebook ### +# v0.19.43 + +using Markdown +using InteractiveUtils + +# ╔═╡ 28975586-853e-4e19-b9eb-65c41fa61a43 +using Pkg + +# ╔═╡ a2d75cfb-feab-4130-8439-30c543618d04 +using DataToFunctions, ImageShow, TestImages + +# ╔═╡ 0ae2da4f-3f75-47bb-a899-9e89c5c3f17c +# Pkg.activate(".") + +# ╔═╡ 4af0c13d-fc42-4fe7-97e6-2248e36b63e2 +# Pkg.add("ImageShow") + +# ╔═╡ 5ac1123d-5df3-4c9d-aff1-ffe91d931497 +data = testimage("resolution_test_512", super_sampling=1) + +# ╔═╡ b0c8d15e-1bd3-4e66-b2b9-57885636eb48 + + +# ╔═╡ 2fce0208-732f-4259-a47d-7f78921bfd87 +f = get_function(data) + +# ╔═╡ 1bed34fb-b29a-4042-a493-4835fdb69a9d + + +# ╔═╡ c287fa80-426b-11ef-125e-5fda207e605c +# ╠═╡ disabled = true +#=╠═╡ +using DataToFunctions + ╠═╡ =# + +# ╔═╡ 87be45c0-8b2e-4d49-abd1-a274b3c1815e + + +# ╔═╡ Cell order: +# ╠═28975586-853e-4e19-b9eb-65c41fa61a43 +# ╠═0ae2da4f-3f75-47bb-a899-9e89c5c3f17c +# ╠═4af0c13d-fc42-4fe7-97e6-2248e36b63e2 +# ╠═a2d75cfb-feab-4130-8439-30c543618d04 +# ╠═5ac1123d-5df3-4c9d-aff1-ffe91d931497 +# ╠═b0c8d15e-1bd3-4e66-b2b9-57885636eb48 +# ╠═2fce0208-732f-4259-a47d-7f78921bfd87 +# ╠═1bed34fb-b29a-4042-a493-4835fdb69a9d +# ╠═c287fa80-426b-11ef-125e-5fda207e605c +# ╠═87be45c0-8b2e-4d49-abd1-a274b3c1815e diff --git a/examples/deformed_fitting.jl b/examples/deformed_fitting.jl new file mode 100644 index 0000000..cdd35d2 --- /dev/null +++ b/examples/deformed_fitting.jl @@ -0,0 +1,42 @@ +using Interpolations +using FourierTools + + +function get_function(data::AbstractArray; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) + new_size = super_sampling.*size(data) + upsampled = fftshift(resample(ifftshift(data), new_size)) + @show upsampled + # return upsampled + # itp = LinearInterpolation(axes(upsampled), upsampled, extrapolation_bc=extrapolation_bc); + interpolation = Interpolations.interpolate(upsampled, interp_type) + interpolation = extrapolate(interpolation, extrapolation_bc) + # center of the original data (too keep the axis and number of datapointsi dentical to the original) + center_orig = (size(data) .÷2 .+1) + # create zero-centered original ranges (== axes) + zero_axes = Tuple(ax .- c for (ax, c) in zip(axes(data), center_orig)) + # center of the upsampled data. This is where to access the upsampled data + function zoomed(shift, zoom, theta) + zoom = zoom .* super_sampling + # careful: The center of the original data is not at the expected position! But rather at: + center_upsamp = new_size .÷2 .+1 # ((center_orig .-1) .*super_sampling .+1) # new_size .÷2 .+1 + scaled_axes = ((ax.-myc) .* z .+ cen for (ax, myc, cen, z) in zip(zero_axes, shift, center_upsamp, zoom)) + # @show Tuple(scaled_axes) + return interpolation[scaled_axes...] + # return extrapolate(scale(interpolation, scaled_axes...), extrapolation_bc) + end + zoomed(p) = zoomed([p[1], p[2]], [p[3], p[4]], [p[5]]) + + return zoomed + + # return (pos) -> interp_linear((center .+ pos)...) + # fitp(t) = interp_linear(t...) + # @time res1 = fitp.(tcoords); # 1 sec + # function my_zoom + +end + + +f_d = get_function(sample_data_d; super_sampling=2, extrapolation_bc=0.0) + + +f_d(true_vals); \ No newline at end of file diff --git a/examples/example_transformators.jl b/examples/example_transformators.jl new file mode 100644 index 0000000..a2c988a --- /dev/null +++ b/examples/example_transformators.jl @@ -0,0 +1,171 @@ +using DataToFunctions +using EvalMultiPoly +using Optim +using NLSolversBase +using Zygote + +using View5D +using TestImages + +# ----------------------------------------------------------------------------- +# Example data +# ----------------------------------------------------------------------------- + +# n = 64 +# data = [ +# sin(0.071 * i) + cos(0.053 * j) + 0.001 * i * j +# for i in 1:n, j in 1:n +# ] +data = Float32.(testimage("resolution_test_512.tif")); +n = size(data, 1) +# ----------------------------------------------------------------------------- +# 1. Affine transform +# ----------------------------------------------------------------------------- + +f_affine = get_function_affine(data; extrapolation_bc=0.0) +f_affine! = get_function_affine_inplace(data; extrapolation_bc=0.0) + +# (shift_x, shift_y, scale_x, scale_y, shear_xy, shear_yx, rotation) +p_affine = (0.25, -0.15, 1.01, 0.99, 0.002, -0.003, 0.01) + +warped_affine = f_affine(p_affine) + +out_affine = similar(data) +f_affine!(out_affine, p_affine) +@assert warped_affine ≈ out_affine + +# Warm up before measuring allocations. +f_affine!(out_affine, p_affine) +println("Affine in-place allocated bytes: ", @allocated(f_affine!(out_affine, p_affine))) + +# ----------------------------------------------------------------------------- +# 2. Fast affine gradient through the custom ChainRules rule +# ----------------------------------------------------------------------------- + +p_affine_vec = collect(p_affine) +interior = 4:n-3 + +affine_loss(p) = begin + y = f_affine(p) + sum(abs2, y[interior, interior] .- data[interior, interior]) +end + +g_affine = Zygote.gradient(affine_loss, p_affine_vec)[1] +println("Affine gradient: ", g_affine) + +# ----------------------------------------------------------------------------- +# 3. Affine registration with Optim.LBFGS +# ----------------------------------------------------------------------------- + +p_true = [0.18, -0.12, 1.006, 0.994, 0.0015, -0.0110, 0.007] +target_affine = f_affine(p_true)[interior, interior] + +registration_loss(p) = begin + y = f_affine(p) + sum(abs2, y[interior, interior] .- target_affine) +end + +function registration_fg!(F, G, p) + if G === nothing + return registration_loss(p) + end + + value, back = Zygote.pullback(registration_loss, p) + gp = back(one(value))[1] + copyto!(G, gp) + return F === nothing ? nothing : value +end + +p0 = [0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0] + +result_affine = Optim.optimize( + NLSolversBase.only_fg!(registration_fg!), + p0, + Optim.LBFGS(), + Optim.Options(iterations=100, g_tol=1e-8, show_trace=true), +) + +println("Affine minimum: ", Optim.minimum(result_affine)) +println("Affine estimate: ", Optim.minimizer(result_affine)) +println("Affine true p: ", p_true) + +@vt data f_affine(p_true) f_affine(Optim.minimizer(result_affine)) +# ----------------------------------------------------------------------------- +# 4. Polynomial transform +# ----------------------------------------------------------------------------- + +order = Val(2) +f_poly = get_function_poly(data, order; extrapolation_bc=0.0) +f_poly! = get_function_poly_inplace(data, order; extrapolation_bc=0.0) + +c_identity = map( + Float64, + EvalMultiPoly.get_identity_multipoly_coeffs(Val(2), order), +) + +M = length(c_identity) +c_poly = ntuple(k -> c_identity[k] + 1e-6 * k, M) + +warped_poly = f_poly(c_poly) +out_poly = similar(data) +f_poly!(out_poly, c_poly) +@assert warped_poly ≈ out_poly + +f_poly!(out_poly, c_poly) # warmup +println("Polynomial in-place allocated bytes: ", @allocated(f_poly!(out_poly, c_poly))) + +# ----------------------------------------------------------------------------- +# 5. Polynomial gradient +# ----------------------------------------------------------------------------- + +c_poly_vec = collect(c_poly) + +poly_loss(c) = begin + y = f_poly(c) + sum(abs2, y[interior, interior] .- data[interior, interior]) +end + +g_poly = Zygote.gradient(poly_loss, c_poly_vec)[1] +println("Polynomial gradient length: ", length(g_poly)) +println("Polynomial gradient: ", g_poly) + +# ----------------------------------------------------------------------------- +# 6. Polynomial registration with Optim.LBFGS +# ----------------------------------------------------------------------------- + +# Generate a target with a small polynomial deformation and recover it locally. +c_true = collect(c_identity) +for k in eachindex(c_true) + c_true[k] += 1e-6 * k +end + +target_poly = f_poly(c_true)[interior, interior] + +poly_registration_loss(c) = begin + y = f_poly(c) + sum(abs2, y[interior, interior] .- target_poly) +end + +function poly_registration_fg!(F, G, c) + if G === nothing + return poly_registration_loss(c) + end + + value, back = Zygote.pullback(poly_registration_loss, c) + gc = back(one(value))[1] + copyto!(G, gc) + return F === nothing ? nothing : value +end + +c0 = collect(c_identity) +result_poly = Optim.optimize( + NLSolversBase.only_fg!(poly_registration_fg!), + c0, + Optim.LBFGS(), + Optim.Options(iterations=100, g_tol=1e-6, show_trace=true), +) + +println("Polynomial minimum: ", Optim.minimum(result_poly)) +println("Polynomial estimate: ", Optim.minimizer(result_poly)) + +@vt data f_poly(c_true) f_poly(Optim.minimizer(result_poly)) diff --git a/examples/perform_fitting.jl b/examples/perform_fitting.jl new file mode 100644 index 0000000..c12925f --- /dev/null +++ b/examples/perform_fitting.jl @@ -0,0 +1,117 @@ +using DataToFunctions +using Optim, StaticArrays, LinearAlgebra +using Zygote +using ForwardDiff, LineSearches, Plots, Printf + + +Base.show(io::IO, f::Float64) = @printf(io, "%.2f", f) + +true_vals = [1.0, 2.0, 1.0, 1.0] +init_x = [0.5, 1.5, 1.0, 1.0] + +sample_data = rand(11,12) + +f = get_function(sample_data; super_sampling=2, extrapolation_bc=0.0); +#f(p0::Vector{Float64}) = f([p0[1], p0[2], p0[3], p0[4]]) + +fitting_data = f(true_vals) .+ rand(11, 12)./10.0 +#f(p2[1], p2[2]) = f(p2::Vector{Tuple{Float64, Float64}}) +loss(p, z) = sum(abs2.(f(p, z) .- fitting_data)) +loss(p2::Vector{Tuple{Float64, Float64}}) = loss(p2[1], p2[2]) +loss(p3) = loss([p3[1], p3[2]], [p3[3], p3[4]]) + + + + + +#Zygote.forwarddiff(loss, init_x) + +ForwardDiff.gradient(loss, init_x)#true_vals .+ [0.0001, 0.0, 0.0, 0.0]) +#conf = ForwardDiff.GradientConfig(f, init_x, chunk::Chunk = Chunk(init_x)) + + +""" +BFGS(; alphaguess = Optim.LineSearches.InitialStatic(), + linesearch = Optim.LineSearches.HagerZhang(), + initial_invH = nothing, + initial_stepnorm = 0.001, + manifold = Optim.Flat() + ) + +GradientDescent(; alphaguess = 0.01, + linesearch = Optim.LineSearches.HagerZhang(), + P = nothing, + precondprep = (P, x) -> nothing +) +""" +lower = [-1*size(fitting_data)[1], -1*size(fitting_data)[2], 0.0, 0.0] +upper = [size(fitting_data)[1], size(fitting_data)[2], size(fitting_data)[1], size(fitting_data)[2]] +#initial_x = [2.0, 2.0] +# requires using LineSearches +inner_optimizer = LBFGS(; m=1, linesearch=LineSearches.BackTracking(order=2)) +res = optimize( + loss, + lower, upper, + init_x, + Fminbox(inner_optimizer), + Optim.Options(store_trace = true, extended_trace = true, iterations=500), + autodiff = :forward +) + + +p00 = heatmap(sample_data, aspect_ratio=1.0, clim=(0.0, 1.0), title="Sample data", legend = :none); +p01 = heatmap(fitting_data, aspect_ratio=1.0, clim=(0.0,1.0), title="Fitting data", legend = :none); +p02 = heatmap(f(Optim.minimizer(res)), aspect_ratio=1.0, clim=(0.0,1.0), title="estimated fit", legend = :none); +p03 = heatmap(fitting_data .- f(Optim.minimizer(res)), aspect_ratio=1.0, clim=(0.0,1.0), title="discrepancy", legend = :none); + +plot(p00, p01, p02, p03, layout=@layout([A B C D]), + framestyle=nothing, showaxis=false, + xticks=false, yticks=false, + size=(700, 300), + plot_title="True vals: $(true_vals)", plot_titlevspan=0.2 +) + +heatmap(fitting_data .- f(Optim.x_trace(res)[end]), aspect_ratio=1, clim=(0.0, 1.0)) + +""" + +res = optimize( + loss, init_x, + LBFGS(), + Optim.Options(store_trace=true, extended_trace=true, iterations=500), + autodiff = :forward + ) +""" + +trace = Optim.trace(res); +trace + + + + +Optim.minimizer(res) +Optim.f_trace(res) +Optim.x_trace(res) +Optim.converged(res) +Optim.g_norm_trace(res) +Optim.g_calls(res) + +loss(Optim.x_trace(res)[end]) +ForwardDiff.gradient(loss, Optim.x_trace(res)[end]) + + + +anim = @animate for i1 in 1:length(Optim.x_trace(res)) + + heatmap(fitting_data .- f(Optim.x_trace(res)[i1]), + aspect_ratio=1, + clim=(0.0, 1.0), + dpi=300 + ) + title!("iteration: $(Int(i1))/$(length(Optim.x_trace(res))), + estimation: $(Optim.x_trace(res)[i1]) + true vals : $(true_vals)") + +end; + +gif(anim, "anim1.mp4", fps=5) diff --git a/examples/perform_random_optim.jl b/examples/perform_random_optim.jl new file mode 100644 index 0000000..b7a1b69 --- /dev/null +++ b/examples/perform_random_optim.jl @@ -0,0 +1,176 @@ +using DataToFunctions +using Optim, StaticArrays, LinearAlgebra +using Zygote +using ForwardDiff, LineSearches, Plots, Printf +using ProfileView, Profile +using Random, Distributions + +Random.seed!(123) + +# to show all the numbers in 2 decimals format +Base.show(io::IO, f::Float64) = @printf(io, "%.2f", f) + +### +# creating the random matrix in size = (11, 12) +sample_data = rand(11, 12) + +# creating the lower and upper bounds of the fitting variables (shift and scaling) +# shift can not be higher than the size of the array, +# scale can not be lower than zero or higher than size of the data, the latter causes the resulting array to be just one pixel +lower = [-1*size(sample_data)[1], -1*size(sample_data)[2], 0.0, 0.0] +upper = [size(sample_data)[1], size(sample_data)[2], size(sample_data)[1], size(sample_data)[2]] + +# preparing the function to fit +f = get_function(sample_data; super_sampling=2, extrapolation_bc=0.0); + +# assigning a scale (multiplier) to the range of true values, wedo not want that the parameters of +# function to be near the limits and cause strange behavior of the optimization +scale_range = 4.0 + +# true values of the fitting are random for the repeatibility +true_vals = (rand(4) .* (upper .- lower)/scale_range ) .+ (lower/scale_range) #[3.0, 5.5, 1.85, 0.6] + +# create the fitting data and adding random noise with scale of 1/10 +d = Normal() +noise = rand(d, 11, 12) + + +fitting_data = f(true_vals) #.+ noise / 10.0 +heatmap(fitting_data, aspect_ratio=1.0) +savefig("fitting_data_without_noise.png") + +# create the loss function and using the Julia's multiple dispatch +# because the gradients function is required just one input (can be vector) +loss(p, z) = sum(abs2.(f(p, z) .- fitting_data)) +loss(p2::Vector{Tuple{Float64, Float64}}) = loss(p2[1], p2[2]) +loss(p3) = loss([p3[1], p3[2]], [p3[3], p3[4]]) + + + + +# initialization of the LBFGS optimizer +inner_optimizer = LBFGS(; m=1, linesearch=LineSearches.BackTracking(order=2)) + + + +function perform_optim_mthr(loss, n_walkers) + """ + defining a function to survey the parameter space to neglect the local + minima and find the global maxima + + loss: the loss function + n_walkers: number of random initial parameter estimation + """ + # allocating the estimation array: + # consists of four parameters [1:4] and the minimum loss function of them [5] + est_m = zeros(n_walkers, 5) + + #x_tr = Array{Any} + #res = Array{Optim.MultivariateOptimizationResults{}}(undef, n_walkers, 1) + #res = Optim.MultivariateOptimizationResults{} + + # defining the random initial parameter values + walkers = (rand(4, n_walkers) .* (upper .- lower)/scale_range ) .+ (lower/scale_range) + + # main loop to do the optimization for each of the initial parameter values + # it uses the Threads to distribute the for loop to each thread + # note that in the settings.json the Julia is started with 16 threads + Threads.@threads for i in 1:size(walkers)[2] + res = optimize( + loss, + lower, upper, # the limits (simple box constraints) + walkers[:, i], + Fminbox(inner_optimizer), # assigning the limits of fitting (simple constraints) along with the LBFGS + Optim.Options(store_trace = true, extended_trace = true, iterations=500), + autodiff = :forward + ); + #x_tr = Optim.x_trace(res) + + # saving each 4 parameters of the fit and the minimum loss function + est_m[i, 1:4] .= Optim.minimizer(res)#x_tr[end] + est_m[i, 5] = minimum(res)#loss(x_tr[end]) + end + # saving the parameters of the minimum of the loss function + ans_m = est_m[argmin(est_m[:, 5]), :] + return est_m, ans_m +end + +#@profile est_m, ans_m = perform_optim_mthr(loss, 10000) + + +# first time: 4.749537 seconds (31.85 M allocations: 3.468 GiB, 8.50% gc time, 88.55% compilation time) +# 2nd time: 0.645307 seconds (12.68 M allocations: 2.525 GiB, 56.30% gc time) +# changing the fitting_data (noise values) : 0.448227 seconds (9.90 M allocations: 1.984 GiB, 27.55% compilation time: 38% of which was recompilation) +@time est_m, ans_m = perform_optim_mthr(loss, 500); + + +println(string(true_vals) * "\n" * string(ans_m)) + +# [2.23, -1.16, 1.26, 2.59] +# [2.26, -1.15, 1.27, 2.59, 0.41] + + +p00 = heatmap(sample_data, aspect_ratio=1.0, clim=(0.0, 1.0), title="Sample data", legend = :none); +p01 = heatmap(fitting_data, aspect_ratio=1.0, clim=(0.0,1.0), title="Fitting data", legend = :none); +p02 = heatmap(f(ans_m[1:4]), aspect_ratio=1.0, clim=(0.0,1.0), title="estimated fit", legend = :none); +p03 = heatmap(fitting_data .- f(ans_m[1:4]), aspect_ratio=1.0, clim=(0.0,1.0), title="discrepancy", legend = :none); + +plot(p00, p01, p02, p03, layout=@layout([A B C D]), + framestyle=nothing, showaxis=false, + xticks=false, yticks=false, + size=(700, 300), + plot_title="True vals: $(true_vals)", plot_titlevspan=0.2 +) +savefig("Output_mth.png") + + +plot( + heatmap(fitting_data .- f(ans_m[1:4]), aspect_ratio=1.0, clim=(0.0, 0.3), title="discrepancy", legend = :none), + heatmap(noise / 10.0 , aspect_ratio=1.0, clim=(0.0, 0.3), title="noise", legend = :none), + layout=@layout([A B]), + framestyle=nothing, showaxis=false, + xticks=false, yticks=false, + size=(700, 300), + plot_title="Noise comparison", plot_titlevspan=0.2 +) +savefig("Discrepancy.png") + +p12 = plot(fitting_data[1, :], color="black", legend=:none); +t = fitting_data[1, :] +for i in 2:size(fitting_data)[1] + t .+= fitting_data[i, :] + plot!(p12, fitting_data[i, :], color="black", legend=:none) +end +display(p12) + +plot(t) + + +function perform_optim_sthr(loss, n_walkers=1000) + est = Array{Float64, 2}(undef, n_walkers, 5) + + walkers = (rand(4, n_walkers) .* (upper .- lower)/scale_range ) .+ (lower/scale_range) + for i in 1:size(walkers)[2] + res = optimize( + loss, + lower, upper, + walkers[:, i], + Fminbox(inner_optimizer), + Optim.Options(store_trace = true, extended_trace = true, iterations=500), + autodiff = :forward + ); + x_tr = Optim.x_trace(res) + est[i, 1:4] .= x_tr[end] + est[i, 5] = loss(x_tr[end]) + end + ans1 = est[argmin(est[:, 5]), :] + return est, ans1 +end + + +@time est_s, ans_s = perform_optim_sthr(loss, 20000); + +println(string(true_vals) * "\n" * string(ans_s)) +# [1.07, -2.63, 0.78, 1.65] +# [1.02, -2.60, 0.77, 1.66, 0.46] + diff --git a/examples/poly_test_exp.jl b/examples/poly_test_exp.jl new file mode 100644 index 0000000..202a7b3 --- /dev/null +++ b/examples/poly_test_exp.jl @@ -0,0 +1,172 @@ +using Images +using DataToFunctions +using FindShift +using Optim, CUDA +using FourierTools +using View5D, Plots, Statistics + +CUDA.allowscalar(false) + +file1 = raw"D:\Hossein\Programming\Julia\DataToFunctions.jl\examples\test_polim1.jpeg" +file2 = raw"D:\Hossein\Programming\Julia\DataToFunctions.jl\examples\test_polim2.jpeg" + +c1 = Float32.(Gray.(load(file1))) +c2 = Float32.(Gray.(load(file2))) + +#TODO increase the size of the images +img11 = c1[30:1519+30, 30:779+30] +img12 = c1[30:1519+30, 800:779+800] +img21 = c2[30:1519+30, 30:779+30] +img22 = c2[30:1519+30, 800:779+800] + +img11_c = CuArray(img11) +img12_c = CuArray(img12) +img21_c = CuArray(img21) +img22_c = CuArray(img22) + +function resize_img(data, scale) + new_size = size(data) .÷ scale + imresize(data, new_size...) +end + +resample_size = 2 +img1_resampled = img11[1:resample_size:end, 1:resample_size:end] #resize_img(img1, 1) +img2_resampled = img12[1:resample_size:end, 1:resample_size:end] #resize_img(img2, 1) + +#im11 = imfilter(img1_resampled, Kernel.gaussian(5)); +#im12 = imfilter(img2_resampled, Kernel.gaussian(5)); + +f = get_interpolated_function(img12_c, PolynomialMode, 2); +loss_p(p1::AbstractArray) = (sum(abs2.(f(Tuple(p1)) .- img11_c))); +loss_updated(p1::AbstractArray) = mapreduce(abs2, +, f(Tuple(p1)) .- img11_c); +#st_vals = [0.0, 1.00, 0.00, 0.0, 0.00, 1.0] #ones(Float64, 6)./10 +st_vals = Float32[0.0, 1.0, 0, 0, 0, 0, 0.0, 0, 0, 1.0, 0, 0]; +# Float64[9.0, 0, 0, 0, 0, 0, 1, 0, 0, 5.0, 0, 0, 1, 0, 0, 0, 0, 0] +# @vv f(Tuple(st_vals)) +# loss_m(st_vals) + +#a = f(Tuple(st_vals)) +# @time a = f(Tuple(st_vals)); +""" +function do_registeration_step(resample_size, gaussian_kernel_size, img1, img2, st_vals=[0f0, 1f0, 0f0, 0f0, 0f0, 1f0]) + img1_resampled = img1[1:resample_size:end, 1:resample_size:end] #resize_img(img1, 1) + img2_resampled = img2[1:resample_size:end, 1:resample_size:end] #resize_img(img2, 1) + + im1 = imfilter(img1_resampled, Kernel.gaussian(gaussian_kernel_size)); + im2 = imfilter(img2_resampled, Kernel.gaussian(gaussian_kernel_size)); + + im1_c = CuArray(im1) + im2_c = CuArray(im2) + + f = get_interpolated_function(im2_c, PolynomialMode, 1); + loss_updated(p1::AbstractArray) = mapreduce(abs2, +, f(Tuple(p1)) .- im1_c); + + #st_vals = Float32[0.0, 1.0, 0, 0, 0, 0, 0.0, 0, 0, 1.0, 0, 0]; + CUDA.@time res = optimize( + loss_updated, + st_vals, + BFGS(),#; initial_stepnorm = 1f-1),#; linesearch=LineSearches.BackTracking(order=2)), + autodiff = :forward + ) + #return [res.minimizer[1]*resample_size, 1.0, 0, 0, 0, 0, res.minimizer[7]*resample_size, 0, 0, 1.0, 0, 0] + return [res.minimizer[1]*resample_size, 1.0, 0, res.minimizer[4]*resample_size, 0, 1.0] +end + +res_step1 = do_registeration_step(10, 5, img11, img12) +@vt img11 get_interpolated_function(img12, PolynomialMode, 1)(Tuple(res_step1)) img12 +""" + + +""" +function g!(G, x) # (G, x) + G .= gradient(loss_updated, x)[1] +end +od = OnceDifferentiable(loss_updated, g!, st_vals) +""" +aligned_imgs = Array(copy(img11_c)) + + +CUDA.@time res = optimize( + loss_updated, + st_vals, + #Newton(), + BFGS(; initial_stepnorm = 1f-2),#; linesearch=LineSearches.BackTracking(order=2)), + #LBFGS(; linesearch=LineSearches.BackTracking(order=3)), + #lower, upper, + #init_x, + #Fminbox(inner_optimizer), + #Optim.Options(store_trace = true, extended_trace = true, iterations=5000), + autodiff = :forward +) +""" +10.271416 seconds (8.78 M CPU allocations: 589.752 MiB, 1.06% gc time) (2.25 k GPU allocations: 52.251 GiB, 0.30% memmgmt time) +* Status: success + +* Candidate solution +Final objective value: 2.966678e+03 + +* Found with +Algorithm: BFGS + +* Convergence measures +|x - x'| = 1.76e-05 ≰ 0.0e+00 +|x - x'|/|x'| = 3.62e-06 ≰ 0.0e+00 +|f(x) - f(x')| = 0.00e+00 ≤ 0.0e+00 +|f(x) - f(x')|/|f(x')| = 0.00e+00 ≤ 0.0e+00 +|g(x)| = 1.17e+05 ≰ 1.0e-01 + +* Work counters +Seconds run: 5 (vs limit Inf) +Iterations: 73 +f(x) calls: 450 +∇f(x) calls: 450 +""" +aligned_imgs = cat(aligned_imgs, Array(f(Tuple(Optim.minimizer(res)))), dims=3); + + +for img_t in [img21_c, img22_c] + f = get_interpolated_function(img_t, PolynomialMode, 2); + loss_p(p1::AbstractArray) = (sum(abs2.(f(Tuple(p1)) .- img11_c))); + loss_updated(p1::AbstractArray) = mapreduce(abs2, +, f(Tuple(p1)) .- img11_c); + #st_vals = Float64[0.0, 1.00, 0.00, size(img11)[2], 0.00, -1.0] #ones(Float64, 6)./10 + st_vals = Float32[0.0, 1.0, 0, 0, 0, 0, size(img11)[2], 0, 0, -1.0, 0, 0]; + # Float64[9.0, 0, 0, 0, 0, 0, 1, 0, 0, 5.0, 0, 0, 1, 0, 0, 0, 0, 0] + # @vv f(Tuple(st_vals)) + # loss_m(st_vals) + + """ + function g!(G, x) # (G, x) + G .= gradient(loss_p, x)[1] + end + od = OnceDifferentiable(loss_p, g!, st_vals) + """ + CUDA.@time res = optimize( + loss_updated, + st_vals, + #Newton(), + BFGS(; initial_stepnorm = 1f-2),#; linesearch=LineSearches.BackTracking(order=2)), + #LBFGS(),#; linesearch=LineSearches.BackTracking(order=3)), + #lower, upper, + #init_x, + #Fminbox(inner_optimizer), + #Optim.Options(store_trace = true, extended_trace = true, iterations=5000), + autodiff = :forward + ) + + + aligned_imgs = cat(aligned_imgs, Array(f(Tuple(Optim.minimizer(res)))), dims=3); +end + + +aligned_imgs = clamp.(aligned_imgs, 0, 1) +@vv aligned_imgs + +aligned_imgs = permutedims(aligned_imgs, [2, 1, 3]) +save("aligned_all_cuda_polyorder2.gif", Gray.(aligned_imgs)) + + + +raw_imgs = cat(img11, img12, img21, img22, dims=3) +raw_imgs = clamp.(raw_imgs, 0, 1) +raw_imgs = permutedims(raw_imgs, [2, 1, 3]) +save("raw_all_imgs.gif", Gray.(raw_imgs)) diff --git a/examples/poly_test_exp_new.jl b/examples/poly_test_exp_new.jl new file mode 100644 index 0000000..b5f8233 --- /dev/null +++ b/examples/poly_test_exp_new.jl @@ -0,0 +1,173 @@ +using Images +using DataToFunctions +using FindShift +using Optim, CUDA +using FourierTools, Zygote +using View5D, Plots, Statistics, LineSearches + +CUDA.allowscalar(false) + +file1 = raw"D:\Hossein\Programming\Julia\DataToFunctions.jl\examples\markerpen_C1_00001.tif" +file2 = raw"D:\Hossein\Programming\Julia\DataToFunctions.jl\examples\markerpen_C2_00001.tif" + +c1 = Float32.(Gray.(load(file1))) +c2 = Float32.(Gray.(load(file2))) + +wide=true +if wide + img11 = c1[180:1990+180, 160:2050+160] + img12 = c1[160:1990+160, 2180:2050+2180] + img21 = c2[130:1990+130, 140:2050+140] + img22 = c2[130:1990+130, 2160:2050+2160] +end + +@vt img11 img12 img21 img22 + + + +img11_c = CuArray(img11./maximum(img11)) +img12_c = CuArray(img12./maximum(img12)) +img21_c = CuArray(img21./maximum(img21)) +img22_c = CuArray(img22./maximum(img22)) + + +resample_size = 2 +img1_resampled = img11[1:resample_size:end, 1:resample_size:end] #resize_img(img1, 1) +img2_resampled = img12[1:resample_size:end, 1:resample_size:end] #resize_img(img2, 1) + +#im11 = imfilter(img1_resampled, Kernel.gaussian(5)); +#im12 = imfilter(img2_resampled, Kernel.gaussian(5)); + + +""" +resample_size=50 +img1_resampled = img11[1:resample_size:end, 1:resample_size:end] #resize_img(img1, 1) +img2_resampled = img12[1:resample_size:end, 1:resample_size:end] #resize_img(img2, 1) + +im1 = imfilter(img1_resampled, Kernel.gaussian(3)); +im2 = imfilter(img2_resampled, Kernel.gaussian(3)); + +im1_c = CuArray(im1./maximum(im1)) +im2_c = CuArray(im2./maximum(im2)) + +f1 = get_interpolated_function(im2_c, PolynomialMode, 2); +loss_updated1(p1::AbstractArray) = mapreduce(abs2, +, f(Tuple(p1)) .- im1_c); + +st_vals = Float32[0.0, 1.0, 0, 0, 0, 0, 0.0, 0, 0, 1.0, 0, 0]; +CUDA.@time res1 = optimize( + loss_updated1, + st_vals, + BFGS(; initial_stepnorm = 1f-2),#; linesearch=LineSearches.BackTracking(order=2)), + autodiff = :forward +) + +@vt Array(im1_c) Array(f1(Tuple(res1.minimizer))) Array(im2_c) + +function do_registeration_step(resample_size, gaussian_kernel_size, img1, img2, st_vals=Float32[0.0, 1.0, 0, 0, 0, 0, 0.0, 0, 0, 1.0, 0, 0]) + img1_resampled = img1[1:resample_size:end, 1:resample_size:end] #resize_img(img1, 1) + img2_resampled = img2[1:resample_size:end, 1:resample_size:end] #resize_img(img2, 1) + + im1 = imfilter(img1_resampled, Kernel.gaussian(gaussian_kernel_size)); + im2 = imfilter(img2_resampled, Kernel.gaussian(gaussian_kernel_size)); + + im1_c = CuArray(im1) + im2_c = CuArray(im2) + + f = get_interpolated_function(im2_c, PolynomialMode, 2); + loss_updated(p1::AbstractArray) = mapreduce(abs2, +, f(Tuple(p1)) .- im1_c); + + #st_vals = Float32[0.0, 1.0, 0, 0, 0, 0, 0.0, 0, 0, 1.0, 0, 0]; + CUDA.@time res = optimize( + loss_updated, + st_vals, + BFGS(),#; initial_stepnorm = 1f-1),#; linesearch=LineSearches.BackTracking(order=2)), + autodiff = :forward + ) + return [res.minimizer[1]*resample_size, 1.0, 0, 0, 0, 0, res.minimizer[7]*resample_size, 0, 0, 1.0, 0, 0] + #return [res.minimizer[1]*resample_size, 1.0, 0, res.minimizer[4]*resample_size, 0, 1.0] +end + +res_step1 = do_registeration_step(10, 5, img11, img12) +@vt img11 get_interpolated_function(img12, PolynomialMode, 2)(Tuple(res_step1)) img12 +""" + + +order=2 +f = get_interpolated_function(img12_c, PolynomialMode, order); +loss_p(p1::AbstractArray) = (sum(abs2.(f(Tuple(p1)) .- img11_c))); +loss_updated(p1::AbstractArray) = mapreduce(abs2, +, f(Tuple(p1)) .- img11_c); +if order ==1 + st_vals = Float32[0.0, 1.00, 0.00, 0.0, 0.00, 1.0] #ones(Float64, 6)./10 +elseif order == 2 + st_vals = Float32[0.0, 1.0, 0, 0, 0, 0, 0.0, 0, 0, 1.0, 0, 0]; +end + + +aligned_imgs = Array(copy(img11_c)) + + +CUDA.@time res = optimize( + loss_updated, + st_vals, + #Newton(), + BFGS(; initial_stepnorm = 1f-1),#; linesearch=LineSearches.BackTracking(order=2)), + #LBFGS(; linesearch=LineSearches.BackTracking(order=3)), + #lower, upper, + #init_x, + #Fminbox(inner_optimizer), + Optim.Options(store_trace = true, extended_trace = true, iterations=5000, g_tol=1f-2), + autodiff = :forward +) + +open("markerpen_new_data.txt", "w") do f + write(f, "\norder $(order) params = $(res.minimizer)") +end + +aligned_imgs = cat(aligned_imgs, Array(f(Tuple(res.minimizer))), dims=3); +@vv aligned_imgs + + +for img_t in [img21_c, img22_c] + f = get_interpolated_function(img_t, PolynomialMode, 2); + loss_p(p1::AbstractArray) = (sum(abs2.(f(Tuple(p1)) .- img11_c))); + loss_updated(p1::AbstractArray) = mapreduce(abs2, +, f(Tuple(p1)) .- img11_c); + #st_vals = Float64[0.0, 1.00, 0.00, size(img11)[2], 0.00, -1.0] #ones(Float64, 6)./10 + #st_vals = Float32[0.0, 1.0, 0, 0, 0, 0, size(img11)[1], 0, 0, -1.0, 0, 0]; + if order ==1 + st_vals = Float32[0.0, 1.00, 0.00, 0.0, 0.00, 1.0] #ones(Float64, 6)./10 + elseif order == 2 + st_vals = Float32[0.0, 1.0, 0, 0, 0, 0, 0.0, 0, 0, 1.0, 0, 0]; + end + + CUDA.@time res = optimize( + loss_updated, + st_vals, + #Newton(), + BFGS(; initial_stepnorm = 1f-1),#; linesearch=LineSearches.BackTracking(order=2)), + #LBFGS(),#; linesearch=LineSearches.BackTracking(order=3)), + #lower, upper, + #init_x, + #Fminbox(inner_optimizer), + Optim.Options(store_trace = true, extended_trace = true, iterations=5000, g_tol=1f-2), + autodiff = :forward + ) + println(res) + open("markerpen_new_data.txt", "a") do f + write(f, "\norder $(order) params = $(res.minimizer)") + end + aligned_imgs = cat(aligned_imgs, Array(f(Tuple(Optim.minimizer(res)))), dims=3); +end + + +aligned_imgs = clamp.(aligned_imgs, 0, 1) +@vv aligned_imgs + +aligned_imgs = permutedims(aligned_imgs, [2, 1, 3]) +save("aligned_all_cuda_polyorder2_markerpen_new.gif", Gray.(aligned_imgs)) + + + +raw_imgs = cat(img11, img12, img21, img22, dims=3) +raw_imgs = clamp.(raw_imgs, 0, 1) +raw_imgs = permutedims(raw_imgs, [2, 1, 3]) +save("raw_all_imgs_markerpen_new.gif", Gray.(raw_imgs)) diff --git a/examples/polynomial_apply.jl b/examples/polynomial_apply.jl new file mode 100644 index 0000000..501502b --- /dev/null +++ b/examples/polynomial_apply.jl @@ -0,0 +1,93 @@ +using DataToFunctions +using Optim, StaticArrays, LinearAlgebra +using PointSpreadFunctions +using Zygote +using ForwardDiff, LineSearches, Plots, Printf +using View5D +using Distributions, Rotations +using Plots +using TestImages +using BenchmarkTools +#using InverseModeling +import Random +using Noise +using IndexFunArrays + + + +function poly_test(;sz=64, dtype=Float64, n_photons=1000) + + sz_psf = (sz, sz, 100) + sampling = (0.040, 0.040, 0.050) + # simulate a confocal PSF + #aberrations = Aberrations([Zernike_HorizontalComa,Zernike_Tip],[0.8,0.7]); + pp = PSFParams(0.5,1.4,1.52, method=MethodPropagateIterative);#, aberrations=aberrations); + + #pp_ex = PSFParams(pp_em; λ=0.488);#, method=MethodPropagateIterative, aplanatic=aplanatic_illumination, aberrations=aberrations); + p_psf_3d = psf(sz_psf, pp, sampling=sampling); + p_psf = p_psf_3d[:, :, 50] + #sample_data = p_psf ./ maximum(p_psf) + + # normalizing the sample data + sample_data = p_psf ./ maximum(p_psf) * n_photons + + #sample_data = make_grid(); + f_1 = get_interpolated_function(Float64.(sample_data), PolynomialMode, 1); # get_function_affine(sample_data); + true_vals = (1.6, 1.05, 0.1, 1.5, 0.01, 1.02) # dtype[2.0, 1.0, 1.01, 1.0, 0.0, 0.0, 0.0] + # true_vals = dtype[rand(-4.0:0.001:4.0), rand(-4.0:0.001:4.0), rand(0.5:0.001:1.5),rand(0.5:0.001:1.5), 0.0, 0.0, rand(0.001:0.001:pi/2.001)] + + dat2 = f_1(true_vals) + + #sample_data = dtype.(TestImages.shepp_logan(sz)) + f = get_interpolated_function(Float64.(sample_data), PolynomialMode, 1); + + + + loss_p(p1::AbstractArray) = (sum(abs2.(f(Tuple(p1)) .- dat2))) + st_vals = [2.1, 1.00, 0.00, 1.5, 0.00, 1.0] #ones(Float64, 6)./10 + #st_vals = Float64[1.0, 0, 0, 0, 0, 0, 1.0, 0, 0, 1.0, 0, 0, 1.0, 0, 0, 0, 0, 0] + # Float64[9.0, 0, 0, 0, 0, 0, 1, 0, 0, 5.0, 0, 0, 1, 0, 0, 0, 0, 0] + # @vv f(Tuple(st_vals)) + # loss_m(st_vals) + + + function g!(G, x) # (G, x) + G .= gradient(loss_p, x)[1] + end + od = OnceDifferentiable(loss_p, g!, st_vals) + res = optimize( + loss_p, + st_vals, + #Newton(), + BFGS(; initial_stepnorm = 1e-2),#; linesearch=LineSearches.BackTracking(order=2)), + #LBFGS(),#; linesearch=LineSearches.BackTracking(order=3)), + #lower, upper, + #init_x, + #Fminbox(inner_optimizer), + #Optim.Options(store_trace = true, extended_trace = true, iterations=5000, g_tol=1e-3), + autodiff = :forward + ) + + # return the estimated parameters + return true_vals, Optim.minimizer(res), res, f, f_1 +end + +#a, b, c = poly_test() +#@vt f(Tuple(b)) f_affine(a) + + +function make_grid!(arr::AbstractArray) + arr[isinteger.(xx(size(arr))./10) .|| isinteger.(yy(size(arr))./10)] .= 1.0 + return arr + +end +function make_grid(sz::NTuple{N, Int}=(64, 64)) where {N} + arr = zeros(Float64, sz) + make_grid!(arr) + return arr +end +function make_grid(::Type{T}, sz::NTuple{N, Int}=(64, 64)) where {N, T} + arr = zeros(T, sz) + make_grid!(arr) + return arr +end \ No newline at end of file diff --git a/examples/star_fitting.jl b/examples/star_fitting.jl new file mode 100644 index 0000000..0bd63c0 --- /dev/null +++ b/examples/star_fitting.jl @@ -0,0 +1,88 @@ +using DataToFunctions +using Optim, StaticArrays, LinearAlgebra +using Zygote +using ForwardDiff, LineSearches, Plots, Printf +using View5D +using Distributions +using Plots + +using Rotations +using CoordinateTransformations + +Base.show(io::IO, f::Float64) = @printf(io, "%.3f", f) + + +# size of the test array to fit +size_arr = 600.0 + +# defining the mean and the variance of the test normal (Gaussian) distribution +μ = [0, 0] +Σ = [size_arr/2.0 0.0; + 0.0 size_arr/2.0] + +Σ_d = [size_arr/0.5 1.5; + 1.5 size_arr/6.0] + +# initializing the multivariate normal distribution +p = MvNormal(μ, Σ) +p_d = MvNormal(μ, Σ_d) + + + +# this part of the code is to define the sample array based on a 2D normal distribution +X = -1*size_arr/2.0:1*size_arr/2.0 +Y = -1*size_arr/2.0:1*size_arr/2.0 + +z = [pdf(p, [x,y]) for y in Y, x in X] +z_d = [pdf(p_d, [x,y]) for y in Y, x in X] + +#@vv z + +heatmap(z, aspect_ratio=1) +heatmap(z_d, aspect_ratio=1) + +# setting a typical values for the shift (1:2) and scale (3:4) +true_vals = [1.1, -1.5, 0.75, 1.5] + +# normalizing the sample data +sample_data = z./maximum(z) +sample_data_d = z_d./maximum(z_d) + +# converting the data to function (DataToFunctions.get_function) +f = get_function(sample_data; super_sampling=2, extrapolation_bc=0.0); +f_d = get_function(sample_data_d; super_sampling=2, extrapolation_bc=0.0); + +# adding some scaled random noise to the fitting data +fitting_data = f(true_vals) .+ rand(size(z)...)./8.0 + +# @vv fitting_data + +# defining the loss function based on the gaussian noise +loss(p) = sum(abs2.(f(p) .- fitting_data)) + +#loss(p3) = loss([p3[1], p3[2]], [p3[3], p3[4]]) + +heatmap(fitting_data, aspect_ratio=1, title="Fitting data") + +# perform the main fit to the fitting data by minimizing the loss function +output = perform_fit(loss, fitting_data) + +# plotting the output of the fitting pocedure for further illustration +begin + p00 = heatmap(sample_data, aspect_ratio=1.0, clim=(0.0, 1.0), title="Sample data", legend = :none); + p01 = heatmap(fitting_data, aspect_ratio=1.0, clim=(0.0,1.0), title="Fitting data", legend = :none); + p02 = heatmap(f(output), aspect_ratio=1.0, clim=(0.0,1.0), title="estimated fit", legend = :none); + p03 = heatmap(fitting_data .- f(output), aspect_ratio=1.0, clim=(0.0, 0.3), title="discrepancy", legend = :none); + + plot(p00, p01, p02, p03, layout=@layout([A B C D]), + framestyle=nothing, showaxis=false, + xticks=false, yticks=false, + size=(700, 300), + plot_title="True vals: $(true_vals) + est vals: $(output)", + plot_titlevspan=0.25 + ) +end + +# comparing the true values to the best fitting parameters +println(string(true_vals) * "\n" * string(output)) diff --git a/examples/star_rotation.jl b/examples/star_rotation.jl new file mode 100644 index 0000000..5704ec0 --- /dev/null +++ b/examples/star_rotation.jl @@ -0,0 +1,90 @@ +using DataToFunctions +using Optim, StaticArrays, LinearAlgebra, FourierTools +using Zygote +using ForwardDiff, LineSearches, Plots, Printf +using View5D +using Distributions +using Plots + +using Rotations +using CoordinateTransformations + +using Interpolations + +function get_function2(data::AbstractArray; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) + new_size = super_sampling.*size(data) + upsampled = fftshift(resample(ifftshift(data), new_size)) + + # return upsampled + # itp = LinearInterpolation(axes(upsampled), upsampled, extrapolation_bc=extrapolation_bc); + interpolation = Interpolations.interpolate(upsampled, interp_type) + interpolation = extrapolate(interpolation, extrapolation_bc) + # center of the original data (too keep the axis and number of datapointsi dentical to the original) + center_orig = (size(data) .÷2 .+1) + # create zero-centered original ranges (== axes) + zero_axes = Tuple(ax .- c for (ax, c) in zip(axes(data), center_orig)) + # center of the upsampled data. This is where to access the upsampled data + cen_vec = SVector(size(data)./ 2.0) + + function transform_axes(t::SVector, mat::SMatrix, cen_vec::SVector, shift_vec::SVector) + + return (mat * (t - cen_vec)) + cen_vec - shift_vec + + end + + function zoomed(shift_vec, zoom, theta) + zoom = zoom .* super_sampling + # # careful: The center of the original data is not at the expected position! But rather at: + # center_upsamp = new_size .÷2 .+1 # ((center_orig .-1) .*super_sampling .+1) # new_size .÷2 .+1 + # scaled_axes = ((ax.-myc) .* z .+ cen for (ax, myc, cen, z) in zip(zero_axes, shift, center_upsamp, zoom)) + + # #sh_x, sh_y = 0.0, 0.0 + s_x, s_y = zoom + theta = theta[1] + + rot_mat = @SMatrix [cos(theta) -1.0*sin(theta); sin(theta) cos(theta)]; + scale_mat = @SMatrix [s_x 0.0; 0.0 s_y]; + mat = rot_mat * scale_mat + + return interpolation.(transform_axes.(SVector.(Tuple.(CartesianIndices(data))), mat, cen_vec, SVector(shift_vec[1], shift_vec[2]))...) + # return extrapolate(scale(interpolation, scaled_axes...), extrapolation_bc) + end + zoomed(p) = zoomed([p[1], p[2]], [p[3], p[4]], p[5]); + # zoomed([p[1], p[2]], [p[3], p[4]], p[5]) = zoomed[p] + return zoomed + +end + +# defining the mean and the varixance of the test normal (Gaussian) distribution +μ = [0, 0] +Σ = [2 0.0; + 0.0 2] + +Σ_d = [2 1.5; + 1.5 2] + +# initializing the multivariate normal distribution +p = MvNormal(μ, Σ) +p_d = MvNormal(μ, Σ_d) + +# size of the test array to fit +size_arr = 12.0 + +# this part of the code is to define the sample array based on a 2D normal distribution +X = -1*size_arr/2.0:1*size_arr/2.0 +Y = -1*size_arr/2.0:1*size_arr/2.0 + +z = [pdf(p, [x,y]) for y in Y, x in X] +heatmap(z, aspect_ratio=1) + + + + +true_vals = [1.1, -1.5, 0.75, 1.5, pi/2] + + +sample_data = z./maximum(z) +f = get_function2(sample_data; super_sampling=2, extrapolation_bc=0.0); +f(true_vals) + +heatmap(f(true_vals), aspect_ratio=1.0, clim=(0.0,1.0), legend = :none); diff --git a/examples/transform_cpu_cuda_ad.jl b/examples/transform_cpu_cuda_ad.jl new file mode 100644 index 0000000..ab962b9 --- /dev/null +++ b/examples/transform_cpu_cuda_ad.jl @@ -0,0 +1,330 @@ +using CUDA +using DataToFunctions +using Interpolations +using EvalMultiPoly +using BenchmarkTools +using Zygote +using ForwardDiff +using Optim + +# ----------------------------------------------------------------------------- +# Synthetic image +# ----------------------------------------------------------------------------- + +N = 256 + +data = Float32[ + sin(0.031f0 * i) + + cos(0.027f0 * j) + + 0.0002f0 * i * j + for i in 1:N, j in 1:N +] + +# ----------------------------------------------------------------------------- +# 1. Affine forward transform +# ----------------------------------------------------------------------------- + +affine = get_function_affine(data) +affine! = get_function_affine_inplace(data) + +p = ( + 0.25f0, # shift x + -0.15f0, # shift y + 1.01f0, # scale x + 0.99f0, # scale y + 0.002f0, # shear xy + -0.003f0, # shear yx + 0.01f0, # rotation [rad] +) + +warped = affine(p) + +out = similar(data) +affine!(out, p) + +@assert warped ≈ out + +# Warm up before measuring allocations. +affine!(out, p) + +println("Affine in-place CPU allocations:") +@show @allocated affine!(out, p) + +println("\nAffine forward benchmark:") +@btime $affine!($out, $p) + +# ----------------------------------------------------------------------------- +# 2. Affine gradient: custom Zygote rrule vs ForwardDiff +# ----------------------------------------------------------------------------- + +p_target = ( + 0.18f0, + -0.11f0, + 1.008f0, + 0.994f0, + 0.0015f0, + -0.002f0, + 0.007f0, +) + +target = affine(p_target) + +affine_loss(q) = + sum(abs2, affine(q) .- target) + +# Warm up. +g_zygote = Zygote.gradient(affine_loss, p)[1] + +println("\nAffine Zygote gradient:") +@show g_zygote + +println("\nAffine Zygote gradient benchmark:") +@btime Zygote.gradient($affine_loss, $p)[1] + +pvec = collect(p) +affine_loss_fd(q) = affine_loss(Tuple(q)) + +g_forward = ForwardDiff.gradient( + affine_loss_fd, + pvec, +) + +println("\nForwardDiff reference:") +@show g_forward +@show maximum(abs.(collect(g_zygote) .- g_forward)) + +# ----------------------------------------------------------------------------- +# 3. Optimize affine parameters with Optim.LBFGS +# ----------------------------------------------------------------------------- + +# Optim works with vectors. DataToFunctions preserves a vector gradient when +# the affine parameters are supplied as a vector. + +p0 = Float32[ + 0.0, + 0.0, + 1.0, + 1.0, + 0.0, + 0.0, + 0.0, +] + +loss_for_optim(q) = + sum(abs2, affine(q) .- target) + +function affine_gradient!(G, q) + G .= Zygote.gradient(loss_for_optim, q)[1] + return G +end + +result = Optim.optimize( + loss_for_optim, + affine_gradient!, + p0, + Optim.LBFGS(), + Optim.Options( + iterations=50, + show_trace=true, + ), +) + +println("\nLBFGS affine fit:") +@show Optim.minimum(result) +@show Optim.minimizer(result) +@show collect(p_target) + +# ----------------------------------------------------------------------------- +# 4. Polynomial transform +# ----------------------------------------------------------------------------- + +order = Val(2) + +poly = get_function_poly( + data, + order, +) + +poly! = get_function_poly_inplace( + data, + order, +) + +c_identity = get_identity_multipoly_coeffs( + Val(2), + Val(2), +) + +# Very small deformation away from the identity. +coeffs = ntuple( + k -> Float32(c_identity[k] + 1e-6 * k), + length(c_identity), +) + +poly_warped = poly(coeffs) + +poly_out = similar(data) +poly!(poly_out, coeffs) + +@assert poly_warped ≈ poly_out + +poly!(poly_out, coeffs) + +println("\nPolynomial in-place CPU allocations:") +@show @allocated poly!(poly_out, coeffs) + +println("\nPolynomial forward benchmark:") +@btime $poly!($poly_out, $coeffs) + +poly_target_coeffs = ntuple( + k -> Float32(c_identity[k] - 7e-7 * k), + length(c_identity), +) + +poly_target = poly(poly_target_coeffs) + +poly_loss(c) = + sum(abs2, poly(c) .- poly_target) + +g_poly = Zygote.gradient( + poly_loss, + coeffs, +)[1] + +println("\nPolynomial Zygote gradient:") +@show g_poly + +println("\nPolynomial Zygote gradient benchmark:") +@btime Zygote.gradient($poly_loss, $coeffs)[1] + +coeffs_vec = collect(coeffs) + +poly_loss_fd(c) = + poly_loss(Tuple(c)) + +g_poly_forward = ForwardDiff.gradient( + poly_loss_fd, + coeffs_vec, +) + +println("\nPolynomial ForwardDiff agreement:") +@show maximum( + abs.( + collect(g_poly) .- + g_poly_forward + ), +) + +# ----------------------------------------------------------------------------- +# 5. Polynomial optimization with Optim.LBFGS +# ----------------------------------------------------------------------------- + +# Optimize a Vector so Optim can update it in place. +poly_c0 = collect(c_identity) + +function poly_loss_optim(c) + return sum( + abs2, + poly(c) .- poly_target, + ) +end + +function poly_gradient!(G, c) + G .= Zygote.gradient( + poly_loss_optim, + c, + )[1] + return G +end + +poly_result = Optim.optimize( + poly_loss_optim, + poly_gradient!, + poly_c0, + Optim.LBFGS(), + Optim.Options( + iterations=30, + show_trace=true, + ), +) + +println("\nLBFGS polynomial fit:") +@show Optim.minimum(poly_result) +@show Optim.minimizer(poly_result) + +# ----------------------------------------------------------------------------- +# 6. Optional CUDA execution +# +# DataToFunctions itself does not import CUDA. If `data` is a CuArray, the +# transform is expressed with GPU-compatible broadcast/reduction abstractions. +# +# The current GPU path is intended first for BSpline(Linear()), which is the +# default and does not require a spline-prefilter solve. +# ----------------------------------------------------------------------------- + + +CUDA.allowscalar(false) + +data_gpu = CuArray(data) + +affine_gpu = get_function_affine(data_gpu) +affine_gpu! = get_function_affine_inplace(data_gpu) + +y_gpu = affine_gpu(p) +out_gpu = similar(data_gpu) +affine_gpu!(out_gpu, p) + +@assert Array(y_gpu) ≈ warped rtol=2e-4 atol=2e-4 +@assert Array(out_gpu) ≈ warped rtol=2e-4 atol=2e-4 + +CUDA.synchronize() + +println("\nCUDA affine forward benchmark:") +@btime begin + $affine_gpu!($out_gpu, $p) + CUDA.synchronize() +end + +# Zygote reaches the custom VJP. The VJP is expressed as a reduction +# over backend-resident Cartesian indices and therefore does not use +# scalar host indexing. +target_gpu = affine_gpu(p_target) + +affine_loss_gpu(q) = + sum(abs2, affine_gpu(q) .- target_gpu) + +g_gpu = Zygote.gradient( + affine_loss_gpu, + p, +)[1] + +println("\nCUDA affine gradient:") +@show g_gpu + +poly_gpu = get_function_poly( + data_gpu, + order, +) + +poly_gpu! = get_function_poly_inplace( + data_gpu, + order, +) + +poly_out_gpu = similar(data_gpu) +poly_gpu!(poly_out_gpu, coeffs) + +@assert Array(poly_out_gpu) ≈ poly_warped rtol=2e-4 atol=2e-4 + +poly_target_gpu = poly_gpu(poly_target_coeffs) + +poly_loss_gpu(c) = + sum(abs2, poly_gpu(c) .- poly_target_gpu) + +g_poly_gpu = Zygote.gradient( + poly_loss_gpu, + coeffs, +)[1] + +println("\nCUDA polynomial gradient:") +@show g_poly_gpu diff --git a/src/DataToFunctions.jl b/src/DataToFunctions.jl index fd58afb..27be1cd 100644 --- a/src/DataToFunctions.jl +++ b/src/DataToFunctions.jl @@ -1,55 +1,13 @@ module DataToFunctions -using Interpolations -using FourierTools -export get_function -""" - get_function(data::AbstractArray; super_sampling=2, extrapolation_bc=Flat(), interp_type=Interpolations.BSpline(Linear())) +# export gpu_or_cpu -returns a function `dat(shift, zoom)` which generates a shifted and scaled version of the original data. -This is useful for fitting with a function which is itself defined by measured data. +# deprecated: +# to include CUDA +# include("requires.jl") -# Arguments -`data`: The data to represent by the function `dat` -`super_sampling`: The factor by which the data is internally represented as a supersampled version (Fourier-based upsampling, see `FourierTools.resample`) -`extrapolation_bc`: The extrapolation boundary condition to select for values outside the range. - By default the value 0.0 is used. Other options are `Flat()`, or `Line()`, See the package `Interpolation` for details. -`interp_type`: The type of interpolation to use. See the package `Interpolation` for details. - -# Example -```jldoctest -``` -""" -function get_function(data::AbstractArray; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) - new_size = super_sampling.*size(data) - upsampled = fftshift(resample(ifftshift(data), new_size)) - @show upsampled - # return upsampled - # itp = LinearInterpolation(axes(upsampled), upsampled, extrapolation_bc=extrapolation_bc); - interpolation = Interpolations.interpolate(upsampled, interp_type) - interpolation = extrapolate(interpolation, extrapolation_bc) - # center of the original data (too keep the axis and number of datapointsi dentical to the original) - center_orig = (size(data) .÷2 .+1) - # create zero-centered original ranges (== axes) - zero_axes = Tuple(ax .- c for (ax, c) in zip(axes(data), center_orig)) - # center of the upsampled data. This is where to access the upsampled data - function zoomed(shift, zoom) - zoom = zoom .* super_sampling - # careful: The center of the original data is not at the expected position! But rather at: - center_upsamp = new_size .÷2 .+1 # ((center_orig .-1) .*super_sampling .+1) # new_size .÷2 .+1 - scaled_axes = ((ax.-myc) .* z .+ cen for (ax, myc, cen, z) in zip(zero_axes, shift, center_upsamp, zoom)) - # @show Tuple(scaled_axes) - return interpolation[scaled_axes...] - # return extrapolate(scale(interpolation, scaled_axes...), extrapolation_bc) - end - - return zoomed - - # return (pos) -> interp_linear((center .+ pos)...) - # fitp(t) = interp_linear(t...) - # @time res1 = fitp.(tcoords); # 1 sec - # function my_zoom - -end - -end # module DataToFunctions +#include("datafunction.jl") +include("transformation_types.jl") +include("transformators.jl") + +end # module DataToFunctions \ No newline at end of file diff --git a/src/datafunction.jl b/src/datafunction.jl new file mode 100644 index 0000000..d0063c1 --- /dev/null +++ b/src/datafunction.jl @@ -0,0 +1,34 @@ +import Base: getindex, setindex!, size, axes, eltype, copy, similar, ndims, iterate, length + + +struct DataFunctionAffine{affineparams{NTuple{N, Number}}, itp, S<:AbstractArray} <: AbstractArray + params::affineparams + interpolation::itp + data::S +end + +function DataFunctionAffine(params::affineparams, interpolation::itp) where {affineparams<:NTuple{N, Number}, itp} + return DataFunctionAffine{affineparams, itp}(params, interpolation) +end + +function getindex(dfa::DataFunctionAffine, I::Vararg{Number, N}) where {N} + return dfa.interpolation(Tuple(I)...) +end + +function copy(s::DataFunctionAffine) + res = similar(s) + res .= s +end + + +function similar(dfa::DataFunctionAffine, ::Type{T}=eltype(dfa)) where {T} + return DataFunctionAffine(dfa.params, similar(dfa.interpolation, T)) +end + +size(dfa::DataFunctionAffine) = size(dfa.data) +axes(dfa::DataFunctionAffine) = axes(dfa.data) +eltype(dfa::DataFunctionAffine) = eltype(dfa.data) +ndims(dfa::DataFunctionAffine) = ndims(dfa.data) +length(dfa::DataFunctionAffine) = length(dfa.data) + + diff --git a/src/transformation_types.jl b/src/transformation_types.jl new file mode 100644 index 0000000..85faf63 --- /dev/null +++ b/src/transformation_types.jl @@ -0,0 +1,21 @@ +abstract type transformation_method end + +""" + AffineMode + +Transformation mode selecting affine coordinate transformations. + +See also [`get_function_affine`](@ref) and +[`get_interpolated_function`](@ref). +""" +struct AffineMode <: transformation_method end + +""" + PolynomialMode + +Transformation mode selecting polynomial coordinate transformations. + +When using this mode with [`get_interpolated_function`](@ref), the polynomial +order must also be specified. +""" +struct PolynomialMode <: transformation_method end \ No newline at end of file diff --git a/src/transformators.jl b/src/transformators.jl new file mode 100644 index 0000000..a3a6f5e --- /dev/null +++ b/src/transformators.jl @@ -0,0 +1,1922 @@ +using Interpolations +using FourierTools +using EvalMultiPoly +using Adapt +import ChainRulesCore + +export get_interpolated_function +export get_function_tuple, get_function_tuple_inplace +export get_function_homogen, get_function_homogen_inplace, get_function_svec +export get_function_affine, get_function_affine_inplace +export get_function_poly, get_function_poly_inplace +# export apply_transform, apply_transform! +# export apply_transform_homogen, apply_transform_homogen! +# export apply_transform_affine, apply_transform_affine! +# export add_dim, red_dim_apply, red_dim, mat_mul +# export func_transform, func_transform_tup +export PolynomialMode, AffineMode + +# ----------------------------------------------------------------------------- +# Legacy implementation -- intentionally left unchanged +# ----------------------------------------------------------------------------- + +""" + get_function(data::AbstractArray; super_sampling=2, extrapolation_bc=Flat(), interp_type=Interpolations.BSpline(Linear())) + +returns a function `dat(shift, zoom)` which generates a shifted and scaled version of the original data. +This is useful for fitting with a function which is itself defined by measured data. + +# Arguments +`data`: The data to represent by the function `dat` +`super_sampling`: The factor by which the data is internally represented as a supersampled version (Fourier-based upsampling, see `FourierTools.resample`) +`extrapolation_bc`: The extrapolation boundary condition to select for values outside the range. + By default the value 0.0 is used. Other options are `Flat()`, or `Line()`, See the package `Interpolation` for details. +`interp_type`: The type of interpolation to use. See the package `Interpolation` for details. + + +""" +function get_function_old(data::AbstractArray; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) + new_size = super_sampling.*size(data) + upsampled = fftshift(resample(ifftshift(data), new_size)) + # @show upsampled + # return upsampled + # itp = LinearInterpolation(axes(upsampled), upsampled, extrapolation_bc=extrapolation_bc); + interpolation = Interpolations.interpolate(upsampled, interp_type) + interpolation = extrapolate(interpolation, extrapolation_bc) + # center of the original data (too keep the axis and number of datapointsi dentical to the original) + center_orig = (size(data) .÷2 .+1) + # create zero-centered original ranges (== axes) + zero_axes = Tuple(ax .- c for (ax, c) in zip(axes(data), center_orig)) + # center of the upsampled data. This is where to access the upsampled data + function zoomed(shift, zoom) + zoom = zoom .* super_sampling + # careful: The center of the original data is not at the expected position! But rather at: + center_upsamp = new_size .÷2 .+1 # ((center_orig .-1) .*super_sampling .+1) # new_size .÷2 .+1 + scaled_axes = ((ax.-myc) .* z .+ cen for (ax, myc, cen, z) in zip(zero_axes, shift, center_upsamp, zoom)) + # @show Tuple(scaled_axes) + return interpolation[scaled_axes...] + # return extrapolate(scale(interpolation, scaled_axes...), extrapolation_bc) + end + + return zoomed + + zoomed(p) = zoomed([p[1], p[2]], [p[3], p[4]]) + # return (pos) -> interp_linear((center .+ pos)...) + # fitp(t) = interp_linear(t...) + # @time res1 = fitp.(tcoords); # 1 sec + # function my_zoom + +end + +# ----------------------------------------------------------------------------- +# Internal utilities +# ----------------------------------------------------------------------------- + +# Interpolations.jl currently supports GPU *usage* of interpolants by first +# constructing the interpolant on CPU and then adapting its storage to the target +# backend. Constructing directly from a GPU array may leave CPU `Array` +# coefficients inside the interpolant, which cannot be passed to a GPU kernel. +# +# CPU path: no copy and no adaptation. +@inline function _make_interpolant( + data::Array, + interp_type, + extrapolation_bc, +) + return extrapolate( + interpolate(data, interp_type), + extrapolation_bc, + ) +end + +# Generic backend path (e.g. CuArray): +# backend array -> CPU values -> CPU interpolant -> backend-adapted interpolant +# +# `typeof(data)` is used as the Adapt target, so DataToFunctions does not need to +# import or depend on CUDA.jl. CUDA.jl registers the corresponding Adapt storage +# rule when the user loads CUDA. +function _make_interpolant( + data::AbstractArray, + interp_type, + extrapolation_bc, +) + cpu_data = Adapt.adapt(Array, data) + + cpu_itp = extrapolate( + interpolate(cpu_data, interp_type), + extrapolation_bc, + ) + + return Adapt.adapt( + typeof(data), + cpu_itp, + ) +end + +""" + add_dim(v::Tuple) + +Append a homogeneous coordinate equal to one to the coordinate tuple `v`. + +The tuple length remains part of the concrete type and the elements are allowed +to have heterogeneous scalar types. This is useful for allocation-free static +coordinate arithmetic and for automatic-differentiation number types. + +# Arguments +- `v`: Cartesian coordinate tuple. + +# Returns +An `(N + 1)`-tuple whose final element is `one(v[1])`. + +# Examples +```julia +julia> add_dim((2.0, 3.0)) +(2.0, 3.0, 1.0) +``` +""" +@inline function add_dim(v::Tuple{Vararg{Any,N}}) where {N} + return ntuple(i -> i <= N ? v[i] : one(v[1]), Val(N + 1)) +end + +""" + red_dim(v::Tuple) + +Convert homogeneous coordinates to Cartesian coordinates. + +For an `N`-component homogeneous coordinate, the first `N-1` entries are divided +by the final homogeneous component. + +# Arguments +- `v`: Homogeneous coordinate tuple. + +# Returns +An `(N - 1)`-tuple containing the normalized Cartesian coordinates. + +# Examples +```julia +julia> red_dim((4.0, 6.0, 2.0)) +(2.0, 3.0) +``` +""" +@inline function red_dim(v::Tuple{Vararg{Any,N}}) where {N} + last_inv = inv(v[N]) + return ntuple(i -> v[i] * last_inv, Val(N - 1)) +end + +""" + red_dim_apply(f, v::Tuple) + +Drop the final component of `v` and splat the remaining values into `f`. + +This helper is retained for compatibility. It does **not** normalize by the +final homogeneous component. Use [`red_dim`](@ref) when homogeneous +normalization is required. +""" +@inline function red_dim_apply(f, v::Tuple{Vararg{Any,N}}) where {N} + return f(ntuple(i -> v[i], Val(N - 1))...) +end + +@inline idx_apply(f, v::Tuple) = f(v...) +@inline func_transform(t, coord_transform_func) = coord_transform_func(Tuple(t)) +@inline func_transform_tup(t, coord_transform_func) = coord_transform_func(Tuple(t)) + +# Matrices are represented as an outer tuple of row tuples: +# +# ((a11, a12, ...), +# (a21, a22, ...), +# ...) +# +# The operations below are fully unrolled for statically-known tuple sizes. + +@inline function _tuple_matvec(A::NTuple{N,<:Tuple}, x::NTuple{N}) where {N} + return ntuple( + i -> sum(ntuple(j -> A[i][j] * x[j], Val(N))), + Val(N), + ) +end + +@inline function _tuple_matmul( + A::NTuple{N,<:Tuple}, + B::NTuple{N,<:Tuple}, +) where {N} + return ntuple( + i -> ntuple( + j -> sum(ntuple(k -> A[i][k] * B[k][j], Val(N))), + Val(N), + ), + Val(N), + ) +end + +"""Multiply tuple matrix `A` by tuple coordinate `v`, returning `A * v`.""" +@inline mat_mul(v::NTuple{N}, A::NTuple{N,<:Tuple}) where {N} = + _tuple_matvec(A, v) + +"""Multiply two square tuple matrices.""" +@inline mat_mul(A::NTuple{N,<:Tuple}, B::NTuple{N,<:Tuple}) where {N} = + _tuple_matmul(A, B) + +# ----------------------------------------------------------------------------- +# Backend-generic indices and transform evaluation +# ----------------------------------------------------------------------------- + +# Plain Arrays use the lazy CartesianIndices object directly. This keeps the CPU +# in-place hot path allocation-free. +@inline _backend_indices(data::Array) = CartesianIndices(axes(data)) + +# Other AbstractArray backends (e.g. CuArray) get an index array on the same +# backend. It is constructed once by get_function_* and reused by every call. +# +# CartesianIndex is an isbits type and can therefore be stored in GPU memory. +function _backend_indices(data::AbstractArray{T,N}) where {T,N} + inds = similar(data, CartesianIndex{N}) + inds .= CartesianIndices(axes(data)) + return inds +end + +@inline function _transform_sample(_, I, itp, coord_transform) + coords = coord_transform(Tuple(I)) + return itp(coords...) +end + +# CPU pure path: only the returned image needs to be allocated. +function _apply_transform( + coord_transform, + data::Array, + itp, + inds::CartesianIndices, +) + return map(inds) do I + coords = coord_transform(Tuple(I)) + itp(coords...) + end +end + +# Generic-array path. A GPU array among the broadcast arguments selects the GPU +# broadcast backend; no CUDA-specific code is required here. +function _apply_transform( + coord_transform, + data::AbstractArray, + itp, + inds, +) + return _transform_sample.( + data, # backend driver; value intentionally unused + inds, + Ref(itp), + Ref(coord_transform), + ) +end + +""" + apply_transform(coord_transform, data, itp) + +Apply a coordinate transform to every output index and sample the interpolation +object `itp` at the transformed coordinates. + +`coord_transform` receives an `N`-tuple and must return an `N`-tuple for +`N`-dimensional `data`. + +# Arguments +- `coord_transform`: Function mapping output coordinates to sampling coordinates. +- `data`: Array defining the output axes and execution backend. +- `itp`: Callable interpolation/extrapolation object. + +# Returns +A newly allocated array on the same backend as `data`. + +# Performance +For `Array`, the implementation uses a specialized Cartesian traversal. Other +`AbstractArray` backends use broadcast, which allows GPU arrays such as +`CuArray` to execute without scalar host indexing when the interpolant is +backend-compatible. + +See also [`apply_transform!`](@ref). +""" +function apply_transform(coord_transform, data::AbstractArray, itp) + return _apply_transform( + coord_transform, + data, + itp, + _backend_indices(data), + ) +end + +# Strict CPU performance path. +function _apply_transform!( + out::Array, + data::Array, + coord_transform, + itp, + inds::CartesianIndices, +) + axes(out) == axes(data) || + throw(DimensionMismatch("input and output axes must match")) + + @inbounds for I in inds + coords = coord_transform(Tuple(I)) + out[I] = itp(coords...) + end + return out +end + +# Backend-generic in-place path. On CuArray this becomes a broadcast kernel. +function _apply_transform!( + out::AbstractArray, + data::AbstractArray, + coord_transform, + itp, + inds, +) + axes(out) == axes(data) || + throw(DimensionMismatch("input and output axes must match")) + + out .= _transform_sample.( + data, + inds, + Ref(itp), + Ref(coord_transform), + ) + return out +end + +""" + apply_transform!(out, coord_transform, itp) + +In-place counterpart of [`apply_transform`](@ref). + +The transformed samples are written into the preallocated array `out`. + +# Arguments +- `out`: Destination array. Its axes define the output coordinate grid. +- `coord_transform`: Function mapping output coordinates to sampling coordinates. +- `itp`: Callable interpolation/extrapolation object. + +# Returns +`out`. + +# Performance +For a plain CPU `Array`, the hot loop is designed to perform no heap allocations +after compilation. Other array backends use an in-place broadcast. +""" +function apply_transform!(out::AbstractArray, coord_transform, itp) + return _apply_transform!( + out, + out, + coord_transform, + itp, + _backend_indices(out), + ) +end + +# ----------------------------------------------------------------------------- +# Homogeneous transforms +# ----------------------------------------------------------------------------- + +@inline function _homogeneous_transform(x, coord_transform) + return red_dim(coord_transform(add_dim(x))) +end + +function _apply_transform_homogen( + coord_transform, + data::AbstractArray, + itp, + inds, +) + return _apply_transform( + x -> _homogeneous_transform(x, coord_transform), + data, + itp, + inds, + ) +end + +""" + apply_transform_homogen(coord_transform, data, itp) + +Apply a transform expressed in homogeneous coordinates. + +For `N`-dimensional data, `coord_transform` receives an `(N + 1)`-tuple and must +return an `(N + 1)`-tuple. The returned homogeneous coordinate is normalized +with [`red_dim`](@ref) before interpolation. + +# Returns +A newly allocated transformed array. + +See also [`apply_transform_homogen!`](@ref) and [`apply_transform_affine`](@ref). +""" +function apply_transform_homogen(coord_transform, data::AbstractArray, itp) + return _apply_transform_homogen( + coord_transform, + data, + itp, + _backend_indices(data), + ) +end + +# CPU specialization keeps the homogeneous operations directly in the hot loop. +function _apply_transform_homogen!( + out::Array, + data::Array, + coord_transform, + itp, + inds::CartesianIndices, +) + axes(out) == axes(data) || + throw(DimensionMismatch("input and output axes must match")) + + @inbounds for I in inds + hcoords = add_dim(Tuple(I)) + coords = red_dim(coord_transform(hcoords)) + out[I] = itp(coords...) + end + return out +end + +# Generic-array path (including GPU arrays). +function _apply_transform_homogen!( + out::AbstractArray, + data::AbstractArray, + coord_transform, + itp, + inds, +) + return _apply_transform!( + out, + data, + x -> _homogeneous_transform(x, coord_transform), + itp, + inds, + ) +end + +""" + apply_transform_homogen!(out, coord_transform, itp) + +In-place counterpart of [`apply_transform_homogen`](@ref). + +For a CPU `Array`, the homogeneous conversion, coordinate transformation, and +interpolation are kept directly in the hot loop to support allocation-free +steady-state execution. + +# Returns +`out`. +""" +function apply_transform_homogen!(out::AbstractArray, coord_transform, itp) + return _apply_transform_homogen!( + out, + out, + coord_transform, + itp, + _backend_indices(out), + ) +end + +# ----------------------------------------------------------------------------- +# Affine / projective transforms +# ----------------------------------------------------------------------------- + +@inline function _check_homogeneous_matrix(matrix, N) + length(matrix) == N + 1 || throw( + DimensionMismatch( + "an $N-dimensional array requires a $(N + 1)×$(N + 1) homogeneous matrix", + ), + ) + return nothing +end + +function _apply_transform_affine( + matrix::NTuple{M,<:Tuple}, + data::AbstractArray{T,N}, + itp, + inds, +) where {M,T,N} + _check_homogeneous_matrix(matrix, N) + return _apply_transform_homogen( + x -> mat_mul(x, matrix), + data, + itp, + inds, + ) +end + +""" + apply_transform_affine(matrix, data, itp) + +Apply an affine or projective transform represented by a homogeneous tuple +matrix. + +For `N`-dimensional `data`, `matrix` must contain `N + 1` row tuples, each of +length `N + 1`. Matrix storage is row-major at the Julia tuple level and the +operation corresponds to `matrix * coordinate`. + +# Arguments +- `matrix`: Homogeneous transformation matrix represented as an outer tuple of row tuples. +- `data`: Array defining output axes and backend. +- `itp`: Callable interpolation/extrapolation object. + +# Returns +A newly allocated transformed array. + +# Examples +```julia +A = ( + (1.0, 0.0, 0.0), + (0.0, 1.0, 0.0), + (0.0, 0.0, 1.0), +) +y = apply_transform_affine(A, data, itp) +``` +""" +function apply_transform_affine( + matrix::NTuple{M,<:Tuple}, + data::AbstractArray{T,N}, + itp, +) where {M,T,N} + return _apply_transform_affine( + matrix, + data, + itp, + _backend_indices(data), + ) +end + +# CPU specialization: keep matrix multiplication, homogeneous reduction, and +# interpolation directly in one loop. This is the strict zero-allocation path. +function _apply_transform_affine!( + out::Array{T,N}, + data::Array, + matrix::NTuple{M,<:Tuple}, + itp, + inds::CartesianIndices, +) where {T,N,M} + _check_homogeneous_matrix(matrix, N) + axes(out) == axes(data) || + throw(DimensionMismatch("input and output axes must match")) + + @inbounds for I in inds + hcoords = add_dim(Tuple(I)) + coords = red_dim(mat_mul(hcoords, matrix)) + out[I] = itp(coords...) + end + return out +end + +# Generic-array path (including GPU arrays). +function _apply_transform_affine!( + out::AbstractArray{T,N}, + data::AbstractArray, + matrix::NTuple{M,<:Tuple}, + itp, + inds, +) where {T,N,M} + _check_homogeneous_matrix(matrix, N) + return _apply_transform_homogen!( + out, + data, + x -> mat_mul(x, matrix), + itp, + inds, + ) +end + +""" + apply_transform_affine!(out, matrix, itp) + +In-place counterpart of [`apply_transform_affine`](@ref). + +# Returns +`out`. + +# Performance +For CPU `Array`s, matrix multiplication, homogeneous reduction, and +interpolation are fused into the explicit hot loop. With tuple matrices and a +preallocated output, the steady-state path is intended to be allocation-free. +""" +function apply_transform_affine!( + out::AbstractArray{T,N}, + matrix::NTuple{M,<:Tuple}, + itp, +) where {T,N,M} + return _apply_transform_affine!( + out, + out, + matrix, + itp, + _backend_indices(out), + ) +end + +# 2-D seven-parameter model: +# +# p = ( +# shift_x, +# shift_y, +# scale_x, +# scale_y, +# shear_xy, +# shear_yx, +# rotation, +# ) + +@inline function _check_affine_params_2d(p) + length(p) == 7 || throw(ArgumentError( + "the 2-D affine parameterization requires 7 parameters: " * + "(shift_x, shift_y, scale_x, scale_y, shear_xy, shear_yx, rotation)", + )) + return nothing +end + +@inline function _affine_matrix_2d(p, x_cen, y_cen) + _check_affine_params_2d(p) + + z = zero(p[1]) + o = one(p[1]) + c = cos(p[7]) + s = sin(p[7]) + + rot_mat = ( + (c, -s, z), + (s, c, z), + (z, z, o), + ) + + shear_mat = ( + (o, p[5], z), + (p[6], o, z), + (z, z, o), + ) + + scale_mat = ( + (o / p[3], z, z), + (z, o / p[4], z), + (z, z, o), + ) + + shift_mat = ( + (o, z, -p[1]), + (z, o, -p[2]), + (z, z, o), + ) + + t_to_origin = ( + (o, z, o * x_cen), + (z, o, o * y_cen), + (z, z, o), + ) + + t_to_center = ( + (o, z, -o * x_cen), + (z, o, -o * y_cen), + (z, z, o), + ) + + matrix = _tuple_matmul(t_to_origin, scale_mat) + matrix = _tuple_matmul(matrix, rot_mat) + matrix = _tuple_matmul(matrix, shear_mat) + matrix = _tuple_matmul(matrix, shift_mat) + return _tuple_matmul(matrix, t_to_center) +end + +""" + _affine_warp_2d_params(p, data, itp, inds) + +Named pure seven-parameter affine warp. + +The named boundary is intentional: the custom ChainRules rule below replaces a +very expensive scalar-by-scalar Zygote reverse pass with a single image-level +vector-Jacobian product. +""" +function _affine_warp_2d_params(p, data::AbstractMatrix, itp, inds) + _check_affine_params_2d(p) + x_cen = size(data, 1) ÷ 2 + 1 + y_cen = size(data, 2) ÷ 2 + 1 + matrix = _affine_matrix_2d(p, x_cen, y_cen) + return _apply_transform_affine(matrix, data, itp, inds) +end + +# Tuple reduction helper used by backend-generic VJPs. +# +# This deliberately does not introduce a custom accumulator type and does not +# extend/overload Base arithmetic. The reduction values are plain isbits tuples. +@inline function _tuple_add( + a::NTuple{N}, + b::NTuple{N}, +) where {N} + return ntuple( + i -> a[i] + b[i], + Val(N), + ) +end + +@inline function _affine_pixel_vjp( + I, + δ, + p, + itp, + x_cen, + y_cen, + ::Type{G}, +) where {G} + shift_x, shift_y = p[1], p[2] + scale_x, scale_y = p[3], p[4] + shear_xy, shear_yx = p[5], p[6] + θ = p[7] + + c = cos(θ) + s = sin(θ) + inv_scale_x = inv(scale_x) + inv_scale_y = inv(scale_y) + + x = I[1] + y = I[2] + + a = x - x_cen - shift_x + b = y - y_cen - shift_y + u = a + shear_xy * b + v = shear_yx * a + b + + r1 = c * u - s * v + r2 = s * u + c * v + + x′ = x_cen + r1 * inv_scale_x + y′ = y_cen + r2 * inv_scale_y + + grad = Interpolations.gradient(itp, x′, y′) + gx = grad[1] + gy = grad[2] + + inv_scale_x2 = inv_scale_x * inv_scale_x + inv_scale_y2 = inv_scale_y * inv_scale_y + + dx_dp1 = (-c + s * shear_yx) * inv_scale_x + dy_dp1 = (-s - c * shear_yx) * inv_scale_y + dx_dp2 = (-c * shear_xy + s) * inv_scale_x + dy_dp2 = (-s * shear_xy - c) * inv_scale_y + + vals = ( + δ * (gx * dx_dp1 + gy * dy_dp1), + δ * (gx * dx_dp2 + gy * dy_dp2), + δ * gx * (-r1 * inv_scale_x2), + δ * gy * (-r2 * inv_scale_y2), + δ * ( + gx * (c * b * inv_scale_x) + + gy * (s * b * inv_scale_y) + ), + δ * ( + gx * (-s * a * inv_scale_x) + + gy * ( c * a * inv_scale_y) + ), + δ * ( + gx * (-r2 * inv_scale_x) + + gy * ( r1 * inv_scale_y) + ), + ) + + return ntuple( + i -> convert(G, vals[i]), + Val(7), + ) +end + +# CPU VJP: one pass, one small preallocated interpolation-gradient buffer. +function _affine_vjp_2d( + p, + data::AbstractMatrix, + itp, + Δ, + inds::CartesianIndices, +) + pt = ntuple(i -> p[i], Val(7)) + P = promote_type(ntuple(i -> typeof(pt[i]), Val(7))...) + G = promote_type(P, eltype(data), eltype(Δ)) + + shift_x, shift_y = pt[1], pt[2] + scale_x, scale_y = pt[3], pt[4] + shear_xy, shear_yx = pt[5], pt[6] + θ = pt[7] + + x_cen = size(data, 1) ÷ 2 + 1 + y_cen = size(data, 2) ÷ 2 + 1 + + c = cos(θ) + s = sin(θ) + inv_scale_x = inv(scale_x) + inv_scale_y = inv(scale_y) + inv_scale_x2 = inv_scale_x * inv_scale_x + inv_scale_y2 = inv_scale_y * inv_scale_y + + dx_dp1 = (-c + s * shear_yx) * inv_scale_x + dy_dp1 = (-s - c * shear_yx) * inv_scale_y + dx_dp2 = (-c * shear_xy + s) * inv_scale_x + dy_dp2 = (-s * shear_xy - c) * inv_scale_y + + g1 = zero(G) + g2 = zero(G) + g3 = zero(G) + g4 = zero(G) + g5 = zero(G) + g6 = zero(G) + g7 = zero(G) + + grad_itp = Vector{G}(undef, 2) + + @inbounds for I in inds + x = I[1] + y = I[2] + + a = x - x_cen - shift_x + b = y - y_cen - shift_y + u = a + shear_xy * b + v = shear_yx * a + b + r1 = c * u - s * v + r2 = s * u + c * v + x′ = x_cen + r1 * inv_scale_x + y′ = y_cen + r2 * inv_scale_y + + Interpolations.gradient!(grad_itp, itp, x′, y′) + gx = grad_itp[1] + gy = grad_itp[2] + + δ = Δ[I] + δgx = δ * gx + δgy = δ * gy + + g1 += δgx * dx_dp1 + δgy * dy_dp1 + g2 += δgx * dx_dp2 + δgy * dy_dp2 + g3 += δgx * (-r1 * inv_scale_x2) + g4 += δgy * (-r2 * inv_scale_y2) + g5 += δgx * (c * b * inv_scale_x) + + δgy * (s * b * inv_scale_y) + g6 += δgx * (-s * a * inv_scale_x) + + δgy * ( c * a * inv_scale_y) + g7 += δgx * (-r2 * inv_scale_x) + + δgy * ( r1 * inv_scale_y) + end + + return (g1, g2, g3, g4, g5, g6, g7) +end + +# Backend-generic VJP. With GPU-backed `inds`/`Δ`, GPUArrays/CUDA can reduce the +# lazy broadcast without scalar host indexing. +function _affine_vjp_2d( + p, + data::AbstractMatrix, + itp, + Δ, + inds::AbstractArray, +) + pt = ntuple(i -> p[i], Val(7)) + P = promote_type(ntuple(i -> typeof(pt[i]), Val(7))...) + G = promote_type(P, eltype(data), eltype(Δ)) + + x_cen = size(data, 1) ÷ 2 + 1 + y_cen = size(data, 2) ÷ 2 + 1 + + bc = Base.broadcasted( + _affine_pixel_vjp, + vec(inds), + vec(Δ), + Ref(pt), + Ref(itp), + Ref(x_cen), + Ref(y_cen), + Ref(G), + ) + + return mapreduce( + identity, + _tuple_add, + bc; + init=ntuple(_ -> zero(G), Val(7)), + ) +end + +@inline _affine_parameter_tangent(::Tuple, g::Tuple) = g +@inline _affine_parameter_tangent(::AbstractVector, g::Tuple) = collect(g) + +function ChainRulesCore.rrule( + ::typeof(_affine_warp_2d_params), + p, + data::AbstractMatrix, + itp, + inds, +) + y = _affine_warp_2d_params(p, data, itp, inds) + + function affine_pullback(Δ_raw) + Δ = ChainRulesCore.unthunk(Δ_raw) + + if Δ isa ChainRulesCore.AbstractZero + return ( + ChainRulesCore.NoTangent(), + ChainRulesCore.ZeroTangent(), + ChainRulesCore.NoTangent(), + ChainRulesCore.NoTangent(), + ChainRulesCore.NoTangent(), + ) + end + + gp = _affine_vjp_2d(p, data, itp, Δ, inds) + + return ( + ChainRulesCore.NoTangent(), + _affine_parameter_tangent(p, gp), + ChainRulesCore.NoTangent(), # source data is treated as constant + ChainRulesCore.NoTangent(), # interpolation object is constant + ChainRulesCore.NoTangent(), # backend indices are constant + ) + end + + return y, affine_pullback +end + +""" + get_function_affine(data; kwargs...) + +Create an interpolation-based affine warp of sampled data. + +The returned function accepts either: + +1. an `(N + 1) × (N + 1)` homogeneous matrix represented as an outer tuple of + row tuples, for arbitrary-dimensional data; or +2. for 2-D data, the seven-parameter representation + `(shift_x, shift_y, scale_x, scale_y, shear_xy, shear_yx, rotation)`. + +# Keyword arguments +- `super_sampling=2`: Retained for API compatibility. The current optimized + affine path operates directly on the interpolation grid. +- `extrapolation_bc=zero(eltype(data))`: Boundary condition or fill value used + outside the interpolation domain. +- `interp_type=Interpolations.BSpline(Linear())`: Interpolation scheme. + +# Returns +A callable `warp(transform)` that allocates and returns the transformed array. + +# Automatic differentiation +The 2-D seven-parameter path has a custom `ChainRulesCore.rrule` that computes +an image-level vector-Jacobian product. Reverse-mode AD therefore differentiates +with respect to the transform parameters without tracing every pixel operation. +The source data and constructed interpolation object are treated as constants +by this rule. + +# GPU support +For non-`Array` backends, the interpolant is constructed on CPU and adapted to +the input backend with Adapt.jl. The forward transform is expressed using +backend-generic broadcast. In particular, a `CuArray` input can execute on CUDA +without DataToFunctions depending directly on CUDA.jl, provided the selected +interpolation/extrapolation combination is GPU-compatible. + +# Examples +```julia +using DataToFunctions + +img = rand(Float32, 64, 64) +warp = get_function_affine(img) + +p = (0.2f0, -0.1f0, 1.01f0, 0.99f0, 0.0f0, 0.0f0, 0.01f0) +warped = warp(p) +``` + +For repeated forward evaluation with a reusable output buffer, see +[`get_function_affine_inplace`](@ref). +""" +function get_function_affine( + data::AbstractArray{T,N}; + super_sampling=2, + extrapolation_bc=zero(T), + interp_type=Interpolations.BSpline(Linear()), +) where {T,N} + _ = super_sampling # retained for API compatibility + + itp = _make_interpolant(data, interp_type, extrapolation_bc) + inds = _backend_indices(data) + + function interpolated(matrix::NTuple{M,<:Tuple}) where {M} + return _apply_transform_affine(matrix, data, itp, inds) + end + + function interpolated(p::AbstractVector) + N == 2 || throw(ArgumentError( + "the 7-parameter affine interface is defined only for 2-D data; " * + "pass an $(N + 1)×$(N + 1) tuple matrix for $N-D data", + )) + return _affine_warp_2d_params(p, data, itp, inds) + end + + function interpolated(p::NTuple{7}) + N == 2 || throw(ArgumentError( + "the 7-parameter affine interface is defined only for 2-D data; " * + "pass an $(N + 1)×$(N + 1) tuple matrix for $N-D data", + )) + return _affine_warp_2d_params(p, data, itp, inds) + end + + return interpolated +end + +""" + get_function_affine_inplace(data; kwargs...) + +Create the in-place affine warp `warp!(out, transform)`. + +The accepted transform representations and keyword arguments are the same as +for [`get_function_affine`](@ref). + +# Returns +A callable that writes transformed samples into `out` and returns `out`. + +# Performance +For CPU `Array`s, a preallocated output and tuple parameters/matrices are +intended to give zero heap allocations in steady state. For GPU-backed arrays, +the same API uses backend-generic in-place broadcast. + +# Automatic differentiation +Use [`get_function_affine`](@ref) for Zygote differentiation. The mutating form +is intended for forward evaluations where output-buffer reuse is important. + +# Examples +```julia +img = rand(Float32, 64, 64) +warp! = get_function_affine_inplace(img) +out = similar(img) +p = (0.2f0, -0.1f0, 1.01f0, 0.99f0, 0.0f0, 0.0f0, 0.01f0) +warp!(out, p) +``` +""" +function get_function_affine_inplace( + data::AbstractArray{T,N}; + super_sampling=2, + extrapolation_bc=zero(T), + interp_type=Interpolations.BSpline(Linear()), +) where {T,N} + _ = super_sampling + + itp = _make_interpolant(data, interp_type, extrapolation_bc) + inds = _backend_indices(data) + + function interpolated!(out, matrix::NTuple{M,<:Tuple}) where {M} + return _apply_transform_affine!(out, data, matrix, itp, inds) + end + + function interpolated!(out, p::AbstractVector) + N == 2 || throw(ArgumentError( + "the 7-parameter affine interface is defined only for 2-D data", + )) + _check_affine_params_2d(p) + x_cen = size(data, 1) ÷ 2 + 1 + y_cen = size(data, 2) ÷ 2 + 1 + matrix = _affine_matrix_2d(p, x_cen, y_cen) + return _apply_transform_affine!(out, data, matrix, itp, inds) + end + + function interpolated!(out, p::NTuple{7}) + N == 2 || throw(ArgumentError( + "the 7-parameter affine interface is defined only for 2-D data", + )) + x_cen = size(data, 1) ÷ 2 + 1 + y_cen = size(data, 2) ÷ 2 + 1 + matrix = _affine_matrix_2d(p, x_cen, y_cen) + return _apply_transform_affine!(out, data, matrix, itp, inds) + end + + return interpolated! +end + +# ----------------------------------------------------------------------------- +# Parameterized user-transform helpers +# ----------------------------------------------------------------------------- + +@inline function _param_transform_sample(_, I, itp, f, params) + coords = f(Tuple(I), params) + return itp(coords...) +end + +function _apply_param_transform( + data::Array, + f, + params, + itp, + inds::CartesianIndices, +) + return map(inds) do I + coords = f(Tuple(I), params) + itp(coords...) + end +end + +function _apply_param_transform( + data::AbstractArray, + f, + params, + itp, + inds, +) + return _param_transform_sample.( + data, + inds, + Ref(itp), + Ref(f), + Ref(params), + ) +end + +function _apply_param_transform!( + out::Array, + data::Array, + f, + params, + itp, + inds::CartesianIndices, +) + axes(out) == axes(data) || + throw(DimensionMismatch("input and output axes must match")) + + @inbounds for I in inds + coords = f(Tuple(I), params) + out[I] = itp(coords...) + end + return out +end + +function _apply_param_transform!( + out::AbstractArray, + data::AbstractArray, + f, + params, + itp, + inds, +) + axes(out) == axes(data) || + throw(DimensionMismatch("input and output axes must match")) + + out .= _param_transform_sample.( + data, + inds, + Ref(itp), + Ref(f), + Ref(params), + ) + return out +end + +@inline function _param_homogeneous_sample(_, I, itp, f, params) + hcoords = add_dim(Tuple(I)) + coords = red_dim(f(hcoords, params)) + return itp(coords...) +end + +function _apply_param_homogeneous( + data::Array, + f, + params, + itp, + inds::CartesianIndices, +) + return map(inds) do I + hcoords = add_dim(Tuple(I)) + coords = red_dim(f(hcoords, params)) + itp(coords...) + end +end + +function _apply_param_homogeneous( + data::AbstractArray, + f, + params, + itp, + inds, +) + return _param_homogeneous_sample.( + data, + inds, + Ref(itp), + Ref(f), + Ref(params), + ) +end + +function _apply_param_homogeneous!( + out::Array, + data::Array, + f, + params, + itp, + inds::CartesianIndices, +) + axes(out) == axes(data) || + throw(DimensionMismatch("input and output axes must match")) + + @inbounds for I in inds + hcoords = add_dim(Tuple(I)) + coords = red_dim(f(hcoords, params)) + out[I] = itp(coords...) + end + return out +end + +function _apply_param_homogeneous!( + out::AbstractArray, + data::AbstractArray, + f, + params, + itp, + inds, +) + axes(out) == axes(data) || + throw(DimensionMismatch("input and output axes must match")) + + out .= _param_homogeneous_sample.( + data, + inds, + Ref(itp), + Ref(f), + Ref(params), + ) + return out +end + +# ----------------------------------------------------------------------------- +# User-defined tuple and homogeneous transforms +# ----------------------------------------------------------------------------- + +""" + get_function_tuple(data, f; kwargs...) + +Create a parameterized warp from a user-supplied Cartesian coordinate function. + +The callback is called as `f(coord, params)`, where `coord` is an `N`-tuple and +must be mapped to another `N`-tuple of sampling coordinates. + +# Arguments +- `data`: Sampled input data. +- `f`: Coordinate transformation callback. + +# Keyword arguments +- `super_sampling=2`: Retained for API compatibility. +- `extrapolation_bc=zero(eltype(data))`: Extrapolation boundary condition. +- `interp_type=Interpolations.BSpline(Linear())`: Interpolation scheme. + +# Returns +A callable `warp(params)` returning a transformed array. + +# Examples +```julia +shift(coord, p) = (coord[1] + p[1], coord[2] + p[2]) +warp = get_function_tuple(data, shift) +y = warp((0.2, -0.1)) +``` + +See also [`get_function_tuple_inplace`](@ref). +""" +function get_function_tuple( + data::AbstractArray{T,N}, + f; + super_sampling=2, + extrapolation_bc=zero(T), + interp_type=Interpolations.BSpline(Linear()), +) where {T,N} + _ = super_sampling + + itp = _make_interpolant(data, interp_type, extrapolation_bc) + inds = _backend_indices(data) + + interpolated(params) = + _apply_param_transform(data, f, params, itp, inds) + + return interpolated +end + +""" + get_function_tuple_inplace(data, f; kwargs...) + +Create the in-place counterpart of [`get_function_tuple`](@ref). + +The returned function is called as `warp!(out, params)` and writes the +transformed samples into `out`. + +For CPU arrays, the callback and interpolation are evaluated directly in the +preallocated hot loop. Other array backends use in-place broadcast. +""" +function get_function_tuple_inplace( + data::AbstractArray{T,N}, + f; + super_sampling=2, + extrapolation_bc=zero(T), + interp_type=Interpolations.BSpline(Linear()), +) where {T,N} + _ = super_sampling + + itp = _make_interpolant(data, interp_type, extrapolation_bc) + inds = _backend_indices(data) + + function interpolated!(out, params) + return _apply_param_transform!( + out, + data, + f, + params, + itp, + inds, + ) + end + + return interpolated! +end + +""" + get_function_homogen(data, f; kwargs...) + +Create a parameterized warp using homogeneous coordinates. + +For `N`-dimensional data, the callback is called as `f(hcoord, params)`, where +`hcoord` is an `(N + 1)`-tuple. The callback must return another homogeneous +`(N + 1)`-tuple, which is normalized before interpolation. + +# Returns +A callable `warp(params)` returning a transformed array. + +# Examples +```julia +shift_h(h, p) = ( + h[1] + p[1] * h[3], + h[2] + p[2] * h[3], + h[3], +) +warp = get_function_homogen(data, shift_h) +y = warp((0.2, -0.1)) +``` + +See also [`get_function_homogen_inplace`](@ref). +""" +function get_function_homogen( + data::AbstractArray{T,N}, + f; + super_sampling=2, + extrapolation_bc=zero(T), + interp_type=Interpolations.BSpline(Linear()), +) where {T,N} + _ = super_sampling + + itp = _make_interpolant(data, interp_type, extrapolation_bc) + inds = _backend_indices(data) + + interpolated(params) = + _apply_param_homogeneous(data, f, params, itp, inds) + + return interpolated +end + +""" + get_function_homogen_inplace(data, f; kwargs...) + +Create the in-place counterpart of [`get_function_homogen`](@ref). + +The returned callable has the form `warp!(out, params)`. For CPU arrays the +homogeneous transform and interpolation are kept in a direct loop; other array +backends use in-place broadcast. +""" +function get_function_homogen_inplace( + data::AbstractArray{T,N}, + f; + super_sampling=2, + extrapolation_bc=zero(T), + interp_type=Interpolations.BSpline(Linear()), +) where {T,N} + _ = super_sampling + + itp = _make_interpolant(data, interp_type, extrapolation_bc) + inds = _backend_indices(data) + + function interpolated!(out, params) + return _apply_param_homogeneous!( + out, + data, + f, + params, + itp, + inds, + ) + end + + return interpolated! +end + +""" + get_function_svec(args...; kwargs...) + +Compatibility alias for [`get_function_homogen`](@ref). + +This function is retained for compatibility with the previous +StaticArrays-based interface. The current implementation uses tuple-based +homogeneous coordinates instead. +""" +# Compatibility name retained from the previous StaticArrays-based API. +get_function_svec(args...; kwargs...) = + get_function_homogen(args...; kwargs...) + +# ----------------------------------------------------------------------------- +# Polynomial transforms +# ----------------------------------------------------------------------------- + +@inline _poly_coeff_tuple(coeffs::Tuple) = coeffs +@inline _poly_coeff_tuple(coeffs::AbstractVector) = Tuple(coeffs) + +@inline function _poly_sample(_, I, itp, poly, coeffs) + coords = poly(Tuple(I), coeffs) + return itp(coords...) +end + +# CPU pure path. +function _poly_warp_params( + coeffs, + data::Array, + itp, + poly, + inds::CartesianIndices, +) + c = _poly_coeff_tuple(coeffs) + + return map(inds) do I + coords = poly(Tuple(I), c) + itp(coords...) + end +end + +# Backend-generic pure path. +function _poly_warp_params( + coeffs, + data::AbstractArray, + itp, + poly, + inds, +) + c = _poly_coeff_tuple(coeffs) + + return _poly_sample.( + data, + inds, + Ref(itp), + Ref(poly), + Ref(c), + ) +end + +# CPU in-place path. Keep polynomial evaluation directly in the loop: this is +# important for the observed zero-allocation EvalMultiPoly hot path. +function _poly_warp_params!( + out::Array, + data::Array, + coeffs, + itp, + poly, + inds::CartesianIndices, +) + axes(out) == axes(data) || + throw(DimensionMismatch("input and output axes must match")) + + c = _poly_coeff_tuple(coeffs) + + @inbounds for I in inds + coords = poly(Tuple(I), c) + out[I] = itp(coords...) + end + return out +end + +# Backend-generic in-place polynomial warp. +function _poly_warp_params!( + out::AbstractArray, + data::AbstractArray, + coeffs, + itp, + poly, + inds, +) + axes(out) == axes(data) || + throw(DimensionMismatch("input and output axes must match")) + + c = _poly_coeff_tuple(coeffs) + + out .= _poly_sample.( + data, + inds, + Ref(itp), + Ref(poly), + Ref(c), + ) + return out +end + +@inline function _poly_basis_coefficients( + coeffs::NTuple{M,T}, +) where {M,T} + return ntuple( + k -> ntuple( + j -> ifelse(j == k, one(T), zero(T)), + Val(M), + ), + Val(M), + ) +end + +@inline function _poly_spatial_vjp(grad_itp, dcoords, ::Val{N}) where {N} + return sum( + ntuple( + d -> grad_itp[d] * dcoords[d], + Val(N), + ), + ) +end + +# CPU VJP. +function _poly_vjp( + coeffs::NTuple{M,C}, + data::AbstractArray{T,N}, + itp, + poly, + Δ, + inds::CartesianIndices, +) where {M,C,T,N} + G = promote_type(C, T, eltype(Δ)) + g = zeros(G, M) + grad_itp = Vector{G}(undef, N) + basis = _poly_basis_coefficients(coeffs) + + @inbounds for I in inds + x = Tuple(I) + coords = poly(x, coeffs) + + Interpolations.gradient!( + grad_itp, + itp, + coords..., + ) + + δ = Δ[I] + + for k in 1:M + dcoords = poly(x, basis[k]) + g[k] += δ * _poly_spatial_vjp( + grad_itp, + dcoords, + Val(N), + ) + end + end + + return ntuple(k -> g[k], Val(M)) +end + +@inline function _poly_pixel_vjp( + I, + δ, + coeffs::NTuple{M,C}, + itp, + poly, + basis, + ::Val{N}, + ::Type{G}, +) where {M,C,N,G} + x = Tuple(I) + coords = poly(x, coeffs) + grad_itp = Interpolations.gradient(itp, coords...) + + vals = ntuple( + k -> begin + dcoords = poly(x, basis[k]) + δ * _poly_spatial_vjp( + grad_itp, + dcoords, + Val(N), + ) + end, + Val(M), + ) + + return ntuple( + k -> convert(G, vals[k]), + Val(M), + ) +end + +# Backend-generic VJP. For GPU-backed arrays this is a reduction of a lazy +# broadcast, avoiding scalar host indexing and avoiding materializing the full +# image-by-coefficient Jacobian. +function _poly_vjp( + coeffs::NTuple{M,C}, + data::AbstractArray{T,N}, + itp, + poly, + Δ, + inds::AbstractArray, +) where {M,C,T,N} + G = promote_type(C, T, eltype(Δ)) + basis = _poly_basis_coefficients(coeffs) + + bc = Base.broadcasted( + _poly_pixel_vjp, + vec(inds), + vec(Δ), + Ref(coeffs), + Ref(itp), + Ref(poly), + Ref(basis), + Ref(Val(N)), + Ref(G), + ) + + return mapreduce( + identity, + _tuple_add, + bc; + init=ntuple(_ -> zero(G), Val(M)), + ) +end + +function _poly_vjp(coeffs::AbstractVector, data, itp, poly, Δ, inds) + return collect( + _poly_vjp( + Tuple(coeffs), + data, + itp, + poly, + Δ, + inds, + ), + ) +end + +@inline _poly_parameter_tangent(::Tuple, g::Tuple) = g +@inline _poly_parameter_tangent(::AbstractVector, g::Tuple) = collect(g) +@inline _poly_parameter_tangent(::AbstractVector, g::AbstractVector) = g + +function ChainRulesCore.rrule( + ::typeof(_poly_warp_params), + coeffs, + data::AbstractArray, + itp, + poly, + inds, +) + y = _poly_warp_params(coeffs, data, itp, poly, inds) + + function poly_pullback(Δ_raw) + Δ = ChainRulesCore.unthunk(Δ_raw) + + if Δ isa ChainRulesCore.AbstractZero + return ( + ChainRulesCore.NoTangent(), + ChainRulesCore.ZeroTangent(), + ChainRulesCore.NoTangent(), + ChainRulesCore.NoTangent(), + ChainRulesCore.NoTangent(), + ChainRulesCore.NoTangent(), + ) + end + + c = _poly_coeff_tuple(coeffs) + gc = _poly_vjp(c, data, itp, poly, Δ, inds) + + return ( + ChainRulesCore.NoTangent(), + _poly_parameter_tangent(coeffs, gc), + ChainRulesCore.NoTangent(), # source data is constant + ChainRulesCore.NoTangent(), # interpolation object is constant + ChainRulesCore.NoTangent(), # generated polynomial is constant + ChainRulesCore.NoTangent(), # backend indices are constant + ) + end + + return y, poly_pullback +end + +""" + get_function_poly(data, ::Val{order}; kwargs...) + get_function_poly(data, order::Integer; kwargs...) + +Create a polynomial coordinate warp using EvalMultiPoly.jl. + +For `N`-dimensional data, the generated polynomial maps each output coordinate +to `N` interpolation coordinates. Coefficients follow the ordering defined by +`EvalMultiPoly.get_multi_poly(Val(N), Val(order))`. + +# Arguments +- `data`: Sampled input data. +- `order`: Polynomial order, preferably supplied as `Val(order)` when the order + is known statically. + +# Keyword arguments +- `super_sampling=1`: Retained for API compatibility. +- `extrapolation_bc=zero(eltype(data))`: Extrapolation boundary condition. +- `interp_type=Interpolations.BSpline(Linear())`: Interpolation scheme. +- `extrapolation=nothing`: Optional explicit extrapolation setting. When given, + it takes precedence over `extrapolation_bc`. + +# Returns +A callable `warp(coeffs)` returning the transformed array. Tuple coefficients +are recommended for the performance-sensitive path. + +# Automatic differentiation +A custom `ChainRulesCore.rrule` computes the vector-Jacobian product with +respect to the polynomial coefficients. It uses the linearity of polynomial +coordinates in their coefficients and the spatial gradient of the interpolation +object, avoiding materialization of the full pixel-by-coefficient Jacobian. +The source data, interpolation object, and generated polynomial evaluator are +treated as constants by the rule. + +# GPU support +For a GPU-backed input, interpolation storage is adapted to the backend and the +forward transform uses backend-generic broadcast. The custom VJP uses a +backend-generic reduction. + +# Examples +```julia +using EvalMultiPoly + +warp = get_function_poly(data, Val(2)) +coeffs = get_identity_multipoly_coeffs(Val(2), Val(2)) +y = warp(coeffs) +``` + +See also [`get_function_poly_inplace`](@ref). +""" +function get_function_poly( + data::AbstractArray{T,N}, + ::Val{N_order}; + super_sampling=1, + extrapolation_bc=zero(T), + interp_type=Interpolations.BSpline(Linear()), + extrapolation=nothing, +) where {T,N,N_order} + _ = super_sampling + + bc = isnothing(extrapolation) ? + extrapolation_bc : + extrapolation + + itp = _make_interpolant(data, interp_type, bc) + poly = get_multi_poly(Val(N), Val(N_order)) + inds = _backend_indices(data) + + interpolated(coeffs) = + _poly_warp_params(coeffs, data, itp, poly, inds) + + return interpolated +end + +function get_function_poly( + data::AbstractArray, + order::Integer; + kwargs..., +) + return get_function_poly( + data, + Val(order); + kwargs..., + ) +end + +""" + get_function_poly_inplace(data, ::Val{order}; kwargs...) + get_function_poly_inplace(data, order::Integer; kwargs...) + +Create the in-place polynomial warp `warp!(out, coeffs)`. + +The polynomial and interpolation options are the same as for +[`get_function_poly`](@ref). + +# Returns +A callable that writes into `out` and returns `out`. + +# Performance +For CPU `Array`s, polynomial evaluation is intentionally kept directly inside +the hot loop. With tuple coefficients and a preallocated output, the +steady-state forward path is intended to perform zero heap allocations. For +GPU-backed arrays, the same API uses an in-place broadcast. + +# Automatic differentiation +Use [`get_function_poly`](@ref) when differentiating with Zygote. This mutating +variant is intended for repeated forward evaluations. +""" +function get_function_poly_inplace( + data::AbstractArray{T,N}, + ::Val{N_order}; + super_sampling=1, + extrapolation_bc=zero(T), + interp_type=Interpolations.BSpline(Linear()), + extrapolation=nothing, +) where {T,N,N_order} + _ = super_sampling + + bc = isnothing(extrapolation) ? + extrapolation_bc : + extrapolation + + itp = _make_interpolant(data, interp_type, bc) + poly = get_multi_poly(Val(N), Val(N_order)) + inds = _backend_indices(data) + + function interpolated!(out, coeffs) + return _poly_warp_params!( + out, + data, + coeffs, + itp, + poly, + inds, + ) + end + + return interpolated! +end + +function get_function_poly_inplace( + data::AbstractArray, + order::Integer; + kwargs..., +) + return get_function_poly_inplace( + data, + Val(order); + kwargs..., + ) +end + +# ----------------------------------------------------------------------------- +# Mode-dispatch convenience API +# ----------------------------------------------------------------------------- + +""" + get_interpolated_function(data, ::Type{AffineMode}; kwargs...) + get_interpolated_function(data, ::Type{PolynomialMode}, order; kwargs...) + get_interpolated_function(data; kwargs...) + +Construct an interpolated transformation function using mode-based dispatch. + +`AffineMode` delegates to [`get_function_affine`](@ref). `PolynomialMode` +delegates to [`get_function_poly`](@ref) and requires a polynomial `order`. +When no mode is supplied, `AffineMode` is selected and a warning is emitted. + +# Keyword arguments +- `super_sampling=2`: Forwarded to the selected constructor. +- `extrapolation_bc=zero(eltype(data))`: Extrapolation boundary condition. +- `interp_type=Interpolations.BSpline(Linear())`: Interpolation scheme. + +# Examples +```julia +f_affine = get_interpolated_function(data, AffineMode) +f_poly = get_interpolated_function(data, PolynomialMode, Val(2)) +``` +""" +function get_interpolated_function( + data::AbstractArray{T,N}, + ::Type{AffineMode}; + super_sampling=2, + extrapolation_bc=zero(T), + interp_type=Interpolations.BSpline(Linear()), +) where {T,N} + return get_function_affine( + data; + super_sampling=super_sampling, + extrapolation_bc=extrapolation_bc, + interp_type=interp_type, + ) +end + +function get_interpolated_function( + data::AbstractArray{T,N}; + super_sampling=2, + extrapolation_bc=zero(T), + interp_type=Interpolations.BSpline(Linear()), +) where {T,N} + @warn "No transformation mode provided. `AffineMode` is used as default" + return get_interpolated_function( + data, + AffineMode; + super_sampling=super_sampling, + extrapolation_bc=extrapolation_bc, + interp_type=interp_type, + ) +end + +function get_interpolated_function( + data::AbstractArray{T,N}, + ::Type{PolynomialMode}, + order=nothing; + super_sampling=2, + extrapolation_bc=zero(T), + interp_type=Interpolations.BSpline(Linear()), +) where {T,N} + isnothing(order) && throw(ArgumentError( + "providing the polynomial order is mandatory for `PolynomialMode`", + )) + + return get_function_poly( + data, + order; + super_sampling=super_sampling, + extrapolation_bc=extrapolation_bc, + interp_type=interp_type, + ) +end diff --git a/src/transformators_old.jl b/src/transformators_old.jl new file mode 100644 index 0000000..ebcbff5 --- /dev/null +++ b/src/transformators_old.jl @@ -0,0 +1,606 @@ +using Interpolations +using FourierTools +using StaticArrays +using EvalMultiPoly + +export get_interpolated_function, get_function_tuple, get_function_svec, get_function_affine +export get_function_poly, get_function_poly_inplace +export add_dim, red_dim_apply, red_dim, mat_mul, func_transform +export extrapolate, interpolate + +export PolynomialMode, AffineMode + +""" + get_function(data::AbstractArray; super_sampling=2, extrapolation_bc=Flat(), interp_type=Interpolations.BSpline(Linear())) + +returns a function `dat(shift, zoom)` which generates a shifted and scaled version of the original data. +This is useful for fitting with a function which is itself defined by measured data. + +# Arguments +`data`: The data to represent by the function `dat` +`super_sampling`: The factor by which the data is internally represented as a supersampled version (Fourier-based upsampling, see `FourierTools.resample`) +`extrapolation_bc`: The extrapolation boundary condition to select for values outside the range. + By default the value 0.0 is used. Other options are `Flat()`, or `Line()`, See the package `Interpolation` for details. +`interp_type`: The type of interpolation to use. See the package `Interpolation` for details. + + +""" +function get_function_old(data::AbstractArray; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) + new_size = super_sampling.*size(data) + upsampled = fftshift(resample(ifftshift(data), new_size)) + # @show upsampled + # return upsampled + # itp = LinearInterpolation(axes(upsampled), upsampled, extrapolation_bc=extrapolation_bc); + interpolation = Interpolations.interpolate(upsampled, interp_type) + interpolation = extrapolate(interpolation, extrapolation_bc) + # center of the original data (too keep the axis and number of datapointsi dentical to the original) + center_orig = (size(data) .÷2 .+1) + # create zero-centered original ranges (== axes) + zero_axes = Tuple(ax .- c for (ax, c) in zip(axes(data), center_orig)) + # center of the upsampled data. This is where to access the upsampled data + function zoomed(shift, zoom) + zoom = zoom .* super_sampling + # careful: The center of the original data is not at the expected position! But rather at: + center_upsamp = new_size .÷2 .+1 # ((center_orig .-1) .*super_sampling .+1) # new_size .÷2 .+1 + scaled_axes = ((ax.-myc) .* z .+ cen for (ax, myc, cen, z) in zip(zero_axes, shift, center_upsamp, zoom)) + # @show Tuple(scaled_axes) + return interpolation[scaled_axes...] + # return extrapolate(scale(interpolation, scaled_axes...), extrapolation_bc) + end + + return zoomed + + zoomed(p) = zoomed([p[1], p[2]], [p[3], p[4]]) + # return (pos) -> interp_linear((center .+ pos)...) + # fitp(t) = interp_linear(t...) + # @time res1 = fitp.(tcoords); # 1 sec + # function my_zoom + +end + +# """ +# add_dim(cind) + +# adds a dimension to a CartesianIndex + +# `cind`: A CartesianIndex +# """ +# function add_dim(cind) +# return SVector.((Tuple(cind))..., 1) +# end + +@inline@inline function add_dim(v::NTuple{N,T}) where {N,T} + SVector{N + 1,T}(v..., one(T)) +end + + +@inline function red_dim(v::StaticVector{N,T}) where {N,T} + last_inv = inv(v[N]) + + return ntuple( + i -> v[i] * last_inv, + Val(N - 1), + ) +end + + +@inline function mat_mul( + v::SVector, + A::SMatrix, +) + A * v +end + +# """ +# red_dim(svec::SVector{S,T}) + +# removes the last dimension of a SVector to convert it from a homogeneous to a Cartesian coordinates + +# `svec::SVector{S,T}`: A SVector +# """ +# # @inline function red_dim(svec::SVector{S,T})::SVector{S-1,T} where {S,T} +# # return @view svec[1:S-1] +# # end +# @inline function red_dim(v::StaticVector{N,T}) where {N,T} +# last_inv = inv(v[N]) +# return ntuple(i -> v[i] * last_inv, Val(N - 1)) +# end + +""" + red_dim_apply(fct, svec::SVector{S,T}) + +applies a function to a SVector by removing the last dimension to convert it from a homogeneous to a Cartesian coordinates + +`fct`: The function to apply +`svec::SVector{S,T}`: A SVector +""" +@inline function red_dim_apply(fct, svec::SVector{S,T}) where {S,T} + return fct((@view svec[1:S-1])...) +end + +""" + red_dim_apply(fct, tup::NTuple{S,T}) + +applies a function to a Tuple by removing the last dimension to convert it from a homogeneous to a Cartesian coordinates + +`fct`: The function to apply +`tup::NTuple{S,T}`: A Tuple +""" +@inline function red_dim_apply(fct, tup::NTuple{S,T}) where {S,T} + return fct(tup[1:S-1]...) +end + +""" + idx_apply(fct, svec::SVector{S,T}) where {S,T} + +applies a function to a SVector + +`fct`: The function to apply +`svec::SVector{S,T}`: A SVector +""" +@inline function idx_apply(fct, svec::SVector{S,T})::Number where {S,T} + return fct(svec...) +end + +""" + idx_apply(fct, tup::NTuple{S,T}) where {S,T} + +applies a function to a Tuple + +`fct`: The function to apply +`tup::NTuple{S,T}`: A Tuple +""" +@inline function idx_apply(fct, tup::NTuple{S,T})::Number where {S,T} + return fct(tup...) +end + +# """ +# mat_mul(t::SVector{N, T2}, matrix_c::SMatrix{N,N,T}) + +# multiplies a SVector with a SMatrix + +# `t::SVector{N, T2}`: The SVector to multiply +# `matrix_c::SMatrix{N,N,T}`: The SMatrix to multiply with +# """ +# @inline function mat_mul(t::SVector{N, T2}, matrix_c::SMatrix{N,N,T})::SVector{N,T} where {N,T, T2} +# return matrix_c * t +# end + +""" + func_transform(t, coord_transform_func::Function)::SVector + +applies a coordinate transformation function to an array or `CartesianIndex` and returns the transformed array + +`t`: The array or `CartesianIndex` to transform +`coord_transform_func::Function`: The function to apply the transformation +""" +@inline function func_transform(t, coord_transform_func::Function)::SVector + return coord_transform_func(Tuple(t)) +end + +""" + func_transform_tup(t, coord_transform_func::Function) + +applies a coordinate transformation function to a Tuple + +`t`: The array or `CartesianIndex` to transform +`coord_transform_func::Function`: The function to apply the transformation +""" +@inline function func_transform_tup(t, coord_transform_func::Function) + return coord_transform_func(Tuple(t)) +end + +@inline function transform_and_sample(itp, coord_transf_func, I) + coords = coord_transf_func(Tuple(I)) + return idx_apply(itp, coords) +end + +""" + apply_transform(coord_transf_func::Function, data::AbstractArray{T}, itp) where {T} + +applies a general coordinate transformation function to the indices of an array and returns the transformed array + +`coord_transf_func:Function`: A function that takes a N-+1 dimensional CartesianIndex and returns a new N+1 dimensional SVector or Tuple +`data::AbstractArray{T}`: The data to transform +`itp`: The interpolation object to use +""" +# function apply_transform(coord_transf_func::Function, data::AbstractArray{T, N}, itp) where {T, N} #, out::AbstractArray{T}) where {T} +# # @info "Applying tuple transformation" + +# # return map((it) -> idx_apply(Interpolations.adapt(gpu_or_cpu(nothing), itp), coord_transf_func(Tuple(it))), CartesianIndices(data)) +# return idx_apply.(Ref(Interpolations.adapt(gpu_or_cpu(nothing), itp)), coord_transf_func.(Tuple.(CartesianIndices(data)))); +# #return idx_apply.(Ref(itp), coord_transf_func.(Tuple.(CartesianIndices(data)))); +# # return idx_apply.(Ref(itp), coord_transf_func.(CartesianIndices(data))); +# end +function apply_transform(coord_transf_func, data, itp) + inds = CartesianIndices(axes(data)) + + sample(I) = transform_and_sample( + itp, + coord_transf_func, + I, + ) + + return sample.(inds) +end + +# """ +# apply_transform_homogen(coord_transf_func::Function, data, itp) +# applies a homogeneous coordinate-based coordinate transformation function to the indices of an array and returns the transformed array + +# `coord_transf_func::Function`: A function that takes a N-+1 dimensional homogeneous SVector returns a new N+1 dimensional SVector +# `data`: The data to transform +# `itp`: The interpolation object to use +# """ +# function apply_transform_homogen(coord_transf_func::Function, data, itp)#, out) +# h_coord_transf_func = (c) -> red_dim(coord_transf_func(add_dim(c))) +# #@info "Applying homogeneous transformation" +# return apply_transform(h_coord_transf_func, data, itp)#, out); +# # out .= itp.(red_dim.(coord_transf_func.(add_dim.(CartesianIndices(data))))); +# # out .= red_dim_apply.(Ref(itp), coord_transf_func.(add_dim.(CartesianIndices(data)))); +# end + +# """ +# apply_transform_affine(mymat::SMatrix{T}, data, itp) where T + +# applies an affine transformation matrix to the indices of an array and returns the transformed array + +# `mymat::SMatrix{T}` The affine transformation matrix to apply +# `data`: The data to transform +# `itp`: The interpolation function (object) to use +# """ +# function apply_transform_affine(mymat::SMatrix{T}, data, itp) where T #, out) where {T} # The SMatrix spec is important to avoid allocations +# # red_dim_apply.(Ref(itp), func_transform.(CartesianIndices(data), Ref(coord_transf_func))); +# # return red_dim_apply.(Ref(itp), mat_mul.(add_dim.(CartesianIndices(data)), Ref(mymat))); +# # @info "Applying affine transformation" +# homogenous_transform = (c) -> mat_mul(c, mymat) +# return apply_transform_homogen(homogenous_transform, data, itp)#, out); +# #return out +# end +@inline mat_mul(v, mymat) = mymat * v + + +function apply_transform_homogen( + coord_transf_func, + data, + itp, +) + apply_transform( + x -> red_dim( + coord_transf_func( + add_dim(x), + ) + ), + data, + itp, + ) +end + + +function apply_transform_affine( + mymat, + data, + itp, +) + apply_transform_homogen( + x -> mat_mul(x, mymat), + data, + itp, + ) +end + +""" + get_function_tuple(data::AbstractArray, fct_tup::Function; super_sampling=2, extrapolation_bc=Flat(), interp_type=Interpolations.BSpline(Linear())) + +returns a function `dat(shift, zoom)` which generates a shifted and scaled version of the original data. +This is useful for fitting with a function which is itself defined by measured data. + +# Arguments +`data`: The data to represent by the function `dat` +`fct_tup`: The function to apply to the data +`super_sampling`: The factor by which the data is internally represented as a supersampled version (Fourier-based upsampling, see `FourierTools.resample`) +`extrapolation_bc`: The extrapolation boundary condition to select for values outside the range. + By default the value 0.0 is used. Other options are `Flat()`, or `Line()`, See the package `Interpolation` for details. +`interp_type`: The type of interpolation to use. See the package `Interpolation` for details. + +# Example + +""" +function get_function_tuple(data::AbstractArray{T, N}, fct_tup::Function; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) where {T, N} + # building the extraplation + interpolation object + itp = extrapolate(interpolate(data, interp_type), extrapolation_bc); + function interpolated(params)#, out = similar(data)) + fct_tup_noparams(ci) = fct_tup(ci, params) + #@show Interpolations.adapt(gpu_or_cpu(1), itp) + return apply_transform(fct_tup_noparams, data, itp); + end + return interpolated +end + +""" + get_function_svec(data::AbstractArray, fct_hom::Function; super_sampling=1, extrapolation_bc=Flat(), interp_type=Interpolations.BSpline(Linear())) + +returns a function `interpolated(params)` which generates a transformed version of the original data parameterized by transform parameters. + +# Arguments +`data`: The data to represent by the function `dat` +`fct_hom`: The function to apply to the data +`super_sampling`: The factor by which the data is internally represented as a supersampled version (Fourier-based upsampling, see `FourierTools.resample`) +`extrapolation_bc`: The extrapolation boundary condition to select for values outside the range. + By default the value 0.0 is used. Other options are `Flat()`, or `Line()`, See the package `Interpolation` for details. +`interp_type`: The type of interpolation to use. See the package `Interpolation` for details. +""" +function get_function_svec(data::AbstractArray{T}, fct_hom::Function; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) where T + # building the extraplation + interpolation object + itp = extrapolate(interpolate(data, interp_type), extrapolation_bc); + function interpolated(params::SVector) #, out = similar(data)) + fct_hom_noparams(c) = fct_hom(c, params) + return apply_transform_homogen(fct_hom_noparams, data, itp)#, out); + # return out; + end + return interpolated +end + +""" + get_function_affine(data::AbstractArray; super_sampling=1, extrapolation_bc=Flat(), interp_type=Interpolations.BSpline(Linear())) + +returns a function `interpolated()` which generates a transformed version of the original data parameterized by transform parameters or by transformation matrix. +This is useful for fitting with a function which is itself defined by measured data. +The returned function supports two ways to be used, with an affine transform matrix `matrix_c` as in input or with a vector `p` of parameters. + + +# Arguments +`data`: The data to represent by the function `dat` +`super_sampling`: The factor by which the data is internally represented as a supersampled version (Fourier-based upsampling, see `FourierTools.resample`) +`extrapolation_bc`: The extrapolation boundary condition to select for values outside the range. + By default the value 0.0 is used. Other options are `Flat()`, or `Line()`, See the package `Interpolation` for details. +`interp_type`: The type of interpolation to use. See the package `Interpolation` for details. + +# Example + +```julia +julia> dat1 = reshape(1:16,(4,4)) +4×4 reshape(::UnitRange{Int64}, 4, 4) with eltype Int64: + 1 5 9 13 + 2 6 10 14 + 3 7 11 15 + 4 8 12 16 + +julia> affine_func = get_function_affine(Float32.(dat1)); + +julia> homogeneous_transform = [1 0 -1; 0 1 1; 0 0 1] # translates by [1,1] +3×3 Matrix{Int64}: + 1 0 -1 + 0 1 1 + 0 0 1 + +julia> affine_func(SMatrix{3,3}(homogeneous_transform)) +4×4 Matrix{Float32}: + 0.0 0.0 0.0 0.0 + 5.0 9.0 13.0 0.0 + 6.0 10.0 14.0 0.0 + 7.0 11.0 15.0 0.0 +``` +""" +function get_function_affine(data::AbstractArray{T}; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) where T + #new_size = super_sampling.*size(data) + #upsampled = fftshift(resample(ifftshift(data), new_size)) + + # building the extraplation + interpolation object + itp = extrapolate(interpolate(data, interp_type), extrapolation_bc); + + function interpolated(matrix_c::SMatrix{T}) where T #, out = similar(data)) where T1 + return apply_transform_affine(matrix_c, data, itp)# , out); + # return out; + end + + function interpolated(p::AbstractVector{T1}) where {T1} #, out = similar(data)) where T1 + x_cen, y_cen = (size(data) .÷ 2.0 .+1) + # x_cen_up, y_cen_up = (size(upsampled) .÷ 2.0 .+ 1.0) + + # creating the matrices of rotation, shear, scale, and shift + rot_mat = @SMatrix T[cos(p[7]) -1.0*sin(p[7]) 0.0; sin(p[7]) cos(p[7]) 0.0; 0.0 0.0 1.0]; + shear_mat = @SMatrix T[1.0 p[5] 0.0; p[6] 1.0 0.0; 0.0 0.0 1.0]; + scale_mat = @SMatrix T[1/p[3] 0.0 0.0; 0.0 1/p[4] 0.0; 0.0 0.0 1.0]; + shift_mat = @SMatrix T[1.0 0.0 -1*p[1]; 0.0 1.0 -1*p[2]; 0.0 0.0 1.0]; + t_to_origin = @SMatrix T[1.0 0.0 1*x_cen; 0.0 1.0 y_cen; 0.0 0.0 1.0]; + t_to_center = @SMatrix T[1.0 0.0 -1.0*x_cen; 0.0 1.0 -1.0*y_cen; 0.0 0.0 1.0]; + # t_orig_upsampled = SMatrix{3, 3}(T[1.0 0.0 -1.0*x_cen_up; 0.0 1.0 -1.0*y_cen_up; 0.0 0.0 1.0]); + + # building the overall transformation matrix + matrix_c = t_to_origin * scale_mat * rot_mat * shear_mat *shift_mat * t_to_center + + return apply_transform_affine(matrix_c, data, itp) #, out); # do not call interolated here for type stability reasons + end + + + function interpolated(p::NTuple{N, T}) where {N, T} #, out = similar(data)) where T1 + x_cen, y_cen = (size(data) .÷ 2.0 .+1) + # x_cen_up, y_cen_up = (size(upsampled) .÷ 2.0 .+ 1.0) + + # creating the matrices of rotation, shear, scale, and shift + rot_mat = @SMatrix T[cos(p[7]) -1.0*sin(p[7]) 0.0; sin(p[7]) cos(p[7]) 0.0; 0.0 0.0 1.0]; + shear_mat = @SMatrix T[1.0 p[5] 0.0; p[6] 1.0 0.0; 0.0 0.0 1.0]; + scale_mat = @SMatrix T[1/p[3] 0.0 0.0; 0.0 1/p[4] 0.0; 0.0 0.0 1.0]; + shift_mat = @SMatrix T[1.0 0.0 -1*p[1]; 0.0 1.0 -1*p[2]; 0.0 0.0 1.0]; + t_to_origin = @SMatrix T[1.0 0.0 1*x_cen; 0.0 1.0 y_cen; 0.0 0.0 1.0]; + t_to_center = @SMatrix T[1.0 0.0 -1.0*x_cen; 0.0 1.0 -1.0*y_cen; 0.0 0.0 1.0]; + # t_orig_upsampled = SMatrix{3, 3}(T[1.0 0.0 -1.0*x_cen_up; 0.0 1.0 -1.0*y_cen_up; 0.0 0.0 1.0]); + + # building the overall transformation matrix + matrix_c = t_to_origin * scale_mat * rot_mat * shear_mat *shift_mat * t_to_center + + return apply_transform_affine(matrix_c, data, itp) #, out); # do not call interolated here for type stability reasons + end + + return interpolated +end + + +""" + get_function_poly(data::AbstractArray, order; super_sampling=1, extrapolation_bc=Flat(), interp_type=Interpolations.BSpline(Linear())) + +returns a function `interpolated(p, [out])` which generates a transformed version of the original data parameterized by transform parameters. +This is useful for fitting with a function which is itself defined by measured data. +The returned function supports two ways to be used, with an affine transform matrix `p` as in input or with a vector `p` of parameters. +The optional argument `out` can be used to store the result of the transformation. + + +# Arguments +`data`: The data to represent by the function `dat` +`order`: The order of the polynomial to use for the transformation +`super_sampling`: The factor by which the data is internally represented as a supersampled version (Fourier-based upsampling, see `FourierTools.resample`) +`extrapolation_bc`: The extrapolation boundary condition to select for values outside the range. + By default the value 0.0 is used. Other options are `Flat()`, or `Line()`, See the package `Interpolation` for details. +`interp_type`: The type of interpolation to use. See the package `Interpolation` for details. +""" +# function get_function_poly(data::AbstractArray{T, N}, order; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) where {T, N} +# pm = get_multi_poly(Val(ndims(data)), Val(order)) +# return get_function_tuple(data, pm; super_sampling= super_sampling, extrapolation_bc=extrapolation_bc, interp_type=interp_type); +# end + +@inline function poly_sample(itp, poly, coeffs, I) + coords = poly(Tuple(I), coeffs) + return itp(coords...) +end + + +""" + get_function_poly(data, ::Val{N_order}) + +Returns a function which warps `data` using a polynomial coordinate +transformation of order `N_order`. + +The returned function is non-mutating and AD-compatible. +Its forward evaluation allocates only the returned output array. +""" +function get_function_poly( + data::AbstractArray{T,N}, + ::Val{N_order}; + extrapolation = zero(T), +) where {T,N,N_order} + + itp = extrapolate( + interpolate(data, BSpline(Linear())), + extrapolation, + ) + + poly = get_multi_poly( + Val(N), + Val(N_order), + ) + + inds = CartesianIndices(data) + + function interpolated(coeffs) + return map(inds) do I + poly_sample( + itp, + poly, + coeffs, + I, + ) + end + end + + return interpolated +end + +""" + get_function_poly_inplace(data, ::Val{N_order}) + +Returns an in-place polynomial warp function + + f!(out, coeffs) + +which writes the warped image into `out`. +""" +function get_function_poly_inplace( + data::AbstractArray{T,N}, + ::Val{N_order}; + extrapolation = zero(T), +) where {T,N,N_order} + + itp = extrapolate( + interpolate(data, BSpline(Linear())), + extrapolation, + ) + + poly = get_multi_poly( + Val(N), + Val(N_order), + ) + + function interpolated!( + out, + coeffs, + ) + @inbounds for I in CartesianIndices(out) + out[I] = poly_sample( + itp, + poly, + coeffs, + I, + ) + end + + return out + end + + return interpolated! +end + + +""" + get_interpolated_function(data::AbstractArray, ::Type{AffineMode}; super_sampling=2, extrapolation_bc=Flat(), interp_type=Interpolations.BSpline(Linear())) + +returns a function `interpolated(p)` which generates a transformed version of the original data parameterized by transform parameters. +This is useful for fitting with a function which is itself defined by measured data. +The returned function supports two ways to be used, with an affine transform matrix `p` as in input or with a vector or tuple `p` of parameters. + +# Arguments +`data`: The data to represent by the function `dat` +`AffineMode`: The transformation mode to use +`super_sampling`: The factor by which the data is internally represented as a supersampled version (Fourier-based upsampling, see `FourierTools.resample`) +`extrapolation_bc`: The extrapolation boundary condition to select for values outside the range. + By default the value 0.0 is used. Other options are `Flat()`, or `Line()`, See the package `Interpolation` for details. +`interp_type`: The type of interpolation to use. See the package `Interpolation` for details. + +# Returns +A function `interpolated(p)` which generates a transformed version of the original data parameterized by transform parameters +""" +function get_interpolated_function(data::AbstractArray{T, N}, ::Type{AffineMode}; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) where {T, N} + return get_function_affine(data; super_sampling=super_sampling, extrapolation_bc=extrapolation_bc, interp_type=interp_type) +end + +function get_interpolated_function(data::AbstractArray{T, N}; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) where {T, N} + @warn "No transformation mode provided. `AffineMode` is used as default" + return get_interpolated_function(data, AffineMode; super_sampling=super_sampling, extrapolation_bc=extrapolation_bc, interp_type=interp_type) +end + +""" + get_interpolated_function(data::AbstractArray, ::Type{PolynomialMode}, order=nothing; super_sampling=2, extrapolation_bc=Flat(), interp_type=Interpolations.BSpline(Linear())) + +returns a function `interpolated(p)` which generates a transformed version of the original data parameterized by transform parameters. +This is useful for fitting with a function which is itself defined by measured data. +The returned function supports polynomial transformations of the data. + +# Arguments +`data`: The data to represent by the function `dat` +`PolynomialMode`: The transformation mode to use +`order`: The order of the polynomial to use for the transformation +`super_sampling`: The factor by which the data is internally represented as a supersampled version (Fourier-based upsampling, see `FourierTools.resample`) +`extrapolation_bc`: The extrapolation boundary condition to select for values outside the range. + By default the value 0.0 is used. Other options are `Flat()`, or `Line()`, See the package `Interpolation` for details. +`interp_type`: The type of interpolation to use. See the package `Interpolation` for details. + +# Returns +A function `interpolated(p)` which generates a transformed version of the original data parameterized by polynomial transform parameters +""" +function get_interpolated_function(data::AbstractArray{T, N}, ::Type{PolynomialMode}, order=nothing; super_sampling=2, extrapolation_bc=zero(eltype(data)), interp_type=Interpolations.BSpline(Linear())) where {T, N} + #TODO from CuArray to CuArray + if isnothing(order) + error("Providing the order of the transformation polynomial is mandatory for the `PolynomialMode`") + end + return get_function_poly(data, order; super_sampling=super_sampling, extrapolation_bc=extrapolation_bc, interp_type=interp_type) +end \ No newline at end of file diff --git a/test/Aqua.jl b/test/Aqua.jl new file mode 100644 index 0000000..e7b0cce --- /dev/null +++ b/test/Aqua.jl @@ -0,0 +1,6 @@ +using Aqua + +Aqua.test_all( + DataToFunctions, + unbound_args=false, + ) \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index ddd940c..8eda64a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,32 +1,38 @@ using Test -# using Zygote +using Zygote using DataToFunctions -@testset "get_function" begin - data = rand(40,41) - for supersamp = 1:5 - f = get_function(data; super_sampling=supersamp); - @test f((0.0,0.0),(1.0,1.0)) ≈ data - end -end +include("Aqua.jl") -@testset "gradient" begin - data = rand(11,10) - f = get_function(data; super_sampling=2); - loss(p,z) = sum(abs2.(f(p, z) .- data)) - @test loss((0.0,0.0),(1.0,1.0)) < 1e-20 - @test loss((0.0,0.001),(1.0,1.0)) > 1e-20 - @test loss((0.0,0.0),(1.0001,1.0)) > 1e-20 +include("test_transformators.jl") - # throws an error... - # Zygote.gradient(loss, (0.0,0.0), (1.0,1.0)) -end +# @testset "get_function_affine" begin +# data = rand(40,41) +# f = get_function_affine(data); +# @test f([0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0]) ≈ data + +# end -@testset "keep center" begin - data = ones(5,4); data[3,3] = 5.0; - f = get_function(data; super_sampling=5); - @test f((0.0,0.0),(2.0,2.0))[3,3] ≈ 5.0 +# @testset "loss" begin +# data = rand(11,10) +# f = get_function_affine(data; super_sampling=2); +# loss(p) = sum(abs2.(f(p) .- data)) +# @test loss([0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0]) < 1e-20 +# @test loss([0.001, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0]) > 1e-20 +# @test loss([0.0, 0.0, 1.001, 1.0, 0.0, 0.0, 0.0]) > 1e-20 +# end - # throws an error... - # Zygote.gradient(loss, (0.0,0.0), (1.0,1.0)) -end +# @testset "gradient" begin +# data = rand(11,10) +# f = get_function_affine(data; super_sampling=2); +# loss(p) = sum(abs2.(f(p) .- data)) +# st_vals = [0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0] +# # throws an error... +# @test Zygote.gradient(loss, st_vals)[1] ≈ zeros(7) +# end + +# @testset "keep center" begin +# data = ones(5,4); data[3,3] = 5.0; +# f = get_function(data); +# @test f([0.0, 0.0, 2.0, 2.0, 0.0, 0.0, 0.0])[3,3] ≈ 5.0 +# end diff --git a/test/test_transformators.jl b/test/test_transformators.jl new file mode 100644 index 0000000..25f2149 --- /dev/null +++ b/test/test_transformators.jl @@ -0,0 +1,410 @@ +using Test +using DataToFunctions +using Interpolations +using EvalMultiPoly +using Zygote +using ForwardDiff + +# ----------------------------------------------------------------------------- +# Helpers +# ----------------------------------------------------------------------------- + +smooth_image(::Type{T}, n=48) where {T} = T[ + sin(T(0.071) * i) + + cos(T(0.053) * j) + + T(0.001) * i * j + for i in 1:n, j in 1:n +] + +function tuple_isapprox(a::Tuple, b::Tuple; rtol=1e-10, atol=1e-10) + length(a) == length(b) || return false + return all(isapprox.(a, b; rtol=rtol, atol=atol)) +end + +# ----------------------------------------------------------------------------- +# Tuple primitives and general transforms +# ----------------------------------------------------------------------------- + +@testset "tuple primitives" begin + h = DataToFunctions.add_dim((2.0, 3.0)) + @test h == (2.0, 3.0, 1.0) + + # Heterogeneous tuple regression: the final coordinate may remain Int. + @test tuple_isapprox( + DataToFunctions.red_dim((4.0, 6.0, 2)), + (2.0, 3.0), + ) + + A = ( + (1.0, 2.0, 0.0), + (0.0, 1.0, 0.0), + (0.0, 0.0, 1.0), + ) + + @test DataToFunctions.mat_mul((2.0, 3.0, 1.0), A) == (8.0, 3.0, 1.0) +end + +@testset "tuple and homogeneous transforms" begin + data = smooth_image(Float64, 24) + + shift_tuple(c, p) = (c[1] + p[1], c[2] + p[2]) + + f = get_function_tuple(data, shift_tuple) + f! = get_function_tuple_inplace(data, shift_tuple) + + p = (0.2, -0.15) + y = f(p) + out = similar(data) + + f!(out, p) + @test y ≈ out + + # Warm first; only the steady-state call is measured. + f!(out, p) + @test @allocated(f!(out, p)) == 0 + + homogeneous_shift(h, p) = ( + h[1] + p[1] * h[3], + h[2] + p[2] * h[3], + h[3], # deliberately leaves this as Int for integer input coordinates + ) + + fh = get_function_homogen(data, homogeneous_shift) + fh! = get_function_homogen_inplace(data, homogeneous_shift) + + yh = fh(p) + fh!(out, p) + @test yh ≈ out + + fh!(out, p) + @test @allocated(fh!(out, p)) == 0 +end + +@testset "arbitrary-dimensional affine matrix" begin + data = reshape( + collect(Float64, 1:(8 * 7 * 6)), + 8, 7, 6, + ) + + f = get_function_affine(data) + f! = get_function_affine_inplace(data) + + A4 = ( + (1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.0), + (0.0, 0.0, 0.0, 1.0), + ) + + @test f(A4) ≈ data + + out = similar(data) + f!(out, A4) + @test out ≈ data + + f!(out, A4) + @test @allocated(f!(out, A4)) == 0 +end + +# ----------------------------------------------------------------------------- +# Affine transform: forward, allocation behavior, AD, optimization +# ----------------------------------------------------------------------------- + +@testset "affine forward and allocation behavior" begin + data = smooth_image(Float64, 48) + + f = get_function_affine(data) + f! = get_function_affine_inplace(data) + + identity_p = (0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0) + @test f(identity_p) ≈ data + + p = (0.25, -0.15, 1.01, 0.99, 0.002, -0.003, 0.01) + + y = f(p) + out = similar(data) + f!(out, p) + + @test y ≈ out + + # Strict zero-allocation CPU hot-path check. + f!(out, p) + @test @allocated(f!(out, p)) == 0 +end + +@testset "affine analytical coordinates" begin + data = smooth_image(Float64, 32) + + p = (0.25, -0.15, 1.01, 0.99, 0.002, -0.003, 0.01) + + x_cen = size(data, 1) ÷ 2 + 1 + y_cen = size(data, 2) ÷ 2 + 1 + + A = DataToFunctions._affine_matrix_2d(p, x_cen, y_cen) + + for (x, y) in ((1, 1), (7, 13), (16, 16), (25, 20)) + matrix_result = DataToFunctions.red_dim(DataToFunctions.mat_mul(DataToFunctions.add_dim((x, y)), A)) + + shift_x, shift_y, + scale_x, scale_y, + shear_xy, shear_yx, + θ = p + + a = x - x_cen - shift_x + b = y - y_cen - shift_y + + u = a + shear_xy * b + v = shear_yx * a + b + + c = cos(θ) + s = sin(θ) + + r1 = c * u - s * v + r2 = s * u + c * v + + analytical_result = ( + x_cen + r1 / scale_x, + y_cen + r2 / scale_y, + ) + + @test all( + isapprox.( + matrix_result, + analytical_result; + rtol=1e-12, + atol=1e-12, + ), + ) + end +end + +@testset "affine Zygote rrule agrees with ForwardDiff" begin + data = smooth_image(Float64, 36) + f = get_function_affine(data) + + p = (0.25, -0.15, 1.01, 0.99, 0.002, -0.003, 0.01) + + target = f(( + 0.18, -0.11, 1.008, 0.994, 0.0015, -0.002, 0.007, + )) + + loss(q) = sum(abs2, f(q) .- target) + + g_zygote = Zygote.gradient(loss, p)[1] + + pvec = collect(p) + loss_fd(q) = loss(Tuple(q)) + g_forward = ForwardDiff.gradient(loss_fd, pvec) + + @test all(isfinite, g_zygote) + @test collect(g_zygote) ≈ g_forward rtol=2e-6 atol=2e-7 +end + +# ----------------------------------------------------------------------------- +# Polynomial transform: forward, allocation behavior, AD +# ----------------------------------------------------------------------------- + +@testset "polynomial forward and allocation behavior" begin + data = smooth_image(Float64, 40) + + order = Val(2) + f = get_function_poly(data, order) + f! = get_function_poly_inplace(data, order) + + c0 = get_identity_multipoly_coeffs(Val(2), Val(2)) + + @test f(c0) ≈ data + + c = ntuple( + k -> c0[k] + 1e-6 * k, + length(c0), + ) + + y = f(c) + out = similar(data) + f!(out, c) + + @test y ≈ out + + f!(out, c) + @test @allocated(f!(out, c)) == 0 +end + +@testset "polynomial Zygote rrule agrees with ForwardDiff" begin + data = smooth_image(Float64, 28) + f = get_function_poly(data, Val(2)) + + c0 = get_identity_multipoly_coeffs(Val(2), Val(2)) + c = ntuple(k -> c0[k] + 5e-7 * k, length(c0)) + + target = f( + ntuple(k -> c0[k] - 3e-7 * k, length(c0)), + ) + + loss(q) = sum(abs2, f(q) .- target) + + g_zygote = Zygote.gradient(loss, c)[1] + + cvec = collect(c) + loss_fd(q) = loss(Tuple(q)) + g_forward = ForwardDiff.gradient(loss_fd, cvec) + + @test all(isfinite, g_zygote) + @test collect(g_zygote) ≈ g_forward rtol=2e-6 atol=2e-7 +end + +# ----------------------------------------------------------------------------- +# Optional CUDA tests +# +# CUDA is intentionally not a dependency of DataToFunctions itself. +# These tests run only when CUDA is installed and a functional device exists. +# ----------------------------------------------------------------------------- + +if Base.find_package("CUDA") !== nothing + @eval using CUDA + + if CUDA.functional() + @testset "CUDA affine and polynomial forward paths" begin + CUDA.allowscalar(false) + + data_cpu = smooth_image(Float32, 48) + data_gpu = CuArray(data_cpu) + + # Regression test for GPU interpolant storage: + # constructing directly from a CuArray can leave CPU coefficients + # inside Interpolations.jl. DataToFunctions must construct on CPU + # and Adapt the completed interpolation object back to the backend. + itp_gpu = DataToFunctions._make_interpolant( + data_gpu, + Interpolations.BSpline(Linear()), + 0f0, + ) + + coords_gpu = CuArray(Float32[1.25, 2.5, 3.75]) + vals_gpu = itp_gpu.(coords_gpu, coords_gpu) + @test vals_gpu isa CuArray + @test all(isfinite, Array(vals_gpu)) + + p = ( + 0.25f0, -0.15f0, + 1.01f0, 0.99f0, + 0.002f0, -0.003f0, + 0.01f0, + ) + + f_cpu = get_function_affine(data_cpu) + f_gpu = get_function_affine(data_gpu) + f_gpu! = get_function_affine_inplace(data_gpu) + + y_cpu = f_cpu(p) + y_gpu = f_gpu(p) + + @test y_gpu isa CuArray + @test Array(y_gpu) ≈ y_cpu rtol=2e-5 atol=2e-5 + + out_gpu = similar(data_gpu) + f_gpu!(out_gpu, p) + @test Array(out_gpu) ≈ y_cpu rtol=2e-5 atol=2e-5 + + c0 = get_identity_multipoly_coeffs(Val(2), Val(2)) + c = ntuple(k -> Float32(c0[k] + 1e-6 * k), length(c0)) + + pf_cpu = get_function_poly(data_cpu, Val(2)) + pf_gpu = get_function_poly(data_gpu, Val(2)) + pf_gpu! = get_function_poly_inplace(data_gpu, Val(2)) + + py_cpu = pf_cpu(c) + py_gpu = pf_gpu(c) + + @test py_gpu isa CuArray + @test Array(py_gpu) ≈ py_cpu rtol=2e-5 atol=2e-5 + + pout_gpu = similar(data_gpu) + pf_gpu!(pout_gpu, c) + @test Array(pout_gpu) ≈ py_cpu rtol=2e-5 atol=2e-5 + end + + @testset "CUDA custom VJPs" begin + CUDA.allowscalar(false) + + data_cpu = smooth_image(Float32, 32) + data_gpu = CuArray(data_cpu) + + affine_gpu = get_function_affine(data_gpu) + + p = ( + 0.15f0, -0.10f0, + 1.006f0, 0.996f0, + 0.001f0, -0.0015f0, + 0.005f0, + ) + + target_affine = affine_gpu(( + 0.10f0, -0.06f0, + 1.003f0, 0.998f0, + 0.0005f0, -0.001f0, + 0.003f0, + )) + + affine_loss_gpu(q) = + sum(abs2, affine_gpu(q) .- target_affine) + + g_affine_gpu = Zygote.gradient( + affine_loss_gpu, + p, + )[1] + + @test all(isfinite, g_affine_gpu) + + affine_cpu = get_function_affine(data_cpu) + target_affine_cpu = Array(target_affine) + affine_loss_cpu(q) = + sum(abs2, affine_cpu(q) .- target_affine_cpu) + g_affine_cpu = Zygote.gradient( + affine_loss_cpu, + p, + )[1] + + @test collect(g_affine_gpu) ≈ collect(g_affine_cpu) rtol=2e-3 atol=2e-3 + + poly_gpu = get_function_poly(data_gpu, Val(2)) + c0 = get_identity_multipoly_coeffs(Val(2), Val(2)) + c = ntuple( + k -> Float32(c0[k] + 5e-7 * k), + length(c0), + ) + + target_poly = poly_gpu( + ntuple( + k -> Float32(c0[k] - 3e-7 * k), + length(c0), + ), + ) + + poly_loss_gpu(q) = + sum(abs2, poly_gpu(q) .- target_poly) + + g_poly_gpu = Zygote.gradient( + poly_loss_gpu, + c, + )[1] + + @test all(isfinite, g_poly_gpu) + + poly_cpu = get_function_poly(data_cpu, Val(2)) + target_poly_cpu = Array(target_poly) + poly_loss_cpu(q) = + sum(abs2, poly_cpu(q) .- target_poly_cpu) + g_poly_cpu = Zygote.gradient( + poly_loss_cpu, + c, + )[1] + + @test collect(g_poly_gpu) ≈ collect(g_poly_cpu) rtol=3e-3 atol=3e-3 + end + else + @info "CUDA.jl is installed, but no functional CUDA device is available; CUDA tests skipped." + end +else + @info "CUDA.jl is not installed; CUDA tests skipped." +end