shnitsel.vis.plot
=================

.. py:module:: shnitsel.vis.plot


Submodules
----------

.. toctree::
   :maxdepth: 1

   /api/shnitsel/vis/plot/common/index
   /api/shnitsel/vis/plot/filtration/index
   /api/shnitsel/vis/plot/kde/index
   /api/shnitsel/vis/plot/p3mhelpers/index
   /api/shnitsel/vis/plot/pca_biplot/index
   /api/shnitsel/vis/plot/polychrom/index
   /api/shnitsel/vis/plot/select/index
   /api/shnitsel/vis/plot/spectra3d/index
   /api/shnitsel/vis/plot/time/index


Functions
---------

.. autoapisummary::

   shnitsel.vis.plot.biplot_kde
   shnitsel.vis.plot.plot_cdf_for_kde
   shnitsel.vis.plot.ski_plots
   shnitsel.vis.plot.pcm_plots


Package Contents
----------------

.. py:function:: biplot_kde(frames, at1 = 0, at2 = 1, at3 = None, at4 = None, geo_filter = None, levels = None, scatter_color = 'time', fill = True, nbins=4, mean=False)

   Generates a biplot that visualizes PCA projections and kernel density estimates (KDE)
   of a property (distance, angle, dihedral angle) describing the geometry of specified
   atoms. The property is chosen based on the number of atoms specified:

   * 2 atoms => distance
   * 3 atoms => angle
   * 4 atoms => dihedral angle

   :param frames: A dataset containing trajectory frames with atomic coordinates.
   :param at1: Indices of the first, second, third and fourth atoms for geometric property calculation.
   :param at2: Indices of the first, second, third and fourth atoms for geometric property calculation.
   :param at3: Indices of the first, second, third and fourth atoms for geometric property calculation.
   :param at4: Indices of the first, second, third and fourth atoms for geometric property calculation.
   :param geo_filter: A list of tuples representing ranges. A KDE is plotted for each range, indicating the distribution of
                      points for which the value of the geometry feature falls in that range.
   :param levels: Contour levels for the KDE plot. Defaults to [0.08, 1]. This parameter is passed to
                  matplotlib.axes.Axes.contour.
   :param scatter_color: Must be one of 'time' or 'geo'. If 'time', the scatter-points will be colored based on the time coordinate;
                         if 'geo', the scatter-points will be colored based on the relevant geometry feature (see above).
   :param fill: Whether to plot filled contours (``fill=True``, uses ``ax.contourf``)
                or just contour lines (``fill=False``, uses ``ax.contour``).
   :param nbins: number of bins to be visualized, must be an integer between 1 and 4

   :returns: The computed KDE data for the atom-atom distance distribution.
   :rtype: kde_dat

   .. rubric:: Notes

   * Computes a geometric property of the specified atoms across all frames.
   * Uses kernel density estimation (KDE) to analyze the distance distributions.
   * Performs PCA on trajectory pairwise distances and visualizes clustering of structural changes.
   * Produces a figure with PCA projection, cluster analysis, and KDE plots.


.. py:function:: plot_cdf_for_kde(z, level, ax=None)

.. py:function:: ski_plots(spectra)

   Plot spectra for different times on top of each other,
   along with a dashed line that tracks the maximum.
   One plot per statecomb; plots stacked vertically.
   Expected to be used on data produced by ``spectra.spectra_all_times``.

   :param spectra: DataArray containing fosc values organized along 'energy', 'time' and
                   'statecomb' dimensions.

   :rtype: Figure object corresponding to plot.

   .. rubric:: Examples

   >>> import shnitsel as st
   >>> from shnitsel.core.plot import spectra3d
   >>> spectra_data = (
           st.io.read(path)
           .st.get_inter_state()
           .st.assign_fosc()
           .st.spectra_all_times())
   >>> spectra3d.ski_plots(spectra_data)


.. py:function:: pcm_plots(spectra)

   Represent fosc as colour in a plot of fosc against time and energy.
   The colour scale is logarithmic.
   One plot per statecomb; plots stacked horizontally.
   Expected to be used on data produced by `spectra.spectra_all_times`.

   :param spectra: DataArray containing fosc values organized along 'energy', 'time' and
                   'statecomb' dimensions.

   :rtype: Figure object corresponding to plot.

   .. rubric:: Examples

   >>> import shnitsel as st
   >>> from shnitsel.core.plot import spectra3d
   >>> spectra_data = (
           st.io.read(path)
           .st.get_inter_state()
           .st.assign_fosc()
           .st.spectra_all_times())
   >>> spectra3d.pcm_plots(spectra_data)