Scales Reference

In Scientific Machine Learning (SciML) and operator learning, physical quantities often span vastly different magnitudes. Feeding raw physical parameters or unscaled grid coordinates directly into a Neural Operator can lead to numerical instability, slow convergence, or poor generalization.

To address this, neojax provides a collection of Scales.

Scales vs. Normalizers: Unlike statistical normalizers (which rely on dataset statistics like mean, variance, or min/max values), Scales apply transformations based on physical invariants, domain characteristics, or grid properties. This ensures that the transformed inputs are non-dimensionalized while retaining their underlying physical relationships.

Using Scales with PhysicsNormalizer

Scales in neojax are highly composable. Instead of applying a single scale manually, you typically pass one or more scales into a PhysicsNormalizer. The PhysicsNormalizer evaluates all provided scales, computes their element-wise product, and uses the combined factor to non-dimensionalize your data.

Here is an example of how you might compose multiple physical scales for a fluid dynamics problem:

import jax.numpy as jnp
from neojax.data.scales import CharacteristicLengthScale, ReynoldsScale
from neojax.data.normalizers import PhysicsNormalizer

# 1. Define physical scales
length_scale = CharacteristicLengthScale(L_ref=1.0)
re_scale = ReynoldsScale(U=10.0, L=1.0, nu=1e-3)

# 2. Compose them into a PhysicsNormalizer
physics_norm = PhysicsNormalizer(length_scale, re_scale)

# 3. Compute the combined scale product on your input data
data = jnp.ones((3, 64, 64))
physics_norm = physics_norm.compute_stats(data)

# 4. Non-dimensionalize the data for model input
non_dimensional_data = physics_norm(data)

# You can safely revert the transformation back to physical units during inference
physical_data = physics_norm.inverse_transform(non_dimensional_data)

Available Scales

Depending on your underlying PDE or physical system, you can choose from several scaling strategies or implement your own:

• PhysicalScale: A Python Protocol that defines the interface for all physical scaling logic. You can implement custom non-dimensionalization by inheriting from equinox.Module and implementing the get_scale(self, x: Array) method.
• CharacteristicLengthScale: Standardizes spatial domains by mapping physical dimensions against a constant reference length scale (\(L\)).
• ReynoldsScale: Specifically tailored for fluid dynamics (e.g., Navier-Stokes equations). It scales parameters according to the Reynolds number (\(Re = UL / \nu\)), effectively capturing the ratio of inertial to viscous forces.
• GridBasedScale: Applies spatially-varying scaling factors. Useful when the scaling depends on the location in the domain, such as heavily varying mesh densities or coordinate-dependent physics.

Below is the detailed API reference for each scale.

Physical Scale

neojax.data.scales.PhysicalScale

Protocol for physical scaling logic.

Physical scales are used to non-dimensionalize inputs or outputs based on domain-specific physics (e.g., fluid dynamics, solid mechanics).

Info

To implement a custom scale, create a class (an equinox.Module) that implements the get_scale method. This method should return an array compatible with the input shape x. The resulting scale object can then be passed to a PhysicsNormalizer.

get_scale(x: Float[Array, ...], **kwargs: typing.Any) -> Array

Returns the scaling factor(s) for the given input.

Parameters:

Name	Type	Description	Default
x	Float[Array, ...]	The input array to be scaled.	required
**kwargs	Any	Additional parameters for the scaling logic.	required

Returns:

Type	Description
Array	The scaling array, which should be broadcast-compatible with x.

Characteristic Length Scale

neojax.data.scales.CharacteristicLengthScale

Scales data using a constant reference length.

Commonly used for simple non-dimensionalization where a single spatial constant (like the chord of an airfoil or the diameter of a pipe) governs the system.

Parameters:

Name	Type	Description	Default
L_ref	float | Float[Array, '...']	The constant reference length scale.	required

Internal Attributes

These fields store the internal state of the scale.

• L_ref (float | Float[Array, "..."]): The stored reference length.

Examples:

scale = CharacteristicLengthScale(L_ref=1.0)
factors = scale.get_scale(x)

get_scale(x: Float[Array, '...']) -> Float[Array, '...']

Returns the reference length as a JAX array.

Parameters:

Name	Type	Description	Default
x	Float[Array, '...']	Input array (unused here but required by protocol).	required

Returns:

Type	Description
Float[Array, '...']	The reference length scale.

Reynolds Scale

neojax.data.scales.ReynoldsScale

Scales data based on the Reynolds number \(Re = UL / \\nu\).

A fundamental scaling in fluid mechanics used to non-dimensionalize velocity and pressure fields.

Parameters:

Name	Type	Description	Default
U	float | Float[Array, '...']	Reference velocity.	required
L	float | Float[Array, '...']	Reference length.	required
nu	float | Float[Array, '...']	Kinematic viscosity.	required

Internal Attributes

These fields store the internal state of the scale.

• U (float | Float[Array, "..."]): Reference velocity.
• L (float | Float[Array, "..."]): Reference length.
• nu (float | Float[Array, "..."]): Kinematic viscosity.

Examples:

# Define physics for a flow problem
scale = ReynoldsScale(U=10.0, L=1.0, nu=1e-3)
re_number = scale.get_scale(x)

get_scale(x: Float[Array, '...']) -> Float[Array, '...']

Computes the Reynolds number.

Parameters:

Name	Type	Description	Default
x	Float[Array, '...']	Input array (unused here but required by protocol).	required

Returns:

Type	Description
Float[Array, '...']	The computed Reynolds number.

GridBasedScale

neojax.data.scales.GridBasedScale

Applies spatially-varying scaling factors from a grid.

Used when the scaling depends on the location in the domain, such as varying mesh density or coordinate-dependent physics.

Parameters:

Name	Type	Description	Default
scale_field	Float[Array, '...']	An array of scaling factors defined on the domain grid.	required

Internal Attributes

These fields store the internal state of the scale.

• scale_field (Float[Array, "..."]): The spatially-varying scaling factors.

Examples:

grid_factors = jnp.linspace(1.0, 2.0, 32)
scale = GridBasedScale(scale_field=grid_factors)
factors = scale.get_scale(x_batch)

get_scale(x: Float[Array, '...'], eps: float = 1e-07) -> Float[Array, '...']

Returns the spatially-varying scale field, broadcast to x.

Parameters:

Name	Type	Description	Default
x	Float[Array, '...']	The input array (e.g., a batch of velocity fields).	required
eps	float	Small value to prevent division by zero. Defaults to 1e-7.	1e-07

Returns:

Type	Description
Float[Array, '...']	The scaling field broadcast to the shape of x.