Skip to content

Conversion Utilties for Normalized and Difficulty Steppers¤

exponax.stepper.generic.normalize_coefficients ¤

normalize_coefficients(
    coefficients: tuple[float, ...],
    *,
    domain_extent: float,
    dt: float
) -> tuple[float, ...]

Normalize the coefficients to a linear time stepper to be used with the normalized linear stepper.

αᵢ = aᵢ Δt / Lⁱ

Warning

A consequence of this normalization is that the normalized coefficients for high order derivatives will be very small.

Arguments:

  • coefficients: coefficients for the linear operator, coefficients[i] is the coefficient for the i-th derivative
  • domain_extent: extent of the domain
  • dt: time step

Returns:

  • normalized_coefficients: normalized coefficients for the linear operator, normalized_coefficients[i] is the coefficient for the i-th derivative
Source code in exponax/stepper/generic/_utils.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def normalize_coefficients(
    coefficients: tuple[float, ...],
    *,
    domain_extent: float,
    dt: float,
) -> tuple[float, ...]:
    """
    Normalize the coefficients to a linear time stepper to be used with the
    normalized linear stepper.

        αᵢ = aᵢ Δt / Lⁱ

    !!! warning
        A consequence of this normalization is that the normalized coefficients
        for high order derivatives will be very small.

    **Arguments:**

    - `coefficients`: coefficients for the linear operator, `coefficients[i]` is
        the coefficient for the `i`-th derivative
    - `domain_extent`: extent of the domain
    - `dt`: time step

    **Returns:**

    - `normalized_coefficients`: normalized coefficients for the linear
        operator, `normalized_coefficients[i]` is the coefficient for the `i`-th
        derivative
    """
    normalized_coefficients = tuple(
        c * dt / (domain_extent**i) for i, c in enumerate(coefficients)
    )
    return normalized_coefficients

exponax.stepper.generic.denormalize_coefficients ¤

denormalize_coefficients(
    normalized_coefficients: tuple[float, ...],
    *,
    domain_extent: float,
    dt: float
) -> tuple[float, ...]

Denormalize the coefficients as they were used in the normalized linear to then be used again in a genric linear stepper with a physical interface.

aᵢ = αᵢ Lⁱ / Δt

Arguments:

  • normalized_coefficients: coefficients for the linear operator, normalized_coefficients[i] is the coefficient for the i-th derivative
  • domain_extent: extent of the domain
  • dt: time step

Returns:

  • coefficients: coefficients for the linear operator, coefficients[i] is the coefficient for the i-th derivative
Source code in exponax/stepper/generic/_utils.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def denormalize_coefficients(
    normalized_coefficients: tuple[float, ...],
    *,
    domain_extent: float,
    dt: float,
) -> tuple[float, ...]:
    """
    Denormalize the coefficients as they were used in the normalized linear to
    then be used again in a genric linear stepper with a physical interface.

        aᵢ = αᵢ Lⁱ / Δt

    **Arguments:**

    - `normalized_coefficients`: coefficients for the linear operator,
        `normalized_coefficients[i]` is the coefficient for the `i`-th
        derivative
    - `domain_extent`: extent of the domain
    - `dt`: time step

    **Returns:**

    - `coefficients`: coefficients for the linear operator, `coefficients[i]` is
        the coefficient for the `i`-th derivative
    """
    coefficients = tuple(
        c_n / dt * domain_extent**i for i, c_n in enumerate(normalized_coefficients)
    )
    return coefficients

exponax.stepper.generic.normalize_convection_scale ¤

normalize_convection_scale(
    convection_scale: float,
    *,
    domain_extent: float,
    dt: float
) -> float

Normalize the scale (=coefficient) in front of the convection term to be used with the normalized generic steppers.

β₁ = b₁ Δt / L

Arguments:

  • convection_scale: scale in front of the convection term, i.e., the b_1 in 𝒩(u) = - b₁ 1/2 (u²)ₓ
  • domain_extent: extent of the domain
  • dt: time step

Returns:

  • normalized_convection_scale: normalized scale in front of the convection
