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.

from __future__ import annotations

import torch

from curvelets.torch import UDCTModule

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.

Setup

shape = (32, 32)
angular_wedges_config = torch.tensor([[3, 3]])
udct_module = UDCTModule(
    shape=shape,
    angular_wedges_config=angular_wedges_config,
)

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.

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

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)

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!

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

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

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

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!