<commit_before>"""Plots radial profiles coming out of stagyy.

Author: <NAME> with inputs from <NAME> and <NAME>
Date: 2015/09/11
"""
from __future__ import print_function

import numpy as np
from scipy import integrate as itg
import f90nml
import os
import sys
import math
import constants
import misc

def _readproffile(args):
    proffile = os.path.join(args.path, args.name+'_rprof.dat')
    if not os.path.isfile(proffile):
        print('No profile file found at', proffile)
        sys.exit()
    timesteps = []
    data0 = []
    lnum = -1
    with open(proffile) as stream:
        for line in stream:
            if line != '\n':
                lnum += 1
                lll = ' '.join(line.split())
                if line[0] == '*':
                    timesteps.append([lnum, int(lll.split(' ')[1]),
                                      float(lll.split(' ')[5])])
                else:
                    llf = np.array(lll.split(' '))
                    data0.append(llf)
    tsteps = np.array(timesteps)
    nsteps = tsteps.shape[0]
    data = np.array(data0)
    # all the processings of timesteps
    # should be in commands.*_cmd
    # instead of main.py
    # since it could be different between
    # the different modules
    istart, ilast, istep = args.timestep
    if ilast == -1:
        ilast = nsteps-1
    if istart == -1:
        istart = nsteps-1
    args.timestep = istart, ilast, istep
    return data, nsteps, tsteps, lnum

def _normprof(rrr, func): # for args.plot_difference
    """Volumetric norm of a profile

    Two arrays: rrr is the radius position and f the function.
    """
    norm = 3./(rrr[-1]**3.-rrr[0]**3.)*itg.trapz(func**2*rrr**2., rrr)
    return norm

def _extrap(xpos, xpoints, ypoints): # for args.plot_difference
    """np.interp function with linear extrapolation.

    Would be best to use degree 3 extrapolation
    """
    ypos = np.interp(xpos, xpoints, ypoints)
    ypos[xpos < xpoints[0]] = ypoints[0]\
            + (xpos[xpos < xpoints[0]]-xpoints[0])\
            * (ypoints[0]-ypoints[1])/(xpoints[0]-xpoints[1])
    ypos[xpos > xpoints[-1]] = ypoints[-1]\
            + (xpos[xpos > xpoints[-1]]-xpoints[-1])\
            * (ypoints[-1]-ypoints[-2])/(xpoints[-1]-xpoints[-2])
    return ypos

def _calc_energy(data, ir0, ir1): # for args.plot_energy
    """Compute energy balance(r)"""
    zgrid = np.array(data[ir0:ir1, 63], float)
    zgrid = np.append(zgrid, 1.)
    dzg = np.array(data[ir0+1:ir1, 0], float) - np.array(data[ir0:ir1-1, 0],
                                                         float)
    qadv = np.array(data[ir0:ir1-1, 60], float)
    qadv = np.insert(qadv, 0, 0.)
    qadv = np.append(qadv, 0.)
    qcond = (np.array(data[ir0:ir1-1, 1], float) -
             np.array(data[ir0+1:ir1, 1], float))/dzg
    qcond0 = (1.-float(data[ir0, 1]))/float(data[ir0, 0])
    qtop = float(data[ir1, 1])/(1.-float(data[ir1, 0]))
    qcond = np.insert(qcond, 0, qcond0)
    qcond = np.append(qcond, qtop)
    qtot = qadv+qcond
    return qtot, qadv, qcond, zgrid