Source code in exponax/stepper/generic/_utils.py 
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
def normalize_convection_scale(
    convection_scale: float,
    *,
    domain_extent: float,
    dt: float,
) -> float:
    """
    Normalize the scale (=coefficient) in front of the convection term to be
    used with the normalized generic steppers.

        β₁ = b₁ Δt / L

    **Arguments:**

    - `convection_scale`: scale in front of the convection term, i.e., the `b_1`
        in `𝒩(u) = - b₁ 1/2 (u²)ₓ`
    - `domain_extent`: extent of the domain
    - `dt`: time step

    **Returns:**

    - `normalized_convection_scale`: normalized scale in front of the convection
    """
    normalized_convection_scale = convection_scale * dt / domain_extent
    return normalized_convection_scale

exponax.stepper.generic.denormalize_convection_scale ¤

denormalize_convection_scale(
    normalized_convection_scale: float,
    *,
    domain_extent: float,
    dt: float
) -> float

Denormalize the scale in front of the convection term as it was used in the normalized generic steppers to then be used again in a generic stepper with a physical interface.

b₁ = β₁ L / Δt

Arguments:

  • normalized_convection_scale: normalized scale in front of the convection
  • domain_extent: extent of the domain
  • dt: time step

Returns:

  • convection_scale: scale in front of the convection term, i.e., the b_1 in 𝒩(u) = - b₁ 1/2 (u²)ₓ
Source code in exponax/stepper/generic/_utils.py 
 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
def denormalize_convection_scale(
    normalized_convection_scale: float,
    *,
    domain_extent: float,
    dt: float,
) -> float:
    """
    Denormalize the scale in front of the convection term as it was used in the
    normalized generic steppers to then be used again in a generic stepper with
    a physical interface.

        b₁ = β₁ L / Δt

    **Arguments:**

    - `normalized_convection_scale`: normalized scale in front of the convection
    - `domain_extent`: extent of the domain
    - `dt`: time step

    **Returns:**

    - `convection_scale`: scale in front of the convection term, i.e., the `b_1`
        in `𝒩(u) = - b₁ 1/2 (u²)ₓ`
    """
    convection_scale = normalized_convection_scale / dt * domain_extent
    return convection_scale

exponax.stepper.generic.normalize_gradient_norm_scale ¤

normalize_gradient_norm_scale(
    gradient_norm_scale: float,
    *,
    domain_extent: float,
    dt: float
) -> float

Normalize the scale in front of the gradient norm term to be used with the normalized generic steppers.

β₂ = b₂ Δt / L²

Arguments:

  • gradient_norm_scale: scale in front of the gradient norm term, i.e., the b_2 in 𝒩(u) = - b₂ 1/2 ‖∇u‖₂²
  • domain_extent: extent of the domain
  • dt: time step

Returns:

  • normalized_gradient_norm_scale: normalized scale in front of the gradient norm term
Source code in exponax/stepper/generic/_utils.py 
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
def normalize_gradient_norm_scale(
    gradient_norm_scale: float,
    *,
    domain_extent: float,
    dt: float,
) -> float:
    """
    Normalize the scale in front of the gradient norm term to be used with the
    normalized generic steppers.

        β₂ = b₂ Δt / L²

    **Arguments:**

    - `gradient_norm_scale`: scale in front of the gradient norm term, i.e., the
        `b_2` in `𝒩(u) = - b₂ 1/2 ‖∇u‖₂²`
    - `domain_extent`: extent of the domain
    - `dt`: time step

    **Returns:**

    - `normalized_gradient_norm_scale`: normalized scale in front of the
        gradient norm term
    """
    normalized_gradient_norm_scale = (
        gradient_norm_scale * dt / jnp.square(domain_extent)
    )
    return normalized_gradient_norm_scale

exponax.stepper.generic.denormalize_gradient_norm_scale ¤

denormalize_gradient_norm_scale(
    normalized_gradient_norm_scale: float,
    *,
    domain_extent: float,
    dt: float
) -> float

Denormalize the scale in front of the gradient norm term as it was used in the normalized generic steppers to then be used again in a generic stepper with a physical interface.

b₂ = β₂ L² / Δt

Arguments:

  • normalized_gradient_norm_scale: normalized scale in front of the gradient norm term
  • domain_extent: extent of the domain
  • dt: time step

