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

.. only:: html

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

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

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

.. _sphx_glr_auto_examples_plot_2dSlab.py:


================================================
NLTE Ca II in a 2D slab with periodic boundaries
================================================

.. GENERATED FROM PYTHON SOURCE LINES 8-13

Lightweaver can solve the NLTE problem in two dimensions. Here we construct
a horizontally homogeneous 2D slab from the FAL C model, with periodic
horizontal boundaries, solve for the Ca II populations, and compare against
the equivalent plane-parallel (1D) solution. Any differences come from the
different angular quadratures and the numerics of the 2D formal solver.

.. GENERATED FROM PYTHON SOURCE LINES 13-21

.. code-block:: Python

    import matplotlib.pyplot as plt
    import numpy as np

    import lightweaver as lw
    from lightweaver.fal import Falc82
    from lightweaver.rh_atoms import CaII_atom, H_6_atom









.. GENERATED FROM PYTHON SOURCE LINES 22-24

Solve for Ca II in statistical equilibrium (H is treated in LTE as a
background species).

.. GENERATED FROM PYTHON SOURCE LINES 24-34

.. code-block:: Python

    def converged_ca(atmos):
        rad_set = lw.RadiativeSet([H_6_atom(), CaII_atom()])
        rad_set.set_active('Ca')
        spect = rad_set.compute_wavelength_grid()
        eq_pops = rad_set.compute_eq_pops(atmos)
        ctx = lw.Context(atmos, spect, eq_pops, Nthreads=2)
        lw.iterate_ctx_se(ctx, quiet=True)
        return eq_pops, ctx









.. GENERATED FROM PYTHON SOURCE LINES 35-38

2D atmospheres are defined on a geometric height grid, with every quantity
provided as a [z, x] array. Here each of the 8 columns is a copy of FAL C,
and the x boundaries are periodic. This uses the 6 rays/octant quadrature of Štěpán et al 2020

.. GENERATED FROM PYTHON SOURCE LINES 38-63

.. code-block:: Python

    fal = Falc82()
    Nx = 8


    def tile(a):
        return np.ascontiguousarray(np.repeat(np.asarray(a)[:, None], Nx, axis=1))


    x = np.arange(Nx) * 50e3
    zero = np.zeros((fal.Nspace, Nx))
    atmos_2d = lw.Atmosphere.make_2d(
        height=np.copy(fal.z),
        x=x,
        temperature=tile(fal.temperature),
        vx=zero,
        vz=zero,
        vturb=tile(fal.vturb),
        ne=tile(fal.ne),
        nh_tot=tile(fal.nh_tot),
        x_lower_bc=lw.PeriodicRadiation(),
        x_upper_bc=lw.PeriodicRadiation(),
    )
    atmos_2d.quadrature(6)
    eq_pops_2d, ctx_2d = converged_ca(atmos_2d)








.. GENERATED FROM PYTHON SOURCE LINES 64-65

The plane-parallel solution for comparison.

.. GENERATED FROM PYTHON SOURCE LINES 65-69

.. code-block:: Python

    atmos_1d = Falc82()
    atmos_1d.quadrature(5)
    eq_pops_1d, ctx_1d = converged_ca(atmos_1d)








.. GENERATED FROM PYTHON SOURCE LINES 70-72

The spatial axis of 2D results is flattened in [z, x] order, so reshape to
[level, z, x] and take a central column.

.. GENERATED FROM PYTHON SOURCE LINES 72-77

.. code-block:: Python

    Nz = atmos_2d.Nz
    ca_2d = np.asarray(eq_pops_2d['Ca']).reshape(-1, Nz, Nx)
    column = ca_2d[:, :, Nx // 2]
    ca_1d = np.asarray(eq_pops_1d['Ca'])








.. GENERATED FROM PYTHON SOURCE LINES 78-81

Plot the relative difference between the 2D column and the 1D solution for
each Ca II level. They agree fairly well, with differences mostly originating
from the angular quadrature (and different formal solver).

.. GENERATED FROM PYTHON SOURCE LINES 81-89

.. code-block:: Python

    z_mm = np.asarray(fal.z) / 1e6
    fig, ax = plt.subplots(figsize=(7, 4), constrained_layout=True)
    for i in range(ca_1d.shape[0]):
        ax.plot(z_mm, column[i] / ca_1d[i] - 1.0, label=f'Level {i}')
    ax.set_xlabel('Height [Mm]')
    ax.set_ylabel(r'$n_{\rm 2D} / n_{\rm 1D} - 1$')
    ax.legend(ncol=2)
    plt.show()



.. image-sg:: /auto_examples/images/sphx_glr_plot_2dSlab_001.png
   :alt: plot 2dSlab
   :srcset: /auto_examples/images/sphx_glr_plot_2dSlab_001.png
   :class: sphx-glr-single-img






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

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


.. _sphx_glr_download_auto_examples_plot_2dSlab.py:

.. only:: html

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

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

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

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

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

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

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


.. only:: html

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

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