jaxrts.experimental.SiiNN.NNModelExpandedZ.expand_Z_features

NNModelExpandedZ.expand_Z_features(x)

Expand the ionization features of the input tensor.

Each normalized ionization value Z_i is transformed into two quantities representing the ionization stage and the local coordinate within that stage.

The transformation is defined as

\[Z_i^{phys} = Z_i \cdot \mathrm{norm\_Z}_i\]

\[n_i = \lfloor Z_i^{phys} \rfloor\]

\[\phi_i = Z_i^{phys} - n_i\]

where norm_Z denotes the normalization factor stored in the model.

The returned feature vector replaces the original Z_i inputs with (φ_i, n_i) while leaving the remaining input parameters unchanged.

Parameters:

x (jax.Array) – Input tensor containing the normalized network inputs.

Returns:

jax.Array – Transformed input tensor with expanded ionization features.