Returns:

  • gradient_norm_scale: scale in front of the gradient norm term, i.e., the b_2 in 𝒩(u) = - b₂ 1/2 ‖∇u‖₂²
Source code in exponax/stepper/generic/_utils.py 
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
def denormalize_gradient_norm_scale(
    normalized_gradient_norm_scale: float,
    *,
    domain_extent: float,
    dt: float,
) -> float:
    """
    Denormalize the scale in front of the gradient norm term as it was used in
    the normalized generic steppers to then be used again in a generic stepper
    with a physical interface.

        b₂ = β₂ L² / Δt

    **Arguments:**

    - `normalized_gradient_norm_scale`: normalized scale in front of the gradient
        norm term
    - `domain_extent`: extent of the domain
    - `dt`: time step

    **Returns:**

    - `gradient_norm_scale`: scale in front of the gradient norm term, i.e., the
        `b_2` in `𝒩(u) = - b₂ 1/2 ‖∇u‖₂²`
    """
    gradient_norm_scale = (
        normalized_gradient_norm_scale / dt * jnp.square(domain_extent)
    )
    return gradient_norm_scale

exponax.stepper.generic.normalize_polynomial_scales ¤

normalize_polynomial_scales(
    polynomial_scales: tuple[float, ...],
    *,
    domain_extent: float = None,
    dt: float
) -> tuple[float, ...]

Normalize the polynomial scales to be used with the normalized polynomial stepper.

Arguments:

  • polynomial_scales: scales for the polynomial operator, polynomial_scales[i] is the scale for the i-th degree polynomial
  • domain_extent: extent of the domain (not needed, kept for compatibility with other normalization APIs)
  • dt: time step

Returns:

  • normalized_polynomial_scales: normalized scales for the polynomial operator, normalized_polynomial_scales[i] is the scale for the i-th degree polynomial
Source code in exponax/stepper/generic/_utils.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
def normalize_polynomial_scales(
    polynomial_scales: tuple[float, ...],
    *,
    domain_extent: float = None,
    dt: float,
) -> tuple[float, ...]:
    """
    Normalize the polynomial scales to be used with the normalized polynomial
    stepper.

    **Arguments:**

    - `polynomial_scales`: scales for the polynomial operator,
        `polynomial_scales[i]` is the scale for the `i`-th degree polynomial
    - `domain_extent`: extent of the domain (not needed, kept for
        compatibility with other normalization APIs)
    - `dt`: time step

    **Returns:**

    - `normalized_polynomial_scales`: normalized scales for the polynomial
        operator, `normalized_polynomial_scales[i]` is the scale for the `i`-th
        degree polynomial
    """
    normalized_polynomial_scales = tuple(c * dt for c in polynomial_scales)
    return normalized_polynomial_scales

exponax.stepper.generic.denormalize_polynomial_scales ¤

denormalize_polynomial_scales(
    normalized_polynomial_scales: tuple[float, ...],
    *,
    domain_extent: float = None,
    dt: float
) -> tuple[float, ...]

Denormalize the polynomial scales as they were used in the normalized polynomial to then be used again in a regular polynomial stepper.

Arguments:

  • normalized_polynomial_scales: scales for the polynomial operator, normalized_polynomial_scales[i] is the scale for the i-th degree polynomial
  • domain_extent: extent of the domain (not needed, kept for compatibility with other normalization APIs)
  • dt: time step

Returns:

  • polynomial_scales: scales for the polynomial operator,
Source code in exponax/stepper/generic/_utils.py 
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
def denormalize_polynomial_scales(
    normalized_polynomial_scales: tuple[float, ...],
    *,
    domain_extent: float = None,
    dt: float,
) -> tuple[float, ...]:
    """
    Denormalize the polynomial scales as they were used in the normalized
    polynomial to then be used again in a regular polynomial stepper.

    **Arguments:**

    - `normalized_polynomial_scales`: scales for the polynomial operator,
        `normalized_polynomial_scales[i]` is the scale for the `i`-th degree
        polynomial
    - `domain_extent`: extent of the domain (not needed, kept for
        compatibility with other normalization APIs)
    - `dt`: time step

    **Returns:**

    - `polynomial_scales`: scales for the polynomial operator,
    """
    polynomial_scales = tuple(c_n / dt for c_n in normalized_polynomial_scales)
    return polynomial_scales

