.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/linops.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_linops.py: Linear Operators and Matrix-Freedom ======================================= In this example we explore some of the functionality from the ``skerch`` linear operator API, specifically: * Creation of a ``skerch``-compatible linear operator (bare-minimum interface) * Matrix-free transposition * Wrapping NumPy-only linear operators into PyTorch * Composition and addition of linear operators * Matrix-free diagonal linear operators * Matrix-free noisy linear operators for sketching This functionality allows us to perform sketches and work with linear operators at scale. .. GENERATED FROM PYTHON SOURCE LINES 19-41 .. code-block:: Python from time import time import matplotlib.pyplot as plt import torch from skerch.algorithms import snorm from skerch.linops import ( CompositeLinOp, DiagonalLinOp, TorchLinOpWrapper, TransposedLinOp, linop_to_matrix, ) from skerch.measurements import ( GaussianNoiseLinOp, PhaseNoiseLinOp, RademacherNoiseLinOp, SsrftNoiseLinOp, ) from skerch.utils import gaussian_noise .. GENERATED FROM PYTHON SOURCE LINES 42-54 ############################################################################## Creating a ``skerch``-compatible linop -------------------------------------- To work with ``skerch``, linear operators must satisfy only 2 requirements, resulting in a *bare-minimum linop interface*: * They must support left- and right-matrix multiplication via ``@`` * They must feature a ``.shape = (height, width)`` attribute This is satisfied by regular matrices like the following ``mat``: .. GENERATED FROM PYTHON SOURCE LINES 55-73 .. code-block:: Python SHAPE = (50, 50) NUM_MEAS = 10 SEED = 12345 DTYPE = torch.complex64 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" mat = gaussian_noise(SHAPE, 0, 10, seed=SEED, dtype=DTYPE, device=DEVICE) ramp = torch.arange(mat.shape[1], dtype=mat.real.dtype, device=mat.device) mat += ramp + 1j * ramp fig, (ax1, ax2) = plt.subplots(ncols=2) ax1.imshow(mat.real.cpu()) ax2.imshow(mat.imag.cpu()) fig.suptitle("Some test matrix") fig.tight_layout() .. image-sg:: /examples/images/sphx_glr_linops_001.png :alt: Some test matrix :srcset: /examples/images/sphx_glr_linops_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 74-85 But matrices also allow for other operations such as arbitrary indexing like ``mat[5, 7]``. This is something that many linear operators don't or cannot provide without substantial overhead. In those restricted cases, we are dealing with a *matrix-free* linear operator, of potentially very large scale, defined mainly by its dimensions and matmul functionality. As an example, the following linop implements the bare-minimum interface. We see that, despite this limitation, we can successfully apply :func:`skerch.algorithms.snorm` to estimate its operator norm: .. GENERATED FROM PYTHON SOURCE LINES 85-125 .. code-block:: Python class SomeLinOp: """Some matrix-free linop to exemplify skerch-compatibility.""" def __init__(self, dims): self.shape = (dims, dims) def __matmul__(self, x): result = x * 0.5 result[0, ...] += x.sum(dim=0) return result def __rmatmul__(self, x): result = x * 0.5 result = (result.transpose(0, -1) + x[..., 0]).transpose(0, -1) return result lop = SomeLinOp(SHAPE[0]) lopmat = linop_to_matrix(lop, dtype=ramp.dtype, device="cpu") lop_S = torch.linalg.svdvals(lopmat) op_norm = snorm( lop, DEVICE, DTYPE, num_meas=NUM_MEAS, seed=SEED + 123, noise_type="gaussian", norm_types=["op"], )[0]["op"] print("Actual norm:", torch.linalg.svdvals(lopmat).max().item()) print("Sketched norm:", op_norm.item()) fig, ax = plt.subplots(figsize=(8, 4)) ax.imshow(lopmat) fig.suptitle("The linear operator, in matrix form") fig.tight_layout() .. image-sg:: /examples/images/sphx_glr_linops_002.png :alt: The linear operator, in matrix form :srcset: /examples/images/sphx_glr_linops_002.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none Actual norm: 7.175588607788086 Sketched norm: 7.088768482208252 .. GENERATED FROM PYTHON SOURCE LINES 126-133 ############################################################################## Matrix-free transposition ------------------------- We can also use :class:`skerch.linops.TransposedLinOp` to transpose linear operators in a matrix-free fashion: .. GENERATED FROM PYTHON SOURCE LINES 134-143 .. code-block:: Python lopT = TransposedLinOp(lop) fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(8, 3)) ax1.imshow(linop_to_matrix(lop, dtype=ramp.dtype, device="cpu")) ax2.imshow(linop_to_matrix(lopT, dtype=ramp.dtype, device="cpu")) fig.suptitle("The linear operator and its transpose") fig.tight_layout() .. image-sg:: /examples/images/sphx_glr_linops_003.png :alt: The linear operator and its transpose :srcset: /examples/images/sphx_glr_linops_003.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 144-152 ############################################################################## PyTorch wrapper --------------- Since ``skerch`` is built on PyTorch, it won't directly work on numpy linops that don't accept PyTorch inputs. This can be solved with the wrapper, which also tracks the ``torch.device``: .. GENERATED FROM PYTHON SOURCE LINES 153-182 .. code-block:: Python class NumpyFlipLinOp: """Numpy-only antidiagonal linop. Doesn't work with torch tensors.""" def __init__(self, shape): self.shape = shape def __matmul__(self, x): return x[::-1].copy() def __rmatmul__(self, x): return self.__matmul__(x) class TorchFlipLinOp(TorchLinOpWrapper, NumpyFlipLinOp): """This wrapper works with tensors and numpy arrays.""" np_lop = NumpyFlipLinOp(SHAPE) torch_lop = TorchFlipLinOp(SHAPE) arr = ramp.cpu().numpy() print("Numpy linop on numpy data:", np_lop @ arr) print("Wrapped linop on numpy data:", torch_lop @ arr) print("Wrapped linop on torch data:", torch_lop @ ramp) .. rst-class:: sphx-glr-script-out .. code-block:: none Numpy linop on numpy data: [49. 48. 47. 46. 45. 44. 43. 42. 41. 40. 39. 38. 37. 36. 35. 34. 33. 32. 31. 30. 29. 28. 27. 26. 25. 24. 23. 22. 21. 20. 19. 18. 17. 16. 15. 14. 13. 12. 11. 10. 9. 8. 7. 6. 5. 4. 3. 2. 1. 0.] Wrapped linop on numpy data: [49. 48. 47. 46. 45. 44. 43. 42. 41. 40. 39. 38. 37. 36. 35. 34. 33. 32. 31. 30. 29. 28. 27. 26. 25. 24. 23. 22. 21. 20. 19. 18. 17. 16. 15. 14. 13. 12. 11. 10. 9. 8. 7. 6. 5. 4. 3. 2. 1. 0.] Wrapped linop on torch data: tensor([49., 48., 47., 46., 45., 44., 43., 42., 41., 40., 39., 38., 37., 36., 35., 34., 33., 32., 31., 30., 29., 28., 27., 26., 25., 24., 23., 22., 21., 20., 19., 18., 17., 16., 15., 14., 13., 12., 11., 10., 9., 8., 7., 6., 5., 4., 3., 2., 1., 0.]) .. GENERATED FROM PYTHON SOURCE LINES 183-191 ############################################################################## Other linear operators ---------------------- We can perform matrix-free compositions and additions of linear operators. Other matrix-free structured linops, such as diagonal and banded, are also available (see :mod:`skerch.linops`). Here we exemplify composition: .. GENERATED FROM PYTHON SOURCE LINES 192-206 .. code-block:: Python k = 2 U, S, Vh = torch.linalg.svd(mat.cpu()) lop_k = CompositeLinOp( [("U_k", U[:, :k]), ("S_k", DiagonalLinOp(S[:k])), ("Vh_k", Vh[:k, :])] ) fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(8, 3)) ax1.imshow(mat.real.cpu()) ax2.imshow(linop_to_matrix(lop_k, dtype=mat.dtype, device="cpu").real) fig.suptitle(f"Matrix and its rank-{k} matrix-free approximation [{lop_k}]") fig.tight_layout() .. image-sg:: /examples/images/sphx_glr_linops_004.png :alt: Matrix and its rank-2 matrix-free approximation [U_k @ S_k @ Vh_k] :srcset: /examples/images/sphx_glr_linops_004.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 207-225 ############################################################################## Matrix-free noisy linear operators for sketching ------------------------------------------------ In order to run the sketches, ``skerch`` provides built-in support for noisy measurements in the form of matrix-free linear operators. In order to facilitate parallelized measurements, these linops have a bit more restrictive requirements than the *bare-minimum* interface discussed above: Besides the ``.shape`` and ``@`` properties required by all ``skerch`` linops, they also must also implement a ``get_blocks`` iterator, that yields blocks of columns with their indices. A good way to satisfy this interface and add a new noise type to ``skerch`` is to extend :class:`skerch.linops.ByBlockLinOp` with ``get_block`` (see :ref:`Extending With Custom Functionality` for an example). The figure below illustrates some of the already supported types of noise: .. GENERATED FROM PYTHON SOURCE LINES 226-247 .. code-block:: Python blocksize = 5 mop1 = RademacherNoiseLinOp(SHAPE, SEED, blocksize=blocksize) mop2 = GaussianNoiseLinOp(SHAPE, SEED, blocksize=blocksize) mop3 = PhaseNoiseLinOp(SHAPE, SEED, blocksize=blocksize) mop4 = SsrftNoiseLinOp(SHAPE, SEED, blocksize=blocksize) fig, (ax1, ax2, ax3, ax4) = plt.subplots(ncols=4, figsize=(8, 2)) ax1.imshow(mop1.to_matrix(DTYPE, "cpu").real) ax2.imshow(mop2.to_matrix(DTYPE, "cpu").real) ax3.imshow(mop3.to_matrix(DTYPE, "cpu").real) ax4.imshow(mop4.to_matrix(DTYPE, "cpu").real) # ax1.set_title("Rademacher") ax2.set_title("Gaussian") ax3.set_title("Phase") ax4.set_title("SSRFT") fig.suptitle("Different types of noise matrices") fig.tight_layout() .. image-sg:: /examples/images/sphx_glr_linops_005.png :alt: Different types of noise matrices, Rademacher, Gaussian, Phase, SSRFT :srcset: /examples/images/sphx_glr_linops_005.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 248-249 And the line below illustrates the behaviour of ``get_blocks``: .. GENERATED FROM PYTHON SOURCE LINES 249-253 .. code-block:: Python [(b.shape, idxs) for b, idxs in mop1.get_blocks(DTYPE)] .. rst-class:: sphx-glr-script-out .. code-block:: none [(torch.Size([50, 5]), range(0, 5)), (torch.Size([50, 5]), range(5, 10)), (torch.Size([50, 5]), range(10, 15)), (torch.Size([50, 5]), range(15, 20)), (torch.Size([50, 5]), range(20, 25)), (torch.Size([50, 5]), range(25, 30)), (torch.Size([50, 5]), range(30, 35)), (torch.Size([50, 5]), range(35, 40)), (torch.Size([50, 5]), range(40, 45)), (torch.Size([50, 5]), range(45, 50))] .. GENERATED FROM PYTHON SOURCE LINES 254-257 To illustrate the necessity of blockwise measurements, consider the following example, where a larger block size results in substantially faster computations: .. GENERATED FROM PYTHON SOURCE LINES 257-278 .. code-block:: Python mop_shape = (mat.shape[1], 100) mop_slow = RademacherNoiseLinOp(mop_shape, SEED, blocksize=1, register=False) mop_fast = RademacherNoiseLinOp(mop_shape, SEED, blocksize=100, register=False) times = [[], []] for _ in range(20): t0 = time() mat @ mop_slow times[0].append(time() - t0) # t0 = time() mat @ mop_fast times[1].append(time() - t0) fig, ax = plt.subplots(figsize=(8, 3)) ax.boxplot(times, label=["blocksize=1", "blocksize=100"]) ax.set_yscale("log") fig.suptitle("Speedup resulting from blockwise measurements") fig.tight_layout() .. image-sg:: /examples/images/sphx_glr_linops_006.png :alt: Speedup resulting from blockwise measurements :srcset: /examples/images/sphx_glr_linops_006.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 279-292 ############################################################################## This concludes the ``skerch`` tour of linear operators! Please refer to the API docs and other examples. for more details. In summary: * We have seen how to create simple matrix-free linops, so that they are compatible with the ``skerch`` routines * We also saw how to manipulate said linops, by transposing them, converting them to matrices and wrapping them to ensure NumPy compatibility in-core sketched methods for a broad class of low-rank matrices * We explored other available linear operators, including compositions and noisy measurement linops, emphasizing the benefit of blockwise measurements .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.733 seconds) .. _sphx_glr_download_examples_linops.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: linops.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: linops.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: linops.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_