Skip to content

Spectral conv

Module: spectral_conv.py

This module provides the implementation of a 1D spectral convolution layer.

Classes:

Name Description
SpectralConv1d

1D spectral convolution layer
Dependencies
  • jax: For array processing
  • equinox: For neural network layers
Key Features
  • Complex multiplication
  • Fourier domain convolution
  • Real and imaginary weights
Authors

Diya Nag Chaudhury

Version Info

29/Dec/2024: Initial version - Diya Nag Chaudhury

References

None

SpectralConv1d

Bases: Module

A 1D spectral convolution layer.

This layer performs a 1D convolution in the Fourier domain.

Attributes: real_weights: jax.Array imag_weights: jax.Array in_channels: int out_channels: int modes: int

Methods: init: Initializes the SpectralConv1d object complex_mult1d: Performs complex multiplication in 1D call: Calls the SpectralConv1d object

Source code in scirex/core/sciml/fno/layers/spectral_conv.py 
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
class SpectralConv1d(eqx.Module):
    """
    A 1D spectral convolution layer.

    This layer performs a 1D convolution in the Fourier domain.

    Attributes:
    real_weights: jax.Array
    imag_weights: jax.Array
    in_channels: int
    out_channels: int
    modes: int

    Methods:
    __init__: Initializes the SpectralConv1d object
    complex_mult1d: Performs complex multiplication in 1D
    __call__: Calls the SpectralConv1d object
    """

    real_weights: jax.Array
    imag_weights: jax.Array
    in_channels: int
    out_channels: int
    modes: int

    def __init__(
        self,
        in_channels,
        out_channels,
        modes,
        *,
        key,
    ):
        """
        Constructor for the SpectralConv1d class.

        Parameters:
        in_channels: int
            Number of input channels
        out_channels: int
            Number of output channels
        modes: int
            Number of modes
        key: jax.random.PRNGKey
            Random key for initialization

        Returns:
        None

        Usage:
        spectral_conv = SpectralConv1d(
            in_channels=1,
            out_channels=1,
            modes=64,
            key=key,
        )
        """

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modes = modes

        scale = 1.0 / (in_channels * out_channels)

        real_key, imag_key = jax.random.split(key)
        self.real_weights = jax.random.uniform(
            real_key,
            (in_channels, out_channels, modes),
            minval=-scale,
            maxval=+scale,
        )
        self.imag_weights = jax.random.uniform(
            imag_key,
            (in_channels, out_channels, modes),
            minval=-scale,
            maxval=+scale,
        )

    def complex_mult1d(
        self,
        x_hat,
        w,
    ):
        """
        Returns the complex multiplication of x_hat and w.

        Parameters:
        x_hat: jax.Array
            Input array in the Fourier domain
        w: jax.Array
            Weights in the Fourier domain

        Returns:
        jax.Array
            Complex multiplication of x_hat and w

        Usage:
        y_hat = spectral_conv.complex_mult1d(x_hat, w)

        """
        return jnp.einsum("iM,ioM->oM", x_hat, w)

    def __call__(
        self,
        x,
    ):
        """
        Forward pass of the SpectralConv1d layer.

        Parameters:
        x: jax.Array
            Input array

        Returns:
        jax.Array
            Output array

        Usage:
        y = spectral_conv(x)
        """
        channels, spatial_points = x.shape

        # shape of x_hat is (in_channels, spatial_points//2+1)
        x_hat = jnp.fft.rfft(x)
        # shape of x_hat_under_modes is (in_channels, self.modes)
        x_hat_under_modes = x_hat[:, : self.modes]
        weights = self.real_weights + 1j * self.imag_weights
        # shape of out_hat_under_modes is (out_channels, self.modes)
        out_hat_under_modes = self.complex_mult1d(x_hat_under_modes, weights)

        # shape of out_hat is (out_channels, spatial_points//2+1)
        out_hat = jnp.zeros((self.out_channels, x_hat.shape[-1]), dtype=x_hat.dtype)
        out_hat = out_hat.at[:, : self.modes].set(out_hat_under_modes)

        out = jnp.fft.irfft(out_hat, n=spatial_points)

        return out