exponax.stepper.generic.reduce_normalized_coefficients_to_difficulty ¤

reduce_normalized_coefficients_to_difficulty(
    normalized_coefficients: tuple[float, ...],
    *,
    num_spatial_dims: int,
    num_points: int
) -> tuple[float, ...]

Reduce the normalized coefficients for a linear operator to a difficulty based interface. This interface is designed to "reduce the intensity of the dynamics" at higher resolutions to make emulator learning across resolutions comparible. Thereby, it resembles the stability numbers of the most compact finite difference scheme of the respective PDE.

γ₀ = α₀

γⱼ = αⱼ Nʲ 2ʲ⁻¹ D

Arguments:

  • normalized_coefficients: normalized coefficients for the linear operator, normalized_coefficients[i] is the coefficient for the i-th derivative
  • num_spatial_dims: number of spatial dimensions d
  • num_points: number of points N used to discretize the domain per dimension

Returns:

  • difficulty_coefficients: difficulty coefficients for the linear operator, difficulty_coefficients[i] is the coefficient for the i-th derivative
Source code in exponax/stepper/generic/_utils.py 
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
def reduce_normalized_coefficients_to_difficulty(
    normalized_coefficients: tuple[float, ...],
    *,
    num_spatial_dims: int,
    num_points: int,
) -> tuple[float, ...]:
    """
    Reduce the normalized coefficients for a linear operator to a difficulty
    based interface. This interface is designed to "reduce the intensity of the
    dynamics" at higher resolutions to make emulator learning across resolutions
    comparible. Thereby, it resembles the stability numbers of the most compact
    finite difference scheme of the respective PDE.

        γ₀ = α₀

        γⱼ = αⱼ Nʲ 2ʲ⁻¹ D

    **Arguments:**

    - `normalized_coefficients`: normalized coefficients for the linear
        operator, `normalized_coefficients[i]` is the coefficient for the `i`-th
        derivative
    - `num_spatial_dims`: number of spatial dimensions `d`
    - `num_points`: number of points `N` used to discretize the domain per
        dimension

    **Returns:**

    - `difficulty_coefficients`: difficulty coefficients for the linear operator,
        `difficulty_coefficients[i]` is the coefficient for the `i`-th derivative
    """
    difficulty_coefficients = list(
        alpha * num_points**j * 2 ** (j - 1) * num_spatial_dims
        for j, alpha in enumerate(normalized_coefficients)
    )
    difficulty_coefficients[0] = normalized_coefficients[0]

    difficulty_coefficients = tuple(difficulty_coefficients)
    return difficulty_coefficients

exponax.stepper.generic.extract_normalized_coefficients_from_difficulty ¤

extract_normalized_coefficients_from_difficulty(
    difficulty_coefficients: tuple[float, ...],
    *,
    num_spatial_dims: int,
    num_points: int
) -> tuple[float, ...]

Extract the normalized coefficients for a linear operator from a difficulty based interface.

α₀ = γ₀

αⱼ = γⱼ / (Nʲ 2ʲ⁻¹ D)

Arguments:

  • difficulty_coefficients: difficulty coefficients for the linear operator, difficulty_coefficients[i] is the coefficient for the i-th derivative
  • num_spatial_dims: number of spatial dimensions d
  • num_points: number of points N used to discretize the domain per dimension

Returns:

  • normalized_coefficients: normalized coefficients for the linear operator, normalized_coefficients[i] is the coefficient for the i-th derivative
