jaxrts.experimental.SiiNN.NNSiiModel

class jaxrts.experimental.SiiNN.NNSiiModel(checkpoint_dir: Path)

A jaxrts.model.IonFeatModel to use a neural network to obtain ion-ion static structure factors.

Initialize the NNSiiModel.

Parameters:

checkpoint_dir (Path) – The Path to the checkpoint directory.

Methods

Rayleigh_weight(plasma_state, setup)	This is the result from Wünsch [Wünsch, 2011], to calculate the Rayleigh weight for a plasma from multiple species.
S_ii(plasma_state, setup)
S_ii_on_grid(plasma_state)	For py:class:jaxrts.models.IonFeatModel that internally calculate S_ii a k-grid independently of the probe setup (e.g. HNC-based models) this method should return (k_grid, S_ab_grid).
__init__(checkpoint_dir)	Initialize the NNSiiModel.
check(plasma_state)	Test if the model is applicable to the PlasmaState.
check_supports_S_ii_grid(plasma_state)	Whether S_ii_on_grid() can be used for this model, given the current plasma_state.
citation([style, comment])	Return bibliographic information for the Model used.
evaluate(plasma_state, setup)
prepare(plasma_state, key)	Modify the plasma_state in place.

Attributes

allowed_keys	A list of keywords where this model is adequate for
cite_keys	A list of bibtex keys.
supports_S_ii_grid	HNC type models calculate the SSF on a k grid, and then interpolate to the scattering vector of the setup class.