.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/plot_09_udct_module_gradcheck.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_auto_examples_plot_09_udct_module_gradcheck.py: PyTorch UDCT Module =================== This example demonstrates the recommended way to use PyTorch bindings for the Uniform Discrete Curvelet Transform (UDCT) via ``UDCTModule``. PyTorch Bindings Architecture ------------------------------ The curvelets library provides two interfaces for PyTorch: 1. **UDCT**: Base class that provides forward and backward transforms but does not integrate with PyTorch's automatic differentiation system. It returns nested coefficient structures (lists of tensors organized by scale and direction). 2. **UDCTModule**: A ``nn.Module`` wrapper around ``UDCT`` that provides full autograd integration. It uses a custom ``torch.autograd.Function`` to enable automatic differentiation through the transform. When called, it returns flattened coefficients as a single tensor, making it PyTorch-friendly for use in neural networks and optimization workflows. **Recommendation**: Use ``UDCTModule`` for all PyTorch workflows. It is the recommended interface because it: - Integrates seamlessly with PyTorch's autograd system - Returns flattened tensors suitable for PyTorch operations - Automatically uses the backward transform for gradient computation - Can be used as a standard ``nn.Module`` in neural network architectures How Autograd Integration Works ------------------------------- ``UDCTModule`` uses a custom autograd function that: - **Forward pass**: Computes the forward curvelet transform and flattens the nested coefficients into a single tensor - **Backward pass**: Automatically uses the backward transform to compute gradients with respect to the input image This means you can use ``UDCTModule`` in any PyTorch computation graph, and gradients will flow correctly through the transform. .. GENERATED FROM PYTHON SOURCE LINES 43-46 .. code-block:: Python from __future__ import annotations .. GENERATED FROM PYTHON SOURCE LINES 47-51 .. code-block:: Python import torch from curvelets.torch import UDCTModule .. GENERATED FROM PYTHON SOURCE LINES 52-78 PyTorch Bindings: UDCT vs UDCTModule #################################### The curvelets library provides two PyTorch interfaces: **UDCT** (not recommended for most PyTorch workflows): - Base transform class without autograd integration - Returns nested coefficient structures (lists of tensors) - Requires manual gradient handling if used in optimization - Useful for inspection/debugging or when you need nested structures **UDCTModule** (recommended for PyTorch workflows): - ``nn.Module`` wrapper with full autograd support - Returns flattened coefficients as a single tensor - Automatically integrates with PyTorch's computation graph - Uses backward transform for gradient computation - Can be used directly in neural networks and optimization loops When to use each: - Use **UDCTModule** for: neural networks, optimization, any workflow requiring gradients, standard PyTorch tensor operations - Use **UDCT** for: inspection of nested structures, debugging, when you specifically need the nested coefficient format This example demonstrates ``UDCTModule``, which is the recommended interface for all PyTorch use cases. .. GENERATED FROM PYTHON SOURCE LINES 80-82 Setup ##### .. GENERATED FROM PYTHON SOURCE LINES 82-89 .. code-block:: Python shape = (32, 32) angular_wedges_config = torch.tensor([[3, 3]]) udct_module = UDCTModule( shape=shape, angular_wedges_config=angular_wedges_config, ) .. GENERATED FROM PYTHON SOURCE LINES 90-95 Forward Transform ################# UDCTModule returns flattened coefficients as a single tensor, making it compatible with standard PyTorch operations. The nested structure is automatically flattened during the forward pass. .. GENERATED FROM PYTHON SOURCE LINES 95-101 .. code-block:: Python input_tensor = torch.randn(*shape, dtype=torch.float64, requires_grad=True) output = udct_module(input_tensor) # Input shape: input_tensor.shape # Output shape: output.shape # Note: Output is a flattened tensor, not a nested structure .. GENERATED FROM PYTHON SOURCE LINES 102-112 Reconstruction via Autograd ############################ UDCTModule's autograd integration works as follows: - Forward pass: Uses forward transform to compute coefficients - Backward pass: Automatically uses backward transform to compute gradients This happens transparently through PyTorch's autograd system. Compute a simple operation on the coefficients (use abs to ensure real scalar) .. GENERATED FROM PYTHON SOURCE LINES 112-134 .. code-block:: Python loss = torch.abs(output).sum() # Backward pass: The custom autograd function automatically applies the # backward transform to compute gradients w.r.t. the input loss.backward() # The gradients in input_tensor.grad demonstrate the backward transform is working grad = input_tensor.grad assert grad is not None # Gradient shape: grad.shape # Gradient norm: grad.norm().item() # The backward transform is automatically used in the autograd graph! # Verify reconstruction matches input # Get nested coefficients and reconstruct using backward transform coeffs_nested = udct_module.struct(output.detach()) reconstructed = udct_module._udct.backward(coeffs_nested) reconstruction_error = torch.abs(input_tensor.detach() - reconstructed).max() # Reconstruction error: reconstruction_error.item() assert torch.allclose(input_tensor.detach(), reconstructed, atol=1e-4), ( f"Reconstruction does not match input! Max error: {reconstruction_error.item():.2e}" ) # Reconstruction matches input tensor! .. GENERATED FROM PYTHON SOURCE LINES 135-143 Using struct() Method ##################### While UDCTModule returns flattened coefficients (PyTorch-friendly), you can convert them back to the nested structure when needed using struct(). This is useful for inspection, visualization, or when working with code that expects nested coefficient structures. Convert flattened coefficients back to nested structure .. GENERATED FROM PYTHON SOURCE LINES 143-148 .. code-block:: Python coeffs_nested_from_struct = udct_module.struct(output.detach()) # Flattened coefficients shape: output.shape # Restructured to nested format with len(coeffs_nested_from_struct) scales # struct() converts flattened coefficients to nested structure using internal state .. GENERATED FROM PYTHON SOURCE LINES 149-156 Gradcheck Verification ###################### PyTorch's gradcheck verifies that the backward pass correctly computes gradients. This confirms that UDCTModule's autograd integration is working correctly - the backward transform is properly used for gradient computation. Clear gradients for gradcheck .. GENERATED FROM PYTHON SOURCE LINES 156-166 .. code-block:: Python input_tensor.grad = None result = torch.autograd.gradcheck( udct_module, input_tensor, fast_mode=True, atol=1e-4, rtol=1e-3, ) assert result, "Gradcheck failed" # This confirms the autograd integration is working correctly! .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 1.612 seconds) .. _sphx_glr_download_auto_examples_plot_09_udct_module_gradcheck.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_09_udct_module_gradcheck.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_09_udct_module_gradcheck.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_09_udct_module_gradcheck.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_