Source code in exponax/stepper/generic/_utils.py 
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
def extract_normalized_coefficients_from_difficulty(
    difficulty_coefficients: tuple[float, ...],
    *,
    num_spatial_dims: int,
    num_points: int,
) -> tuple[float, ...]:
    """
    Extract the normalized coefficients for a linear operator from a difficulty
    based interface.

        α₀ = γ₀

        αⱼ = γⱼ / (Nʲ 2ʲ⁻¹ D)

    **Arguments:**

    - `difficulty_coefficients`: difficulty coefficients for the linear operator,
        `difficulty_coefficients[i]` is the coefficient for the `i`-th derivative
    - `num_spatial_dims`: number of spatial dimensions `d`
    - `num_points`: number of points `N` used to discretize the domain per
        dimension

    **Returns:**

    - `normalized_coefficients`: normalized coefficients for the linear operator,
        `normalized_coefficients[i]` is the coefficient for the `i`-th derivative
    """
    normalized_coefficients = list(
        gamma / (num_points**j * 2 ** (j - 1) * num_spatial_dims)
        for j, gamma in enumerate(difficulty_coefficients)
    )
    normalized_coefficients[0] = difficulty_coefficients[0]

    normalized_coefficients = tuple(normalized_coefficients)
    return normalized_coefficients

exponax.stepper.generic.reduce_normalized_convection_scale_to_difficulty ¤

reduce_normalized_convection_scale_to_difficulty(
    normalized_convection_scale: float,
    *,
    num_spatial_dims: int,
    num_points: int,
    maximum_absolute: float
) -> float

Reduce the normalized convection scale to a difficulty based interface.

δ₁ = β₁ * M * N * D

Arguments:

  • normalized_convection_scale: normalized convection scale, see also exponax.stepper.generic.normalize_convection_scale
  • num_spatial_dims: number of spatial dimensions d
  • num_points: number of points N used to discretize the domain per dimension
  • maximum_absolute: maximum absolute value of the input state the resulting stepper is applied to

Returns:

  • difficulty_convection_scale: difficulty convection scale
Source code in exponax/stepper/generic/_utils.py 
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
def reduce_normalized_convection_scale_to_difficulty(
    normalized_convection_scale: float,
    *,
    num_spatial_dims: int,
    num_points: int,
    maximum_absolute: float,
) -> float:
    """
    Reduce the normalized convection scale to a difficulty based interface.

        δ₁ = β₁ * M * N * D

    **Arguments:**

    - `normalized_convection_scale`: normalized convection scale, see also
        `exponax.stepper.generic.normalize_convection_scale`
    - `num_spatial_dims`: number of spatial dimensions `d`
    - `num_points`: number of points `N` used to discretize the domain per
        dimension
    - `maximum_absolute`: maximum absolute value of the input state the
        resulting stepper is applied to

    **Returns:**

    - `difficulty_convection_scale`: difficulty convection scale
    """
    difficulty_convection_scale = (
        normalized_convection_scale * maximum_absolute * num_points * num_spatial_dims
    )
    return difficulty_convection_scale

exponax.stepper.generic.extract_normalized_convection_scale_from_difficulty ¤

extract_normalized_convection_scale_from_difficulty(
    difficulty_convection_scale: float,
    *,
    num_spatial_dims: int,
    num_points: int,
    maximum_absolute: float
) -> float

Extract the normalized convection scale from a difficulty based interface.

β₁ = δ₁ / (M * N * D)

Arguments:

  • difficulty_convection_scale: difficulty convection scale
  • num_spatial_dims: number of spatial dimensions d
  • num_points: number of points N used to discretize the domain per dimension
  • maximum_absolute: maximum absolute value of the input state the resulting stepper is applied to

Returns:

  • normalized_convection_scale: normalized convection scale, see also exponax.stepper.generic.normalize_convection_scale
Source code in exponax/stepper/generic/_utils.py 
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
def extract_normalized_convection_scale_from_difficulty(
    difficulty_convection_scale: float,
    *,
    num_spatial_dims: int,
    num_points: int,
    maximum_absolute: float,
) -> float:
    """
    Extract the normalized convection scale from a difficulty based interface.

        β₁ = δ₁ / (M * N * D)

    **Arguments:**

    - `difficulty_convection_scale`: difficulty convection scale
    - `num_spatial_dims`: number of spatial dimensions `d`
    - `num_points`: number of points `N` used to discretize the domain per
        dimension
    - `maximum_absolute`: maximum absolute value of the input state the
        resulting stepper is applied to

    **Returns:**

    - `normalized_convection_scale`: normalized convection scale, see also
        `exponax.stepper.generic.normalize_convection_scale`
    """
    normalized_convection_scale = difficulty_convection_scale / (
        maximum_absolute * num_points * num_spatial_dims
    )
    return normalized_convection_scale

