.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "gen_examples/plot_reproduce_Gregori2003_fig6.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_plot_reproduce_Gregori2003_fig6.py>`
        to download the full example code

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

.. _sphx_glr_gen_examples_plot_reproduce_Gregori2003_fig6.py:


Calculate the full structure factors for various plasma conditions
==================================================================

This code reproduces Fig. 6 of :cite:`Gregori.2003`, showing the impact of
changing temperature, ionization and electron density on the scattering
spectra.

It also shows how to attach :py:class:`jaxrts.models.Model` objects to a
:py:class`jaxrts.plasma_state.PlasmaState`, to easily calculate the full S_ee
with :py:meth:`jaxrts.plasma_state.PlasmaState.probe`.

.. GENERATED FROM PYTHON SOURCE LINES 13-144



.. image-sg:: /gen_examples/images/sphx_glr_plot_reproduce_Gregori2003_fig6_001.svg
   :alt: $Z=$ 2, $T=$ 1eV, $n_e=$ 3e+23/cm$^3$, $T=$ 1eV, $Z=$ 2, $T=$ 10eV, $n_e=$ 3e+23/cm$^3$, $T=$ 10eV, $Z=$ 2, $T=$ 40eV, $n_e=$ 3e+23/cm$^3$, $T=$ 40eV
   :srcset: /gen_examples/images/sphx_glr_plot_reproduce_Gregori2003_fig6_001.svg
   :class: sphx-glr-single-img





.. code-block:: Python


    from functools import partial

    import jax.numpy as jnp
    import matplotlib.pyplot as plt
    import scienceplots  # noqa: F401

    import jaxrts
    from jaxrts.models import (
        Gregori2003IonFeat,
        Neglect,
        RPA,
        SchumacherImpulse,
    )

    ureg = jaxrts.ureg

    element = jaxrts.elements.Element("Be")


    def set_state(state, Z_f, n_e, T):
        state.Z_free = jnp.array([Z_f])
        state.mass_density = (
            jnp.array([n_e])
            / (1 * ureg.centimeter**3)
            * element.atomic_mass
            / state.Z_free[0]
        )
        state.T_e = jnp.array([T]) * ureg.electron_volt / ureg.k_B


    state = jaxrts.PlasmaState(
        ions=[element],
        Z_free=jnp.array([2]),
        mass_density=jnp.array([3e23])
        / (1 * ureg.centimeter**3)
        * element.atomic_mass
        / 2,
        T_e=jnp.array([10]) * ureg.electron_volt / ureg.k_B,
    )

    setup = jaxrts.setup.Setup(
        ureg("160°"),
        ureg("4750 eV"),
        ureg("4750 eV") + jnp.linspace(-250, 100, 500) * ureg.electron_volt,
        partial(
            jaxrts.instrument_function.instrument_gaussian,
            sigma=ureg("50.0eV") / ureg.hbar / (2 * jnp.sqrt(2 * jnp.log(2))),
        ),
    )

    state["ionic scattering"] = Gregori2003IonFeat()
    state["free-free scattering"] = RPA()
    state["bound-free scattering"] = SchumacherImpulse()
    state["free-bound scattering"] = Neglect()

    plt.style.use("science")
    fig, ax = plt.subplots(3, 2, sharex=True, sharey=True, figsize=(8, 8))

    state.Z_free = jnp.array([2.5])
    state.mass_density = (
        jnp.array([3e23])
        / (1 * ureg.centimeter**3)
        * element.atomic_mass
        / state.Z_free[0]
    )
    state.probe(setup).m_as(ureg.second)

    for Z, ne, T in [(2, 0.5e23, 1), (2, 1e23, 1), (2, 5e23, 1)]:
        set_state(state, Z, ne, T)
        ax[0, 0].plot(
            (setup.measured_energy - setup.energy).m_as(ureg.electron_volt),
            state.probe(setup).m_as(ureg.second),
            label=f"$n_e = $ {ne}/cm$^3$",
        )
        ax[0, 0].set_title(f"$Z=$ {Z}, $T=$ {T}eV")

    for Z, ne, T in [(2, 0.5e23, 10), (2, 1e23, 10), (2, 5e23, 10)]:
        set_state(state, Z, ne, T)
        ax[1, 0].plot(
            (setup.measured_energy - setup.energy).m_as(ureg.electron_volt),
            state.probe(setup).m_as(ureg.second),
            label=f"$n_e = $ {ne}/cm$^3$",
        )
        ax[1, 0].set_title(f"$Z=$ {Z}, $T=$ {T}eV")

    for Z, ne, T in [(2, 0.5e23, 40), (2, 1e23, 40), (2, 5e23, 40)]:
        set_state(state, Z, ne, T)
        ax[2, 0].plot(
            (setup.measured_energy - setup.energy).m_as(ureg.electron_volt),
            state.probe(setup).m_as(ureg.second),
            label=f"$n_e = $ {ne}/cm$^3$",
        )
        ax[2, 0].set_title(f"$Z=$ {Z}, $T=$ {T}eV")

    for Z, ne, T in [(2.0, 3e23, 1), (2.5, 3e23, 1), (3.0, 3e23, 1)]:
        set_state(state, Z, ne, T)
        ax[0, 1].plot(
            (setup.measured_energy - setup.energy).m_as(ureg.electron_volt),
            state.probe(setup).m_as(ureg.second),
            label=f"$Z=$ {Z}",
        )
        ax[0, 1].set_title(f"$n_e=$ {ne}/cm$^3$, $T=$ {T}eV")

    for Z, ne, T in [(2.0, 3e23, 10), (2.5, 3e23, 10), (3.0, 3e23, 10)]:
        set_state(state, Z, ne, T)
        ax[1, 1].plot(
            (setup.measured_energy - setup.energy).m_as(ureg.electron_volt),
            state.probe(setup).m_as(ureg.second),
            label=f"$Z=$ {Z}",
        )
        ax[1, 1].set_title(f"$n_e=$ {ne}/cm$^3$, $T=$ {T}eV")

    for Z, ne, T in [(2.0, 3e23, 40), (2.5, 3e23, 40), (3.0, 3e23, 40)]:
        set_state(state, Z, ne, T)
        ax[2, 1].plot(
            (setup.measured_energy - setup.energy).m_as(ureg.electron_volt),
            state.probe(setup).m_as(ureg.second),
            label=f"$Z=$ {Z}",
        )
        ax[2, 1].set_title(f"$n_e=$ {ne}/cm$^3$, $T=$ {T}eV")

    for a in ax.flatten():
        a.legend()

    for a in ax[2, :]:
        a.set_xlabel("Energy shift [eV]")
    for a in ax[:, 0]:
        a.set_ylabel("$S_{ee}^{tot}$ [s]")

    plt.show()


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

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


.. _sphx_glr_download_gen_examples_plot_reproduce_Gregori2003_fig6.py:

.. only:: html

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

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

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

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

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


.. only:: html

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

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