utax 0.0.2__tar.gz

utax-0.0.2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Aymeric Galan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

utax-0.0.2/PKG-INFO

@@ -0,0 +1,32 @@
+Metadata-Version: 2.4
+Name: utax
+Version: 0.0.2
+Summary: Utility functions for signal processing, compatible with the differentiable programming library JAX.
+Home-page: https://github.com/aymgal/utax
+Author: Austin Peel, Martin Millon, Frederic Dux, Kevin Michalewicz
+Author-email: Aymeric Galan <aymeric.galan@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/aymgal/utax
+Project-URL: Repository, https://github.com/aymgal/utax
+Keywords: jax,signal processing,wavelet,convolution,interpolation
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.5.0
+Requires-Dist: jaxlib>=0.5.0
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pytest-pep8; extra == "dev"
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-python
+
+![License](https://img.shields.io/github/license/aymgal/utax)
+![PyPi python support](https://img.shields.io/badge/Python-3.10-blue)
+[![Tests](https://github.com/aymgal/utax/actions/workflows/ci_tests.yml/badge.svg?branch=main)](https://github.com/aymgal/utax/actions/workflows/ci_tests.yml)
+[![Coverage Status](https://coveralls.io/repos/github/aymgal/utax/badge.svg?branch=main)](https://coveralls.io/github/aymgal/utax?branch=main)
+
+# `utax`
+
+Utility functions for applications in signal processing problems, compatible with the differentable programming library `JAX`.

utax-0.0.2/README.md

@@ -0,0 +1,8 @@
+![License](https://img.shields.io/github/license/aymgal/utax)
+![PyPi python support](https://img.shields.io/badge/Python-3.10-blue)
+[![Tests](https://github.com/aymgal/utax/actions/workflows/ci_tests.yml/badge.svg?branch=main)](https://github.com/aymgal/utax/actions/workflows/ci_tests.yml)
+[![Coverage Status](https://coveralls.io/repos/github/aymgal/utax/badge.svg?branch=main)](https://coveralls.io/github/aymgal/utax?branch=main)
+
+# `utax`
+
+Utility functions for applications in signal processing problems, compatible with the differentable programming library `JAX`.

utax-0.0.2/pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = ["setuptools>=64"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "utax"
+version = "0.0.2"
+description = "Utility functions for signal processing, compatible with the differentiable programming library JAX."
+readme = "README.md"
+license = {text = "MIT"}
+authors = [
+    {name = "Aymeric Galan", email = "aymeric.galan@gmail.com"},
+    {name = "Austin Peel"},
+    {name = "Martin Millon"},
+    {name = "Frederic Dux"},
+    {name = "Kevin Michalewicz"},
+]
+keywords = ["jax", "signal processing", "wavelet", "convolution", "interpolation"]
+requires-python = ">=3.10"
+dependencies = [
+    "jax>=0.5.0",
+    "jaxlib>=0.5.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "pytest-cov",
+    "pytest-pep8",
+]
+
+[project.urls]
+Homepage = "https://github.com/aymgal/utax"
+Repository = "https://github.com/aymgal/utax"
+
+[tool.setuptools.packages.find]
+include = ["utax*"]

utax-0.0.2/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

utax-0.0.2/setup.py

@@ -0,0 +1,31 @@
+import setuptools
+import os
+
+__name__ = 'utax'
+
+release_info = {}
+infopath = os.path.abspath(os.path.join(os.path.dirname(__file__),
+                           __name__, 'info.py'))
+with open(infopath) as open_file:
+    exec(open_file.read(), release_info)
+
+this_directory = os.path.abspath(os.path.dirname(__file__))
+with open(os.path.join(this_directory, 'README.md'), encoding='utf-8') as f:
+    long_description = f.read()
+
+setuptools.setup(
+    name=__name__,
+    author=release_info['__author__'],
+    author_email=release_info['__email__'],
+    version=release_info['__version__'],
+    url=release_info['__url__'],
+    packages=setuptools.find_packages(),
+    python_requires=release_info['__python__'],
+    install_requires=release_info['__requires__'],
+    license=release_info['__license__'],
+    description=release_info['__about__'],
+    long_description=long_description,
+    long_description_content_type='text/markdown',
+    setup_requires=release_info['__setup_requires__'],
+    tests_require=release_info['__tests_require__']
+)

utax-0.0.2/utax/__init__.py

@@ -0,0 +1,9 @@
+from importlib.metadata import version, PackageNotFoundError
+
+__author__ = 'Aymeric Galan, Austin Peel, Martin Millon, Frederic Dux, Kevin Michalewicz'
+__email__ = 'aymeric.galan@gmail.com'
+
+try:
+    __version__ = version("utax")
+except PackageNotFoundError:
+    __version__ = "unknown"

utax-0.0.2/utax/convolution/__init__.py

@@ -0,0 +1,5 @@
+from .functions import (convolve_separable_dilated,
+                        build_convolution_matrix,
+                        convolve_dilated1D)
+from .classes import (GaussianFilter, 
+                      BlurringOperator)

utax-0.0.2/utax/convolution/classes.py

@@ -0,0 +1,195 @@
+from functools import partial
+import numpy as np
+from scipy.linalg import toeplitz
+import jax.numpy as jnp
+from jax.scipy import stats as jstats
+from jax import jit
+
+from utax.convolution import convolve_separable_dilated
+
+
+class GaussianFilter(object):
+    """JAX-friendly Gaussian filter."""
+    def __init__(self, sigma, truncate=4.0, mode='edge'):
+        """Convolve an image by a gaussian filter.
+
+        Parameters
+        ----------
+        sigma : float
+            Standard deviation of the Gaussian kernel.
+        truncate : float, optional
+            Truncate the filter at this many standard deviations.
+            Default is 4.0.
+
+        Note
+        ----
+        Reproduces `scipy.ndimage.gaussian_filter` with high accuracy.
+
+        """
+        if sigma <= 0:
+            self.kernel = None
+        else:
+            self.kernel = self.gaussian_kernel(sigma, truncate)
+        self.mode = mode
+
+    def gaussian_kernel(self, sigma, truncate):
+        # Determine the kernel pixel size (rounded up to an odd int)
+        self.radius = int(jnp.ceil(2 * truncate * sigma)) // 2
+        npix = self.radius * 2 + 1  # always at least 1
+
+        # Return the identity if sigma is not a positive number
+        if sigma <= 0:
+            return jnp.ones(1)
+
+        # Compute the kernel
+        x = jnp.ravel(jnp.indices((npix,)))  # pixel coordinates
+        kernel = jstats.norm.pdf((x-self.radius) / sigma)
+        kernel /= kernel.sum()
+
+        return kernel
+
+    @partial(jit, static_argnums=(0,))
+    def __call__(self, image):
+        """Jit-compiled convolution an image by a gaussian filter.
+
+        Parameters
+        ----------
+        image : array_like
+            Image to filter.
+        """
+        # Convolve
+        # pad_mode = ['constant', 'edge'][mode == 'nearest']
+        # image_padded = jnp.pad(image, pad_width=radius, mode=pad_mode)
+        if self.kernel is None:
+            return image
+        return convolve_separable_dilated(image, self.kernel, boundary=self.mode)
+
+
+class BlurringOperator(object):
+
+    def __init__(self, nx, ny, kernel):
+        self.target_shape = nx, ny
+        self.kernel = kernel
+        conv_matrix, self.temp_shape = self.toeplitz_matrix(self.target_shape, self.kernel)
+        self.conv_matrix = jnp.array(conv_matrix)
+        # get the indices for cropping the output of the full convolution
+        nxk, nyk = self.kernel.shape
+        if nxk % 2 == 0:
+            self.i1, self.i2 = nxk//2-1, -nxk//2
+        else:
+            self.i1, self.i2 = nxk//2, -nxk//2+1
+        if nyk % 2 == 0:
+            self.j1, self.j2 = nyk//2-1, -nyk//2
+        else:
+            self.j1, self.j2 = nyk//2, -nyk//2+1
+
+    @partial(jit, static_argnums=(0, 2))
+    def convolve(self, image, out_padding='same'):
+        image_conv = self.v2m(self.conv_matrix.dot(self.m2v(image)), self.temp_shape)
+        if out_padding == 'same':
+            return image_conv[self.i1:self.i2, self.j1:self.j2]
+        elif out_padding == 'full':
+            return image_conv
+        else:
+            raise ValueError(f"padding model '{out_padding}' is not supported.")
+
+    @partial(jit, static_argnums=(0, 2))
+    def convolve_transpose(self, image, in_padding='same'):
+        if in_padding == 'full':
+            image_padded = image
+        elif in_padding == 'same':
+            image_padded = jnp.pad(image, ((self.i1, -self.i2), (self.j1, -self.j2)), 
+                                   'constant', constant_values=0)
+        else:
+            raise ValueError(f"padding model '{in_padding}' is not supported.")
+        image_conv_t = self.v2m(self.conv_matrix.T.dot(self.m2v(image_padded)), self.target_shape)
+        return image_conv_t
+
+    @staticmethod
+    def toeplitz_matrix(input_shape, kernel, verbose=False):
+        """
+        Performs 2D convolution between input I and filter F by converting the F to a toeplitz matrix and multiply it
+          with vectorizes version of I
+          By : AliSaaalehi@gmail.com
+          
+        Arg:
+        input shape of the image to be convolved (I) -- 2D numpy matrix
+        convolution kernel (F) -- numpy 2D matrix
+        verbose -- if True, all intermediate resutls will be printed after each step of the algorithms
+        
+        Returns: 
+        output -- 2D numpy matrix, result of convolving I with F
+        """
+        # number of columns and rows of the input 
+        I_row_num, I_col_num = input_shape 
+
+        # number of columns and rows of the filter
+        F_row_num, F_col_num = kernel.shape
+
+        #  calculate the output dimensions
+        output_row_num = I_row_num + F_row_num - 1
+        output_col_num = I_col_num + F_col_num - 1
+        if verbose: print('output dimension:', output_row_num, output_col_num)
+
+        # zero pad the filter
+        F_zero_padded = np.pad(kernel, ((output_row_num - F_row_num, 0),
+                                        (0, output_col_num - F_col_num)), 
+                               'constant', constant_values=0)
+        if verbose: print('F_zero_padded: ', F_zero_padded)
+
+        # use each row of the zero-padded F to creat a toeplitz matrix. 
+        #  Number of columns in this matrices are same as numbe of columns of input signal
+        toeplitz_list = []
+        for i in range(F_zero_padded.shape[0]-1, -1, -1): # iterate from last row to the first row
+            c = F_zero_padded[i, :] # i th row of the F 
+            r = np.r_[c[0], np.zeros(I_col_num-1)] # first row for the toeplitz fuction should be defined otherwise
+                                                                # the result is wrong
+            toeplitz_m = toeplitz(c,r) # this function is in scipy.linalg library
+            toeplitz_list.append(toeplitz_m)
+            if verbose: print('F '+ str(i)+'\n', toeplitz_m)
+
+            # doubly blocked toeplitz indices: 
+        #  this matrix defines which toeplitz matrix from toeplitz_list goes to which part of the doubly blocked
+        c = range(1, F_zero_padded.shape[0]+1)
+        r = np.r_[c[0], np.zeros(I_row_num-1, dtype=int)]
+        doubly_indices = toeplitz(c, r)
+        if verbose: print('doubly indices \n', doubly_indices)
+
+        ## creat doubly blocked matrix with zero values
+        toeplitz_shape = toeplitz_list[0].shape # shape of one toeplitz matrix
+        h = toeplitz_shape[0]*doubly_indices.shape[0]
+        w = toeplitz_shape[1]*doubly_indices.shape[1]
+        doubly_blocked_shape = [h, w]
+        doubly_blocked = np.zeros(doubly_blocked_shape)
+
+        # tile toeplitz matrices for each row in the doubly blocked matrix
+        b_h, b_w = toeplitz_shape # hight and withs of each block
+        for i in range(doubly_indices.shape[0]):
+            for j in range(doubly_indices.shape[1]):
+                start_i = i * b_h
+                start_j = j * b_w
+                end_i = start_i + b_h
+                end_j = start_j + b_w
+                doubly_blocked[start_i: end_i, start_j:end_j] = toeplitz_list[doubly_indices[i,j]-1]
+
+        if verbose: print('doubly_blocked: ', doubly_blocked)
+        
+        out_shape = (output_row_num, output_col_num)
+        
+        return doubly_blocked, out_shape
+
+    @staticmethod
+    def m2v(mat):
+        return jnp.flipud(mat).flatten(order='C')
+
+    @staticmethod
+    def v2m(vec, output_shape):
+        return jnp.flipud(vec.reshape(output_shape, order='C'))
+
+    # @staticmethod
+    # def m2v_t(mat):
+    #     return jnp.flipud(mat.flatten(order='C'))
+
+    # @staticmethod
+    # def v2m_t(vec, output_shape):
+    #     return jnp.flipud(vec).reshape(output_shape, order='C')

utax-0.0.2/utax/convolution/functions.py

@@ -0,0 +1,211 @@
+from functools import partial
+import numpy as np
+from scipy import sparse
+import jax.numpy as jnp
+from jax.scipy import stats as jstats
+from jax import jit
+from jax.lax import conv_general_dilated, conv_dimension_numbers
+
+
+@partial(jit, static_argnums=(2, 3))
+def convolve_separable_dilated(image2D, kernel1D, dilation=1, boundary='edge'):
+    """
+     
+     Convolves an image contained in image2D with the 1D kernel kernel1D.
+     The operation is basically the following:
+         blured2D = image2D * (kernel1D ∧ kernel1D )
+     where ∧ is a wedge product, here a tensor product. 
+     
+     
+
+    Parameters
+    ----------
+    image2D : 2D array
+        imaged to be convolved with the kernel.
+    kernel1D : 1D array
+        kernel to convolve the image with..
+    dilation : TYPE, optional
+        makes the spacial extent of the kernel bigger. The default is 1.
+
+    Returns
+    -------
+    2D array
+        image convoluted by the kernel.
+
+    """
+    
+    # padding
+    b = int(kernel1D.size // 2) * dilation
+    padded = jnp.pad(image2D, ((b, b), (b, b)), mode=boundary)
+    # Fred D.: THIS PADDING IS DANGEROUS AS IT WILL OVERFLOW MEMORY VERY QUICKLY
+    # I LEAVE IT AS ORIGINALLY IMPLEMENTED AS I DO NOT WANT TO CHANGE THE
+    # OUTPUT OF THE WAVELET TRANSFORM (this could have impact on the science)
+    
+    
+    # specify the row and column operations for the jax convolution routine:
+    image = jnp.expand_dims(padded, (2,))
+    # shape (Nx, Ny, 1) -- (N, W, C)
+    # we treat the Nx as the batch number!! (because it is a 1D convolution 
+    # over the rows)
+    kernel = jnp.expand_dims(kernel1D, (0,2,))
+    # here we have kernel shape ~(I,W,O)
+    # so: 
+    # (Nbatch, Width, Channel) * (Inputdim, Widthkernel, Outputdim) 
+    #                                            -> (Nbatch, Width, Channel)
+    # where Nbatch is our number of rows.
+    dimension_numbers = ('NWC', 'IWO', 'NWC')
+    dn = conv_dimension_numbers(image.shape, 
+                                kernel.shape, 
+                                dimension_numbers)
+    # with these conv_general_dilated knows how to handle the different
+    # axes:
+    rowblur = conv_general_dilated(image, kernel,
+                                   window_strides=(1,),
+                                   padding='VALID',
+                                   rhs_dilation=(dilation,),
+                                   dimension_numbers=dn)
+    
+    # now we do the same for the columns, hence this time we have
+    # (Height, Nbatch, Channel) * (Inputdim, Widthkernel, Outputdim) 
+    #                                            -> (Height, Nbatch, Channel)
+    # where Nbatch is our number of columns.
+    dimension_numbers = ('HNC', 'IHO', 'HNC')
+    dn = conv_dimension_numbers(image.shape, 
+                                kernel.shape, 
+                                dimension_numbers)
+    
+    rowcolblur = conv_general_dilated(rowblur, kernel,
+                                      window_strides=(1,),
+                                      padding='VALID',
+                                      rhs_dilation=(dilation,),
+                                      dimension_numbers=dn)
+    
+    return rowcolblur[:,:,0]
+
+@partial(jit, static_argnums=(2, 3))
+def convolve_dilated1D(signal1D, kernel1D, dilation=1, boundary='edge'):
+    """
+
+     Convolves a vector contained in signal1D with the 1D kernel.
+     The operation is basically the following:
+         blured1D = signal1D * kernel1D
+
+    Parameters
+    ----------
+    signal1D : 1D array
+        vector to be convolved with the kernel.
+    kernel1D : 1D array
+        kernel to convolve signal1D with..
+    dilation : TYPE, optional
+        makes the spacial extent of the kernel bigger. The default is 1.
+
+    Returns
+    -------
+    1D array convoluted by the kernel.
+    """
+
+    # padding
+    b = int(kernel1D.size // 2) * dilation
+    padded = jnp.pad(signal1D, ((b, b)), mode=boundary)
+
+    shape = kernel1D.shape
+    strides = tuple(1 for s in shape)
+    rowblur = conv_general_dilated(padded[None, None], kernel1D[None, None],
+                                   window_strides=strides,
+                                   padding='VALID',
+                                   rhs_dilation=(dilation,),
+                                   )
+
+    return rowblur[0, 0]
+
+
+
+def build_convolution_matrix(psf_kernel_2d, image_shape):
+    """
+    Build a sparse matrix to convolve an image via matrix-vector product.
+    Ported from C++ code in VKL from Vernardos & Koopmans 2022.
+
+    Note: only works with square kernel with odd number of pixels on the side, 
+    lower than the number of pixels on the side of the image to be convolved.
+
+    Authors: @gvernard, @aymgal
+    """
+    Ni, Nj = image_shape
+    Ncropx, Ncropy = psf_kernel_2d.shape
+
+    def setCroppedLimitsEven(k, Ncrop, Nimg, Nquad):
+        if k < (Nquad - 1):
+            Npre   = k
+            Npost  = Nquad
+            offset = Nquad - k
+        elif k > (Nimg - Nquad - 1):
+            Npre   = Nquad
+            Npost  = Nimg - k
+            offset = 0
+        else:
+            Npre   = Nquad
+            Npost  = Nquad
+            offset = 0
+        return Npre, Npost, offset
+
+    def setCroppedLimitsOdd(k, Ncrop, Nimg, Nquad):
+        if k < (Nquad - 1):
+            Npre   = k
+            Npost  = Nquad
+            offset = Nquad - 1 - k
+        elif k > (Nimg - Nquad - 1):
+            Npre   = Nquad - 1
+            Npost  = Nimg - k
+            offset = 0
+        else:
+            Npre   = Nquad-1
+            Npost  = Nquad
+            offset = 0
+        return Npre, Npost, offset
+
+    # get the correct method to offset the PSF kernel from the above
+    if Ncropx % 2 == 0:
+        # Warning: this might be broken in certain cases
+        func_limits_x = setCroppedLimitsEven
+        Nquadx = Ncropx//2
+    else:
+        func_limits_x = setCroppedLimitsOdd
+        Nquadx = int(np.ceil(Ncropx/2.))
+    if Ncropx % 2 == 0:
+        # Warning: this might be broken in certain cases
+        func_limits_y = setCroppedLimitsEven
+        Nquady = Ncropy//2
+    else:
+        func_limits_y = setCroppedLimitsOdd
+        Nquady = int(np.ceil(Ncropy/2.))
+
+    # create the blurring matrix in a sparse form
+    blur = psf_kernel_2d.flatten()
+    sparse_B_rows, sparse_B_cols = [], []
+    sparse_B_values  = []
+    for i in range(Ni):  # loop over image rows
+        for j in range(Nj):  # loop over image columns
+
+            Nleft, Nright, crop_offsetx = func_limits_x(j, Ncropx, Nj, Nquadx)
+            Ntop, Nbottom, crop_offsety = func_limits_y(i, Ncropy, Ni, Nquady)
+            
+            crop_offset = crop_offsety*Ncropx + crop_offsetx
+
+            for ii in range(i-Ntop, i+Nbottom):  # loop over PSF rows
+                ic = ii - i + Ntop
+
+                for jj in range(j-Nleft, j+Nright):  # loop over PSF columns
+                    jc = jj - j + Nleft;
+
+                    val = blur[crop_offset + ic*Ncropx + jc]
+
+                    # save entries 
+                    # (note: rows and cols were inverted from the VKL code)
+                    sparse_B_rows.append(ii*Nj + jj)
+                    sparse_B_cols.append(i*Nj + j)
+                    sparse_B_values.append(val)
+
+    # populate the sparse matrix
+    blurring_matrix = sparse.csr_matrix((sparse_B_values, (sparse_B_rows, sparse_B_cols)), 
+                                        shape=(Ni**2, Nj**2))
+    return blurring_matrix

utax-0.0.2/utax/info.py

@@ -0,0 +1,29 @@
+"""PACKAGE INFO
+
+This module provides some basic information about the package.
+
+"""
+
+# Set the package release version
+version_info = (0, 0, 2)
+__version__ = '.'.join(str(c) for c in version_info)
+
+# Set the package details
+__author__ = 'Aymeric Galan, Austin Peel, Martin Millon, Frederic Dux, Kevin Michalewicz'
+__email__ = 'aymeric.galan@gmail.com'
+__year__ = '2022'
+__url__ = 'https://github.com/aymgal/utax'
+__description__ = 'Utility functions for signal processing, compatible with the differentable programming library JAX.'
+__python__ = '>=3.10'
+__requires__ = [
+    'jax>=0.5.0', 
+    'jaxlib>=0.5.0', 
+]  # Package dependencies
+
+# Default package properties
+__license__ = 'MIT'
+__about__ = ('{} Author: {}, Email: {}, Year: {}, {}'
+             ''.format(__name__, __author__, __email__, __year__,
+                       __description__))
+__setup_requires__ = ['pytest-runner', ]
+__tests_require__ = ['pytest', 'pytest-cov', 'pytest-pep8']

utax-0.0.2/utax/interpolation.py

@@ -0,0 +1,235 @@
+import numpy as np
+import jax.numpy as jnp
+from jax import lax, vmap  #, jit
+
+
+
+class BilinearInterpolator(object):
+    """Bilinear interpolation of a 2D field.
+
+    Functionality is modelled after scipy.interpolate.RectBivariateSpline
+    when `kx` and `ky` are both equal to 1. Results match the scipy version when
+    interpolated values lie within the x and y domain (boundaries included).
+    Returned values can be significantly different outside the natural domain,
+    as the scipy version does not extrapolate. Evaluation of this jax version
+    is MUCH SLOWER as well.
+
+    """
+    def __init__(self, x, y, z, allow_extrapolation=True):
+        self.z = jnp.array(z)
+
+        # Sort x if not increasing
+        x = jnp.array(x)
+        x_sorted = jnp.sort(x)
+        flip_x = ~jnp.all(jnp.diff(x) >= 0)
+
+        def x_keep_fn(_):
+            return x, self.z
+
+        def x_sort_fn(_):
+            return x_sorted, jnp.flip(self.z, axis=0)
+
+        self.x, self.z = lax.cond(flip_x, x_sort_fn, x_keep_fn, operand=None)
+
+        # Sort y if not increasing
+        y = jnp.array(y)
+        y_sorted = jnp.sort(y)
+        flip_y = ~jnp.all(jnp.diff(y) >= 0)
+
+        def y_keep_fn(_):
+            return y, self.z
+
+        def y_sort_fn(_):
+            return y_sorted, jnp.flip(self.z, axis=1)
+
+        self.y, self.z = lax.cond(flip_y, y_sort_fn, y_keep_fn, operand=None)
+        self._extrapol_bool = allow_extrapolation
+
+    def __call__(self, x, y, dx=0, dy=0):
+        """Vectorized evaluation of the interpolation or its derivatives.
+
+        Parameters
+        ----------
+        x, y : array_like
+            Position(s) at which to evaluate the interpolation.
+        dx, dy : int, either 0 or 1
+            If 1, return the first partial derivative of the interpolation
+            with respect to that coordinate. Only one of (dx, dy) should be
+            nonzero at a time.
+
+        """
+        x = jnp.atleast_1d(x)
+        y = jnp.atleast_1d(y)
+
+        error_msg_type = "dx and dy must be integers"
+        error_msg_value = "dx and dy must only be either 0 or 1"
+        assert isinstance(dx, int) and isinstance(dy, int), error_msg_type
+        assert dx in (0, 1) and dy in (0, 1), error_msg_value
+        if dx == 1: dy = 0
+
+        return vmap(self._evaluate, in_axes=(0, 0, None, None))(x, y, dx, dy)
+
+    # @partial(jit, static_argnums=(0,))
+    def _compute_coeffs(self, x, y):
+        # Find the pixel that the point (x, y) falls in
+        # x_ind = jnp.digitize(x, self.x_padded) - 1
+        # y_ind = jnp.digitize(y, self.y_padded) - 1
+        x_ind = jnp.searchsorted(self.x, x, side='right') - 1
+        x_ind = jnp.clip(x_ind, a_min=0, a_max=(len(self.x) - 2))
+        y_ind = jnp.searchsorted(self.y, y, side='right') - 1
+        y_ind = jnp.clip(y_ind, a_min=0, a_max=(len(self.y) - 2))
+
+        # Determine the coordinates and dimensions of this pixel
+        x1 = self.x[x_ind]
+        x2 = self.x[x_ind + 1]
+        y1 = self.y[y_ind]
+        y2 = self.y[y_ind + 1]
+        area = (x2 - x1) * (y2 - y1)
+
+        # Compute function values at the four corners
+        # Edge padding is implicitly constant
+        v11 = self.z[x_ind, y_ind]
+        v12 = self.z[x_ind, y_ind + 1]
+        v21 = self.z[x_ind + 1, y_ind]
+        v22 = self.z[x_ind + 1, y_ind + 1]
+
+        # Compute the coefficients
+        a0_ = v11 * x2 * y2 - v12 * x2 * y1 - v21 * x1 * y2 + v22 * x1 * y1
+        a1_ = -v11 * y2 + v12 * y1 + v21 * y2 - v22 * y1
+        a2_ = -v11 * x2 + v12 * x2 + v21 * x1 - v22 * x1
+        a3_ = v11 - v12 - v21 + v22
+
+        return a0_ / area, a1_ / area, a2_ / area, a3_ / area
+
+    def _evaluate(self, x, y, dx=0, dy=0):
+        """Single-point evaluation of the interpolation."""
+        a0, a1, a2, a3 = self._compute_coeffs(x, y)
+        if (dx, dy) == (0, 0):
+            result = a0 + a1 * x + a2 * y + a3 * x * y
+        elif (dx, dy) == (1, 0):
+            result = a1 + a3 * y
+        else:
+            result = a2 + a3 * x
+        # if extrapolation is not allowed, then we mask out values outside the original bounding box
+        result = lax.cond(self._extrapol_bool, 
+                          lambda _: result, 
+                          lambda _: result * (x >= self.x[0]) * (x <= self.x[-1]) * (y >= self.y[0]) * (y <= self.y[-1]), 
+                          operand=None)
+        return result
+
+
+class BicubicInterpolator(object):
+    """Bicubic interpolation of a 2D field.
+
+    Functionality is modelled after scipy.interpolate.RectBivariateSpline
+    when `kx` and `ky` are both equal to 3.
+
+    """
+    def __init__(self, x, y, z, zx=None, zy=None, zxy=None, allow_extrapolation=True):
+        self.z = jnp.array(z)
+        if np.all(np.diff(x) >= 0):  # check if sorted in increasing order
+            self.x = jnp.array(x)
+        else:
+            self.x = jnp.array(np.sort(x))
+            self.z = jnp.flip(self.z, axis=1)
+        if np.all(np.diff(y) >= 0):  # check if sorted in increasing order
+            self.y = jnp.array(y)
+        else:
+            self.y = jnp.array(np.sort(y))
+            self.z = jnp.flip(self.z, axis=0)
+
+        # Assume uniform coordinate spacing
+        self.dx = self.x[1] - self.x[0]
+        self.dy = self.y[1] - self.y[0]
+
+        # Compute approximate partial derivatives if not provided
+        if zx is None:
+            self.zx = jnp.gradient(z, axis=0) / self.dx
+        else:
+            self.zx = zy
+        if zy is None:
+            self.zy = jnp.gradient(z, axis=1) / self.dy
+        else:
+            self.zy = zx
+        if zxy is None:
+            self.zxy = jnp.gradient(self.zx, axis=1) / self.dy
+        else:
+            self.zxy = zxy
+
+        # Prepare coefficients for function evaluations
+        self._A = jnp.array([[1., 0., 0., 0.],
+                            [0., 0., 1., 0.],
+                            [-3., 3., -2., -1.],
+                            [2., -2., 1., 1.]])
+        self._B = jnp.array([[1., 0., -3., 2.],
+                            [0., 0., 3., -2.],
+                            [0., 1., -2., 1.],
+                            [0., 0., -1., 1.]])
+        row0 = [self.z[:-1,:-1], self.z[:-1,1:], self.dy * self.zy[:-1,:-1], self.dy * self.zy[:-1,1:]]
+        row1 = [self.z[1:,:-1], self.z[1:,1:], self.dy * self.zy[1:,:-1], self.dy * self.zy[1:,1:]]
+        row2 = self.dx * jnp.array([self.zx[:-1,:-1], self.zx[:-1,1:],
+                                   self.dy * self.zxy[:-1,:-1], self.dy * self.zxy[:-1,1:]])
+        row3 = self.dx * jnp.array([self.zx[1:,:-1], self.zx[1:,1:],
+                                   self.dy * self.zxy[1:,:-1], self.dy * self.zxy[1:,1:]])
+        self._m = jnp.array([row0, row1, row2, row3])
+
+        self._m = jnp.transpose(self._m, axes=(2, 3, 0, 1))
+
+        self._extrapol_bool = allow_extrapolation
+
+    def __call__(self, x, y, dx=0, dy=0):
+        """Vectorized evaluation of the interpolation or its derivatives.
+
+        Parameters
+        ----------
+        x, y : array_like
+            Position(s) at which to evaluate the interpolation.
+        dx, dy : int, either 0, 1, or 2
+            Return the nth partial derivative of the interpolation
+            with respect to the specified coordinate. Only one of (dx, dy)
+            should be nonzero at a time.
+
+        """
+        x = jnp.atleast_1d(x)
+        y = jnp.atleast_1d(y)
+        if x.ndim == 1:
+            vmap_call = vmap(self._evaluate, in_axes=(0, 0, None, None))
+        elif x.ndim == 2:
+            vmap_call = vmap(vmap(self._evaluate, in_axes=(0, 0, None, None)),
+                             in_axes=(0, 0, None, None))
+        return vmap_call(x, y, dx, dy)
+
+    def _evaluate(self, x, y, dx=0, dy=0):
+        """Evaluate the interpolation at a single point."""
+        # Determine which pixel (i, j) the point (x, y) falls in
+        i = jnp.maximum(0, jnp.searchsorted(self.x, x) - 1)
+        j = jnp.maximum(0, jnp.searchsorted(self.y, y) - 1)
+
+        # Rescale coordinates into (0, 1)
+        u = (x - self.x[i]) / self.dx
+        v = (y - self.y[j]) / self.dy
+
+        # Compute interpolation coefficients
+        a = jnp.dot(self._A, jnp.dot(self._m[i, j], self._B))
+
+        if dx == 0:
+            uu = jnp.asarray([1., u, u**2, u**3])
+        if dx == 1:
+            uu = jnp.asarray([0., 1., 2. * u, 3. * u**2]) / self.dx
+        if dx == 2:
+            uu = jnp.asarray([0., 0., 2., 6. * u]) / self.dx**2
+        if dy == 0:
+            vv = jnp.asarray([1., v, v**2, v**3])
+        if dy == 1:
+            vv = jnp.asarray([0., 1., 2. * v, 3. * v**2]) / self.dy
+        if dy == 2:
+            vv = jnp.asarray([0., 0., 2., 6. * v]) / self.dy**2
+        result = jnp.dot(uu, jnp.dot(a, vv))
+
+        # if extrapolation is not allowed, then we mask out values outside the original bounding box
+        result = lax.cond(self._extrapol_bool, 
+                          lambda _: result, 
+                          lambda _: result * (x >= self.x[0]) * (x <= self.x[-1]) * (y >= self.y[0]) * (y <= self.y[-1]), 
+                          operand=None)
+        return result
+        

utax-0.0.2/utax/wavelet.py

@@ -0,0 +1,243 @@
+from functools import partial
+import jax.numpy as jnp
+from jax import jit
+
+from utax.convolution import convolve_separable_dilated, convolve_dilated1D
+
+
+class WaveletTransform(object):
+    """
+    Class that handles wavelet transform using JAX, using the 'a trous' algorithm
+
+    Parameters
+    ----------
+    nscales : number of scales in the decomposition
+    dim : dimensionality of the convolution (2 for image, 1 for 1D vector)
+    self._type : supported types are 'starlet', 'battle-lemarie-1', 'battle-lemarie-3'
+
+    """
+    def __init__(self, nscales, wavelet_type='starlet', second_gen=False, dim=2):
+        self._n_scales = nscales
+        self._second_gen = second_gen
+        if wavelet_type == 'starlet':
+            self._h = jnp.array([1., 4., 6., 4., 1.]) / 16.
+        elif wavelet_type == 'battle-lemarie-1':  # (order 1 = 2 vanishing moments)
+            self._h = jnp.array([-0.000122686, -0.000224296, 0.000511636, 
+                        0.000923371, -0.002201945, -0.003883261, 0.009990599, 
+                        0.016974805, -0.051945337, -0.06910102, 0.39729643, 
+                        0.817645956, 0.39729643, -0.06910102, -0.051945337, 
+                        0.016974805, 0.009990599, -0.003883261, -0.002201945,
+                        0.000923371, 0.000511636, -0.000224296, -0.000122686])
+        elif wavelet_type == 'battle-lemarie-3':  # (order 2 = 4 vanishing moments)
+            self._h = jnp.array([0.000146098, -0.000232304, -0.000285414, 
+                           0.000462093, 0.000559952, -0.000927187, -0.001103748, 
+                           0.00188212, 0.002186714, -0.003882426, -0.00435384, 
+                           0.008201477, 0.008685294, -0.017982291, -0.017176331, 
+                           0.042068328, 0.032080869, -0.110036987, -0.050201753, 
+                           0.433923147, 0.766130398, 0.433923147, -0.050201753, 
+                          -0.110036987, 0.032080869, 0.042068328, -0.017176331, 
+                          -0.017982291, 0.008685294, 0.008201477, -0.00435384, 
+                          -0.003882426, 0.002186714, 0.00188212, -0.001103748, 
+                          -0.000927187, 0.000559952, 0.000462093, -0.000285414, 
+                          -0.000232304, 0.000146098])
+            
+        elif wavelet_type == 'battle-lemarie-5':  # (order 5 = 6 vanishing moments)
+            self._h = jnp.array([1.4299532e-004,
+  1.5656611e-004,
+ -2.2509746e-004,
+ -2.4421337e-004,
+  3.5556002e-004,
+  3.8161407e-004,
+ -5.6393016e-004,
+ -5.9748378e-004,
+  8.9882146e-004,
+  9.3739129e-004,
+ -1.4412528e-003,
+ -1.4737089e-003,
+  2.3286290e-003,
+  2.3211761e-003,
+ -3.7992267e-003,
+ -3.6602095e-003,
+  6.2791340e-003,
+  5.7683203e-003,
+ -1.0562022e-002,
+ -9.0493510e-003,
+  1.8208558e-002,
+  1.4009689e-002,
+ -3.2519969e-002,
+ -2.1006296e-002,
+  6.1312356e-002,
+  2.9474179e-002,
+ -1.2926869e-001,
+ -3.7019995e-002,
+  4.4246341e-001,
+  7.4723338e-001,
+  4.4246341e-001,
+ -3.7019995e-002,
+ -1.2926869e-001,
+  2.9474179e-002,
+  6.1312356e-002,
+ -2.1006296e-002,
+ -3.2519969e-002,
+  1.4009689e-002,
+  1.8208558e-002,
+ -9.0493510e-003,
+ -1.0562022e-002,
+  5.7683203e-003,
+  6.2791340e-003,
+ -3.6602095e-003,
+ -3.7992267e-003,
+  2.3211761e-003,
+  2.3286290e-003,
+ -1.4737089e-003,
+ -1.4412528e-003,
+  9.3739129e-004,
+  8.9882146e-004,
+ -5.9748378e-004,
+ -5.6393016e-004,
+  3.8161407e-004,
+  3.5556002e-004,
+ -2.4421337e-004,
+ -2.2509746e-004,
+  1.5656611e-004,
+  1.4299532e-004])
+        else:
+            raise ValueError(f"'{wavelet_type}' starlet transform is not supported")
+            
+        self._h /= jnp.sum(self._h)
+        self._fac = len(self._h) // 2
+
+        if self._second_gen:
+            self.decompose = self._decompose_2nd_gen 
+            self.reconstruct = self._reconstruct_2nd_gen
+        else:
+            self.decompose = self._decompose_1st_gen
+            self.reconstruct = self._reconstruct_1st_gen
+
+        self._dim = dim
+        if self._dim == 1:
+            self.convolve = convolve_dilated1D
+            if self._second_gen:
+                raise NotImplementedError('2nd generation of wavelet not yet implemented for 1D decomposition.')
+        elif self._dim == 2:
+            self.convolve = convolve_separable_dilated
+        else:
+            raise ValueError(f"Dimensionality {dim} not supported.")
+
+    @property
+    def scale_norms(self):
+        if not hasattr(self, '_norms'):
+            npix_dirac = 2**(self._n_scales + 2)
+            if self._dim == 1 :
+                dirac = (jnp.arange(npix_dirac) == int(npix_dirac / 2)).astype(float)
+                wt_dirac = self.decompose(dirac)
+                self._norms = jnp.sqrt(jnp.sum(wt_dirac**2, axis=(1,)))
+            else:
+                dirac = jnp.diag((jnp.arange(npix_dirac) == int(npix_dirac / 2)).astype(float))
+                wt_dirac = self.decompose(dirac)
+                self._norms = jnp.sqrt(jnp.sum(wt_dirac**2, axis=(1, 2,)))
+        return self._norms
+
+    
+    @partial(jit, static_argnums=(0,))
+    def _decompose_1st_gen(self, image):
+        """Decompose an image into the chosen wavelet basis"""
+        # Validate input
+        assert self._n_scales >= 0, "nscales must be a non-negative integer"
+        if self._n_scales == 0:
+            return image
+
+        # Preparations
+        # image = jnp.copy(image)
+        kernel = self._h.copy()
+
+        # Compute the first scale:
+        c1 = self.convolve(image, kernel)
+        # Wavelet coefficients:
+        w0 = (image - c1)  
+        result = jnp.expand_dims(w0, 0)
+        cj = c1
+        
+        # Compute the remaining scales
+        # at each scale, the kernel becomes larger ( a trou ) using the
+        # dilation argument in the jax wrapper for convolution.
+        for step in range(1, self._n_scales):
+            cj1 = self.convolve(cj, kernel, dilation=self._fac**step)
+            # wavelet coefficients
+            wj = (cj - cj1)
+            result = jnp.concatenate((result, jnp.expand_dims(wj, 0)), axis=0)
+            cj = cj1
+        
+        # Append final coarse scale
+        result = jnp.concatenate((result, jnp.expand_dims(cj, axis=0)), axis=0)
+        return result
+    
+    @partial(jit, static_argnums=(0,))
+    def _decompose_2nd_gen(self, image):
+        """Decompose an image into the chosen wavelet basis"""
+        # Validate input
+        assert self._n_scales >= 0, "nscales must be a non-negative integer"
+        if self._n_scales == 0:
+            return image
+
+        # Preparations
+        # image = jnp.copy(image)
+        kernel = self._h.copy()
+
+        # Compute the first scale:
+        c1 = self.convolve(image, kernel)
+        c1p = self.convolve(c1, kernel)
+        # Wavelet coefficients:
+        w0 = (image - c1p)  
+        result = jnp.expand_dims(w0, 0)
+        cj = c1
+        
+        # Compute the remaining scales
+        # at each scale, the kernel becomes larger ( a trou ) using the
+        # dilation argument in the jax wrapper for convolution.
+        for step in range(1, self._n_scales):
+            cj1  = self.convolve(cj, kernel, dilation=self._fac**step)
+            cj1p = self.convolve(cj1, kernel, dilation=self._fac**step)
+            # wavelet coefficients
+            wj = (cj - cj1p)
+            result = jnp.concatenate((result, jnp.expand_dims(wj, 0)), axis=0)
+            cj = cj1
+        
+        # Append final coarse scale
+        result = jnp.concatenate((result, jnp.expand_dims(cj, axis=0)), axis=0)
+        return result
+
+
+    @partial(jit, static_argnums=(0,))
+    def _reconstruct_1st_gen(self, coeffs):
+        return jnp.sum(coeffs, axis=0)
+    
+    
+    @partial(jit, static_argnums=(0,))
+    def _reconstruct_2nd_gen(self, coeffs):
+        # Validate input
+        assert coeffs.shape[-3] == self._n_scales+1, \
+               "Wavelet coefficients are not consistent with number of scales"
+        if self._n_scales == 0:
+            return coeffs[0, :, :]
+        
+        kernel = self._h
+        
+        # Start with the last scale 'J-1'
+        cJ = coeffs[self._n_scales, :, :]
+        cJp = self.convolve(cJ, kernel,
+                                       dilation=self._fac**(self._n_scales-1))
+        
+
+        wJ = coeffs[self._n_scales-1, :, :]
+        cj = cJp + wJ
+        
+        # Compute the remaining scales
+        for ii in range(self._n_scales-2, -1, -1):
+            cj1 = cj
+            cj1p = self.convolve(cj1, kernel, dilation=self._fac**ii)
+            wj1 = coeffs[ii, :, :]
+            cj = cj1p + wj1
+
+        result = cj
+        return result

utax-0.0.2/utax.egg-info/PKG-INFO

@@ -0,0 +1,32 @@
+Metadata-Version: 2.4
+Name: utax
+Version: 0.0.2
+Summary: Utility functions for signal processing, compatible with the differentiable programming library JAX.
+Home-page: https://github.com/aymgal/utax
+Author: Austin Peel, Martin Millon, Frederic Dux, Kevin Michalewicz
+Author-email: Aymeric Galan <aymeric.galan@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/aymgal/utax
+Project-URL: Repository, https://github.com/aymgal/utax
+Keywords: jax,signal processing,wavelet,convolution,interpolation
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.5.0
+Requires-Dist: jaxlib>=0.5.0
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pytest-pep8; extra == "dev"
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-python
+
+![License](https://img.shields.io/github/license/aymgal/utax)
+![PyPi python support](https://img.shields.io/badge/Python-3.10-blue)
+[![Tests](https://github.com/aymgal/utax/actions/workflows/ci_tests.yml/badge.svg?branch=main)](https://github.com/aymgal/utax/actions/workflows/ci_tests.yml)
+[![Coverage Status](https://coveralls.io/repos/github/aymgal/utax/badge.svg?branch=main)](https://coveralls.io/github/aymgal/utax?branch=main)
+
+# `utax`
+
+Utility functions for applications in signal processing problems, compatible with the differentable programming library `JAX`.

utax-0.0.2/utax.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+LICENSE
+README.md
+pyproject.toml
+setup.py
+utax/__init__.py
+utax/info.py
+utax/interpolation.py
+utax/wavelet.py
+utax.egg-info/PKG-INFO
+utax.egg-info/SOURCES.txt
+utax.egg-info/dependency_links.txt
+utax.egg-info/requires.txt
+utax.egg-info/top_level.txt
+utax/convolution/__init__.py
+utax/convolution/classes.py
+utax/convolution/functions.py

utax-0.0.2/utax.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

utax-0.0.2/utax.egg-info/requires.txt

@@ -0,0 +1,7 @@
+jax>=0.5.0
+jaxlib>=0.5.0
+
+[dev]
+pytest
+pytest-cov
+pytest-pep8

utax-0.0.2/utax.egg-info/top_level.txt

@@ -0,0 +1 @@
+utax