exponax.stepper.generic.reduce_normalized_gradient_norm_scale_to_difficulty ¤

reduce_normalized_gradient_norm_scale_to_difficulty(
    normalized_gradient_norm_scale: float,
    *,
    num_spatial_dims: int,
    num_points: int,
    maximum_absolute: float
) -> float

Reduce the normalized gradient norm scale to a difficulty based interface.

δ₂ = β₂ * M * N² * D

Arguments:

  • normalized_gradient_norm_scale: normalized gradient norm scale, see also exponax.stepper.generic.normalize_gradient_norm_scale
  • num_spatial_dims: number of spatial dimensions d
  • num_points: number of points N used to discretize the domain per dimension
  • maximum_absolute: maximum absolute value of the input state the resulting stepper is applied to

Returns:

  • difficulty_gradient_norm_scale: difficulty gradient norm scale
Source code in exponax/stepper/generic/_utils.py 
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
def reduce_normalized_gradient_norm_scale_to_difficulty(
    normalized_gradient_norm_scale: float,
    *,
    num_spatial_dims: int,
    num_points: int,
    maximum_absolute: float,
) -> float:
    """
    Reduce the normalized gradient norm scale to a difficulty based interface.

        δ₂ = β₂ * M * N² * D

    **Arguments:**

    - `normalized_gradient_norm_scale`: normalized gradient norm scale, see also
        `exponax.stepper.generic.normalize_gradient_norm_scale`
    - `num_spatial_dims`: number of spatial dimensions `d`
    - `num_points`: number of points `N` used to discretize the domain per
        dimension
    - `maximum_absolute`: maximum absolute value of the input state the
        resulting stepper is applied to

    **Returns:**

    - `difficulty_gradient_norm_scale`: difficulty gradient norm scale
    """
    difficulty_gradient_norm_scale = (
        normalized_gradient_norm_scale
        * maximum_absolute
        * jnp.square(num_points)
        * num_spatial_dims
    )
    return difficulty_gradient_norm_scale

exponax.stepper.generic.extract_normalized_gradient_norm_scale_from_difficulty ¤

extract_normalized_gradient_norm_scale_from_difficulty(
    difficulty_gradient_norm_scale: float,
    *,
    num_spatial_dims: int,
    num_points: int,
    maximum_absolute: float
) -> float

Extract the normalized gradient norm scale from a difficulty based interface.

β₂ = δ₂ / (M * N² * D)

Arguments:

  • difficulty_gradient_norm_scale: difficulty gradient norm scale
  • num_spatial_dims: number of spatial dimensions d
  • num_points: number of points N used to discretize the domain per dimension
  • maximum_absolute: maximum absolute value of the input state the resulting stepper is applied to

Returns:

  • normalized_gradient_norm_scale: normalized gradient norm scale, see also exponax.stepper.generic.normalize_gradient_norm_scale
Source code in exponax/stepper/generic/_utils.py 
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
def extract_normalized_gradient_norm_scale_from_difficulty(
    difficulty_gradient_norm_scale: float,
    *,
    num_spatial_dims: int,
    num_points: int,
    maximum_absolute: float,
) -> float:
    """
    Extract the normalized gradient norm scale from a difficulty based interface.

        β₂ = δ₂ / (M * N² * D)

    **Arguments:**

    - `difficulty_gradient_norm_scale`: difficulty gradient norm scale
    - `num_spatial_dims`: number of spatial dimensions `d`
    - `num_points`: number of points `N` used to discretize the domain per
        dimension
    - `maximum_absolute`: maximum absolute value of the input state the
        resulting stepper is applied to

    **Returns:**

    - `normalized_gradient_norm_scale`: normalized gradient norm scale, see also
        `exponax.stepper.generic.normalize_gradient_norm_scale`
    """
    normalized_gradient_norm_scale = difficulty_gradient_norm_scale / (
        maximum_absolute * jnp.square(num_points) * num_spatial_dims
    )
    return normalized_gradient_norm_scale