__call__(x)

Forward pass of the SpectralConv1d layer.

x: jax.Array Input array

jax.Array Output array

Usage: y = spectral_conv(x)

Source code in scirex/core/sciml/fno/layers/spectral_conv.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
def __call__(
    self,
    x,
):
    """
    Forward pass of the SpectralConv1d layer.

    Parameters:
    x: jax.Array
        Input array

    Returns:
    jax.Array
        Output array

    Usage:
    y = spectral_conv(x)
    """
    channels, spatial_points = x.shape

    # shape of x_hat is (in_channels, spatial_points//2+1)
    x_hat = jnp.fft.rfft(x)
    # shape of x_hat_under_modes is (in_channels, self.modes)
    x_hat_under_modes = x_hat[:, : self.modes]
    weights = self.real_weights + 1j * self.imag_weights
    # shape of out_hat_under_modes is (out_channels, self.modes)
    out_hat_under_modes = self.complex_mult1d(x_hat_under_modes, weights)

    # shape of out_hat is (out_channels, spatial_points//2+1)
    out_hat = jnp.zeros((self.out_channels, x_hat.shape[-1]), dtype=x_hat.dtype)
    out_hat = out_hat.at[:, : self.modes].set(out_hat_under_modes)

    out = jnp.fft.irfft(out_hat, n=spatial_points)

    return out

__init__(in_channels, out_channels, modes, *, key)

Constructor for the SpectralConv1d class.

in_channels: int Number of input channels out_channels: int Number of output channels modes: int Number of modes key: jax.random.PRNGKey Random key for initialization

Returns: None

Usage: spectral_conv = SpectralConv1d( in_channels=1, out_channels=1, modes=64, key=key, )

Source code in scirex/core/sciml/fno/layers/spectral_conv.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def __init__(
    self,
    in_channels,
    out_channels,
    modes,
    *,
    key,
):
    """
    Constructor for the SpectralConv1d class.

    Parameters:
    in_channels: int
        Number of input channels
    out_channels: int
        Number of output channels
    modes: int
        Number of modes
    key: jax.random.PRNGKey
        Random key for initialization

    Returns:
    None

    Usage:
    spectral_conv = SpectralConv1d(
        in_channels=1,
        out_channels=1,
        modes=64,
        key=key,
    )
    """

    self.in_channels = in_channels
    self.out_channels = out_channels
    self.modes = modes

    scale = 1.0 / (in_channels * out_channels)

    real_key, imag_key = jax.random.split(key)
    self.real_weights = jax.random.uniform(
        real_key,
        (in_channels, out_channels, modes),
        minval=-scale,
        maxval=+scale,
    )
    self.imag_weights = jax.random.uniform(
        imag_key,
        (in_channels, out_channels, modes),
        minval=-scale,
        maxval=+scale,
    )

complex_mult1d(x_hat, w)

Returns the complex multiplication of x_hat and w.

x_hat: jax.Array Input array in the Fourier domain w: jax.Array Weights in the Fourier domain

jax.Array Complex multiplication of x_hat and w

Usage: y_hat = spectral_conv.complex_mult1d(x_hat, w)

Source code in scirex/core/sciml/fno/layers/spectral_conv.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
def complex_mult1d(
    self,
    x_hat,
    w,
):
    """
    Returns the complex multiplication of x_hat and w.

    Parameters:
    x_hat: jax.Array
        Input array in the Fourier domain
    w: jax.Array
        Weights in the Fourier domain

    Returns:
    jax.Array
        Complex multiplication of x_hat and w

    Usage:
    y_hat = spectral_conv.complex_mult1d(x_hat, w)

    """
    return jnp.einsum("iM,ioM->oM", x_hat, w)