def rprof_cmd(args):
    """Plot radial profiles"""
    plt = args.plt


    if not (args.plot_conctheo and args.plot_temperature
            and args.plot_concentration):
        args.plot_difference = False

    if args.plot_difference:
        # plot time series of difference profiles
        # initialize the plot here
        figd, axax = plt.subplots(3, sharex=True)

    linestyles = ('-', '--', '-.', ':')
    lwdth = args.linewidth
    ftsz = args.fontsize

    # parameters for the theoretical composition profiles

    if not args.par_nml: # too much parameter to define arbitrarily
        print('No par file found. Abort.')
        sys.exit()

    spherical = args.par_nml['geometry']['shape'].lower() == 'spherical'
    if spherical:
        rcmb = args.par_nml['geometry']['r_cmb']
    else:
        rcmb = 0.

    rmin = rcmb
    rmax = rcmb + 1.
    if args.plot_conctheo:
        if 'fe_eut' in args.par_nml['tracersin']:
            xieut = args.par_nml['tracersin']['fe_eut']
        else:
            xieut = 0.8
        if 'k_fe' in args.par_nml['tracersin']:
            k_fe = args.par_nml['tracersin']['k_fe']
        else:
            k_fe = 0.85
        if 'fe_cont' in args.par_nml['tracersin']:
            xi0l = args.par_nml['tracersin']['fe_cont']
        else:
            xi0l = 0.1
        xi0s = k_fe*xi0l
        xired = xi0l/xieut
        rsup = (rmax**3-xired**(1/(1-k_fe))*(rmax**3-rmin**3))**(1./3.)
        print('rmin, rmax, rsup=', rmin, rmax, rsup)

        def initprof(rpos):
            """Theoretical initial profile."""
            if rpos < rsup:
                return xi0s*((rmax**3-rmin**3)/(rmax**3-rpos**3))**(1-k_fe)
            else:
                return xieut

    data, nsteps, tsteps, lnum = _readproffile(args)

    istart, ilast, istep = args.timestep

    nzp = []
    for iti in range(0, nsteps-1):
        nzp = np.append(nzp, tsteps[iti+1, 0]-tsteps[iti, 0]-1)

    nzp = np.append(nzp, lnum-tsteps[nsteps-1, 0])

    nzs = [[0, 0, 0]]
    nzc = 0
    for iti in range(1, nsteps):
        if nzp[iti] != nzp[iti-1]:
            nzs.append([iti, iti-nzc, int(nzp[iti-1])])
            nzc = iti
    if nzp[nsteps-1] != nzs[-1][1]:
        nzs.append([nsteps, nsteps-nzc, int(nzp[nsteps-1])])

    nzi = np.array(nzs)

    def plotprofiles(quant, vartuple, integrate=False):
        """Plot the chosen profiles for the chosen timesteps

        quant holds the strings for the x axis annotation and
        the legends for the additional profiles

        vartuple contains the numbers of the column to be plotted
        """

        if integrate:
            integ = lambda f, r: f * (r/rmax)**2

        if quant[0] == 'Grid':
            fig, axe = plt.subplots(2, sharex=True)
        else:
            fig, axe = plt.subplots()

        timename = str(istart) + "_" + str(ilast) + "_" + str(istep)
        if args.plot_difference:
            concdif = []
            tempdif = []
            wmax = []

        for step in range(istart, ilast, istep):
            step = step +1# starts at 0=> 15 is the 16th
            # find the indices
            ann = sorted(np.append(nzi[:, 0], step))
            inn = ann.index(step)
            nnz = np.multiply(nzi[:, 1], nzi[:, 2])
<commit_msg>Move nzi calculation to rprof._readproffile<commit_after>"""Plots radial profiles coming out of stagyy.

Author: <NAME> with inputs from <NAME> and <NAME>
Date: 2015/09/11
"""
from __future__ import print_function

import numpy as np
from scipy import integrate as itg
import f90nml
import os
import sys
import math
import constants
import misc

def _readproffile(args):
    proffile = os.path.join(args.path, args.name+'_rprof.dat')
    if not os.path.isfile(proffile):
        print('No profile file found at', proffile)
        sys.exit()
    timesteps = []
    data0 = []
    lnum = -1
    with open(proffile) as stream:
        for line in stream:
            if line != '\n':
                lnum += 1
                lll = ' '.join(line.split())
                if line[0] == '*':
                    timesteps.append([lnum, int(lll.split(' ')[1]),
                                      float(lll.split(' ')[5])])
                else:
                    llf = np.array(lll.split(' '))
                    data0.append(llf)
    tsteps = np.array(timesteps)
    nsteps = tsteps.shape[0]
    data = np.array(data0)
    # all the processings of timesteps
    # should be in commands.*_cmd
    # instead of main.py
    # since it could be different between
    # the different modules
    istart, ilast, istep = args.timestep
    if ilast == -1:
        ilast = nsteps-1
    if istart == -1:
        istart = nsteps-1
    args.timestep = istart, ilast, istep

    nzp = []
    for iti in range(0, nsteps-1):
        nzp = np.append(nzp, tsteps[iti+1, 0]-tsteps[iti, 0]-1)

    nzp = np.append(nzp, lnum-tsteps[nsteps-1, 0])
    nzs = [[0, 0, 0]]
    nzc = 0
    for iti in range(1, nsteps):
        if nzp[iti] != nzp[iti-1]:
            nzs.append([iti, iti-nzc, int(nzp[iti-1])])
            nzc = iti
    if nzp[nsteps-1] != nzs[-1][1]:
        nzs.append([nsteps, nsteps-nzc, int(nzp[nsteps-1])])
    nzi = np.array(nzs)
    return data, tsteps, nzi

def _normprof(rrr, func): # for args.plot_difference
    """Volumetric norm of a profile

    Two arrays: rrr is the radius position and f the function.
    """
    norm = 3./(rrr[-1]**3.-rrr[0]**3.)*itg.trapz(func**2*rrr**2., rrr)
    return norm

