<commit_before>#!/usr/bin/python3
# -*- coding: utf-8 -*-
import getopt
import numpy as np
import os
import sys
from os.path import isfile, join, dirname
from scipy import interpolate
from scipy.optimize import fmin

import matplotlib.pyplot as plt
from matplotlib import gridspec

from pystella import velocity
from pystella.model.stella import Stella
from pystella.rf import band
from pystella.rf import light_curve_func as lcf
from pystella.util.phys_var import phys

__author__ = 'bakl'

ROOT_DIRECTORY = dirname(dirname(os.path.abspath(__file__)))

_colors = ["blue", "cyan", "brown", 'darkseagreen', 'tomato', 'olive', 'orange',
           'skyblue', 'darkviolet']
# colors = {"B-V": "blue", 'B-V-I': "cyan", 'V-I': "brown"}
lntypes = {"B-V": "-", 'B-V-I': "-.", 'V-I': "--",
           "u-g": "-", 'g-r-i': "-.", 'r-i': "--"}
markers = {u'D': u'diamond', 6: u'caretup', u's': u'square', u'x': u'x',
           5: u'caretright', u'^': u'triangle_up', u'd': u'thin_diamond', u'h': u'hexagon1',
           u'+': u'plus', u'*': u'star', u'o': u'circle', u'p': u'pentagon', u'3': u'tri_left',
           u'H': u'hexagon2', u'v': u'triangle_down', u'8': u'octagon', u'<': u'triangle_left'}
markers = list(markers.keys())


