Butterfly module
----------------

.. automodule:: lazylinop.butterfly

Preamble
~~~~~~~~

This module gathers utilities to manipulate and create a particular family of :py:class:`.LazyLinOp`
associated to matrices with the so-called *butterfly* structure, *i.e.*, which are expressed as a
product of a few compatible *Kronecker-Sparse factors*.
Classical examples of such operators are the Hadamard transform and the Discrete-Fourier-Transform.

**Construction of a butterfly** :py:class:`.LazyLinOp` **from a given array.**

The main function of the module is :py:func:`ksd` (*Kronecker-sparse decomposition*), which computes a near best approximation of a given array or of a :py:class:`.LazyLinOp` by an operator with a prescribed butterfly structure :ref:`[1] <references>`.

*Example:*
Consider an $N\times N$ Hadamard matrix (with $N$ a power of two):

.. code-block:: python

   >>> import scipy as sp
   >>> H = sp.linalg.hadamard(N)

Turning this NumPy array into an efficient butterfly implementation requires a simple call ``L = ksd(H, chain)`` where ``chain`` is a :py:class:`Chain` object specifying the nature of the Kronecker-sparse factors involved in the decomposition.
The classical chain associated to the Hadamard matrix is square-dyadic, and can be retrieved as follows:

.. code-block:: python

   >>> chain = Chain.square_dyadic(H.shape)
   >>> L = ksd(H, chain)

It is also possible to use other chains (see below), leading to implementations with different tradeoffs between memory consumption, energy consumption and computational speed.

.. code-block:: python

   >>> chain2 = Chain.monarch(H.shape, ...)
   >>> L2 = ksd(H, chain2)

.. FIXME: add timeit and memory consumption.

**Construction of individual or multiple Kronecker-sparse factors.**

Under the hood, the operator ``L`` resulting from a call to :py:func:`ksd` is a (lazy) product of Kronecker-sparse factors, each of them benefitting from an optimized implementation on various backends via the function :py:func:`ksm` (**Kronecker-sparse multiplication** operator).
A single Kronecker-sparse factor is a structure determined by a 4D array $T$ of shape ``(a, b, c, d)``, and is fully determined as ``L = ksm(T)`` given a 4-dimensional array of values $T_{i,j,k,l}$ of size $a\times b\times c\times d$.
The shape of ``L`` is then $abd\times acd$, and converting it to an array would yield a structured sparse matrix with support (the set of possible nonzero entries) $I_a\otimes 1_{b\times c}\otimes I_d$.
See :ref:`[1] <references>` and the documentation of :py:func:`ksm` for further details and illustrations.

*Implementation parameters.*

Options of :py:func:`ksm` (and of :py:func:`ksd`) notably allow to use either the CPU or the GPU with various implementation parameters, or to directly build a lazy product of Kronecker-sparse factors ``L = ksm([T1, ..., Tn], ...)`` to minimize memory movements between CPU and GPU when computing ``L @ x``.

**Saving and loading butterfly operators.**

The :py:class:`.LazyLinOp` resulting from a call to :py:func:`ksm` or :py:func:`ksd` can be saved to disk (as a ``.json`` file) and load using:

.. code-block:: python

   >>> lazylinop.butterfly.save(L, 'filename')
   >>> L = lazylinop.butterfly.load('filename')

**Plot butterfly operators.**

The :py:class:`.LazyLinOp` resulting from a call to :py:func:`ksm` or :py:func:`ksd` can be plotted or saved to disk (as a ``.png`` file and a ``.svg`` file) using:

.. code-block:: python

   >>> N = 16
   >>> L = lazylinop.butterfly.dft(N)
   >>> lazylinop.butterfly.plot(L)
   >>> lazylinop.butterfly.plot(L, "dft_16x16")
It is also possible to save to disk as a MATLAB-style file and load using:

.. code-block:: python

   >>> lazylinop.butterfly.savemat(L, 'filename')
   >>> L = lazylinop.butterfly.loadmat('filename')

**Pre-built Discrete-Fourier-Transform.**

We provide :py:func:`dft` a pre-built :py:class:`.LazyLinOp` ``L = B @ P`` with the Butterfly structure ``B`` multiplied by a bit-reversal permutation ``P`` corresponding to the Discrete-Fourier-Transform :ref:`[2] <references>`.
The ``L`` corresponding to a DFT of a signal of size $256$ with $8$ factors (square-dyadic decomposition) is given by:

.. code-block:: python

   >>> L = lazylinop.butterfly.dft(256)

**Pre-built Fast-Walsh-Hadamard-Transform.**

We provide :py:func:`fwht` a pre-built :py:class:`.LazyLinOp` ``L`` with the Butterfly structure corresponding to the Fast-Walsh-Hadamard-Transform (FWHT).
The ``L`` corresponding to a FWHT of a signal of size $256$ with $8$ factors (square-dyadic decomposition) is given by:

.. code-block:: python

   >>> L = lazylinop.butterfly.fwht(256)

**Chains and their manipulation.**

A *chain* encodes the structure of the Butterfly operator ``L`` and corresponds to a sequence of tuples

.. math::

   \begin{equation}
   \left(\left(a_l,~b_l,~c_l,~d_l\right)\right)_{l=1}^n
   \end{equation}

such that Kronecker-sparse operators $L_l$ with the corresponding structure have *compatible sizes*, *i.e.* such that the product $L = L_1\cdots L_n$ is well-defined. The shape of $L$ is given by the attribute ``chain.shape``.
The concatenation of compatible chains is implemented via ``chain = chain1 @ chain2``.
Specifying it explicitly is simple:

