jaxrts.experimental.SiiNN.NNModelExpandedZ.__init__

NNModelExpandedZ.__init__(din, dhid, dout, rngs, no_of_atoms=None)[source]

Construct an expanded neural network model.

The dimensionality of the first layer is increased because each ionization variable Z_i is replaced by two features: its integer ionization stage and its fractional coordinate within that stage.

Parameters:

  • din (int) – Number of input nodes in the original model.

  • dhid (list[int]) – List containing the number of neurons in each hidden layer.

  • dout (int) – Number of output nodes.

  • rngs (nnx.Rngs) – Random number generator used for parameter initialization.

  • no_of_atoms (int or None, optional) – Number of ion species in the plasma. If not given, this quantity is inferred from din as din - 3.