def plot_scm(models_data, bands, z, is_fit=True, xlim=None, ylim=None):
    # setup figure
    plt.matplotlib.rcParams.update({'font.size': 14})
    # plt.rc('text', usetex=True)
    # plt.rc('font', family='serif')
    fig = plt.figure(num=len(bands), figsize=(9, 9), dpi=100, facecolor='w', edgecolor='k')
    gs1 = gridspec.GridSpec(len(bands) // 2 + len(bands) % 2, 2)
    gs1.update(wspace=0., hspace=0., left=0.1, right=0.9)

    ax_cache = {}

    # create the grid of figures
    for ib, bname in enumerate(bands):
        icol = ib % 2
        irow = ib // 2
        ax = fig.add_subplot(gs1[irow, icol])
        ax_cache[bname] = ax
        # ax.legend(prop={'size': 6})
        # x
        if icol > 0:
            ax.yaxis.tick_right()
            ax.yaxis.set_label_position("right")
        ax.set_ylabel(r'$V_{ph}$ [1000 km/s]')

        if irow == 0:
            ax.set_xlabel(r'$M_{%s}$' % bname)

        if xlim is None:
            ax.set_xlim(max(models_data[bname]) + 2., min(models_data[bname]) - 2.)
        else:
            ax.set_xlim(xlim)
        # xstart, xend = 0, 20000.
        # ax.xaxis.set_ticks(np.arange(5000, xend, (xend - xstart) / 4.))
        # ax.xaxis.set_major_formatter(ticker.FormatStrFormatter('%d'))
        # y
        if ylim is not None:
            ax.set_ylim(ylim)
        # ax.set_yticks(np.arange(0.5, ylim[1], 0.5))
        # ax.text(.5, .9, bset, horizontalalignment='center', transform=ax.transAxes)
        # ax.set_title(bset)
        ax.set_title(bname, x=0.5, y=0.9)

    # plot data
    mi = 0
    vel = models_data['v'] / 1e8  # convert to 1000 km/c
    for ib, bname in enumerate(bands):
        ax = ax_cache[bname]
        x = models_data[bname]
        bcolor = _colors[ib % (len(_colors) - 1)]
        ax.plot(x, vel, marker=markers[mi % (len(markers) - 1)], label='Models',
                markersize=5, color=bcolor, ls="", linewidth=1.5)
        # выбросы todo сделать относительно фита
        # bakl
        if is_fit:
            ax = ax_cache[bname]
            mag = models_data[bname]
            Av = 0.
            a = scm_fit_bakl(mag, vel, z=z, Av=Av)
            mag_bakl = hamuy_fit(a, vel, z, Av)
            # sort
            sorti = np.argsort(mag_bakl)
            xb = mag_bakl[sorti]
            yb = vel[sorti]
            # if mag_bakl is not None:
            ax.plot(xb, yb, label='Baklanov', color="orange", ls="-", lw=2)
            # ax.plot(mag_bakl, vel, marker='o', label='Baklanov', markersize=5, color="blue", ls="")
            print("Bakl fit for %s: %4.2f, %4.2f" % (bname, a[0], a[1]))

    # plt.title('; '.join(set_bands) + ' filter response')
    # plt.grid()
    return fig, ax_cache


def plot_scm_fit(ax_cache, bands, models_data, z):
    # hamuy
    for bname in bands:
        ax = ax_cache[bname]
        ylim = ax.get_ylim()
        yh = np.linspace(ylim[0], ylim[1], num=50)
        xx = scm_fit(yh, Av=0, bname=bname, z=z, src='hamuy')
        if yh is not None:
            ax.plot(xx, yh, color="darkviolet", ls="--", linewidth=2.5, label='Hamuy')
    # kasen
    for bname in bands:
        ax = ax_cache[bname]
        xx = scm_fit(yh, Av=0, bname=bname, z=z, src='kasen')
        if yh is not None:
            ax.plot(xx, yh, color="red", ls="--", linewidth=2.5, label='Kasen')
    # nugent
    yn = np.linspace(ylim[0], ylim[1], num=len(models_data['V']))
    for bname in bands:
        ax = ax_cache[bname]
        V_I = models_data['V'] - models_data['I']  # todo make for any bands
        xn = scm_fit(yn, Av=V_I, src='nugent')
        if yn is not None:
            ax.plot(xn, yn, label='Nugent', color="blue", ls="--")
            # ax.plot(xn, yn, marker='x', label='Nugent', markersize=5, color="blue", ls="")
            # ax.plot(xx, yy, color="orange", ls="--", linewidth=2.5, label='Nugent')


def scm_fit(v, Av=0, bname=None, z=0.003, src='hamuy'):
    """
    Magnitude V from Hamuy & Pinto 2002, Nugent 2006
    :param bname:  band name
    :param v:  expansion velocity [1000 km/c]
    :param Av:  host and Galaxy extinction; Av  => V-I for Nugent method<commit_msg>- FIX: Nugent fit only for band I<commit_after>#!/usr/bin/python3
# -*- coding: utf-8 -*-
import getopt
import numpy as np
import os
import sys
from os.path import isfile, join, dirname
from scipy import interpolate
from scipy.optimize import fmin

import matplotlib.pyplot as plt
from matplotlib import gridspec

from pystella import velocity
from pystella.model.stella import Stella
from pystella.rf import band
from pystella.rf import light_curve_func as lcf
from pystella.util.phys_var import phys

__author__ = 'bakl'

ROOT_DIRECTORY = dirname(dirname(os.path.abspath(__file__)))

_colors = ["blue", "cyan", "brown", 'darkseagreen', 'tomato', 'olive', 'orange',
           'skyblue', 'darkviolet']
# colors = {"B-V": "blue", 'B-V-I': "cyan", 'V-I': "brown"}
lntypes = {"B-V": "-", 'B-V-I': "-.", 'V-I': "--",
           "u-g": "-", 'g-r-i': "-.", 'r-i': "--"}
markers = {u'D': u'diamond', 6: u'caretup', u's': u'square', u'x': u'x',
           5: u'caretright', u'^': u'triangle_up', u'd': u'thin_diamond', u'h': u'hexagon1',
           u'+': u'plus', u'*': u'star', u'o': u'circle', u'p': u'pentagon', u'3': u'tri_left',
           u'H': u'hexagon2', u'v': u'triangle_down', u'8': u'octagon', u'<': u'triangle_left'}
markers = list(markers.keys())


def plot_scm(models_data, bands, z, is_fit=True, xlim=None, ylim=None):
    # setup figure
    plt.matplotlib.rcParams.update({'font.size': 14})
    # plt.rc('text', usetex=True)
    # plt.rc('font', family='serif')
    fig = plt.figure(num=len(bands), figsize=(9, 9), dpi=100, facecolor='w', edgecolor='k')
    gs1 = gridspec.GridSpec(len(bands) // 2 + len(bands) % 2, 2)
    gs1.update(wspace=0., hspace=0., left=0.1, right=0.9)

    ax_cache = {}

    # create the grid of figures
    for ib, bname in enumerate(bands):
        icol = ib % 2
        irow = ib // 2
        ax = fig.add_subplot(gs1[irow, icol])
        ax_cache[bname] = ax
        # ax.legend(prop={'size': 6})
        # x
        if icol > 0:
            ax.yaxis.tick_right()
            ax.yaxis.set_label_position("right")
        ax.set_ylabel(r'$V_{ph}$ [1000 km/s]')

        if irow == 0:
            ax.set_xlabel(r'$M_{%s}$' % bname)

        if xlim is None:
            ax.set_xlim(max(models_data[bname]) + 2., min(models_data[bname]) - 2.)
        else:
            ax.set_xlim(xlim)
        # xstart, xend = 0, 20000.
        # ax.xaxis.set_ticks(np.arange(5000, xend, (xend - xstart) / 4.))
        # ax.xaxis.set_major_formatter(ticker.FormatStrFormatter('%d'))
        # y
        if ylim is not None:
            ax.set_ylim(ylim)
        # ax.set_yticks(np.arange(0.5, ylim[1], 0.5))
        # ax.text(.5, .9, bset, horizontalalignment='center', transform=ax.transAxes)
        # ax.set_title(bset)
        ax.set_title(bname, x=0.5, y=0.9)

    # plot data
    mi = 0
    vel = models_data['v'] / 1e8  # convert to 1000 km/c
    for ib, bname in enumerate(bands):
        ax = ax_cache[bname]
        x = models_data[bname]
        bcolor = _colors[ib % (len(_colors) - 1)]
        ax.plot(x, vel, marker=markers[mi % (len(markers) - 1)], label='Models',
                markersize=5, color=bcolor, ls="", linewidth=1.5)
        # выбросы todo сделать относительно фита
        # bakl
        if is_fit:
            ax = ax_cache[bname]
            mag = models_data[bname]
            Av = 0.
            a = scm_fit_bakl(mag, vel, z=z, Av=Av)
            mag_bakl = hamuy_fit(a, vel, z, Av)
            # sort
            sorti = np.argsort(mag_bakl)
            xb = mag_bakl[sorti]
            yb = vel[sorti]
            # if mag_bakl is not None:
            ax.plot(xb, yb, label='Baklanov', color="orange", ls="-", lw=2)
            # ax.plot(mag_bakl, vel, marker='o', label='Baklanov', markersize=5, color="blue", ls="")
            print("Bakl fit for %s: %4.2f, %4.2f" % (bname, a[0], a[1]))

    # plt.title('; '.join(set_bands) + ' filter response')
    # plt.grid()
    return fig, ax_cache


def plot_scm_fit(ax_cache, bands, models_data, z):
    # hamuy
    for bname in bands:
        ax = ax_cache[bname]
        ylim = ax.get_ylim()
        yh = np.linspace(ylim[0], ylim[1], num=50)
        xx = scm_fit(yh, Av=0, bname=bname, z=z, src='hamuy')
        if yh is not None:
            ax.plot(xx, yh, color="darkviolet", ls="--", linewidth=2.5, label='Hamuy')
    # kasen
    for bname in bands:
        ax = ax_cache[bname]
        xx = scm_fit(yh, Av=0, bname=bname, z=z, src='kasen')
        if yh is not None:
            ax.plot(xx, yh, color="red", ls="--", linewidth=2.5, label='Kasen')
    # nugent
    bname = 'I'
    if bname in bands:
        yn = np.linspace(ylim[0], ylim[1], num=len(models_data['V']))
        ax = ax_cache[bname]
        V_I = models_data['V'] - models_data['I']  # todo make for any bands
        xn = scm_fit(yn, Av=V_I, src='nugent')
        if yn is not None:
            ax.plot(xn, yn, label='Nugent', color="blue", ls="--")
            # ax.plot(xn, yn, marker='x', label='Nugent', markersize=5, color="blue", ls="")
            # ax.plot(xx, yy, color="orange", ls="--", linewidth=2.5, label='Nugent')


def scm_fit(v, Av=0, bname=None, z=0.003, src='hamuy'):
    """
    Magnitude V from Hamuy & Pinto 2002, Nugent 2006
    :param bname:  band name
    :param v:  expansion velocity [1000 km/c]
    :param Av:  host and Galaxy extinction; Av  => V-I for Nugent method
{'fraction_non_alphanumeric': 0.17505368312949304, 'fraction_numerical': 0.021099803939874893, 'mean_word_length': 2.509829619921363, 'pattern_counts': {'":': 6, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '426918', 'n_tokens_mistral': 4377, 'n_tokens_neox': 4187, 'n_words': 950}