.. code-block:: python

   >>> chain = Chain([(a1, b1, c1, d1), ..., (an, bn, cn, dn)])

The most standard chain is probably the so-called *square-dyadic chain*, which is behind usual fast implementations of the Hadamard transform and of DFTs.
It is only defined when $\mathtt{shape}=\left(N,~N\right)$ with $N$ a power of two as:

.. code-block:: python

   >>> sd_chain = Chain.square_dyadic(shape)

Another common chain with arbitrary non-prime shape is Monarch chain with two factors.

.. code-block:: python

   >>> chain = Chain.monarch(shape, ...)

You can visualize the structure of all Kronecker-sparse factors in these chains as follows:

.. code-block:: python

   >>> sd_chain.plot()

.. image:: _static/square_dyadic.svg
   :height: 150
   :align: center
   :alt: Output file of ``sd_chain.plot("square_dyadic")``.

.. code-block:: python

   >>> chain.plot()

.. image:: _static/monarch.svg
   :height: 150
   :align: center
   :alt: Output file of ``chain.plot("monarch")``.

and anticipate the memory footprint of the corresponding operators:

.. code-block:: python

   >>> sd_chain = Chain.square_dyadic((32, 32))
   >>> sd_chain.mem(np.float32)
   1280
   >>> chain = Chain.monarch((32, 32))
   >>> chain.mem(np.float32)
   1536

As you can see, the memory footprint of a *monarch* chain is higher than that of a *square-dyadic* one, but when it fits in the memory of a GPUs it can lead to a faster implementation.

*Chainability.*

This above chains satisfy a "chainability" property, which is crucial to ensure that :py:func:`ksd` can be run, with approximation guarantees :ref:`[1] <references>`.
Chainability can be checked using the attribute ``chain.chainable``, and :py:func:`ksd` will generate an error if the provided chain is not chainable.

.. _references:

References
~~~~~~~~~~

[1] Butterfly Factorization with Error Guarantees.
Leon Zheng, Quoc-Tung Le, Elisa Riccietti, and Remi Gribonval
https://hal.science/hal-04763712v1/document

[2] Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations.
Dao T, Gu A, Eichhorn M, Rudra A, Re C.
Proc Mach Learn Res. 2019 Jun;97:1517-1527.
PMID: 31777847; PMCID: PMC6879380.

Butterfly construction
~~~~~~~~~~~~~~~~~~~~~~

1. :meth:`lazylinop.butterfly.ksd`
2. :meth:`lazylinop.butterfly.ksm`
3. :meth:`lazylinop.butterfly.dft`
4. :meth:`lazylinop.butterfly.fwht`
5. :meth:`lazylinop.butterfly.hadamard`
6. :meth:`lazylinop.butterfly.suksd`
7. :meth:`lazylinop.butterfly.fnt`
8. :meth:`lazylinop.butterfly.dct`
9. :meth:`lazylinop.butterfly.dst`

.. autofunction:: lazylinop.butterfly.ksd
.. autofunction:: lazylinop.butterfly.ksm
.. autofunction:: lazylinop.butterfly.dft
.. autofunction:: lazylinop.butterfly.fwht
.. autofunction:: lazylinop.butterfly.hadamard
.. autofunction:: lazylinop.butterfly.suksd
.. autofunction:: lazylinop.butterfly.fnt
.. autofunction:: lazylinop.butterfly.dct
.. autofunction:: lazylinop.butterfly.dst

Chain class
~~~~~~~~~~~

.. autoclass:: lazylinop.butterfly.Chain
.. autofunction:: lazylinop.butterfly.Chain.square_dyadic
.. autofunction:: lazylinop.butterfly.Chain.monarch
.. autofunction:: lazylinop.butterfly.Chain.wip_non_redundant

Chain methods
~~~~~~~~~~~~~

.. autofunction:: lazylinop.butterfly.Chain.mem
.. autofunction:: lazylinop.butterfly.Chain.plot

Chain attributes
~~~~~~~~~~~~~~~~

.. autoattribute:: lazylinop.butterfly.Chain.T

.. .. autoattribute:: lazylinop.butterfly.Chain.shape
.. .. autoattribute:: lazylinop.butterfly.Chain.n_patterns
.. .. autoattribute:: lazylinop.butterfly.Chain.ks_patterns
.. .. autoattribute:: lazylinop.butterfly.Chain.chainable

Utilities
~~~~~~~~~

1. :meth:`lazylinop.butterfly.fuse`
2. :meth:`lazylinop.butterfly.optimize`

.. autofunction:: lazylinop.butterfly.fuse
.. autofunction:: lazylinop.butterfly.optimize

I/O (load/save)
~~~~~~~~~~~~~~~

1. :meth:`lazylinop.butterfly.load`
2. :meth:`lazylinop.butterfly.save`
3. :meth:`lazylinop.butterfly.loadmat`
4. :meth:`lazylinop.butterfly.savemat`

.. autofunction:: lazylinop.butterfly.load
.. autofunction:: lazylinop.butterfly.save
.. autofunction:: lazylinop.butterfly.loadmat
.. autofunction:: lazylinop.butterfly.savemat

Plot
~~~~

1. :meth:`lazylinop.butterfly.plot`

.. autofunction:: lazylinop.butterfly.plot