# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Model atmospheric emission and absorption for spectroscopic simulations.

The atmosphere model is responsible for calculating the spectral flux density
arriving at the telescope given a source flux entering the atmosphere. The
calculation is either performed as:

.. math::

    f(\\lambda) = 10^{-e(\\lambda) X / 2.5} s(\\lambda) + a b(\\lambda)

if ``extinct_emission`` is False, or else as:

.. math::

    f(\\lambda) = 10^{-e(\\lambda) X / 2.5} \\left[
    s(\\lambda) + a b(\\lambda)\\right]

where :math:`s(\\lambda)` is the source flux entering the atmosphere,
:math:`e(\\lambda)` is the zenith extinction, :math:`X` is the airmass,
:math:`a` is the fiber entrance face area, and :math:`b(\\lambda)` is the
sky emission surface brightness.  The sky brightness can optionally include
a scattered moonlight component.

An atmosphere model is usually initialized from a configuration used to create
a simulator and then accessible via its ``atmosphere`` attribute, for example:

    >>> import specsim.simulator
    >>> simulator = specsim.simulator.Simulator('test')  # doctest: +IGNORE_OUTPUT
    >>> simulator.atmosphere.airmass
    1.0

See :doc:`/api` for examples of changing model parameters defined in the
configuration.  Certain parameters can also be changed after a model has
been initialized, for example:

    >>> simulator.atmosphere.airmass = 1.5
    >>> simulator.atmosphere.moon.moon_phase = 0.25
    >>> simulator.atmosphere.moon.moon_zenith = 25 * u.deg

