Skip to content

Utils Reference

This page contains the API reference for utility functions in neojax.

neojax.utils

Functions

count_array_params(arr, dims=None)

Returns the number of parameters (elements) in a single array, optionally, along certain dimensions only

Parameters

array : Array dims : int list or None, default is None if not None, the dimensions to consider when counting the number of parameters (elements)

Notes

One complex number is counted as two parameters (we count real and imaginary parts)'

Source code in neojax/utils.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def count_array_params(arr: Array, dims=None) -> int:
    """Returns the number of parameters (elements) in a single array, optionally, along certain dimensions only

    Parameters
    ----------
    array : Array
    dims : int list or None, default is None
        if not None, the dimensions to consider when counting the number of parameters (elements)

    Notes
    -----
    One complex number is counted as two parameters (we count real and imaginary parts)'
    """
    if dims is None:
        dims = list(arr.shape)
    else:
        dims = [arr.shape[d] for d in dims]
    n_params = prod(dims)
    if jnp.iscomplex(arr):
        return 2 * n_params
    return n_params

count_model_params(model)

Returns the total number of parameters of an equinox model

Notes

One complex number is counted as two parameters (we count real and imaginary parts)'

Source code in neojax/utils.py 
19
20
21
22
23
24
25
26
27
28
def count_model_params(model: eqx.Module) -> int:
    """Returns the total number of parameters of an equinox model

    Notes
    -----
    One complex number is counted as two parameters (we count real and imaginary parts)'
    """
    trainable_params = eqx.filter(model, eqx.is_inexact_array)
    total_params = jax.vmap(count_array_params)(trainable_params)
    return int(jnp.sum(total_params))

spectrum_2d(signal, n_observations, normalize=True)

This function computes the spectrum of a 2D signal using the Fast Fourier Transform (FFT).

Parameters

signal : an array of shape (T * n_observations * n_observations) A 2D discretized signal represented as a 1D array with shape (T * n_observations * n_observations), where T is the number of time steps and n_observations is the spatial size of the signal.

T can be any number of channels that we reshape into and
n_observations * n_observations is the spatial resolution.

n_observations: an integer Number of discretized points. Basically the resolution of the signal. normalize: bool whether to apply normalization to the output of the 2D FFT. If True, normalizes the outputs by 1/n_observations (actually 1/sqrt(n_observations * n_observations)). Returns

spectrum: a array A 1D array of shape (s,) representing the computed spectrum. The spectrum is computed using a square approximation to radial binning, meaning that the wavenumber 'bin' into which a particular coefficient is the coefficient's location along the diagonal, indexed from the top-left corner of the 2d FFT output.

Source code in neojax/utils.py 
 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
def spectrum_2d(signal: Array, n_observations: int, normalize: bool = True) -> Array:
    """This function computes the spectrum of a 2D signal using the Fast Fourier Transform (FFT).

    Parameters
    ----------
    signal : an array of shape (T * n_observations * n_observations)
        A 2D discretized signal represented as a 1D array with shape
        (T * n_observations * n_observations), where T is the number of time
        steps and n_observations is the spatial size of the signal.

        T can be any number of channels that we reshape into and
        n_observations * n_observations is the spatial resolution.
    n_observations: an integer
        Number of discretized points. Basically the resolution of the signal.
    normalize: bool
        whether to apply normalization to the output of the 2D FFT.
        If True, normalizes the outputs by ``1/n_observations``
        (actually ``1/sqrt(n_observations * n_observations)``).
    Returns
    --------
    spectrum: a array
        A 1D array of shape (s,) representing the computed spectrum.
        The spectrum is computed using a square approximation to radial
        binning, meaning that the wavenumber 'bin' into which a particular
        coefficient is the coefficient's location along the diagonal, indexed
        from the top-left corner of the 2d FFT output.
    """
    T = signal.shape[0]
    signal = jnp.reshape(signal, (T, n_observations, n_observations))

    if normalize:
        signal = jnp.fft.fft2(signal, norm="ortho")
    else:
        signal = jnp.fft.rfft2(
            signal, s=(n_observations, n_observations), norm="backward"
        )

    # 2d wavenumbers following numpy fft convention
    k_max = n_observations // 2
    wavenumers = jnp.repeat(
        jnp.concat(
            (
                jnp.arange(start=0, stop=k_max, step=1),
                jnp.arange(start=-k_max, stop=0, step=1),
            ),
            axis=0,
        ),
        n_observations,
        1,
    )
    k_x = jnp.transpose(wavenumers, (0, 1))
    k_y = wavenumers

    # Sum wavenumbers
    sum_k = jnp.sqrt(k_x**2 + k_y**2)
    sum_k = sum_k

    # Remove symmetric components from wavenumbers
    index = -1.0 * jnp.ones((n_observations, n_observations))
    k_max1 = k_max + 1
    index[0:k_max1, 0:k_max1] = sum_k[0:k_max1, 0:k_max1]

    spectrum = jnp.zeros((T, n_observations))
    for j in range(1, n_observations + 1):
        ind = jnp.where(index == j)
        spectrum[:, j - 1] = jnp.sum(jnp.abs(signal[:, ind[0], ind[1]]) ** 2, axis=1)

    spectrum = jnp.mean(spectrum, axis=0)
    return spectrum

validate_scaling_factor(scaling_factor, n_dim, n_layers=None)

Parameters

scaling_factor : None OR float OR list[float] Or list[list[float]] n_dim : int n_layers : int or None; defaults to None If None, return a single list (rather than a list of lists) with factor repeated dim times.

Source code in neojax/utils.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
195
196
197
198
199
200
201
202
203
204
205
206
def validate_scaling_factor(
    scaling_factor: Union[None, Number, list[Number], list[list[Number]]],
    n_dim: int,
    n_layers: Optional[int] = None,
) -> Union[None, list[float], list[list[float]]]:
    """
    Parameters
    ----------
    scaling_factor : None OR float OR list[float] Or list[list[float]]
    n_dim : int
    n_layers : int or None; defaults to None
        If None, return a single list (rather than a list of lists)
        with `factor` repeated `dim` times.
    """
    if scaling_factor is None:
        return None
    if isinstance(scaling_factor, (float, int)):
        if n_layers is None:
            return [float(scaling_factor)] * n_dim

        return [[float(scaling_factor)] * n_dim] * n_layers

    if (
        isinstance(scaling_factor, list)
        and len(scaling_factor) > 0
        and all([isinstance(s, (float, int)) for s in scaling_factor])
    ):
        if n_layers is None and len(scaling_factor) == n_dim:
            # this is a dim-wise scaling
            return [float(s) for s in scaling_factor]
        return [[float(s)] * n_dim for s in scaling_factor]

    if (
        isinstance(scaling_factor, list)
        and len(scaling_factor) > 0
        and all([isinstance(s, (list)) for s in scaling_factor])
    ):
        s_sub_pass = True
        for s in scaling_factor:
            if all([isinstance(s_sub, (int, float)) for s_sub in s]):
                pass
            else:
                s_sub_pass = False
            if s_sub_pass:
                return scaling_factor

    return None