def _extrap(xpos, xpoints, ypoints): # for args.plot_difference
    """np.interp function with linear extrapolation.

    Would be best to use degree 3 extrapolation
    """
    ypos = np.interp(xpos, xpoints, ypoints)
    ypos[xpos < xpoints[0]] = ypoints[0]\
            + (xpos[xpos < xpoints[0]]-xpoints[0])\
            * (ypoints[0]-ypoints[1])/(xpoints[0]-xpoints[1])
    ypos[xpos > xpoints[-1]] = ypoints[-1]\
            + (xpos[xpos > xpoints[-1]]-xpoints[-1])\
            * (ypoints[-1]-ypoints[-2])/(xpoints[-1]-xpoints[-2])
    return ypos

def _calc_energy(data, ir0, ir1): # for args.plot_energy
    """Compute energy balance(r)"""
    zgrid = np.array(data[ir0:ir1, 63], float)
    zgrid = np.append(zgrid, 1.)
    dzg = np.array(data[ir0+1:ir1, 0], float) - np.array(data[ir0:ir1-1, 0],
                                                         float)
    qadv = np.array(data[ir0:ir1-1, 60], float)
    qadv = np.insert(qadv, 0, 0.)
    qadv = np.append(qadv, 0.)
    qcond = (np.array(data[ir0:ir1-1, 1], float) -
             np.array(data[ir0+1:ir1, 1], float))/dzg
    qcond0 = (1.-float(data[ir0, 1]))/float(data[ir0, 0])
    qtop = float(data[ir1, 1])/(1.-float(data[ir1, 0]))
    qcond = np.insert(qcond, 0, qcond0)
    qcond = np.append(qcond, qtop)
    qtot = qadv+qcond
    return qtot, qadv, qcond, zgrid

def rprof_cmd(args):
    """Plot radial profiles"""
    plt = args.plt


    if not (args.plot_conctheo and args.plot_temperature
            and args.plot_concentration):
        args.plot_difference = False

    if args.plot_difference:
        # plot time series of difference profiles
        # initialize the plot here
        figd, axax = plt.subplots(3, sharex=True)

    linestyles = ('-', '--', '-.', ':')
    lwdth = args.linewidth
    ftsz = args.fontsize

    # parameters for the theoretical composition profiles

    if not args.par_nml: # too much parameter to define arbitrarily
        print('No par file found. Abort.')
        sys.exit()

    spherical = args.par_nml['geometry']['shape'].lower() == 'spherical'
    if spherical:
        rcmb = args.par_nml['geometry']['r_cmb']
    else:
        rcmb = 0.

    rmin = rcmb
    rmax = rcmb + 1.
    if args.plot_conctheo:
        if 'fe_eut' in args.par_nml['tracersin']:
            xieut = args.par_nml['tracersin']['fe_eut']
        else:
            xieut = 0.8
        if 'k_fe' in args.par_nml['tracersin']:
            k_fe = args.par_nml['tracersin']['k_fe']
        else:
            k_fe = 0.85
        if 'fe_cont' in args.par_nml['tracersin']:
            xi0l = args.par_nml['tracersin']['fe_cont']
        else:
            xi0l = 0.1
        xi0s = k_fe*xi0l
        xired = xi0l/xieut
        rsup = (rmax**3-xired**(1/(1-k_fe))*(rmax**3-rmin**3))**(1./3.)
        print('rmin, rmax, rsup=', rmin, rmax, rsup)

        def initprof(rpos):
            """Theoretical initial profile."""
            if rpos < rsup:
                return xi0s*((rmax**3-rmin**3)/(rmax**3-rpos**3))**(1-k_fe)
            else:
                return xieut

    data, tsteps, nzi = _readproffile(args)

    istart, ilast, istep = args.timestep

    def plotprofiles(quant, vartuple, integrate=False):
        """Plot the chosen profiles for the chosen timesteps

        quant holds the strings for the x axis annotation and
        the legends for the additional profiles

        vartuple contains the numbers of the column to be plotted
        """

        if integrate:
            integ = lambda f, r: f * (r/rmax)**2

        if quant[0] == 'Grid':
            fig, axe = plt.subplots(2, sharex=True)
        else:
            fig, axe = plt.subplots()

        timename = str(istart) + "_" + str(ilast) + "_" + str(istep)
        if args.plot_difference:
            concdif = []
            tempdif = []
            wmax = []

        for step in range(istart, ilast, istep):
            step = step +1# starts at 0=> 15 is the 16th
            # find the indices
            ann = sorted(np.append(nzi[:, 0], step))
            inn = ann.index(step)
            nnz = np.multiply(nzi[:, 1], nzi[:, 2])

{'fraction_non_alphanumeric': 0.14431602279924002, 'fraction_numerical': 0.023749208359721344, 'mean_word_length': 2.0827232796486093, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1383414', 'n_tokens_mistral': 4748, 'n_tokens_neox': 4487, 'n_words': 1078}