.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "gen_examples/ionic/plot_Sii_grid_interpolation.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_gen_examples_ionic_plot_Sii_grid_interpolation.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_gen_examples_ionic_plot_Sii_grid_interpolation.py:


S_ii grid interpolation
=======================

More expensive calculation methods like DFT-MD or Average Atom can yield more
accurate calculations for the static structure factors than the simpler models
shipped with ``jaxrts``. However, such simulations are normally only available
on a grid. Hence, we provide a :py:class:`jaxrts.models.GridInterpolationSii`
class, interpolating with over such a grid. Centerpiece is the
:py:class:`jaxrts.weissker_interpolator.SiiInterpolator`, which implements the
interpolation scheme proposed by :cite:`Weissker.2009`, lineraly interpolating
the SSF's integral over k, rather than the function ifself. However, as no sum
rule exists for the SSFs, we interpolate the function's norm, rather than
putting it to a known value.

.. GENERATED FROM PYTHON SOURCE LINES 16-115



.. image-sg:: /gen_examples/ionic/images/sphx_glr_plot_Sii_grid_interpolation_001.svg
   :alt: plot Sii grid interpolation
   :srcset: /gen_examples/ionic/images/sphx_glr_plot_Sii_grid_interpolation_001.svg
   :class: sphx-glr-single-img





.. code-block:: Python


    import jaxrts
    from jaxrts.weissker_interpolator import AutoNormInterpolator
    import jax.numpy as jnp
    import matplotlib.pyplot as plt
    from copy import deepcopy

    ureg = jaxrts.ureg

    state = jaxrts.PlasmaState(
        ions=[jaxrts.Element("Al")],
        Z_free=jnp.array([3]),
        mass_density=jnp.array([10]) * ureg.gram / ureg.centimeter**3,
        T_e=100 * ureg.electron_volt / ureg.k_B,
    )
    state["ion-ion Potential"] = jaxrts.hnc_potentials.YukawaShortRangeRepulsion(
        ureg("250pm")
    )
    state["ionic scattering"] = jaxrts.models.OnePotentialHNCIonFeat()

    target_k, target_Sii = state["ionic scattering"].S_ii_on_grid(state)

    # Create a 2x2 grid around the target state
    T_range = 100 * ureg.electron_volt / ureg.k_B
    rho_range = 8.0 * ureg.gram / ureg.centimeter**3

    grid_Sii = jnp.zeros([*target_Sii.shape, 2, 2])
    for idx in ([0, 0], [0, 1], [1, 0], [1, 1]):
        t, r = idx
        new_state = deepcopy(state)
        new_state.T_e += (t - 0.5) * T_range
        new_state.T_i = jaxrts.units.to_array([new_state.T_e])
        new_state.mass_density += (r - 0.5) * rho_range
        _, calc_Sii = new_state["ionic scattering"].S_ii_on_grid(new_state)
        grid_Sii = grid_Sii.at[:, :, :, *idx].set(
            calc_Sii.m_as(ureg.dimensionless)
        )

    # Cut the grid at k == 10/a0
    cutoff = 10 / ureg.a0
    grid_Sii = grid_Sii[:, :, target_k < cutoff, :, :]
    target_Sii = target_Sii[:, :, target_k < cutoff]
    target_k = target_k[target_k < cutoff]


    interpolator = AutoNormInterpolator(
        grid_Sii * ureg.dimensionless,
        target_k,
        state.T_e + (T_range * jnp.array([-0.5, 0.5])),
        state.mass_density[0] + (rho_range * jnp.array([-0.5, 0.5])),
    )
    # The second argument specifies the attributes of a plasma state that should be
    # used when calling the `interpolator`.
    # The syntax is either a plain attribute (`T_e`) or a jnpu function (`sum`) and
    # then the :py:class:`jaxrts.plasmastate.PlasmaState` attributes that should be
    # passed to that function (`mass_density`), separated by a '%'
    Sii_model = jaxrts.models.GridInterpolationSii(
        interpolator, ["T_e", "sum%mass_density"]
    )

    k, Sii = Sii_model.S_ii_on_grid(state)
    fig, ax = plt.subplots()
    for idx in ([0, 1], [1, 0], [0, 0], [1, 1]):
        ax.plot(
            k.m_as(1 / ureg.a0),
            grid_Sii[0, 0, :, *idx],
            color="gray",
            label="grid input" if idx == [0, 0] else None,
            alpha=0.7,
        )
    ax.plot(
        k.m_as(1 / (1 * ureg.a0)),
        Sii[0, 0, :].m_as(ureg.dimensionless),
        label="inferred",
    )
    ax.plot(
        k.m_as(1 / (1 * ureg.a0)),
        target_Sii[0, 0, :].m_as(ureg.dimensionless),
        ls="dashed",
        label="target",
    )
    ax.plot(
        k.m_as(1 / ureg.a0),
        (
            grid_Sii[0, 0, :, 0, 1]
            + grid_Sii[0, 0, :, 1, 0]
            + grid_Sii[0, 0, :, 0, 0]
            + grid_Sii[0, 0, :, 1, 1]
        )
        / 4,
        ls="dotted",
        label="direct mean of the grid input",
    )
    ax.set_xlabel("$k$ (1/a0)")
    ax.set_ylabel("$S_{ii}$ (dimensionless)")
    ax.set_xlim(0, 10)
    plt.legend()
    plt.tight_layout()
    plt.show()


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 2.965 seconds)


.. _sphx_glr_download_gen_examples_ionic_plot_Sii_grid_interpolation.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: plot_Sii_grid_interpolation.ipynb <plot_Sii_grid_interpolation.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: plot_Sii_grid_interpolation.py <plot_Sii_grid_interpolation.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: plot_Sii_grid_interpolation.zip <plot_Sii_grid_interpolation.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_