See :class:`Atmosphere` and :class:`Moon` for details.
"""
from __future__ import print_function, division

import numpy as np

import astropy.units as u

import speclite.filters

import specsim.config

class Atmosphere(object):
    """Model atmospheric surface brightness and extinction.

    A simulation uses only our read-only :attr:`surface_brightness` and
    :attr:`extinction` attributes.  Use the :attr:`condition` and
    :attr:`airmass` attributes to update this model.  Refer to the
    :class:`moon model <Moon>` for details on updating the optional
    scattered moon model.

    Parameters
    ----------
    wavelength : astropy.units.Quantity
        Array of wavelengths with units where data is tabulated.
    surface_brightness_dict : dict
        Dictionary of tabulated sky emission surface brightness values. Each
        dictionary key defines a possible sky condition.
    extinction_coefficient : array
        Array of extinction coefficients tabulated on ``wavelength``.
    extinct_emission : bool
        If set, atmospheric extinction is applied to sky emission.
    condition : str
        Sky emission condition to use, which must be one of the keys
        of ``surface_brightness_dict``.
    seeing : dict or None
        Dictionary of seeing PSF parameters to use which must contain keys
        "fwhm_ref", "wlen_ref" and "moffat_beta".  Seeing is used to define
        the atmospheric PSF, which is only used when
        :attr:`instrument.fiberloss_method` equals "galsim".
    airmass : float
        Airmass of the observation.
    moon : :class:`Moon` or None
        Model to use for scattered moonlight.
    """
    def __init__(self, wavelength, surface_brightness_dict,
                 extinction_coefficient, extinct_emission, condition, airmass,
                 seeing, moon):
        self._wavelength = wavelength
        self._surface_brightness_dict = surface_brightness_dict
        self._extinction_coefficient = extinction_coefficient
        self._extinct_emission = extinct_emission
        self._condition_names = surface_brightness_dict.keys()
        self._moon = moon
        self.condition = condition
        self.airmass = airmass
        if seeing is not None:
            for required in ('fwhm_ref', 'wlen_ref', 'moffat_beta'):
                if required not in seeing:
                    raise ValueError('Missing required seeing key "{0}"'
                                     .format(required))
        self._seeing = seeing

    @property
    def moon(self):
        """Moon or None: Model of scattered moonlight.

        See :class:`Moon` for details on changing scattered moon simulation
        parameters via this attribute.
        """
        return self._moon

    @property
    def surface_brightness(self):
        """astropy.units.Quantity: Total sky surface brightness.

        Includes both dark sky emission and (if configured) scattered moonlight.
        Changes to :attr:`condition` or :attr:`airmass` are reflected here.
        """
        sky = self._surface_brightness_dict[self.condition].copy()
        if self._extinct_emission:
            sky *= self.extinction
        if self.moon is not None and self.moon.visible:
            sky += self.moon.surface_brightness
        return sky

    @property
    def extinction(self):
        """numpy.ndarray: The extinction factor for the current model airmass.

        Tabulated as a function of wavelength. Changes to :attr:`airmass`
        automatically update these values.
        """
        return self._extinction

    @property
    def condition(self):
        """str: Sky emission condition.

        Must be one of the predefined names in :attr:`condition_names`.
        """
        return self._condition

    @condition.setter
    def condition(self, name):
        if name not in self._condition_names:
            raise ValueError(
                "Invalid condition '{0}'. Pick one of {1}."
                .format(name, self._condition_names))
        self._condition = name

    @property
    def condition_names(self):
        """list: The list of valid sky condition names.

        The valid names are keys of the ``atmosphere.sky.table.paths`` node,
        or "default" if only a single path is specified via a
        ``atmosphere.sky.table.path`` node.
        """
        return self._condition_names

    @property
    def airmass(self):
        """float: Observing airmass.

        Changes to this value automatically propagate to our scattered
        moon model, if there is one.
        """
        return self._airmass

    @airmass.setter
    def airmass(self, airmass):
        self._airmass = airmass
        self._extinction = 10 ** (-self._extinction_coefficient * airmass / 2.5)
        if self.moon is not None:
            self.moon.airmass = airmass

    @property
    def seeing_moffat_beta(self):
        """float: Beta parameter for atmospheric Moffat profile.

        Returns None if no seeing has been specified.
        """
        return self._seeing['moffat_beta'] if self._seeing else None

    @property
    def seeing_wlen_ref(self):
        """float: Reference wavelength for :attr:`seeing_fwhm_ref`

        Returns None if no seeing has been specified.
        """
        return self._seeing['wlen_ref'] if self._seeing else None

    @property
    def seeing_fwhm_ref(self):
        """float: FWHM zenith seeing at :attr:`seeing_wlen_ref`.

        Returns None if no seeing has been specified.
        """
        return self._seeing['fwhm_ref'] if self._seeing else None

    @seeing_fwhm_ref.setter
    def seeing_fwhm_ref(self, fwhm_ref):
        try:
            self._seeing['fwhm_ref'] = fwhm_ref.to(u.arcsec)
        except TypeError:
            raise ValueError('Seeing has not been initialized.')
        except (u.UnitConversionError, AttributeError):
            raise ValueError('Invalid units for seeing_fwhm_ref.')

    def get_seeing_fwhm(self, wavelength):
        """Calculate the seeing FWHM at the specified wavelength.

        Assumes that seeing scales with wavelength with a power -1/5, as
        predicted by Kolmogorov turbulence theory.

        Parameters
        ----------
        wavelength : astropy.units.Quantity
            Wavelength in units convertible to Angstroms.

        Returns
        -------
        astropy.units.Quantity
            Full-width half maximum of seeing distribution at the specified
            wavelength, in on-sky angular units.
        """
        wlen_ratio = (wavelength.to(u.Angstrom).value /
                      self._seeing['wlen_ref'].to(u.Angstrom).value)
        return self._seeing['fwhm_ref'] * wlen_ratio ** (-0.2)

    def plot(self):
        """Plot a summary of this atmosphere model.

        Requires that the matplotlib package is installed.
        """
        import matplotlib.pyplot as plt

        fig, ax1 = plt.subplots(figsize=(8, 4))
        ax1_rhs = ax1.twinx()

        wave = self._wavelength.to(u.Angstrom).value
        wave_unit = u.Angstrom

        sky_unit = 1e-17 * u.erg / (u.cm**2 * u.s * u.Angstrom * u.arcsec**2)
        sky = self.surface_brightness.to(sky_unit).value
        sky_min, sky_max = np.percentile(sky, (1, 99))

        ext = self._extinction_coefficient
        ext_min, ext_max = np.percentile(ext, (1, 99))

        ax1.scatter(wave, sky, color='g', lw=0, s=1.)
        if self.moon is not None and self.moon.visible:
            moon = self.moon.surface_brightness.to(sky_unit).value
            ax1.scatter(wave, moon, color='b', lw=0, s=1.)
            # Adjust the vertical limits to include the moon.
            moon_min, moon_max = np.percentile(moon, (1, 99))
            sky_min = min(moon_min, sky_min)
            sky_max = max(moon_max, sky_max)
        ax1_rhs.scatter(wave, ext, color='r', lw=0, s=1.)

        ax1.set_yscale('log')
        ax1_rhs.set_yscale('log')

        ax1.set_ylabel(
            'Surface Brightness [$10^{-17}\\mathrm{erg}/(\\mathrm{cm}^2' +
            '\\mathrm{s} \\AA)/\\mathrm{arcsec}^2$]')
        ax1.set_ylim(0.5 * sky_min, 1.5 * sky_max)
        ax1_rhs.set_ylabel('Zenith Extinction')
        ax1_rhs.set_ylim(0.5 * ext_min, 1.5 * ext_max)

        ax1.set_xlabel('Wavelength [$\\AA$]')
        ax1.set_xlim(wave[0], wave[-1])

        ncol = 2
        ax1.plot([], [], 'g-',
                 label='Total Emission ({0})'.format(self.condition))
        if self.moon is not None and self.moon.visible:
            ax1.plot([], [], 'b-', label='Scattered Moon')
            ncol += 1
        ax1.plot([], [], 'r-', label='Extinction')
        ax1.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
                   ncol=ncol, mode='expand', borderaxespad=0.)

class Moon(object):
    """Model of scattered moonlight.

    Most of the work is performed by :func:`krisciunas_schaefer`, which
    implements the model of their 1991 paper. This class uses the predicted
    V-band surface brightness to normalize an input lunar spectrum (or solar
    spectrum if you assume the moon is grey).

    The predicted :attr:`surface_brightness` is automatically updated to
    reflect changes in the following attributes: :attr:`airmass`,
    :attr:`moon_zenith`, :attr:`moon_phase` and :attr:`separation_angle`.

    This implementation is loosely based on and tested against [IDL code]
    (https://desi.lbl.gov/svn/code/desimodel/tags/0.4.2/pro/lunarmodel.pro)
    from Connie Rockosi.

    Parameters
    ----------
    wavelength : astropy.units.Quantity
        Array of wavelengths with units where data is tabulated.
    moon_spectrum : astropy.units.Quantity
        Tabulated spectrum of scattered moonlight with units of flux density.
        The normalization does not matter since it will be fixed by
        :meth:`get_lunar_surface_brightness`.  A solar spectrum can be used,
        which effectively assumes that the moon's reflectance is wavelength
        independent.
    extinction_coefficient : array
        Array of extinction coefficients tabulated on ``wavelength``.
    airmass : float
        Airmass of the observation.
    moon_zenith : astropy.units.Quantity
        See :func:`krisciunas_schaefer`.
    separation_angle : astropy.units.Quantity
        See :func:`krisciunas_schaefer`.
    moon_phase : float
        See :func:`krisciunas_schaefer`.
    """
    def __init__(self, wavelength, moon_spectrum, extinction_coefficient,
                 airmass, moon_zenith, separation_angle, moon_phase):
        self._wavelength = wavelength
        self._moon_spectrum = moon_spectrum
        self._extinction_coefficient = extinction_coefficient

        # Calculate the V-band extinction of the moon spectrum.
        self._vband = speclite.filters.load_filter('bessell-V')
        V = self._vband.get_ab_magnitude(moon_spectrum, wavelength)
        extinction = 10 ** (-extinction_coefficient / 2.5)
        Vstar = self._vband.get_ab_magnitude(
            moon_spectrum * extinction, wavelength)
        self._vband_extinction = Vstar - V

        # Initialize the model parameters.
        self.airmass = airmass
        self.moon_zenith = moon_zenith
        self.separation_angle = separation_angle
        self.moon_phase = moon_phase

    def _update(self):
        """Update the model based on the current parameter values.
        """
        self._update_required = False

        if not self.visible:
            self._surface_brightness = (
                np.zeros_like(self._moon_spectrum) / (u.arcsec ** 2))
            self._scattered_V = None
            return

        # Calculate the V-band surface brightness of scattered moonlight.
        self._scattered_V = krisciunas_schaefer(
            self.obs_zenith, self.moon_zenith, self.separation_angle,
            self.moon_phase, self.vband_extinction)

        # Calculate the wavelength-dependent extinction of moonlight
        # scattered once into the observed field of view.
        scattering_airmass = (
            1 - 0.96 * np.sin(self.moon_zenith) ** 2) ** (-0.5)
        extinction = (
            10 ** (-self._extinction_coefficient * scattering_airmass / 2.5) *
            (1 - 10 ** (-self._extinction_coefficient * self.airmass / 2.5)))
        self._surface_brightness = self._moon_spectrum * extinction

        # Renormalized the extincted spectrum to the correct V-band magnitude.
        raw_V = self._vband.get_ab_magnitude(
            self._surface_brightness, self._wavelength) * u.mag
        area = 1 * u.arcsec ** 2
        self._surface_brightness *= 10 ** (
            -(self._scattered_V * area - raw_V) / (2.5 * u.mag)) / area

    @property
    def scattered_V(self):
        """V-band surface brightness of scattered moonlight.

        This is a read-only attribute whose value depends
        on the current values of :attr:`airmass`, :attr:`moon_zenith`,
        :attr:`moon_phase` and :attr:`separation_angle`.  Returns None if
        the moon is below the horizon.
        """
        if self._update_required:
            self._update()
        return self._scattered_V

    @property
    def surface_brightness(self):
        """astropy.units.Quantity: Tabulated scattered moon surface brightness.

        This is the only model attribute used for simulation. Its value depends
        on the current values of :attr:`airmass`, :attr:`moon_zenith`,
        :attr:`moon_phase` and :attr:`separation_angle`.
        """
        if self._update_required:
            self._update()
        return self._surface_brightness

    @property
    def airmass(self):
        """Airmass of observation used for lunar scattering model.

        Changes to this value will update :attr:`obs_zenith` and
        :attr:`surface_brightness`.

        This should normally be the same airmass that is used in the
        :class:`Atmosphere` model to calculate source extinction, but this
        is not checked here.
        """
        return self._airmass

    @airmass.setter
    def airmass(self, airmass):
        # Remove any dimensionless astropy.units.Quantity wrapper since
        # np.arcsin(Quantity(1)) has u.rad added automatically, but we
        # add it explicitly below.
        self._airmass = np.float(airmass)
        # Estimate the zenith angle corresponding to this observing airmass.
        # We invert eqn.3 of KS1991 for this (instead of eqn.14).
        self._obs_zenith = np.arcsin(
            np.sqrt((1 - self._airmass ** -2) / 0.96)) * u.rad
        self._update_required = True

    @property
    def visible(self):
        """bool: Read-only visibility of the moon.

        The visibility criterion is :attr:`moon_zenith` < 90 degrees.
        """
        return self._visible

    @property
    def moon_phase(self):
        """Phase of the moon.

        See :func:`krisciunas_schaefer`. Changes to this value will update
        :attr:`surface_brightness`.
        """
        return self._moon_phase

    @moon_phase.setter
    def moon_phase(self, moon_phase):
        self._moon_phase = moon_phase
        self._update_required = True

    @property
    def obs_zenith(self):
        """Read-only value of the observing zenith angle.

        This attribute is calculated from :attr:`airmass` by inverting
        Eqn.3 of Krisciunas & Schaefer 1991:

        .. math::

            X = (1 - 0.96 \\sin^2 Z)^{-0.5}
        """
        return self._obs_zenith

    @property
    def moon_zenith(self):
        """Moon zenith angle.

        See :func:`krisciunas_schaefer`. Changes to this value will update
        :attr:`surface_brightness` and :attr:`visible`.
        """
        return self._moon_zenith
        self._update_required = True

    @moon_zenith.setter
    def moon_zenith(self, moon_zenith):
        self._moon_zenith = moon_zenith
        self._visible = self._moon_zenith < 90 * u.deg

    @property
    def separation_angle(self):
        """Read-only value of the observation-moon separation angle.

        See :func:`krisciunas_schaefer`. Changes to this value will update
        :attr:`surface_brightness`.
        """
        return self._separation_angle

    @separation_angle.setter
    def separation_angle(self, separation_angle):
        self._separation_angle = separation_angle
        self._update_required = True

    @property
    def vband_extinction(self):
        """Read-only value of the V-band extinction of the moon spectrum.

        Calculated as V* - V where V is the Bessell-V magnitude of the
        input lunar spectrum and V* is calculated with airmass 1.0
        extinction applied.
        """
        return self._vband_extinction

def krisciunas_schaefer(obs_zenith, moon_zenith, separation_angle, moon_phase,
                        vband_extinction):
    """Calculate the scattered moonlight surface brightness in V band.

    Based on Krisciunas and Schaefer, "A model of the brightness of moonlight",
    PASP, vol. 103, Sept. 1991, p. 1033-1039 (http://dx.doi.org/10.1086/132921).
    Equation numbers in the code comments refer to this paper.

    The function :func:`plot_lunar_brightness` provides a convenient way to
    plot this model's predictions as a function of observation pointing.

    Units are required for the angular inputs and the result has units of
    surface brightness, for example:

    >>> sb = krisciunas_schaefer(20*u.deg, 70*u.deg, 50*u.deg, 0.25, 0.15)
    >>> print(np.round(sb, 3))
    19.855 mag / arcsec2

    The output is automatically broadcast over input arrays following the usual
    numpy rules.

    This method has several caveats but the authors find agreement with data at
    the 8% - 23% level.  See the paper for details.

    Parameters
    ----------
    obs_zenith : astropy.units.Quantity
        Zenith angle of the observation in angular units.
    moon_zenith : astropy.units.Quantity
        Zenith angle of the moon in angular units.
    separation_angle : astropy.units.Quantity
        Opening angle between the observation and moon in angular units.
    moon_phase : float
        Phase of the moon from 0.0 (full) to 1.0 (new), which can be calculated
        as abs((d / D) - 1) where d is the time since the last new moon
        and D = 29.5 days is the period between new moons.  The corresponding
        illumination fraction is ``0.5*(1 + cos(pi * moon_phase))``.
    vband_extinction : float
        V-band extinction coefficient to use.

    Returns
    -------
    astropy.units.Quantity
        Observed V-band surface brightness of scattered moonlight.
    """
    moon_phase = np.asarray(moon_phase)
    if np.any((moon_phase < 0) | (moon_phase > 1)):
        raise ValueError(
            'Invalid moon phase {0}. Expected 0-1.'.format(moon_phase))
    # Calculate the V-band magnitude of the moon (eqn. 9).
    abs_alpha = 180. * moon_phase
    m = -12.73 + 0.026 * abs_alpha + 4e-9 * abs_alpha ** 4
    # Calculate the illuminance of the moon outside the atmosphere in
    # foot-candles (eqn. 8).
    Istar = 10 ** (-0.4 * (m + 16.57))
    # Calculate the scattering function (eqn.21).
    rho = separation_angle.to(u.deg).value
    f_scatter = (10 ** 5.36 * (1.06 + np.cos(separation_angle) ** 2) +
                 10 ** (6.15 - rho / 40.))
    # Calculate the scattering airmass along the lines of sight to the
    # observation and moon (eqn. 3).
    X_obs = (1 - 0.96 * np.sin(obs_zenith) ** 2) ** (-0.5)
    X_moon = (1 - 0.96 * np.sin(moon_zenith) ** 2) ** (-0.5)
    # Calculate the V-band moon surface brightness in nanoLamberts.
    B_moon = (f_scatter * Istar *
        10 ** (-0.4 * vband_extinction * X_moon) *
        (1 - 10 ** (-0.4 * (vband_extinction * X_obs))))
    # Convert from nanoLamberts to to mag / arcsec**2 using eqn.19 of
    # Garstang, "Model for Artificial Night-Sky Illumination",
    # PASP, vol. 98, Mar. 1986, p. 364 (http://dx.doi.org/10.1086/131768)
    return ((20.7233 - np.log(B_moon / 34.08)) / 0.92104 *
            u.mag / (u.arcsec ** 2))

def plot_lunar_brightness(moon_zenith, moon_azimuth, moon_phase,
                          vband_extinction=0.162, ngrid=250,
                          cmap='YlGnBu', figure_size=(8, 6)):
    """Create a polar plot of the scattered moon brightness in V band.

    Evaluates the model of :func:`krisciunas_schaefer` on a polar grid of
    observation pointings, for a fixed moon position and phase.

    This method requires that matplotlib is installed.

    Parameters
    ----------
    moon_zenith : astropy.units.Quantity
        See :func:`krisciunas_schaefer`.
    moon_azimuth : astropy.units.Quantity
        Aziumuthal angle of the moon in angular units.  Azimuth is measured
        clockwize from zero (North).
    moon_phase : float
        See :func:`krisciunas_schaefer`.
    vband_extinction : float
        See :func:`krisciunas_schaefer`.
    ngrid : int
        Size of observing location zenith and azimuth grids to use.
    cmap : str
        Name of the matplotlib color map to use.
    figure_size : tuple or None
        Tuple (width, height) giving the figure dimensions in inches.

    Returns
    -------
    tuple
        Tuple (fig, ax, cax) of matplotlib objects created for this plot. You
        can ignore these unless you want to make further changes to the plot.
    """
    import matplotlib.pyplot as plt

    # Build a grid in observation (zenith, azimuth).
    # Build a grid in observation (zenith, azimuth).
    obs_zenith = np.linspace(0., 90., ngrid, endpoint=False) * u.deg
    obs_az = (np.linspace(0., 360., ngrid) * u.deg)[:, np.newaxis]

    # Calculate the separation angles.
    cos_sep = (np.cos(moon_zenith) * np.cos(obs_zenith) +
               np.cos(moon_azimuth - obs_az) * np.sin(moon_zenith) *
               np.sin(obs_zenith))
    sep = np.arccos(cos_sep)

    # Calculate the V-band moon brightness.
    moon_V = krisciunas_schaefer(
        obs_zenith, moon_zenith, sep, moon_phase, vband_extinction)

    # Initialize the plot. We are borrowing from:
    # http://blog.rtwilson.com/producing-polar-contour-plots-with-matplotlib/
    fig, ax = plt.subplots(
        figsize=figure_size, subplot_kw=dict(projection='polar'))
    r, theta = np.meshgrid(
        obs_zenith.to(u.deg).value, obs_az.to(u.rad).value[:,0], copy=False)
    ax.set_theta_zero_location('N')
    ax.set_theta_direction(-1)
    ax.set_ylim(0., 90.)

    # Draw a polar contour plot.
    cax = ax.contourf(theta, r, moon_V.value, 50, cmap=cmap)
    fig.colorbar(cax).set_label('Scattered Moon V [mag/arcsec2]')

    # Draw a point indicating the moon position.
    plt.scatter(moon_azimuth.to(u.rad).value, moon_zenith.to(u.deg).value,
                s=150., marker='o', color='w', lw=0.5, edgecolor='k')

    # Add labels.
    xy, coords = (1., 0.), 'axes fraction'
    plt.annotate('$k_V$ = {0:.3f}'.format(vband_extinction),
                 xy, xy, coords, coords,
                 horizontalalignment='right', verticalalignment='top',
                 size='x-large', color='k')
    xy, coords = (0., 0.), 'axes fraction'
    plt.annotate('$\\phi$ = {0:.1f}%'.format(100. * moon_phase),
                 xy, xy, coords, coords,
                 horizontalalignment='left', verticalalignment='top',
                 size='x-large', color='k')

    plt.tight_layout()
    return fig, ax, cax

def initialize(config):
    """Initialize the atmosphere model from configuration parameters.

    After an atmosphere model has been initialized, further changes to the
    input configuration will no effect unless this method is called to
    initialize a new model. However, certain model attributes can be
    varied after a model is initialized.  See :class:`Atmosphere` and
    :class:`Moon` for details.

    Parameters
    ----------
    config : :class:`specsim.config.Configuration`
        The configuration parameters to use.

    Returns
    -------
    Atmosphere
        An initialized :class:`atmosphere model <Atmosphere>`, possibly
        containing a :class:`scattered moonlight model <Moon>`.
    """
    atm_config = config.atmosphere

    # Load tabulated data.
    surface_brightness_dict = config.load_table(
        atm_config.sky, 'surface_brightness', as_dict=True)
    extinction_coefficient = config.load_table(
        atm_config.extinction, 'extinction_coefficient')

    # Initialize an optional atmospheric seeing PSF.
    psf_config = getattr(atm_config, 'seeing', None)
    if psf_config:
        seeing = dict(
            fwhm_ref=specsim.config.parse_quantity(psf_config.fwhm_ref),
            wlen_ref=specsim.config.parse_quantity(psf_config.wlen_ref),
            moffat_beta=float(psf_config.moffat_beta))
    else:
        seeing = None

    # Initialize an optional lunar scattering model.
    moon_config = getattr(atm_config, 'moon', None)
    if moon_config:
        moon_spectrum = config.load_table(moon_config, 'flux')
        c = config.get_constants(moon_config,
            ['moon_zenith', 'separation_angle', 'moon_phase'])
        moon = Moon(
            config.wavelength, moon_spectrum, extinction_coefficient,
            atm_config.airmass, c['moon_zenith'], c['separation_angle'],
            c['moon_phase'])
    else:
        moon = None

    atmosphere = Atmosphere(
        config.wavelength, surface_brightness_dict, extinction_coefficient,
        atm_config.extinct_emission, atm_config.sky.condition,
        atm_config.airmass, seeing, moon)

    if config.verbose:
        print(
            "Atmosphere initialized with condition '{0}' from {1}."
            .format(atmosphere.condition, atmosphere.condition_names))
        if seeing:
            print('Seeing is {0} at {1} with Moffat beta {2}.'
                  .format(seeing['fwhm_ref'], seeing['wlen_ref'],
                          seeing['moffat_beta']))
        if moon:
            print(
                'Lunar V-band extinction coefficient is {0:.5f}.'
                .format(moon.vband_extinction))

    return atmosphere


#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Built-in imports
import warnings
import itertools

# 3rd party imports
import numpy as np
import xarray as xr

from scipy import constants

# Local imports
from ..pyrf import iso86012datetime64, time_clip, ts_scalar, ts_skymap

from .psd_rebin import psd_rebin

__author__ = "Louis Richard"
__email__ = "louisr@irfu.se"
__copyright__ = "Copyright 2020-2021"
__license__ = "MIT"
__version__ = "2.3.7"
__status__ = "Prototype"

def _coord_sys(coord_sys):
    x_vec = coord_sys[0, :] / np.linalg.norm(coord_sys[0, :])
    y_vec = coord_sys[1, :] / np.linalg.norm(coord_sys[1, :])

    z_vec = np.cross(x_vec, y_vec) / np.linalg.norm(np.cross(x_vec, y_vec))
    y_vec = np.cross(z_vec, x_vec) / np.linalg.norm(np.cross(z_vec, x_vec))

    changed_xyz = [False, False, False]

    for i, vec, c in zip([0, 1, 2], [x_vec, y_vec, z_vec], ["x", "y", "z"]):
        if abs(np.rad2deg(np.arccos(np.dot(vec, coord_sys[:, i])))) > 1.:
            msg = " ".join(["In making 'xyz' a right handed orthogonal"
                            "coordinate system, {}".format(c, i),
                            "(in-plane {:d}) was changed from",
                            np.array2string(coord_sys[:, i]),
                            "to",
                            np.array2string(x_vec),
                            "Please verify that this is according to your"
                            "intentions."])
            warnings.warn(msg, UserWarning)
            changed_xyz[i] = True

    return x_vec, y_vec, z_vec, changed_xyz

def _init(vdf, tint):
    assert isinstance(vdf, xr.Dataset), "vdf must be a xarray.Dataset"

    len_e = 32

    if vdf.phi.data.ndim == 1:
        phi = xr.DataArray(np.tile(vdf.phi.data, (len(vdf.data), 1)),
                           coords=[vdf.time.data,
                                   np.arange(len(vdf.phi.data))],
                           dims=["time", "idx"])
    else:
        phi = vdf.phi

    theta = vdf.theta
    polar = np.deg2rad(theta)
    azimuthal = np.deg2rad(phi)
    step_table = vdf.attrs.get("esteptable", np.zeros(len(vdf.time)))

    energy0 = vdf.attrs.get("energy0", vdf.energy.data[0, :])
    energy1 = vdf.attrs.get("energy1", vdf.energy.data[1, :])

    diff_energ = np.median(np.diff(np.log10(energy0))) / 2

    energy0_edges = np.hstack([10 ** (np.log10(energy0) - diff_energ),
                               10 ** (np.log10(energy0[-1]) + diff_energ)])
    energy1_edges = np.hstack([10 ** (np.log10(energy1) - diff_energ),
                               10 ** (np.log10(energy1[-1]) + diff_energ)])

    if tint is not None and len(tint) == 1:
        t_id = np.argmin(
            np.abs(vdf.time.data - iso86012datetime64(np.array(tint))[0]))

        dist = vdf.data.data[t_id, ...]
        dist = dist[None, ...]
        step_table = step_table[t_id]
        azimuthal = azimuthal[t_id, ...]

        if step_table.data:
            energy_edges = energy1_edges
        else:
            energy_edges = energy1_edges

    elif tint is not None and len(tint) == 2:
        dist = time_clip(vdf.data, tint)
        step_table = ts_scalar(vdf.time.data, step_table)
        step_table = time_clip(step_table, tint)
        azimuthal = time_clip(azimuthal, tint)

        if len(dist.time) > 1 and list(energy0) != list(energy1):
            print("notice: Rebinning distribution.")
            temp = ts_skymap(dist.time.data, dist, time_clip(vdf.energy, tint),
                             np.rad2deg(azimuthal), theta)
            newt, dist, energy, phi = psd_rebin(temp, phi, energy0, energy1,
                                                step_table)
            dist = ts_skymap(newt, dist, np.tile(energy, (len(newt), 1)), phi,
                             theta)
            dist = time_clip(dist.data, tint).data
            azimuthal = xr.DataArray(phi,
                                     coords=[newt, np.arange(phi.shape[1])],
                                     dims=["time", "odx"])
            len_e = dist.shape[1]
            energy_edges = np.hstack(
                [10 ** (np.log10(energy) - diff_energ / 2),
                 10 ** (np.log10(energy[-1]) + diff_energ / 2)])
        else:
            if all(step_table.data):
                energy_edges = energy1_edges
            else:
                energy_edges = energy0_edges
    else:
        raise ValueError("Invalid time interval")

    return dist, polar.data, azimuthal.data, energy_edges, len_e

def _cotrans(dist, polar, azimuthal, x_vec, y_vec, z_vec, e_lim,
             bin_corr):
    # Construct polar and azimuthal angle matrices
    polar = np.ones((len(dist), 1)) * polar

    f_mat = np.zeros((len(dist), dist.shape[2], dist.shape[1]))  #
    # azimuthal, energy
    edges_az = np.linspace(0, 2 * np.pi, azimuthal.shape[1] + 1)

    for i in range(len(dist)):
        pol_mat, azm_mat = np.meshgrid(polar[i, :], azimuthal[i, :])

        # '-' because the data shows which direction the particles were
        # coming from
        x_mat = -np.sin(pol_mat) * np.cos(azm_mat)
        y_mat = -np.sin(pol_mat) * np.sin(azm_mat)
        z_mat = -np.cos(pol_mat)

        # Transform into different coordinate system
        xx_mat = np.reshape(x_mat, (x_mat.shape[0] * x_mat.shape[1], 1))
        yy_mat = np.reshape(y_mat, (y_mat.shape[0] * y_mat.shape[1], 1))
        zz_mat = np.reshape(z_mat, (z_mat.shape[0] * z_mat.shape[1], 1))

        new_tmp_x = np.dot(np.hstack([xx_mat, yy_mat, zz_mat]), x_vec)
        new_tmp_y = np.dot(np.hstack([xx_mat, yy_mat, zz_mat]), y_vec)
        new_tmp_z = np.dot(np.hstack([xx_mat, yy_mat, zz_mat]), z_vec)

        new_x_mat = np.reshape(new_tmp_x, (x_mat.shape[0], x_mat.shape[1]))
        new_y_mat = np.reshape(new_tmp_y, (x_mat.shape[0], x_mat.shape[1]))
        new_z_mat = np.reshape(new_tmp_z, (x_mat.shape[0], x_mat.shape[1]))

        elevation_angle = np.arctan(
            new_z_mat / np.sqrt(new_x_mat ** 2 + new_y_mat ** 2))
        plane_az = np.arctan2(new_y_mat, new_x_mat) + np.pi

        # gets velocity in direction normal to 'z'-axis
        geo_factor_elev = np.cos(elevation_angle)

        # geoFactorBinSize - detector bins in 'equator' plane are bigger and
        # get a larger weight. I think this is not good for the
        # implementation in this function
        if bin_corr:
            geo_factor_bin_size = np.sin(pol_mat)
        else:
            geo_factor_bin_size = np.ones(pol_mat.shape)

        f_mat[i, ...] = _cotrans_jit(dist.data[i, ...], elevation_angle, e_lim,
                                     plane_az, edges_az, geo_factor_elev,
                                     geo_factor_bin_size)

    return f_mat

def _cotrans_jit(dist, elevation_angle, elevation_lim, plane_az, edges_az,
                 geo_factor_elev, geo_factor_bin_size):
    out = np.zeros((dist.shape[1], dist.shape[0]))  # azimuthal, energy

    for ie, iaz in itertools.product(range(dist.shape[0]),
                                     range(dist.shape[1])):
        # dist.data has dimensions nT x nE x nAz x nPol
        c_mat = dist[ie, ...].copy()
        c_mat = c_mat * geo_factor_elev * geo_factor_bin_size
        c_mat[np.abs(elevation_angle) > np.deg2rad(elevation_lim)] = np.nan
        # use 0.1 deg to fix Az angle edges bug
        c_mat[plane_az < edges_az[iaz] - np.deg2rad(.1)] = np.nan
        # use 0.1 deg to fix Az angle edges bug
        c_mat[plane_az > edges_az[iaz + 1] + np.deg2rad(.1)] = np.nan

        out[iaz, ie] = np.nanmean(c_mat)

    return out

def vdf_projection(vdf, tint, coord_sys: np.ndarray = np.eye(3),
                   sc_pot: xr.DataArray = None, e_lim: float = 20,
                   bins_correction: bool = False):
    r"""Computes projection of the velocity distribution onto a specified
    plane.

    Parameters
    ----------
    vdf : xarray.Dataset
        Electron or ion 3D skymap velocity distribution function.
    tint : list of str
        Computes data for time interval if len(tint) = 2 or closest time if
        len(tint) = 1. For tint includes two or more distributions the
        energies are rebinned into 64 channels.
    coord_sys : ndarray, Optional
        3x3 matrix with 1st column is x, 2nd column is y and 3rd column is z.
        z is normal to the projection plane and x and y are made orthogonal to
        z and each other if they are not already. Default is np.eye(3)
        (project onto spacecraft spin plane).
    sc_pot : xarray.DataArray, Optional
        Spacecraft potential to correct velocities. For a single value of tint
        the closest potential is used. For an interval the spacecraft
        potential is average over that interval. Default is None (no
        correction).
    e_lim : float, Optional
        Elevation angle limit in degrees above/below projection plane to
        include in projection. Default is e_lim = 20.
    bins_correction : bool, Optional
        Flag to correction elevation bins. Default is False.

    Returns
    -------
    v_x : ndarray
        2D grid of the velocities in the x direction.
    v_y : ndarray
        2D grid of the velocities in the y direction.
    f_mat : ndarray
        2D projection of the velocity distribution onto the specified plane

    """

    specie = vdf.attrs.get("species", "electrons")
    is_des = specie.lower() == "electrons"

    dist, polar, azimuthal, energy_edges, len_e = _init(vdf, tint)
    x_vec, y_vec, z_vec, changed_xyz = _coord_sys(coord_sys)

    if azimuthal.ndim == 1:
        azimuthal = np.ones((len(dist), 1)) * azimuthal

    f_mat = _cotrans(dist, polar, azimuthal, x_vec, y_vec, z_vec, e_lim,
                     bins_correction)
    if len(dist) == 1:
        f_mat = np.squeeze(f_mat)
    else:
        f_mat = np.squeeze(np.nanmean(f_mat, axis=0))

    if sc_pot is not None:
        if len(tint) == 1:
            time_datetime64 = iso86012datetime64(np.array(tint))[0]
            t_id = np.argmin(np.abs(sc_pot.time.data - time_datetime64))
            sc_pot = sc_pot.data[t_id]
        else:
            sc_pot = time_clip(sc_pot, tint)
            sc_pot = np.nanmean(sc_pot.data)
    else:
        sc_pot = 0.

    if is_des:
        mass = constants.electron_mass
    else:
        mass = constants.proton_mass
        sc_pot *= -1

    q_e = constants.elementary_charge

    speed_table = np.sqrt((energy_edges - sc_pot) * q_e * 2 / mass)
    speed_table = np.real(speed_table * 1e-3)  # km/s

    r_en = speed_table
    v_x = np.matmul(r_en[:, None],
                    np.cos(np.linspace(0, 2 * np.pi, azimuthal.shape[1] + 1)
                           + np.pi)[None, :])
    v_y = np.matmul(r_en[:, None],
                    np.sin(np.linspace(0, 2 * np.pi, azimuthal.shape[1] + 1)
                           + np.pi)[None, :])

    f_mat[f_mat <= 0] = np.nan

    return v_x, v_y, f_mat


#!/usr/bin/env python
# encoding: utf-8

################################################################################
#
#   RPMDrate - Bimolecular reaction rates via ring polymer molecular dynamics
#
#   Copyright (c) 2012 by Joshua W. Allen (jwallen@mit.edu)
#                         William H. Green (whgreen@mit.edu)
#                         Yury V. Suleimanov (ysuleyma@mit.edu, ysuleyma@princeton.edu)
#
#   Permission is hereby granted, free of charge, to any person obtaining a
#   copy of this software and associated documentation files (the "Software"),
#   to deal in the Software without restriction, including without limitation
#   the rights to use, copy, modify, merge, publish, distribute, sublicense,
#   and/or sell copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
#   THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#   DEALINGS IN THE SOFTWARE.
#
################################################################################

"""
This module contains representations of various thermostats available in RPMD.
"""

import os.path
import numpy

import rpmdrate.constants as constants
import rpmdrate.quantity as quantity

################################################################################

class AndersenThermostat:
    """
    A representation of an Andersen thermostat, a simple method of sampling
    the NVT ensemble which periodically replaces the momenta with a fresh
    sampling from a Gaussian distribution at the temperature of interest, as
    if resulting from a collision with a heat bath.
    """

    def __init__(self, samplingTime=None):
        if samplingTime is not None:
            self.samplingTime = float(quantity.convertTime(samplingTime, "ps") / 2.418884326505e-5)
        else:
            self.samplingTime = 0.0

    def activate(self, module, Natoms, Nbeads):
        """
        Set the thermostat as active in the Fortran layer of the given
        `module`.
        """
        module.thermostat = 1
        module.andersen_sampling_time = self.samplingTime

################################################################################

class GLEThermostat(object):
    """
    A representation of a colored-noise, generalized Langevin equation
    thermostat. The GLE thermostat offers significantly faster convergence
    when compared to the simpler Andersen thermostat, but requires more effort
    to set up. In particular, the thermostat requires two matrices be specified
    as input; these can be generated at

        http://gle4md.berlios.de/compose.php?page=matrix

    and supplied either as a file on disk or as numpy arrays. The A and C
    matrices are assumed to have units of s^-1 and K, respectively.
    """

    def __init__(self, A=None, C=None):
        self.A = A
        self.C = C

    @property
    def A(self):
        return self._A
    @A.setter
    def A(self, value):
        if isinstance(value, (list,tuple)) and len(value) == 2 and os.path.exists(value[0]):
            # value is the path of a file on disk to load from and the corresponding units
            path, units = value
            _A = []
            f = open(path, 'r')
            for line in f:
                # Remove comment
                if '#' in line: line = line[0:line.index('#')].strip()
                tokens = line.split()
                if len(tokens) > 0:
                    _A.append([float(t) for t in tokens])
            f.close()
            self._A = numpy.array(quantity.convertFrequency((_A,units),'s^-1'))
        elif isinstance(value, numpy.ndarray):
            self._A = value
        elif value is None:
            self._A = None
        else:
            raise ValueError('Unexpected value {0!r} for A attribute.'.format(value))

    @property
    def C(self):
        return self._C
    @C.setter
    def C(self, value):
        if isinstance(value, (list,tuple)) and len(value) == 2 and os.path.exists(value[0]):
            # value is the path of a file on disk to load from and the corresponding units
            path, units = value
            _C = []
            f = open(path, 'r')
            for line in f:
                # Remove comment
                if '#' in line: line = line[0:line.index('#')].strip()
                tokens = line.split()
                if len(tokens) > 0:
                    _C.append([float(t) for t in tokens])
            f.close()
            self._C = numpy.array(quantity.convertTemperature((_C,units),'K'))
        elif isinstance(value, numpy.ndarray):
            self._C = value
        elif value is None:
            self._C = None
        else:
            raise ValueError('Unexpected value {0!r} for C attribute.'.format(value))

    def activate(self, module, Natoms, Nbeads):
        """
        Set the thermostat as active in the Fortran layer of the given
        `module`.
        """
        Ns = self._A.shape[0] - 1
        module.thermostat = 2
        module.gle_ns = Ns
        module.gle_a[0:Ns+1,0:Ns+1] = self._A * 2.418884326505e-17  # s^-1 to atomic units of inverse time
        if self._C is None:
            module.gle_c[0:Ns+1,0:Ns+1] = numpy.zeros((Ns+1,Ns+1))
            for s in range(Ns+1):
                module.gle_c[s,s] = Nbeads / module.beta
        else:
            module.gle_c[0:Ns+1,0:Ns+1] = self._C * constants.kB / 4.35974417e-18  # K to atomic units of energy


def selection_3():

    # Library import
    import numpy
    import matplotlib
    import matplotlib.pyplot   as plt
    import matplotlib.gridspec as gridspec

    # Library version
    matplotlib_version = matplotlib.__version__
    numpy_version      = numpy.__version__

    # Histo binning
    xBinning = numpy.linspace(0.0,1000.0,101,endpoint=True)

    # Creating data sequence: middle of each bin
    xData = numpy.array([5.0,15.0,25.0,35.0,45.0,55.0,65.0,75.0,85.0,95.0,105.0,115.0,125.0,135.0,145.0,155.0,165.0,175.0,185.0,195.0,205.0,215.0,225.0,235.0,245.0,255.0,265.0,275.0,285.0,295.0,305.0,315.0,325.0,335.0,345.0,355.0,365.0,375.0,385.0,395.0,405.0,415.0,425.0,435.0,445.0,455.0,465.0,475.0,485.0,495.0,505.0,515.0,525.0,535.0,545.0,555.0,565.0,575.0,585.0,595.0,605.0,615.0,625.0,635.0,645.0,655.0,665.0,675.0,685.0,695.0,705.0,715.0,725.0,735.0,745.0,755.0,765.0,775.0,785.0,795.0,805.0,815.0,825.0,835.0,845.0,855.0,865.0,875.0,885.0,895.0,905.0,915.0,925.0,935.0,945.0,955.0,965.0,975.0,985.0,995.0])

    # Creating weights for histo: y4_PT_0
    y4_PT_0_weights = numpy.array([0.0,0.0,77.1570947131,70.4141404187,65.4685046035,60.1625490931,55.8433127479,50.5660372133,46.0298010516,41.5918048069,37.0432726556,32.3227926499,28.6667757435,24.3311394121,21.9893253936,19.1603117874,16.4500300807,15.2954990576,13.2934927516,11.4020223521,10.0878234641,8.15132110272,7.54539761542,6.67745034984,5.77675111197,5.035723739,4.43389224824,3.9712614397,3.51272422769,3.22204407365,2.68162493093,2.31725163925,2.00200710599,1.8546200307,1.64991580391,1.43702358405,1.23231935727,1.03989752009,0.998956354728,0.77378174526,0.777876141796,0.626394669973,0.544513139257,0.528136753114,0.417596446649,0.368467568219,0.266115454825,0.233362802539,0.262021378289,0.180139687574,0.126916612609,0.192421957181,0.151481111823,0.10644618993,0.0614112680365,0.0409408453577,0.0573171915008,0.0655053445723,0.0573171915008,0.0655053445723,0.0573171915008,0.0286585877504,0.0409408453577,0.0368467568219,0.0204704186788,0.0163763341431,0.0163763341431,0.0,0.0,0.0122822536073,0.00818816907154,0.00818816907154,0.0122822536073,0.0,0.00409408453577,0.0,0.00409408453577,0.00409408453577,0.00409408453577,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_1
    y4_PT_1_weights = numpy.array([0.0,0.0,0.012170493784,0.0242554668822,0.0121753353338,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_2
    y4_PT_2_weights = numpy.array([0.0,0.0,0.120516882096,0.190667418918,0.170688939005,0.230969775691,0.220865995678,0.150608273824,0.0702865634698,0.0100459438961,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_3
    y4_PT_3_weights = numpy.array([0.0,0.0,0.198079494608,0.313504754258,0.445485992003,0.467609163722,0.555495714045,0.550065266778,0.66560578175,0.500527416221,0.484060164817,0.434396780617,0.4346267225,0.291517882011,0.280484693568,0.176024330435,0.132027673006,0.0659114595513,0.0274955403792,0.0219792671012,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_4
    y4_PT_4_weights = numpy.array([0.0,0.0,0.107569860232,0.152926540276,0.210197820351,0.240799506276,0.261535469359,0.272370115988,0.26647304052,0.260512233252,0.234839540748,0.220064946095,0.240789926465,0.217085664782,0.206244003646,0.217111558331,0.199345296751,0.209240079692,0.167774646956,0.169731052908,0.152957845016,0.107567695756,0.0967098411313,0.0769771127922,0.0513190906281,0.0404724592134,0.0256560981852,0.0197415465757,0.00691344537533,0.0039493642969,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_5
    y4_PT_5_weights = numpy.array([0.0,0.0,0.026469421877,0.0368111391992,0.0463847330278,0.0582429397943,0.0610034742498,0.063025161425,0.0627723705057,0.0703289704757,0.0632682696279,0.0693279296382,0.0647784533032,0.0657827750611,0.0617574057619,0.0683043624811,0.0620074359067,0.0607485227243,0.057973624239,0.0564609998791,0.0546985854417,0.0549544172141,0.0451169373535,0.0486532894611,0.0456235994952,0.0441150962915,0.0405815849808,0.0403250730175,0.037557744656,0.0373196338551,0.0347905563338,0.0274836224064,0.0184009863138,0.0173888783715,0.0121017470906,0.00957741890171,0.00781076327444,0.00352860713483,0.00302506466826,0.00125954976101,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_6
    y4_PT_6_weights = numpy.array([0.0,0.0,0.0117456805932,0.0160510745363,0.0160202634536,0.0145996546028,0.017471013581,0.0166051341817,0.0226137751752,0.0246249031193,0.0197479046086,0.0211858684365,0.0191812086204,0.0177440945465,0.01545980566,0.0200310426635,0.0171768487175,0.0166043644045,0.0177452941993,0.0146041333066,0.0148982281904,0.0146048331041,0.013461813914,0.0103127253157,0.0143026005755,0.0140366175557,0.00887486445792,0.0117502692651,0.0157392547825,0.0105858062812,0.0148863416305,0.0134570153028,0.0105856663217,0.0117342838916,0.0131758866663,0.0128871402346,0.010307316881,0.00800864615675,0.0120300482921,0.00773851133852,0.00744036662685,0.00944570224741,0.0065943574768,0.00543734833563,0.00544087731428,0.00458387734516,0.00343941757198,0.00286485386084,0.00372256662366,0.00229384012227,0.00315330237875,0.00114373298353,0.000570390518938,0.000570257457448,0.000855298861446,0.00142800071144,0.0,0.000855811113191,0.000283957617664,0.000287252564049,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_7
    y4_PT_7_weights = numpy.array([0.0,0.0,0.00125284526549,0.00168499598199,0.00185749861251,0.00183575971121,0.00146715015491,0.00174627318844,0.00142449880858,0.00129467728017,0.00131802299351,0.00125308205443,0.00125139351877,0.00136074097602,0.000950247018111,0.00109966712725,0.00107859542594,0.0011877844651,0.000993808220713,0.000734331544785,0.000993849292069,0.000907425938116,0.000864211746567,0.000690913672835,0.000539714503417,0.000756157617944,0.000518467620202,0.000777672303167,0.000583256427627,0.000561812150443,0.000389169458214,0.000538299218003,0.00068936804873,0.000821239373105,0.00058142204674,0.000669569140375,0.000539869568742,0.000777473651913,0.000799001748191,0.00086378552647,0.000475164161317,0.000475416875887,0.000712993717852,0.000734323581971,0.000496547669656,0.000561526327331,0.000820794712808,0.000583353657777,0.000777707926283,0.000842200852303,0.000734339088503,0.000821016833409,0.000583154587427,0.000647766117466,0.000842340411096,0.000777750674021,0.000604580005314,0.000583402691947,0.000604679750036,0.000647765698371,0.000496820919905,0.000604629877675,0.000388622915807,0.000453413106247,0.00038890961902,0.000518156232265,0.000151311277457,0.000216013835034,0.000151211826102,0.000172851066499,0.000129724256356,4.32307039111e-05,2.15259505091e-05,4.32537122525e-05,2.1598366016e-05,8.64288860134e-05,0.0,0.0,2.1625766478e-05,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_8
    y4_PT_8_weights = numpy.array([0.0,0.0,0.000254018618774,0.000340757810034,0.000227098125768,0.000170284616993,0.000225875351156,0.000170440229354,0.00014128029926,0.000255240799445,0.000198750691169,0.000142064062893,8.51446694092e-05,0.000142025798288,0.000198913282328,5.66653206676e-05,8.46910920435e-05,5.68266794011e-05,5.68161963546e-05,0.000140559537847,8.50190955793e-05,5.68978186016e-05,8.47994119074e-05,8.43635042659e-05,8.51739803635e-05,2.84292647251e-05,2.83498697542e-05,5.6797754508e-05,2.84292647251e-05,0.000112168537726,0.000113373508977,2.84080907531e-05,5.67658896104e-05,2.84292647251e-05,0.0,2.84489093008e-05,5.68570000539e-05,0.0,0.0,5.68781740259e-05,0.0,2.84292647251e-05,2.83973998275e-05,2.84292647251e-05,0.0,2.83498697542e-05,5.66865985791e-05,8.52544147299e-05,0.0,5.6611405736e-05,2.84292647251e-05,5.68055054291e-05,5.68266794011e-05,0.0,5.68173990838e-05,2.83684897829e-05,0.000113596533563,0.0,5.68266794011e-05,5.67658896104e-05,0.0,2.84489093008e-05,0.0,0.000113114179789,2.84489093008e-05,8.5052133509e-05,8.52240792286e-05,2.83684897829e-05,0.0,0.0,5.67183446887e-05,2.83498697542e-05,0.000141993933391,0.0,0.0,5.6797754508e-05,2.83498697542e-05,0.0,5.64019972314e-05,0.0,0.000113494539163,0.0,2.83498697542e-05,0.0,5.67658896104e-05,2.83973998275e-05,0.0,2.83684897829e-05,0.0,2.84292647251e-05,2.84292647251e-05,2.84489093008e-05,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_9
    y4_PT_9_weights = numpy.array([0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_10
    y4_PT_10_weights = numpy.array([0.0,0.0,1.05462838872,0.0,1.0521138287,0.0,1.0529581672,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_11
    y4_PT_11_weights = numpy.array([0.0,0.0,1.15201585626,4.14563454124,2.07136934156,2.30323550525,2.53414099414,1.38301356993,0.689926667059,0.690363197688,0.921327864146,0.460933689557,1.38374637267,0.230645222875,0.69068560015,0.230673159298,0.230128533541,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_12
    y4_PT_12_weights = numpy.array([0.0,0.0,2.35389244344,1.27419718775,1.57835356252,1.21847771158,0.858756115524,0.526086439007,0.443056541367,0.581372385888,0.387378994831,0.387488627348,0.470541990192,0.276966356723,0.249345271318,0.249169359214,0.221480994065,0.166164465482,0.249253910888,0.193844252367,0.138586309862,0.0829313822454,0.0553826158335,0.221369784379,0.027658683588,0.0553746530507,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_13
    y4_PT_13_weights = numpy.array([0.0,0.0,0.746043904483,0.574843087491,0.353046902195,0.201630148184,0.1511673845,0.110770130708,0.131139960181,0.131117201951,0.0705911738551,0.0906791654103,0.121036095087,0.0503486326784,0.0705673839189,0.050436461239,0.0704838156989,0.0806136555195,0.0604845717045,0.030176514568,0.0403463421156,0.0403272737534,0.0402882449065,0.0604881341261,0.0907026519035,0.0302662912331,0.0302700660648,0.0100533236174,0.0302963806474,0.0100733690662,0.0302270014251,0.0302662912331,0.0,0.0,0.010103349241,0.0302004744324,0.0100953565507,0.0,0.0,0.0100853125852,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_14
    y4_PT_14_weights = numpy.array([0.0,0.0,0.316899125281,0.195206783525,0.118805873014,0.0820486906046,0.0509245525742,0.0594061683611,0.0424593269034,0.0311218026312,0.0339577659611,0.0339541878371,0.022619179794,0.0169751894332,0.0226083838631,0.0169668443248,0.0254663584264,0.0197985254483,0.019781338911,0.0169776517981,0.0283019100796,0.0169629545578,0.0141450472878,0.0113239851129,0.0113149051424,0.0141430658536,0.0198120107433,0.00564490605272,0.0113192642977,0.0226471237878,0.0169683294386,0.0169794370126,0.0113288367412,0.00282347952472,0.00565764494347,0.0141606602201,0.0141401071683,0.0169860046014,0.0169683217437,0.0141465593338,0.0113152975818,0.00849572052068,0.00849465477839,0.0084750289609,0.0,0.0,0.0,0.0,0.00283250409196,0.00283041377503,0.0113201607524,0.0,0.0,0.0,0.00282747702018,0.0,0.0,0.0,0.00282544018274,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_15
    y4_PT_15_weights = numpy.array([0.0,0.0,0.0197527910952,0.0137030901086,0.0121841674005,0.003067188295,0.00761977251509,0.00455442442883,0.00150849604052,0.0,0.00458426252112,0.00152094965608,0.00153153115177,0.0,0.0,0.0,0.00150849604052,0.0,0.0,0.0,0.0,0.0,0.0,0.00457388076549,0.0,0.0,0.0,0.0,0.0,0.00154541013132,0.0,0.00153219773991,0.0,0.0,0.00152449651954,0.0,0.00458437007346,0.0,0.0,0.00152305579093,0.0,0.00152192826419,0.0,0.0,0.0,0.0,0.0,0.0,0.00305829572562,0.00152162924505,0.0,0.0,0.0,0.0,0.00151727396619,0.00304183194421,0.0,0.0,0.00154541013132,0.00150849604052,0.0015126539431,0.0,0.00153629536591,0.00152094965608,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating weights for histo: y4_PT_16
    y4_PT_16_weights = numpy.array([0.0,0.0,0.0021680746957,0.00162497960558,0.00126313090509,0.000901888554933,0.000360688396201,0.00018116843031,0.00036127033852,0.000722436192292,0.0,0.000180533822134,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.000180203005096,0.000180626141377,0.0,0.0,0.0,0.0,0.000180970240375,0.0,0.0,0.0,0.0,0.0,0.000180657132616,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.000180734321975,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.000180626141377,0.0,0.0,0.0,0.0,0.0,0.0,0.000180626141377,0.0,0.0,0.0,0.0,0.000180626141377,0.0,0.0,0.000180553032852,0.000180503369873,0.000180003621709,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0])

    # Creating a new Canvas
    fig   = plt.figure(figsize=(12,6),dpi=80)
    frame = gridspec.GridSpec(1,1,right=0.7)
    pad   = fig.add_subplot(frame[0])

    # Creating a new Stack
    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights+y4_PT_8_weights+y4_PT_9_weights+y4_PT_10_weights+y4_PT_11_weights+y4_PT_12_weights+y4_PT_13_weights+y4_PT_14_weights+y4_PT_15_weights+y4_PT_16_weights,\
             label="$bg\_dip\_1600\_inf$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#e5e5e5", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights+y4_PT_8_weights+y4_PT_9_weights+y4_PT_10_weights+y4_PT_11_weights+y4_PT_12_weights+y4_PT_13_weights+y4_PT_14_weights+y4_PT_15_weights,\
             label="$bg\_dip\_1200\_1600$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#f2f2f2", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights+y4_PT_8_weights+y4_PT_9_weights+y4_PT_10_weights+y4_PT_11_weights+y4_PT_12_weights+y4_PT_13_weights+y4_PT_14_weights,\
             label="$bg\_dip\_800\_1200$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#ccc6aa", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights+y4_PT_8_weights+y4_PT_9_weights+y4_PT_10_weights+y4_PT_11_weights+y4_PT_12_weights+y4_PT_13_weights,\
             label="$bg\_dip\_600\_800$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#ccc6aa", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights+y4_PT_8_weights+y4_PT_9_weights+y4_PT_10_weights+y4_PT_11_weights+y4_PT_12_weights,\
             label="$bg\_dip\_400\_600$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#c1bfa8", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights+y4_PT_8_weights+y4_PT_9_weights+y4_PT_10_weights+y4_PT_11_weights,\
             label="$bg\_dip\_200\_400$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#bab5a3", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights+y4_PT_8_weights+y4_PT_9_weights+y4_PT_10_weights,\
             label="$bg\_dip\_100\_200$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#b2a596", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights+y4_PT_8_weights+y4_PT_9_weights,\
             label="$bg\_dip\_0\_100$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#b7a39b", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights+y4_PT_8_weights,\
             label="$bg\_vbf\_1600\_inf$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#ad998c", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights,\
             label="$bg\_vbf\_1200\_1600$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#9b8e82", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights,\
             label="$bg\_vbf\_800\_1200$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#876656", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights,\
             label="$bg\_vbf\_600\_800$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#afcec6", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights,\
             label="$bg\_vbf\_400\_600$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#84c1a3", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights,\
             label="$bg\_vbf\_200\_400$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#89a8a0", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights,\
             label="$bg\_vbf\_100\_200$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#829e8c", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights+y4_PT_1_weights,\
             label="$bg\_vbf\_0\_100$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#adbcc6", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    pad.hist(x=xData, bins=xBinning, weights=y4_PT_0_weights,\
             label="$signal$", histtype="step", rwidth=1.0,\
             color=None, edgecolor="#7a8e99", linewidth=1, linestyle="solid",\
             bottom=None, cumulative=False, normed=False, align="mid", orientation="vertical")

    # Axis
    plt.rc('text',usetex=False)
    plt.xlabel(r"p_{T} [ j_{2} ]   ( GeV ) ",\
               fontsize=16,color="black")
    plt.ylabel(r"$\mathrm{Events}$ $(\mathcal{L}_{\mathrm{int}} = 40.0\ \mathrm{fb}^{-1})$ ",\
               fontsize=16,color="black")

    # Boundary of y-axis
    ymax=(y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights+y4_PT_8_weights+y4_PT_9_weights+y4_PT_10_weights+y4_PT_11_weights+y4_PT_12_weights+y4_PT_13_weights+y4_PT_14_weights+y4_PT_15_weights+y4_PT_16_weights).max()*1.1
    ymin=0 # linear scale
    #ymin=min([x for x in (y4_PT_0_weights+y4_PT_1_weights+y4_PT_2_weights+y4_PT_3_weights+y4_PT_4_weights+y4_PT_5_weights+y4_PT_6_weights+y4_PT_7_weights+y4_PT_8_weights+y4_PT_9_weights+y4_PT_10_weights+y4_PT_11_weights+y4_PT_12_weights+y4_PT_13_weights+y4_PT_14_weights+y4_PT_15_weights+y4_PT_16_weights) if x])/100. # log scale
    plt.gca().set_ylim(ymin,ymax)

    # Log/Linear scale for X-axis
    plt.gca().set_xscale("linear")
    #plt.gca().set_xscale("log",nonposx="clip")

    # Log/Linear scale for Y-axis
    plt.gca().set_yscale("linear")
    #plt.gca().set_yscale("log",nonposy="clip")

    # Legend
    plt.legend(bbox_to_anchor=(1.05,1), loc=2, borderaxespad=0.)

    # Saving the image
    plt.savefig('../../HTML/MadAnalysis5job_0/selection_3.png')
    plt.savefig('../../PDF/MadAnalysis5job_0/selection_3.png')
    plt.savefig('../../DVI/MadAnalysis5job_0/selection_3.eps')

# Running!
if __name__ == '__main__':
    selection_3()


#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Analytical nuclear gradients for domain decomposition COSMO

See also

[1] Fast Domain Decomposition Algorithm for Continuum Solvation Models: Energy and First Derivatives.
F. Lipparini, B. Stamm, E. Cances, Y. Maday, B. Mennucci
J. Chem. Theory Comput., 9, 3637-3648 (2013)
http://dx.doi.org/10.1021/ct400280b

[2] Quantum, classical, and hybrid QM/MM calculations in solution: General implementation of the ddCOSMO linear scaling strategy.
F. Lipparini, G. Scalmani, L. Lagardere, B. Stamm, E. Cances, Y. Maday, J.-P.Piquemal, M. J. Frisch, B. Mennucci
J. Chem. Phys., 141, 184108 (2014)
http://dx.doi.org/10.1063/1.4901304
'''

import ctypes
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf import gto
from pyscf import df
from pyscf.dft import gen_grid, numint
from pyscf.symm import sph
from pyscf.solvent import ddcosmo
from pyscf.grad import rhf as rhf_grad
from pyscf.grad import rks as rks_grad

# TODO: Define attribute grad_method.base to point to the class of the 0th
# order calculation for all gradients class. Then this function can be
# extended and used as the general interface to initialize solvent gradients.
def ddcosmo_grad(grad_method, pcmobj=None):
    grad_method_class = grad_method.__class__
    class WithSolventGrad(grad_method.__class__):
        def __init__(self, pcmobj):
            self.__dict__.update(grad_method.__dict__)
            self.with_solvent = pcmobj
            self.de_solvent = None
            self.de_solute = None
            self._keys = self._keys.union(['with_solvent', 'de_solvent', 'de_solute'])

        def kernel(self, dm=None, atmlst=None):
            if dm is None:
                dm = grad_method.base.make_rdm1(ao_repr=True)

            # de_solvent needs to be called first because _finalize method
            # is called in the grad_method.kernel function.  de_solvent is
            # required by the _finalize method.
            self.de_solvent = kernel(self.with_solvent, dm)
            self.de_solute = grad_method_class.kernel(self, atmlst=atmlst)
            self.de = self.de_solute + self.de_solvent

            if self.verbose >= logger.NOTE:
                logger.note(self, '--------------- %s (%s) gradients ---------------',
                            grad_method.base.__class__.__name__,
                            self.with_solvent.__class__.__name__)
                rhf_grad._write(self, self.mol, self.de, self.atmlst)
                logger.note(self, '----------------------------------------------')
            return self.de

        def _finalize(self):
            # disable _finalize. It is called in grad_method.kernel method
            # where self.de was not yet initialized.
            pass

    if pcmobj is None:
        pcmobj = ddcosmo.DDCOSMO(mf.mol)
    return WithSolventGrad(pcmobj)

# Inject DDCOSMO gradients into other modules
try:
    from pyscf import grad
    for mod in dir(grad):
        if 1 or hasattr(mod, 'Gradients'):
            mod.Gradients.DDCOSMO = ddcosmo_grad
    del(mod, grad)
except Exception as e:
    print('Error for registering ddcosmo gradients: ' + str(e))

def kernel(pcmobj, dm, verbose=None):
    mol = pcmobj.mol
    natm = mol.natm
    lmax = pcmobj.lmax
    if pcmobj.grids.coords is None:
        pcmobj.grids.build(with_non0tab=True)

    if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
        # UHF density matrix
        dm = dm[0] + dm[1]

    r_vdw = ddcosmo.get_atomic_radii(pcmobj)
    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, lmax, True))

    fi = ddcosmo.make_fi(pcmobj, r_vdw)
    ui = 1 - fi
    ui[ui<0] = 0

    cached_pol = ddcosmo.cache_fake_multipoles(pcmobj.grids, r_vdw, lmax)

    nlm = (lmax+1)**2
    L0 = ddcosmo.make_L(pcmobj, r_vdw, ylm_1sph, fi)
    L0 = L0.reshape(natm*nlm,-1)
    L1 = make_L1(pcmobj, r_vdw, ylm_1sph, fi)

    phi0 = ddcosmo.make_phi(pcmobj, dm, r_vdw, ui)
    phi1 = make_phi1(pcmobj, dm, r_vdw, ui)
    L0_X = numpy.linalg.solve(L0, phi0.ravel()).reshape(natm,-1)
    psi0, vmat, L0_S = \
            ddcosmo.make_psi_vmat(pcmobj, dm, r_vdw, ui, pcmobj.grids, ylm_1sph,
                                  cached_pol, L0_X, L0)
    e_psi1 = make_e_psi1(pcmobj, dm, r_vdw, ui, pcmobj.grids, ylm_1sph,
                         cached_pol, L0_X, L0)
    dielectric = pcmobj.eps
    if dielectric > 0:
        f_epsilon = (dielectric-1.)/dielectric
    else:
        f_epsilon = 1
    de = .5 * f_epsilon * e_psi1
    de+= .5 * f_epsilon * numpy.einsum('jx,azjx->az', L0_S, phi1)
    de-= .5 * f_epsilon * numpy.einsum('aziljm,il,jm->az', L1, L0_S, L0_X)
    return de

def make_L1(pcmobj, r_vdw, ylm_1sph, fi):
    # See JCTC, 9, 3637, Eq (18)
    mol = pcmobj.mol
    natm = mol.natm
    lmax = pcmobj.lmax
    eta = pcmobj.eta
    nlm = (lmax+1)**2

    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    ngrid_1sph = weights_1sph.size
    atom_coords = mol.atom_coords()
    ylm_1sph = ylm_1sph.reshape(nlm,ngrid_1sph)

    Lmat = numpy.zeros((natm,3,natm,nlm,natm,nlm))
    fi1 = make_fi1(pcmobj, pcmobj.get_atomic_radii())

    for ja in range(natm):
        part_weights = weights_1sph.copy()
        part_weights[fi[ja]>1] /= fi[ja,fi[ja]>1]

        part_weights1 = numpy.zeros((natm,3,ngrid_1sph))
        tmp = part_weights[fi[ja]>1] / fi[ja,fi[ja]>1]
        part_weights1[:,:,fi[ja]>1] = -tmp * fi1[:,:,ja,fi[ja]>1]

        for ka in ddcosmo.atoms_with_vdw_overlap(ja, atom_coords, r_vdw):
            vjk = r_vdw[ja] * coords_1sph + atom_coords[ja] - atom_coords[ka]
            rv = lib.norm(vjk, axis=1)
            tjk = rv / r_vdw[ka]
            wjk0 = pcmobj.regularize_xt(tjk, eta, r_vdw[ka])
            wjk1 = regularize_xt1(tjk, eta*r_vdw[ka])
            sjk = vjk.T / rv
            wjk1 = 1./r_vdw[ka] * wjk1 * sjk

            wjk01 = wjk0 * part_weights1
            wjk0 *= part_weights
            wjk1 *= part_weights

            pol0 = sph.multipoles(vjk, lmax)
            pol1 = multipoles1(vjk, lmax)
            p1 = 0
            for l in range(lmax+1):
                fac = 4*numpy.pi/(l*2+1) / r_vdw[ka]**(l+1)
                p0, p1 = p1, p1 + (l*2+1)
                a = numpy.einsum('xn,zn,mn->zxm', ylm_1sph, wjk1, pol0[l])
                a+= numpy.einsum('xn,n,zmn->zxm', ylm_1sph, wjk0, pol1[l])
                Lmat[ja,:,ja,:,ka,p0:p1] += -fac * a
                Lmat[ka,:,ja,:,ka,p0:p1] -= -fac * a
                a = numpy.einsum('xn,azn,mn->azxm', ylm_1sph, wjk01, pol0[l])
                Lmat[:,:,ja,:,ka,p0:p1] += -fac * a
    return Lmat

def multipoles1(r, lmax, reorder_dipole=True):
    ngrid = r.shape[0]
    xs = numpy.ones((lmax+1,ngrid))
    ys = numpy.ones((lmax+1,ngrid))
    zs = numpy.ones((lmax+1,ngrid))
    for i in range(1,lmax+1):
        xs[i] = xs[i-1] * r[:,0]
        ys[i] = ys[i-1] * r[:,1]
        zs[i] = zs[i-1] * r[:,2]
    ylms = []
    for l in range(lmax+1):
        nd = (l+1)*(l+2)//2
        c = numpy.empty((nd,3,ngrid))
        k = 0
        for lx in reversed(range(0, l+1)):
            for ly in reversed(range(0, l-lx+1)):
                lz = l - lx - ly
                c[k,0] = lx * xs[lx-1] * ys[ly] * zs[lz]
                c[k,1] = ly * xs[lx] * ys[ly-1] * zs[lz]
                c[k,2] = lz * xs[lx] * ys[ly] * zs[lz-1]
                k += 1
        ylm = gto.cart2sph(l, c.reshape(nd,3*ngrid).T)
        ylm = ylm.reshape(3,ngrid,l*2+1).transpose(0,2,1)
        ylms.append(ylm)

# when call libcint, p functions are ordered as px,py,pz
# reorder px,py,pz to p(-1),p(0),p(1)
    if (not reorder_dipole) and lmax >= 1:
        ylms[1] = ylms[1][:,[1,2,0]]
    return ylms

def regularize_xt1(t, eta):
    xt = numpy.zeros_like(t)
    # no response if grids are inside the cavity
    # inner = t <= 1-eta
    # xt[inner] = 0
    on_shell = (1-eta < t) & (t < 1)
    ti = t[on_shell]
    xt[on_shell] = -30./eta**5 * (1-ti)**2 * (1-eta-ti)**2
    return xt

def make_fi1(pcmobj, r_vdw):
    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    mol = pcmobj.mol
    eta = pcmobj.eta
    natm = mol.natm
    atom_coords = mol.atom_coords()
    ngrid_1sph = coords_1sph.shape[0]
    fi1 = numpy.zeros((natm,3,natm,ngrid_1sph))
    for ia in range(natm):
        for ja in ddcosmo.atoms_with_vdw_overlap(ia, atom_coords, r_vdw):
            v = r_vdw[ia]*coords_1sph + atom_coords[ia] - atom_coords[ja]
            rv = lib.norm(v, axis=1)
            t = rv / r_vdw[ja]
            xt1 = regularize_xt1(t, eta*r_vdw[ja])
            s_ij = v.T / rv
            xt1 = 1./r_vdw[ja] * xt1 * s_ij
            fi1[ia,:,ia] += xt1
            fi1[ja,:,ia] -= xt1

    fi = ddcosmo.make_fi(pcmobj, r_vdw)
    fi1[:,:,fi<1e-20] = 0
    return fi1

def make_phi1(pcmobj, dm, r_vdw, ui):
    mol = pcmobj.mol
    natm = mol.natm
    nlm = (pcmobj.lmax+1)**2

    if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
        dm = dm[0] + dm[1]
    tril_dm = lib.pack_tril(dm+dm.T)
    nao = dm.shape[0]
    diagidx = numpy.arange(nao)
    diagidx = diagidx*(diagidx+1)//2 + diagidx
    tril_dm[diagidx] *= .5

    atom_coords = mol.atom_coords()
    atom_charges = mol.atom_charges()

    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, pcmobj.lmax, True))
    extern_point_idx = ui > 0

    fi1 = make_fi1(pcmobj, pcmobj.get_atomic_radii())
    fi1[:,:,ui==0] = 0
    ui1 = -fi1

    ngrid_1sph = weights_1sph.size
    v_phi0 = numpy.empty((natm,ngrid_1sph))
    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        d_rs = atom_coords.reshape(-1,1,3) - cav_coords
        v_phi0[ia] = numpy.einsum('z,zp->p', atom_charges, 1./lib.norm(d_rs,axis=2))
    phi1 = -numpy.einsum('n,ln,azjn,jn->azjl', weights_1sph, ylm_1sph, ui1, v_phi0)

    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        for ja in range(natm):
            rs = atom_coords[ja] - cav_coords
            d_rs = lib.norm(rs, axis=1)
            v_phi = atom_charges[ja] * numpy.einsum('px,p->px', rs, 1./d_rs**3)
            tmp = numpy.einsum('n,ln,n,nx->xl', weights_1sph, ylm_1sph, ui[ia], v_phi)
            phi1[ja,:,ia] += tmp  # response of the other atoms
            phi1[ia,:,ia] -= tmp  # response of cavity grids

    int3c2e = mol._add_suffix('int3c2e')
    int3c2e_ip1 = mol._add_suffix('int3c2e_ip1')
    aoslices = mol.aoslice_by_atom()
    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        #fakemol = gto.fakemol_for_charges(cav_coords[ui[ia]>0])
        fakemol = gto.fakemol_for_charges(cav_coords)
        v_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e, aosym='s1')
        v_phi = numpy.einsum('ij,ijk->k', dm, v_nj)
        phi1[:,:,ia] += numpy.einsum('n,ln,azn,n->azl', weights_1sph, ylm_1sph, ui1[:,:,ia], v_phi)

        v_e1_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e_ip1, comp=3, aosym='s1')
        v_e2_nj = v_e1_nj + v_e1_nj.transpose(0,2,1,3)
        phi1_e2_nj = numpy.einsum('ji,xijr->xr', dm, v_e2_nj)
        phi1[ia,:,ia] += numpy.einsum('n,ln,n,xn->xl', weights_1sph, ylm_1sph, ui[ia], phi1_e2_nj)

        for ja in range(natm):
            shl0, shl1, p0, p1 = aoslices[ja]
            phi1_nj  = numpy.einsum('ij,xijr->xr', dm[p0:p1  ], v_e1_nj[:,p0:p1])
            phi1_nj += numpy.einsum('ji,xijr->xr', dm[:,p0:p1], v_e1_nj[:,p0:p1])
            phi1[ja,:,ia] -= numpy.einsum('n,ln,n,xn->xl', weights_1sph, ylm_1sph, ui[ia], phi1_nj)
    return phi1

def make_e_psi1(pcmobj, dm, r_vdw, ui, grids, ylm_1sph, cached_pol, L_X, L):
    mol = pcmobj.mol
    natm = mol.natm
    lmax = pcmobj.lmax
    nlm = (lmax+1)**2

    if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
        dm = dm[0] + dm[1]
    ni = numint.NumInt()
    max_memory = pcmobj.max_memory - lib.current_memory()[0]
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm)
    den = numpy.empty((4,grids.weights.size))

    ao_loc = mol.ao_loc_nr()
    vmat = numpy.zeros((3,nao,nao))
    psi1 = numpy.zeros((natm,3))
    i1 = 0
    for ia, (coords, weight, weight1) in enumerate(rks_grad.grids_response_cc(grids)):
        i0, i1 = i1, i1 + weight.size
        ao = ni.eval_ao(mol, coords, deriv=1)
        mask = gen_grid.make_mask(mol, coords)
        den[:,i0:i1] = make_rho(0, ao, mask, 'GGA')

        fak_pol, leak_idx = cached_pol[mol.atom_symbol(ia)]
        eta_nj = 0
        p1 = 0
        for l in range(lmax+1):
            fac = 4*numpy.pi/(l*2+1)
            p0, p1 = p1, p1 + (l*2+1)
            eta_nj += fac * numpy.einsum('mn,m->n', fak_pol[l], L_X[ia,p0:p1])
        psi1 -= numpy.einsum('n,n,zxn->zx', den[0,i0:i1], eta_nj, weight1)
        psi1[ia] -= numpy.einsum('xn,n,n->x', den[1:4,i0:i1], eta_nj, weight)

        vtmp = numpy.zeros((3,nao,nao))
        aow = numpy.einsum('pi,p->pi', ao[0], weight*eta_nj)
        rks_grad._d1_dot_(vtmp, mol, ao[1:4], aow, mask, ao_loc, True)
        vmat += vtmp

    aoslices = mol.aoslice_by_atom()
    for ia in range(natm):
        shl0, shl1, p0, p1 = aoslices[ia]
        psi1[ia] += numpy.einsum('xij,ij->x', vmat[:,p0:p1], dm[p0:p1]) * 2
    return psi1

if __name__ == '__main__':
    from pyscf import scf
    mol = gto.M(atom='H 0 0 0; H 0 1 1.2; H 1. .1 0; H .5 .5 1', unit='B')
    mf = ddcosmo.ddcosmo_for_scf(scf.RHF(mol))
    mf.kernel()
    de = mf.nuc_grad_method().kernel()
    de_cosmo = kernel(mf.with_solvent, mf.make_rdm1())
    dm1 = mf.make_rdm1()

    mol = gto.M(atom='H 0 0 -0.001; H 0 1 1.2; H 1. .1 0; H .5 .5 1', unit='B')
    mf = ddcosmo.ddcosmo_for_scf(scf.RHF(mol))
    e1 = mf.kernel()
    e1_cosmo = mf.with_solvent.energy(dm1)

    mol = gto.M(atom='H 0 0 0.001; H 0 1 1.2; H 1. .1 0; H .5 .5 1', unit='B')
    mf = ddcosmo.ddcosmo_for_scf(scf.RHF(mol))
    e2 = mf.kernel()
    e2_cosmo = mf.with_solvent.energy(dm1)
    print(abs((e2-e1)/0.002 - de[0,2]).max())
    print(abs((e2_cosmo-e1_cosmo)/0.002 - de_cosmo[0,2]).max())