Exemple #1
0
def gas_alpha(sim=False, t_range=[0, -1], OVERWRITE=False):
    """This calculates the global alpha-value from ts object via
            alphah=(2./3.)*(ts3.uxuym)/(ts3.rhom*ts3.csm**2)
    for the lastNentries of the time series.
    Input:
        t_range         use this time range of the timeseries
        OVERWRITE       overwrite alpha file in sim.pc_datadir/alpha_<lastNentries>

    return:
        alpha dictionary
    """

    from pencil import get_sim
    from pencil.sim import sim
    from pencil import io
    from os.path import exists, join
    import numpy as np

    def empirical_std_deviation(x):
        """
        (Geschaetzte) Streuung der Messwerte x um den (unbekannten) wahren Wert
        (estimated) stray of the measurments x around the (unknown) true value

        s(x) = SQRT( 1./(M-1) * SUM( (x-<x>)**2 ) )"""
        import numpy as np

        x = np.array(x)
        M = np.size(x)
        xm = np.mean(x)

        # return np.sqrt(1./(M-1.)*np.sum((x-xm)**2))
        return np.sqrt(M / (M - 1.0) * ((1.0 / M * np.sum(x**2)) - xm**2))

    def std_deviation_of_mean_value(x):
        """
        Messunsicherheit des Mittelwertes <x>
        Empirical standarddeviation of the arithmetical mean value

        s(<x>) = s(x)/SQRT( M ) = SQRT( 1./(M-1) * SUM( (x-<x>)**2 ) ) / SQRT( M )"""

        import numpy as np

        x = np.array(x)
        M = np.size(x)

        if M == 1:
            return 0

        return empirical_std_deviation(x) / np.sqrt(M)

    if type(sim) == sim.__Simulation__:
        SIM = sim
    else:
        SIM = get_sim()

    filename = "alpha_" + str(t_range[0]) + "_" + str(t_range[1])

    ## skip if nothing is new
    if not OVERWRITE and io.exists(name=filename, sim=SIM):
        print('~ Alpha for "' + SIM.name + '" already exists. Loading file...')
        return io.load(name=filename, sim=SIM)

    print('~ Calculating alpha for "' + SIM.name + '" in "' + SIM.path + '"')

    ## import time series object
    try:
        print("~ reading time series..")
        ts = SIM.get_ts()
    except:
        print("! ERROR: Couldnt read time series!")
        return False

    ## correct if csm quantity is not correctly exported
    try:
        csm = ts.csm
        if csm[0] == 0:
            csm = 1.0
    except:
        print(
            "? WARNING: Couldnt find >csm< in time series, set it to csm = 1. This may be incorrect!"
        )
        csm = 1.0

    if t_range[1] == -1:
        t_range[1] = ts.t[-1]
    id_min = np.argmin(np.abs(ts.t - t_range[0]))
    id_max = np.argmin(np.abs(ts.t - t_range[1]))

    alpha_dict = {}

    ## calculate alpha
    print("~ calculating alpha, its mean value and standard deviation for " +
          str(t_range))
    alpha_dict["alpha"] = -(2.0 / 3.0) * (ts.uxuym) / (ts.rhom * csm**2)
    alpha_dict["alpha_mean"] = np.mean(alpha_dict["alpha"][id_min:id_max])
    alpha_dict["alpha_stddev"] = std_deviation_of_mean_value(
        alpha_dict["alpha"][id_min:id_max])
    print("~ alpha_mean = " + str(alpha_dict["alpha_mean"]) +
          " and alpha_stddev = " + str(alpha_dict["alpha_stddev"]))

    ## calculate alpha minus mean_flux
    print(
        "~ calculating alpha minus mean_flux (alpha_mmf), its mean value and standard deviation"
    )
    alpha_dict["alpha_mmf"] = (-(2.0 / 3.0) * (ts.uxuym - ts.uxm * ts.uym) /
                               (ts.rhom * csm**2))
    alpha_dict["alpha_mmf_mean"] = np.mean(
        alpha_dict["alpha_mmf"][id_min:id_max])
    alpha_dict["alpha_mmf_stddev"] = std_deviation_of_mean_value(
        alpha_dict["alpha_mmf"][id_min:id_max])
    print("~ alpha_mmf_mean = " + str(alpha_dict["alpha_mmf_mean"]) +
          " and alpha_mmf_stddev = " + str(alpha_dict["alpha_mmf_stddev"]))

    import math

    for v in alpha_dict.values():
        if type(v) == type(1.1) and math.isnan(v):
            io.debug_breakpoint()

    ## save alpha in plk
    print("~ saving alpha values in " + SIM.pc_datadir + "/" + filename)
    io.save(alpha_dict, filename, folder=SIM.pc_datadir)

    return alpha_dict
Exemple #2
0
    def copy(
        self,
        path_root=".",
        name=False,
        start_optionals=False,
        optionals=True,
        quiet=True,
        rename_submit_script=False,
        OVERWRITE=False,
    ):
        """
        This method does a copy of the simulation object by creating a new
        directory 'name' in 'path_root' and copy all simulation components and
        optiona)
                ls to its directory.
        This method neither links/compiles the simulation.
        If start_optionals it creates data dir.
        It does not overwrite anything, unless OVERWRITE is True.

        Submit Script Rename:
            Name in submit scripts will be renamed if possible!
            Submit scripts will be identified by submit* plus appearenace of old
            simulation name inside, latter will be renamed!

        Parameters
        ----------
        path_root : string
            Path to new sim.-folder(sim.-name). This folder will
            be created if not existing! Relative paths are
            thought to be relative to the python current workdir

        name : string
            Name of new simulation, will be used as folder name.
            Rename will also happen in submit script if found.
            Simulation folders is not allowed to preexist!

        optionals : bool
            Add list of further files to be copied. Wildcasts
            allowed according to glob module!
            Set True to use self.optionals.

        start optionals :
            Add list of further files to be copied.
            Wildcasts allowed according to glob module!
            Set True to use self.optionals.

        quiet : bool
            Set True to suppress output.

        rename_submit_script : bool
            Set False if no renames shall be performed in submit* files

        OVERWRITE : bool
            Set True to overwrite no matter what happens!
        """
        from glob import glob
        from numpy import size
        from os import listdir, symlink
        from shutil import copyfile

        from pencil import get_sim
        from pencil.io import (
            mkdir,
            get_systemid,
            rename_in_submit_script,
            debug_breakpoint,
        )
        from pencil import is_sim_dir

        # set up paths
        if path_root == False or type(path_root) != type("string"):
            print("! ERROR: No path_root specified to copy the simulation to.")
            return False
        path_root = abspath(path_root)  # simulation root dir

        # name and folder of new simulation but keep name of old if sim with old
        # name is NOT existing in NEW directory
        if name == False:
            name = self.name
        if exists(join(path_root, name)) and OVERWRITE == False:
            name = name + "_copy"
            if exists(join(path_root, name)):
                name = name + str(
                    size([f
                          for f in listdir(path_root) if f.startswith(name)]))
            print(
                "? Warning: No name specified and simulation with that name " +
                "already found! New simulation name now " + name)
        path_newsim = join(path_root, name)  # simulation abspath
        path_newsim_src = join(path_newsim, "src")
        if islink(join(path_root, self.name, "data")):
            link_data = True
            oldtmp = os.path.realpath(join(path_root, self.name, "data"))
            newtmp = join(str.strip(str.strip(oldtmp, "data"), self.name),
                          name, "data")
            if exists(newtmp) and OVERWRITE == False:
                raise ValueError(
                    "Data directory {0} already exists".format(newtmp))
            else:
                path_newsim_data = newtmp
                path_newsim_data_link = join(path_newsim, "data")
        else:
            link_data = False
            path_newsim_data = join(path_newsim, "data")

        path_initial_condition = join(self.path, "initial_condition")
        if exists(path_initial_condition):
            has_initial_condition_dir = True
            path_newsim_initcond = join(path_newsim, "initial_condition")
        else:
            has_initial_condition_dir = False

        if type(optionals) == type(["list"]):
            optionals = self.optionals + optionals  # optional files to be copied
        if optionals == True:
            optionals = self.optionals
        if type(optionals) == type("string"):
            optionals = [optionals]
        if type(optionals) != type(["list"]):
            print("! ERROR: optionals must be of type list!")

        tmp = []
        for opt in optionals:
            files = glob(join(self.path, opt))
            for f in files:
                tmp.append(basename(f))
        optionals = tmp

        # optional files to be copied
        if type(start_optionals) == type(["list"]):
            start_optionals = self.start_optionals + start_optionals
        if start_optionals == False:
            start_optionals = self.start_optionals
        if type(start_optionals) == type("string"):
            start_optionals = [start_optionals]
        if type(start_optionals) != type(["list"]):
            print("! ERROR: start_optionals must be of type list!")

        tmp = []
        for opt in start_optionals:
            files = glob(join(self.datadir, opt))
            for f in files:
                tmp.append(basename(f))
        start_optionals = tmp
        ## check if the copy was already created
        if is_sim_dir(path_newsim) and OVERWRITE == False:
            if not quiet:
                print("? WARNING: Simulation already exists." +
                      " Returning with existing simulation.")
            return get_sim(path_newsim, quiet=quiet)

        ## expand list of optionals wildcasts

        # check existence of path_root+name, a reason to stop and not overwrite
        if OVERWRITE == False and exists(path_newsim):
            print(
                "! ERROR: Folder to copy simulation to already exists!\n! -> "
                + path_newsim)
            return False

        # check existance of self.components
        for comp in self.components:
            if not exists(join(self.path, comp)):
                print("! ERROR: Couldnt find component " + comp +
                      " from simulation " + self.name + " at location " +
                      join(self.path, comp))
                return False

        # check existance of optionals
        for opt in optionals:
            if not exists(join(self.path, opt)):
                print("! ERROR: Couldnt find optional component " + opt +
                      " from simulation " + self.name + " at location " +
                      join(self.path, opt))
                return False

        # check existance of self.start_components
        for comp in self.start_components:
            if not exists(join(self.datadir, comp)):
                print("! ERROR: Couldnt find component " + comp +
                      " from simulation " + self.name + " at location " +
                      join(self.path, comp))
                return False

        # check existance of start_optionals
        for opt in start_optionals:
            if not exists(join(self.datadir, opt)):
                print("! ERROR: Couldnt find optional component " + opt +
                      " from simulation " + self.name + " at location " +
                      join(self.datadir, opt))
                return False

        # create folders
        if mkdir(path_newsim) == False and OVERWRITE == False:
            print("! ERROR: Couldnt create new simulation directory " +
                  path_newsim + " !!")
            return False

        if mkdir(path_newsim_src) == False and OVERWRITE == False:
            print("! ERROR: Couldnt create new simulation src directory " +
                  path_newsim_src + " !!")
            return False

        if mkdir(path_newsim_data) == False and OVERWRITE == False:
            print("! ERROR: Couldnt create new simulation data directory " +
                  path_newsim_data + " !!")
            return False
        if link_data:
            symlink(path_newsim_data, path_newsim_data_link)

        # copy files
        files_to_be_copied = []
        for f in self.components + optionals:
            f_path = abspath(join(self.path, f))
            copy_to = abspath(join(path_newsim, f))
            if f_path == copy_to:
                print("!! ERROR: file path f_path equal to destination " +
                      "copy_to. Debug this line manually!")
                debug_breakpoint()
            copyfile(f_path, copy_to)

        for f in self.start_components + start_optionals:
            f_path = abspath(join(self.datadir, f))
            copy_to = abspath(join(path_newsim_data, f))
            if f_path == copy_to:
                print("!! ERROR: file path f_path equal to destination " +
                      "copy_to. Debug this line manually!")
                debug_breakpoint()
            copyfile(f_path, copy_to)

        # Organizes any personalized initial conditions
        if has_initial_condition_dir:
            if mkdir(path_newsim_initcond) == False and OVERWRITE == False:
                print(
                    "! ERROR: Couldnt create new simulation initial_condition"
                    + " directory " + path_newsim_initcond + " !!")
                return False

            for f in listdir(path_initial_condition):
                f_path = abspath(join(path_initial_condition, f))
                copy_to = abspath(join(path_newsim_initcond, f))

                if f_path == copy_to:
                    print("!! ERROR: file path f_path equal to destination " +
                          "copy_to. Debug this line manually!")
                    debug_breakpoint()
                copyfile(f_path, copy_to)

        # modify name in submit script files
        if rename_submit_script != False:
            if type(rename_submit_script) == type("STRING"):
                rename_in_submit_script(new_name=rename_submit_script,
                                        sim=get_sim(path_newsim))
            else:
                print("!! ERROR: Could not understand rename_submit_script=" +
                      str(rename_submit_script))

        # done
        return get_sim(path_newsim)
def gas_velo_at_particle_pos(varfiles='last4',
                             sim=False,
                             scheme='tsc',
                             use_IDL=False,
                             OVERWRITE=False):
    """This script calulates the gas velocity at the particle position and stores this together
  with particle position, containing grid cell idicies, particle velocities, and particle index
  in a gas_velo_at_particle_pos file.

  Args:
    varfiles:       specifiy varfiles for calculation, e.g. 'last', 'first',
                        'all', 'VAR###', 'last4', 'first3'
    scheme:         possible are:
                        - ngp: nearest grid point
                        - cic: cloud in cell
                        - tsc: triangular shaped cloud
    OVERWRITE:		set to True to overwrite already calculated results
  """

    from pencil import get_sim
    from pencil import read
    from pencil import diag
    from pencil.io import mkdir
    from os import listdir
    from os.path import exists, join, dirname
    import numpy as np

    GAS_VELO_TAG = 'gas_velo_at_particle_pos'

    if sim == False:
        sim = get_sim()
        if sim == False:
            print('! ERROR: Specify simulation object!')
            return False
    SIM = sim

    if use_IDL:
        print(
            '? WARNING: IDL VERSION OF THIS SCRIPT BY JOHANSEN, not recommended for 2D data'
        )
        from ...backpack import pidly
        print('## starting IDL engine..')
        IDL = pidly.IDL(long_delay=0.05)  # start IDL engine

        ## skip if nothing is new
        if (not OVERWRITE) and (exists(join(
                SIM.pc_datadir, 'sigma.pkl'))) and (exists(
                    join(SIM.pc_datadir, 'zeta.pkl'))):
            print('~ ' + SIM.name +
                  ' is already calculated and up-to-date! -> skipping it!')

        else:
            ## start calculations
            print('~ Calculating gas_velo_at_particle_pos for "' + SIM.name +
                  '" in "' + SIM.path + '"')
            IDL('.COMPILE ' + str(
                join(dirname(diag.particle.__file__),
                     'gas_velo_at_particle_pos.pro')))
            IDL.pro('gas_velo_at_particle_pos',
                    datadir=SIM.datadir,
                    destination=GAS_VELO_TAG,
                    doforthelastNvar=varfiles[4:])
            files = [
                i.split('_')[-1].split('.sav')[0]
                for i in listdir(join(SIM.pc_datadir, GAS_VELO_TAG))
                if i.startswith(GAS_VELO_TAG) and i.endswith('.sav')
                or i.endswith('.pkl')
            ]
            if files == []:
                print(
                    '!! ERROR: No calc_gas_speed_at_particle_position-files found for '
                    + SIM.name + '! Use idl script to produce them first!')

        IDL.close()
        return True

    else:
        print('~ Calculating gas_velo_at_particle_pos for "' + SIM.name +
              '" in "' + SIM.path + '"')
        save_destination = join(SIM.pc_datadir, GAS_VELO_TAG)
        mkdir(save_destination)
        varlist = SIM.get_varlist(pos=varfiles, particle=False)
        pvarlist = SIM.get_varlist(pos=varfiles, particle=True)

        for f, p in zip(varlist, pvarlist):
            save_filename = GAS_VELO_TAG + '_' + scheme + '_' + f[3:]
            if not OVERWRITE and exists(save_filename,
                                        folder=save_destination):
                continue

            print('## Reading ' + f + ' ...')
            ff = read.var(datadir=SIM.datadir,
                          varfile=f,
                          quiet=True,
                          trimall=False)
            pp = read.pvar(datadir=SIM.datadir, varfile=p)

            ## remove ghost zones from grid, call the reduced grid the "real grid"
            realgridx = ff.x[ff.l1:ff.l2]
            realgridy = ff.y[ff.m1:ff.m2]
            realgridz = ff.z[ff.n1:ff.n2]
            nx = ff.l2 - ff.l1
            ny = ff.m2 - ff.m1
            nz = ff.n2 - ff.n1

            ## prepare list for all quantities
            l_ipars = pp.ipars  # particle number   KNOWN
            l_px = pp.xp
            l_py = pp.yp
            l_pz = pp.zp  # particle absolut position KNOWN
            l_vx = pp.vpx
            l_vy = pp.vpy
            l_vz = pp.vpz  # particle velocity KNOWN
            l_rix = []
            l_riy = []
            l_riz = [
            ]  # particle untrimmed realgrid index (grid index = l/m/n + readgrid index ???)
            l_ix = []
            l_iy = []
            l_iz = []  # particle grid index (in untrimmed grid)
            l_ux = []
            l_uy = []
            l_uz = []  # underlying gas velocity at position of particle

            ## get index of realgrid cell for each particle
            for i in range(len(l_ipars)):
                l_rix.append(np.abs(realgridx - l_px[i]).argmin())
                l_riy.append(np.abs(realgridy - l_py[i]).argmin())
                l_riz.append(np.abs(realgridz - l_pz[i]).argmin())

            ## convert into untrimmed grid
            l_ix = np.array(l_rix) + ff.l1
            l_iy = np.array(l_riy) + ff.m1
            l_iz = np.array(l_riz) + ff.n1

            ## NGP
            if scheme == 'ngp' or scheme == 'NGP':
                print('## Calculating gas velocities via ' + scheme)
                l_ux = ff.ux[l_iz, l_iy, l_ix]
                l_uy = ff.uy[l_iz, l_iy, l_ix]
                l_uz = ff.uz[l_iz, l_iy, l_ix]

            ## CIC
            if scheme == 'cic' or scheme == 'CIC':
                print('## Calculating gas velocities via ' + scheme)
                for ix0, iy0, iz0, px, py, pz in zip(
                        l_ix, l_iy, l_iz, l_px, l_py,
                        l_pz):  # for each particle
                    if ff.x[ix0] > px:
                        ix0 = ix0 - 1  # ix0 must be left to particle
                    if ff.y[iy0] > py:
                        iy0 = iy0 - 1  # iy0 must be below the particle
                    if ff.z[iz0] > pz:
                        iz0 = iz0 - 1  # iz0 must be under particle

                    ix1 = ix0
                    iy1 = iy0
                    iz1 = iz0  # if a dim. is zero, this is default, else:
                    if nx > 1:
                        ix1 = ix0 + 1
                        dx_1 = 1. / ff.dx  # if a dim is non-zero, ajust ix1 to right cell
                    if ny > 1:
                        iy1 = iy0 + 1
                        dy_1 = 1. / ff.dy  # if a dim is non-zero, ajust iy1 to above cell
                    if nz > 1:
                        iz1 = iz0 + 1
                        dz_1 = 1. / ff.dz  # if a dim is non-zero, ajust iz1 to above cell

                    ux = 0.
                    uy = 0.
                    uz = 0.
                    for ix in [ix0, ix1]:
                        for iy in [iy0, iy1]:
                            for iz in [iz0, iz1]:
                                weight = 1.
                                if nx > 1:
                                    weight = weight * (
                                        1. - abs(px - ff.x[ix]) * dx_1)
                                if ny > 1:
                                    weight = weight * (
                                        1. - abs(py - ff.y[iy]) * dy_1)
                                if nz > 1:
                                    weight = weight * (
                                        1. - abs(pz - ff.z[iz]) * dz_1)

                                ux = ux + weight * ff.ux[iz, iy, ix]
                                uy = uy + weight * ff.uy[iz, iy, ix]
                                uz = uz + weight * ff.uz[iz, iy, ix]

                                if iz0 == iz1: break  # beware of degeneracy:
                            if iy0 == iy1: break  # beware of degeneracy:
                        if ix0 == ix1: break  # beware of degeneracy:

                    l_ux.append(ux)
                    l_uy.append(uy)
                    l_uz.append(uz)

            ## TSC
            if scheme == 'tsc' or scheme == 'TSC':
                for ix0, iy0, iz0, px, py, pz in zip(
                        l_ix, l_iy, l_iz, l_px, l_py,
                        l_pz):  # for each particle
                    ixx0 = ix0
                    ixx1 = ix0  # beware of degeneracy
                    iyy0 = iy0
                    iyy1 = iy0
                    izz0 = iz0
                    izz1 = iz0

                    if nx > 1:
                        ixx0 = ix0 - 1
                        ixx1 = ix0 + 1
                        dx_1 = 1. / ff.dx
                        dx_2 = 1. / ff.dx**2
                    if ny > 1:
                        iyy0 = iy0 - 1
                        iyy1 = iy0 + 1
                        dy_1 = 1. / ff.dy
                        dy_2 = 1. / ff.dy**2
                    if nz > 1:
                        izz0 = iz0 - 1
                        izz1 = iz0 + 1
                        dz_1 = 1. / ff.dz
                        dz_2 = 1. / ff.dz**2

                    ux = 0.
                    uy = 0.
                    uz = 0.
                    for ix in [ix0, ixx0, ixx1]:
                        weight_x = 0.
                        if ix - ix0 == -1 or ix - ix0 == 1:
                            weight_x = 1.125 - 1.5 * abs(
                                px - ff.x[ix]) * dx_1 + 0.5 * abs(
                                    px - ff.x[ix])**2 * dx_2
                        elif nx != 1:
                            weight_x = 0.75 - (px - ff.x[ix])**2 * dx_2

                        for iy in [iy0, iyy0, iyy1]:
                            weight_y = 0.
                            if iy - iy0 == -1 or iy - iy0 == 1:
                                weight_y = 1.125 - 1.5 * abs(
                                    py - ff.y[iy]) * dy_1 + 0.5 * abs(
                                        py - ff.y[iy])**2 * dy_2
                            elif ny != 1:
                                weight_y = 0.75 - (py - ff.y[iy])**2 * dy_2

                            for iz in [iz0, izz0, izz1]:
                                weight_z = 0.
                                if iz - iz0 == -1 or iz - iz0 == 1:
                                    weight_z = 1.125 - 1.5 * abs(
                                        pz - ff.z[iz]) * dz_1 + 0.5 * abs(
                                            pz - ff.z[iz])**2 * dz_2
                                elif nz != 1:
                                    weight_z = 0.75 - (pz - ff.z[iz])**2 * dz_2

                                weight = 1.
                                if nx > 1: weight = weight * weight_x
                                if ny > 1: weight = weight * weight_y
                                if nz > 1: weight = weight * weight_z

                                ux = ux + weight * ff.ux[iz, iy, ix]
                                uy = uy + weight * ff.uy[iz, iy, ix]
                                uz = uz + weight * ff.uz[iz, iy, ix]

                                if izz0 == izz1: break  # beware of degeneracy:
                            if iyy0 == iyy1: break  # beware of degeneracy:
                        if ixx0 == ixx1: break  # beware of degeneracy:

                    l_ux.append(ux)
                    l_uy.append(uy)
                    l_uz.append(uz)

            ## Convert all information into a single record array
            data_set = np.core.records.fromarrays(
                [
                    l_ipars.astype('int'), l_px, l_py, l_pz, l_vx, l_vy, l_vz,
                    l_rix, l_riy, l_riz, l_ix, l_iy, l_iz, l_ux, l_uy, l_uz
                ],
                names=
                'ipar, ipx, ipy, ipz, vx, vy, vz, rix, riy, riz, ix, iy, iz, ux, uy, uz',
                formats=
                'int, float, float, float, float, float, float, int, int, int, int, int, int, float, float, float'
            )
            gas_velo_at_particle_pos = np.sort(data_set,
                                               order=['ix', 'iy', 'iz'])

            Nix = int(gas_velo_at_particle_pos['rix'].max() + 1)
            Niy = int(gas_velo_at_particle_pos['riy'].max() + 1)
            Niz = int(gas_velo_at_particle_pos['riz'].max() + 1)

            Npar_arr = np.array([
                gas_velo_at_particle_pos['rix'],
                gas_velo_at_particle_pos['riy'],
                gas_velo_at_particle_pos['riz']
            ])
            #rgrid_edges = (grid.x[1:]-(grid.x[1:]-grid.x[:-1])/2)[2:-2]
            xrange = np.arange(
                0,
                float(gas_velo_at_particle_pos['rix'].max()) + 2)
            xrange = xrange - 0.5
            yrange = np.arange(
                0,
                float(gas_velo_at_particle_pos['riy'].max()) + 2)
            zrange = np.arange(
                0,
                float(gas_velo_at_particle_pos['riz'].max()) + 2)

            Npar_hist, edges = np.histogramdd(Npar_arr.T,
                                              bins=(xrange, yrange, zrange))
            Npar_hist, edges = np.histogramdd(Npar_arr.T, bins=(Nix, Niy, Niz))

            gas_velo_at_particle_pos = {
                'time':
                ff.t,
                'par_pos':
                np.array([
                    gas_velo_at_particle_pos['ipx'],
                    gas_velo_at_particle_pos['ipy'],
                    gas_velo_at_particle_pos['ipz']
                ]),
                'par_velo':
                np.array([
                    gas_velo_at_particle_pos['vx'],
                    gas_velo_at_particle_pos['vy'],
                    gas_velo_at_particle_pos['vz']
                ]),
                'par_idx':
                np.array([
                    gas_velo_at_particle_pos['rix'],
                    gas_velo_at_particle_pos['riy'],
                    gas_velo_at_particle_pos['riz']
                ]),
                'npar':
                np.array(Npar_hist[gas_velo_at_particle_pos['rix'],
                                   gas_velo_at_particle_pos['riy'],
                                   gas_velo_at_particle_pos['riz']]),
                'gas_velo':
                np.array([
                    gas_velo_at_particle_pos['ux'],
                    gas_velo_at_particle_pos['uy'],
                    gas_velo_at_particle_pos['uz']
                ])
            }

            print('## Saving dataset into ' + save_destination + '...')
            pkl_save(
                {
                    'gas_velo_at_particle_pos': gas_velo_at_particle_pos,
                    't': ff.t
                },
                save_filename,
                folder=save_destination)
        print('## Done!')
# coding: utf-8

# In[1]:

import pencil as pc

# You may need to compile and run simlation object from below first!

# In[2]:

sim = pc.get_sim('../sample_simulations/2d_streaming_instability/')

# ## Reading simulation parameters the easy way
# Use typ what quantitiy you need, parser does the rest!

# In[3]:

sim.get_value('beta_glnrho_global')

# ## Getting an overview of VAR files and PVAR files

# In[4]:

sim = pc.get_sim('../sample_simulations/2d_streaming_instability/')
varlist = sim.get_varlist()
pvarlist = sim.get_varlist(particle=True)
print(varlist[-1])
print(pvarlist[-1])

# ## Get last 10 VAR files
# Also check out all abilities of get_varlist and come up with your own!
def gas_velo_at_particle_pos(
    varfiles="last4", sim=False, scheme="tsc", use_IDL=False, OVERWRITE=False
):
    """This script calulates the gas velocity at the particle position and stores this together
    with particle position, containing grid cell idicies, particle velocities, and particle index
    in a gas_velo_at_particle_pos file.

    Args:
      varfiles:       specifiy varfiles for calculation, e.g. 'last', 'first',
                          'all', 'VAR###', 'last4', 'first3'
      scheme:         possible are:
                          - ngp: nearest grid point
                          - cic: cloud in cell
                          - tsc: triangular shaped cloud
      OVERWRITE:		set to True to overwrite already calculated results
    """

    import os
    from pencil import io
    from pencil import read
    from pencil import get_sim
    from os.path import exists
    import numpy as np

    GAS_VELO_TAG = "gas_velo_at_particle_pos"

    if sim == False:
        sim = get_sim()
        if sim == False:
            print("! ERROR: Specify simulation object!")
            return False
    SIM = sim

    if use_IDL:
        print(
            "? WARNING: IDL VERSION OF THIS SCRIPT BY JOHANSEN, not recommended for 2D data"
        )
        from pencil.backpack import pidly

        print("## starting IDL engine..")
        IDL = pidly.IDL(long_delay=0.05)  # start IDL engine

        ## skip if nothing is new
        if (
            (not OVERWRITE)
            and (exists(os.path.join(SIM.pc_datadir, "sigma.pkl")))
            and (exists(os.path.join(SIM.pc_datadir, "zeta.pkl")))
        ):
            print(
                "~ "
                + SIM.name
                + " is already calculated and up-to-date! -> skipping it!"
            )

        else:
            ## start calculations
            print(
                '~ Calculating gas_velo_at_particle_pos for "'
                + SIM.name
                + '" in "'
                + SIM.path
                + '"'
            )
            IDL.pro(
                "gas_velo_at_particle_pos",
                datadir=SIM.datadir,
                destination=GAS_VELO_TAG,
                doforthelastNvar=varfiles[4:],
            )
            files = [
                i.split("_")[-1].split(".sav")[0]
                for i in os.listdir(os.path.join(SIM.pendatadir, GAS_VELO_TAG))
                if i.startswith(GAS_VELO_TAG)
                and i.endswith(".sav")
                or i.endswith(".pkl")
            ]
            if files == []:
                print(
                    "!! ERROR: No calc_gas_speed_at_particle_position-files found for "
                    + SIM.name
                    + "! Use idl script to produce them first!"
                )

        IDL.close()
        return True

    else:
        print(
            '~ Calculating gas_velo_at_particle_pos for "'
            + SIM.name
            + '" in "'
            + SIM.path
            + '"'
        )
        save_destination = os.path.join(SIM.pc_datadir, GAS_VELO_TAG)
        io.mkdir(save_destination)
        varlist = SIM.get_varlist(pos=varfiles, particle=False)
        pvarlist = SIM.get_varlist(pos=varfiles, particle=True)

        for f, p in zip(varlist, pvarlist):
            save_filename = GAS_VELO_TAG + "_" + scheme + "_" + f[3:]
            if not OVERWRITE and exists(save_filename, folder=save_destination):
                continue

            print("## Reading " + f + " ...")
            ff = read.var(datadir=SIM.datadir, varfile=f, quiet=True, trimall=False)
            pp = read.pvar(datadir=SIM.datadir, varfile=p)

            ## remove ghost zones from grid, call the reduced grid the "real grid"
            realgridx = ff.x[ff.l1 : ff.l2]
            realgridy = ff.y[ff.m1 : ff.m2]
            realgridz = ff.z[ff.n1 : ff.n2]
            nx = ff.l2 - ff.l1
            ny = ff.m2 - ff.m1
            nz = ff.n2 - ff.n1

            ## prepare list for all quantities
            l_ipars = pp.ipars  # particle number   KNOWN
            l_px = pp.xp
            l_py = pp.yp
            l_pz = pp.zp  # particle absolut position KNOWN
            l_vx = pp.vpx
            l_vy = pp.vpy
            l_vz = pp.vpz  # particle velocity KNOWN
            l_rix = []
            l_riy = []
            l_riz = (
                []
            )  # particle untrimmed realgrid index (grid index = l/m/n + readgrid index ???)
            l_ix = []
            l_iy = []
            l_iz = []  # particle grid index (in untrimmed grid)
            l_ux = []
            l_uy = []
            l_uz = []  # underlying gas velocity at position of particle

            ## get index of realgrid cell for each particle
            for i in range(len(l_ipars)):
                l_rix.append(np.abs(realgridx - l_px[i]).argmin())
                l_riy.append(np.abs(realgridy - l_py[i]).argmin())
                l_riz.append(np.abs(realgridz - l_pz[i]).argmin())

            ## convert into untrimmed grid
            l_ix = np.array(l_rix) + ff.l1
            l_iy = np.array(l_riy) + ff.m1
            l_iz = np.array(l_riz) + ff.n1

            ## NGP
            if scheme == "ngp" or scheme == "NGP":
                print("## Calculating gas velocities via " + scheme)
                l_ux = ff.ux[l_iz, l_iy, l_ix]
                l_uy = ff.uy[l_iz, l_iy, l_ix]
                l_uz = ff.uz[l_iz, l_iy, l_ix]

            ## CIC
            if scheme == "cic" or scheme == "CIC":
                print("## Calculating gas velocities via " + scheme)
                for ix0, iy0, iz0, px, py, pz in zip(
                    l_ix, l_iy, l_iz, l_px, l_py, l_pz
                ):  # for each particle
                    if ff.x[ix0] > px:
                        ix0 = ix0 - 1  # ix0 must be left to particle
                    if ff.y[iy0] > py:
                        iy0 = iy0 - 1  # iy0 must be below the particle
                    if ff.z[iz0] > pz:
                        iz0 = iz0 - 1  # iz0 must be under particle

                    ix1 = ix0
                    iy1 = iy0
                    iz1 = iz0  # if a dim. is zero, this is default, else:
                    if nx > 1:
                        ix1 = ix0 + 1
                        dx_1 = (
                            1.0 / ff.dx
                        )  # if a dim is non-zero, ajust ix1 to right cell
                    if ny > 1:
                        iy1 = iy0 + 1
                        dy_1 = (
                            1.0 / ff.dy
                        )  # if a dim is non-zero, ajust iy1 to above cell
                    if nz > 1:
                        iz1 = iz0 + 1
                        dz_1 = (
                            1.0 / ff.dz
                        )  # if a dim is non-zero, ajust iz1 to above cell

                    ux = 0.0
                    uy = 0.0
                    uz = 0.0
                    for ix in [ix0, ix1]:
                        for iy in [iy0, iy1]:
                            for iz in [iz0, iz1]:
                                weight = 1.0
                                if nx > 1:
                                    weight = weight * (1.0 - abs(px - ff.x[ix]) * dx_1)
                                if ny > 1:
                                    weight = weight * (1.0 - abs(py - ff.y[iy]) * dy_1)
                                if nz > 1:
                                    weight = weight * (1.0 - abs(pz - ff.z[iz]) * dz_1)

                                ux = ux + weight * ff.ux[iz, iy, ix]
                                uy = uy + weight * ff.uy[iz, iy, ix]
                                uz = uz + weight * ff.uz[iz, iy, ix]

                                if iz0 == iz1:
                                    break  # beware of degeneracy:
                            if iy0 == iy1:
                                break  # beware of degeneracy:
                        if ix0 == ix1:
                            break  # beware of degeneracy:

                    l_ux.append(ux)
                    l_uy.append(uy)
                    l_uz.append(uz)

            ## TSC
            if scheme == "tsc" or scheme == "TSC":
                for ix0, iy0, iz0, px, py, pz in zip(
                    l_ix, l_iy, l_iz, l_px, l_py, l_pz
                ):  # for each particle
                    ixx0 = ix0
                    ixx1 = ix0  # beware of degeneracy
                    iyy0 = iy0
                    iyy1 = iy0
                    izz0 = iz0
                    izz1 = iz0

                    if nx > 1:
                        ixx0 = ix0 - 1
                        ixx1 = ix0 + 1
                        dx_1 = 1.0 / ff.dx
                        dx_2 = 1.0 / ff.dx ** 2
                    if ny > 1:
                        iyy0 = iy0 - 1
                        iyy1 = iy0 + 1
                        dy_1 = 1.0 / ff.dy
                        dy_2 = 1.0 / ff.dy ** 2
                    if nz > 1:
                        izz0 = iz0 - 1
                        izz1 = iz0 + 1
                        dz_1 = 1.0 / ff.dz
                        dz_2 = 1.0 / ff.dz ** 2

                    ux = 0.0
                    uy = 0.0
                    uz = 0.0
                    for ix in [ix0, ixx0, ixx1]:
                        weight_x = 0.0
                        if ix - ix0 == -1 or ix - ix0 == 1:
                            weight_x = (
                                1.125
                                - 1.5 * abs(px - ff.x[ix]) * dx_1
                                + 0.5 * abs(px - ff.x[ix]) ** 2 * dx_2
                            )
                        elif nx != 1:
                            weight_x = 0.75 - (px - ff.x[ix]) ** 2 * dx_2

                        for iy in [iy0, iyy0, iyy1]:
                            weight_y = 0.0
                            if iy - iy0 == -1 or iy - iy0 == 1:
                                weight_y = (
                                    1.125
                                    - 1.5 * abs(py - ff.y[iy]) * dy_1
                                    + 0.5 * abs(py - ff.y[iy]) ** 2 * dy_2
                                )
                            elif ny != 1:
                                weight_y = 0.75 - (py - ff.y[iy]) ** 2 * dy_2

                            for iz in [iz0, izz0, izz1]:
                                weight_z = 0.0
                                if iz - iz0 == -1 or iz - iz0 == 1:
                                    weight_z = (
                                        1.125
                                        - 1.5 * abs(pz - ff.z[iz]) * dz_1
                                        + 0.5 * abs(pz - ff.z[iz]) ** 2 * dz_2
                                    )
                                elif nz != 1:
                                    weight_z = 0.75 - (pz - ff.z[iz]) ** 2 * dz_2

                                weight = 1.0
                                if nx > 1:
                                    weight = weight * weight_x
                                if ny > 1:
                                    weight = weight * weight_y
                                if nz > 1:
                                    weight = weight * weight_z

                                ux = ux + weight * ff.ux[iz, iy, ix]
                                uy = uy + weight * ff.uy[iz, iy, ix]
                                uz = uz + weight * ff.uz[iz, iy, ix]

                                if izz0 == izz1:
                                    break  # beware of degeneracy:
                            if iyy0 == iyy1:
                                break  # beware of degeneracy:
                        if ixx0 == ixx1:
                            break  # beware of degeneracy:

                    l_ux.append(ux)
                    l_uy.append(uy)
                    l_uz.append(uz)

            ## Convert all information into a single record array
            data_set = np.core.records.fromarrays(
                [
                    l_ipars.astype("int"),
                    l_px,
                    l_py,
                    l_pz,
                    l_vx,
                    l_vy,
                    l_vz,
                    l_rix,
                    l_riy,
                    l_riz,
                    l_ix,
                    l_iy,
                    l_iz,
                    l_ux,
                    l_uy,
                    l_uz,
                ],
                names="ipar, ipx, ipy, ipz, vx, vy, vz, rix, riy, riz, ix, iy, iz, ux, uy, uz",
                formats="int, float, float, float, float, float, float, int, int, int, int, int, int, float, float, float",
            )
            gas_velo_at_particle_pos = np.sort(data_set, order=["ix", "iy", "iz"])

            Nix = int(gas_velo_at_particle_pos["rix"].max() + 1)
            Niy = int(gas_velo_at_particle_pos["riy"].max() + 1)
            Niz = int(gas_velo_at_particle_pos["riz"].max() + 1)

            Npar_arr = np.array(
                [
                    gas_velo_at_particle_pos["rix"],
                    gas_velo_at_particle_pos["riy"],
                    gas_velo_at_particle_pos["riz"],
                ]
            )
            # rgrid_edges = (grid.x[1:]-(grid.x[1:]-grid.x[:-1])/2)[2:-2]
            xrange = np.arange(0, float(gas_velo_at_particle_pos["rix"].max()) + 2)
            xrange = xrange - 0.5
            yrange = np.arange(0, float(gas_velo_at_particle_pos["riy"].max()) + 2)
            zrange = np.arange(0, float(gas_velo_at_particle_pos["riz"].max()) + 2)

            Npar_hist, edges = np.histogramdd(Npar_arr.T, bins=(xrange, yrange, zrange))
            Npar_hist, edges = np.histogramdd(Npar_arr.T, bins=(Nix, Niy, Niz))

            gas_velo_at_particle_pos = {
                "time": ff.t,
                "par_pos": np.array(
                    [
                        gas_velo_at_particle_pos["ipx"],
                        gas_velo_at_particle_pos["ipy"],
                        gas_velo_at_particle_pos["ipz"],
                    ]
                ),
                "par_velo": np.array(
                    [
                        gas_velo_at_particle_pos["vx"],
                        gas_velo_at_particle_pos["vy"],
                        gas_velo_at_particle_pos["vz"],
                    ]
                ),
                "par_idx": np.array(
                    [
                        gas_velo_at_particle_pos["rix"],
                        gas_velo_at_particle_pos["riy"],
                        gas_velo_at_particle_pos["riz"],
                    ]
                ),
                "npar": np.array(
                    Npar_hist[
                        gas_velo_at_particle_pos["rix"],
                        gas_velo_at_particle_pos["riy"],
                        gas_velo_at_particle_pos["riz"],
                    ]
                ),
                "gas_velo": np.array(
                    [
                        gas_velo_at_particle_pos["ux"],
                        gas_velo_at_particle_pos["uy"],
                        gas_velo_at_particle_pos["uz"],
                    ]
                ),
            }

            print("## Saving dataset into " + save_destination + "...")
            io.pkl_save(
                {"gas_velo_at_particle_pos": gas_velo_at_particle_pos, "t": ff.t},
                save_filename,
                folder=save_destination,
            )
        print("## Done!")
def part_to_grid(
    xp,
    yp,
    zp=False,
    quantity=False,
    Nbins=[1024, 1024, 1024],
    sim=False,
    extent=False,
    fill_gaps=False,
):
    """Bins quantity based on position data xp and yp to 1024^2 bins like a histrogram.
    This method is not using TSC.

    Args:
        - xp, yp:       array of x and y positions
        - zp:           specify if 3D run, set False to have zp == 0
        - quantity:     array of same shape as xp and yp, but with quantity to bin, set it False to count number of occurrences/histrogram2d
        - Nbins:        number of histrogram bins for each direction. if 2d only the first two entries in Nbins are used
        - sim:          to extract extent from by loading ghostzone-free grid
        - extent:       [[xmin, xmax],[ymin, ymax]] or set false and instead give a sim
                        set extent manually e.g. if you want to include ghost zones
        - fill_gaps     interpolate empty grid cells

    Returns: arr, xgrid, ygrid
        - arr:          2d array with binned values
        - x-/ygrid:     linspace of used x/y grid
        - zgrid:        if zp != False

    Example:
        vpx = part_to_grid_2d(pvar.xp, pvar.yp, pvar.vpx), notice that this will execute pc.get_sim() internally to get the extent
    """

    import numpy as np
    from pencil import get_sim
    from pencil.calc import fill_gaps_in_grid

    if not xp.shape == yp.shape and yp.shape == xp.shape and yp.shape == zp.shape:
        print("! ERROR: Shape of xp, yp, zp and quantity needs to be equal!")

    if extent == False and sim == False:
        sim = get_sim()

    if quantity == False:
        quantity = xp / xp

    if extent == False:
        grid = sim.grid
        if type(zp) == type(False) and zp == False:
            extent = [
                [grid.x[0] - grid.dx / 2, grid.x[-1] + grid.dx / 2],
                [grid.y[0] - grid.dy / 2, grid.y[-1] + grid.dy / 2],
            ]
        else:
            extent = [
                [grid.x[0] - grid.dx / 2, grid.x[-1] + grid.dx / 2],
                [grid.y[0] - grid.dy / 2, grid.y[-1] + grid.dy / 2],
                [grid.z[0] - grid.dz / 2, grid.z[-1] + grid.dz / 2],
            ]

    if type(zp) == type(False) and zp == False:
        arr = np.zeros((2, Nbins[0], Nbins[1]))
    else:
        arr = np.zeros((3, Nbins[0], Nbins[1], Nbins[2]))

    arr[:] = np.NAN
    xgrid = (np.linspace(extent[0][0], extent[0][1], num=Nbins[0] + 1)[:-1] +
             np.linspace(extent[0][0], extent[0][1], num=Nbins[0] + 1)[1:]) / 2
    ygrid = (np.linspace(extent[1][0], extent[1][1], num=Nbins[1] + 1)[:-1] +
             np.linspace(extent[1][0], extent[1][1], num=Nbins[1] + 1)[1:]) / 2
    if type(zp) == type(False) and zp == False:
        for x, y, q in zip(xp, yp, quantity):
            idx = np.argmin(np.abs(x - xgrid))
            idy = np.argmin(np.abs(y - ygrid))
            if np.isnan(arr[0, idx, idy]):
                arr[0, idx, idy] = 0
                arr[1, idx, idy] = 0
            arr[0, idx, idy] += q
            arr[1, idx, idy] += 1

    else:
        zgrid = (
            np.linspace(extent[2][0], extent[2][1], num=Nbins[2] + 1)[:-1] +
            np.linspace(extent[2][0], extent[2][1], num=Nbins[2] + 1)[1:]) / 2
        for x, y, z, q in zip(xp, yp, zp, quantity):
            idx = np.argmin(np.abs(x - xgrid))
            idy = np.argmin(np.abs(y - ygrid))
            idz = np.argmin(np.abs(z - zgrid))
            if np.isnan(arr[0, idx, idy, idz]):
                arr[0, idx, idy, idz] = 0
                arr[1, idx, idy, idz] = 0
            arr[0, idx, idy, idz] += q
            arr[1, idx, idy, idz] += 1

    arr = arr[0] / arr[1]

    if fill_gaps == True:
        arr = fill_gaps_in_grid(arr, key=np.NAN)

    if type(zp) == type(False) and zp == False:
        return arr, xgrid, ygrid
    else:
        return arr, xgrid, ygrid, zgrid
def get_value_from_file(filename,
                        quantity,
                        change_quantity_to=None,
                        sim=False,
                        filepath=False,
                        DEBUG=False,
                        silent=False):
    """ Use to read in a quantity from
        - *.in
        - *.local
        - submit*, i.e. submit.sh, submit.csh, files, only works if computer is readily specified in pc.io.get_systemid

    Please add further functionallity by yourself!

    Args:
        filename:   can be "run.in", "start.in", "cparam.local", path to that file is extraced from filepath or sim object
        quantity:   variable to read in from file
        sim:        put simulation object here, file will be found by filename automatically
        filepath:   normally not needed, specify here where to find the file with filename, can be a list of paths if unshure
        DEBUG:      make dry run, tell me what you would do but dont change anything!
        silent:     suppress certain output by setting True

    Return:
        Returns None if not successful
    """

    import os
    import numpy as np
    from os.path import join, abspath, exists, split, isfile
    from pencil import get_sim
    from pencil.math import is_number, is_float, is_int
    from pencil.io import timestamp, debug_breakpoint, mkdir
    import re
    import copy

    def string_to_tuple(s):
        q = s.split(',')

        if is_number(q[0]):
            q = np.array([float(t) for t in q])
            q_type = 'TUPLE_FLOAT'
            return q, q_type

        if q[0] == 'T' or q[0] == 'F':
            q = np.array([bool(t == 'T') for t in q])
            q_type = 'TUPLE_BOOL'
            return q, q_type

        if type(q[0]) == type('string'):
            q = [t.replace('"', '').replace("'", '') for t in q]
            q_type = 'TUPLE_STRING'
            return q, q_type

        print('! ERROR: Could not parse string ' + s + ' into a tuple!')
        print(
            '! DEBUG_BREAKPOINT AKTIVATED - check out the following variables: string s, tuple q, first entry in tuple q[0]'
        )
        debug_breakpoint()
        return None, None

    def tuple_to_string(t, q_type):
        return ','.join([str(a) for a in t])

    ######## prepare filename and quantity
    filename = filename.strip()  # get rid of whitespaces
    quantity = quantity.strip()
    q_type = False  # q_type will store the type of the quantity value once found and identified

    split_filename = split(filename)
    if sim == False and split_filename[0] != '' and filepath == False:
        filepath = split_filename[0]
        filename = split_filename[1]

    ######## find correct file
    # prepare search_path list to search filename in
    if filepath == False:
        if sim == False:
            sim = get_sim()
        else:
            filepath = sim.path
        search_paths = [sim.path, join(sim.path,
                                       'src')]  # add other search paths here!!

    elif type(filepath) == type('string'):
        if filepath.endswith(filename):
            filepath = filepath[:-len(
                filename
            )]  # clean filepath if filename occures to be in there at the end
        search_paths = [abspath(filepath.strip())]  # correct path format

    elif type(filepath) == type(['list']):
        search_paths = filepath

    else:
        print('! ERROR: Filename ' + str(filename) +
              ' could not be interprated or found!')
        return None

    absolute_filepath = None
    for search_path in search_paths:
        tmp_path = join(search_path, filename)
        if os.path.isfile(tmp_path):
            absolute_filepath = tmp_path
            break

    # Traps the case of not being able to find the file
    if absolute_filepath is None:
        if DEBUG:
            print('~ DEBUG: File {0} not found in {1}!'.format(
                filename, search_paths))
        return None

    ######## open file
    # now having absolute filepath to file, lets check that file and find quantity inside!
    if DEBUG:
        print('~ DEBUG: Found file {0} in {1}'.format(filename, filepath))

    with open(absolute_filepath, 'r') as f:
        data_raw = f.readlines()

    ######## find line in file which quantity in
    line_matches = []
    # scan through file for differently for different files
    if filename.endswith(
            '.in') or 'cparam.local' or 'Makefile.local' in filename:
        FILE_IS = 'IN_LOCAL'
        SYM_COMMENT = '!'
        SYM_ASSIGN = '='
        SYM_SEPARATOR = ','

        for ii, line in enumerate(data_raw):
            if line.strip().startswith('&'):
                continue  # filter out lines with &something, e.g. &density_run_pars
            quantity_match_tmp = re.search(
                '[^0-9a-zA-Z_]*{0}[^0-9a-zA-Z_]'.format(quantity),
                line.split(SYM_COMMENT)[0])
            # Check if this substring occurs as a string.
            if quantity_match_tmp:
                quantity_match = quantity_match_tmp
                if str.count(line[0:quantity_match.start()], "'") % 2 == 0 and \
                   str.count(line[0:quantity_match.start()], '"') % 2 == 0:
                    if line_matches:
                        line_matches[0] = ii
                    else:
                        line_matches.append(ii)

    elif filename.startswith('submit') and filename.split('.')[-1] in [
            'csh', 'sh'
    ]:
        FILE_IS = 'SUBMIT'
        SYM_COMMENT = False
        SYM_ASSIGN = '='
        SYM_SEPARATOR = ','

        for ii, line in enumerate(data_raw):
            if line.replace(' ', '').startswith('#@') and quantity in line:
                quantity_match_tmp = re.search(
                    '[^0-9a-zA-Z_]*{0}[^0-9a-zA-Z_]'.format(quantity),
                    line.split(SYM_COMMENT)[0])
                if quantity_match_tmp:
                    quantity_match = quantity_match_tmp
                    if line_matches:
                        line_matches[0] = ii
                    else:
                        line_matches.append(ii)
    else:
        print(
            '! ERROR: Filename unknown! No parsing possible! Please enhance this function to work with '
            + filename)

    if len(line_matches) > 1:
        print('! ERROR: Found more than one line with keyword "' + quantity +
              '" inside!')
        return None
    if len(line_matches) == 0:
        if silent == False:
            print('! ERROR: Found no line with keyword "' + quantity +
                  '" inside ' + join(filepath, filename) + '!')
        return None

    filename = os.path.basename(filename)

    ######## get line with quantity inside
    line = data_raw[line_matches[0]]

    ######## do separation of quantity from rest of line, i.e. get rid of comments and other quantities defined in this line
    comment = ''
    if SYM_COMMENT:
        tmp = line.partition(SYM_COMMENT)  # strip away comment
        line = tmp[0]
        if tmp[-1] != '':
            comment = SYM_COMMENT + tmp[-1]  # and store for later

#    line = line.replace(' ','').replace('\n', '')                               # do cleanup in this line

# Find the position where the quantity is stored.
    pos_equal_sign_left = quantity_match.end() + str.find(
        line[quantity_match.end() - 1:], '=')
    #    pos_equal_sign_right = pos_equal_sign_left + str.find(line[pos_equal_sign_left:], ', *[0-9a-zA-Z][0-9a-zA-Z]* *= *[0-9a-zA-Z][0-9a-zA-Z]*')
    pos_equal_sign_right = pos_equal_sign_left + str.find(
        line[pos_equal_sign_left:], '=')
    if pos_equal_sign_right < pos_equal_sign_left:
        pos_equal_sign_right = -1
        pos_right_comma = -1
    else:
        pos_right_comma = str.rfind(line[:pos_equal_sign_right], ',')

    # Change the quantity in the line string.
    q = copy.copy(line[pos_equal_sign_left:pos_right_comma])
    #    qs = line.partition(quantity+SYM_ASSIGN)
    #    if SYM_ASSIGN in qs[-1]:
    #        qs = qs[:2]+qs[-1].partition(SYM_ASSIGN)
    #        #qs = qs[:2]+qs[-1].partition(SYM_ASSIGN)
    #        qs = qs[:2]+qs[2].rpartition(',')+qs[3:]
    #
    #    qs = list(qs)
    #    q = qs[2]

    while q.endswith('\t'):
        q = q[:-1]
        comment = '\t' + comment  # take care of trailing tabulator
    while q.endswith(','):
        q = q[:-1]  # remove trailing ,

    ######## do a cleanup of quantity value q and convert into string, float, int or array, also remember data type of q
    if q.startswith("'") and q.endswith(
            "'"):  # quantity q is string in form of 'STRING'
        q = q[1:-1]
        q_type = 'STRING'

    elif q.startswith('"') and q.endswith(
            '"'):  # quantity q is string in form of "STRING"
        q = q[1:-1]
        q_type = 'STRING'

    elif not is_number(
            q[0]):  # quantity q is string in form of not beeing a number
        q = q.strip().replace('"', '').replace("'", '')
        q_type = 'STRING'

    try:
        float(q)
        q_type = 'FLOAT'
        if is_int(q):
            q = int(q)
            q_type = 'INT'
    except:
        if type(q) == type('string') and ',' in q:
            q, q_type = string_to_tuple(q)  # q is a TULPE_something
            print('q = {0}, q_type = {1}'.format(q, q_type))

        if type(q) == type('string') and q in ['F', 'f']:  # q is BOOL
            q = False
            q_type = 'BOOL'

        if type(q) == type('string') and q in ['T', 't']:
            q = True
            q_type = 'BOOL'

        if type(q) == type('string'):
            if is_number(q[0]):
                q_type = 'STRING'

    if q_type == False:  # catch if type of q was not recognized
        print(
            '! ERROR: Couldnt identify the data type of the quantity value: ' +
            str(q))
        DEBUG = True
        debug_breakpoint()
    elif DEBUG:
        print('~ DEBUG: q_type = ' + q_type)
    if q_type == 'FLOAT':
        q = float(q)
    elif q_type == 'INT':
        q = int(q)

    ######## if value of quantity has to be changed do:
    if change_quantity_to != None:

        ####### prepare change_quantity_to for string injection
        if q_type == 'STRING':
            if not FILE_IS == 'SUBMIT':
                change_quantity_to = "'" + change_quantity_to + "'"

        elif q_type == 'BOOL':
            change_quantity_to = bool(change_quantity_to in ['T', 't', True])
            if change_quantity_to == True:
                change_quantity_to = 'T'
            elif change_quantity_to == False:
                change_quantity_to = 'F'
            else:
                print(
                    '! ERROR: There is something deeply wrong here! change_quantity_to should be bool...'
                )
                debug_breakpoint()
                return None

        elif q_type == 'FLOAT':
            change_quantity_to = '%e' % change_quantity_to

        elif q_type.startswith('TUPLE'):
            if q_type.endswith('BOOL'):
                if type(change_quantity_to) == type(
                    ['list', 'of', 'bool', 'or', 'strings']):
                    for ii, val in enumerate(change_quantity_to):
                        if val in ['T', 't', True]:
                            change_quantity_to[ii] = 'T'
                        elif val in ['F', 'f', False]:
                            change_quantity_to[ii] = 'F'
                        else:
                            print(
                                '! ERROR: There is something deeply wrong here! change_quantity_to['
                                + str(ii) +
                                '] should be bool or string representation, but it is '
                                + str(change_quantity_to[ii]))
                            debug_breakpoint()
                            return None
                change_quantity_to = ','.join(
                    [str(t) for t in change_quantity_to])
            if q_type.endswith('FLOAT'):
                change_quantity_to = str(list(change_quantity_to))[1:-1]
            if q_type.endswith('STRING'):
                change_quantity_to = str(list(change_quantity_to))[1:-1]

        if DEBUG:
            print('~ DEBUG: Would change quantity ' + quantity + ' from ' +
                  str(q) + ' to ' + str(change_quantity_to))
        q = str(change_quantity_to)

        ######## further formatting
        new_line = line[:pos_equal_sign_left] + q + line[
            pos_right_comma:] + '\t' + comment  # create new line and add comment stripped away before
        #        new_line = ''.join(qs).replace(SYM_SEPARATOR, SYM_SEPARATOR+' ')+'\t'+comment    # create new line and add comment stripped away before
        new_line = new_line.rstrip(
        )  # clean empty spaces on the right, no one needs that...
        if new_line[-1] != '\n': new_line = new_line + '\n'
        if FILE_IS == 'SUBMIT':
            new_line = new_line.replace('#@', '#@ ').replace(
                '=', ' = ')  # optimizing format of submit script

        if DEBUG:
            print('~ DEBUG: old line: ' + str(data_raw[line_matches[0]])[:-1])
            print('~ DEBUG: new line: ' + str(new_line)[:-1])

        if not DEBUG:
            ####### do backup of file before changing it
            from shutil import copyfile
            target = join(sim.path, 'pc/backups/' + timestamp())
            mkdir(target)
            target = join(target, filename)
            copyfile(absolute_filepath, target)

            # replace line in raw data
            data_raw[line_matches[0]] = new_line

            # save on drive
            f.close()
            with open(absolute_filepath, 'w') as f:
                for l in data_raw:
                    f.write(l)

    ######## DONE!
    return q
Exemple #8
0
    def __init__(self,
                 datadir=False,
                 sim=False,
                 tmin=0,
                 tmax=-1,
                 noutmax='-1',
                 swap_endian=False,
                 quiet=False,
                 use_existing_pstalk_sav=False):
        """
        Read PSTALK files from Pencil Code using IDL.
        Uses IDL<->Python Bridge, this must be activated manually!

        Args:
            - datadir      specify datadir, default False
            - sim           specify simulation from which you want to read
            - swap_endian   change if needed to True, default False
            - quiet         verbosity, default False
            - use_existing_pstalk_sav
                            use existing <sim.datadir>/data/pc/tmp/pstalk.sav for speed up


        """

        import numpy as np
        import os
        from os.path import join
        from pencil import get_sim

        if datadir == False:
            if sim == False:
                sim = get_sim()
            datadir = sim.datadir

        if quiet == False:
            quiet = '0'
        else:
            quiet = '1'

        if swap_endian == False:
            swap_endian = '0'
        else:
            swap_endian = '1'

        if use_existing_pstalk_sav == True:
            from scipy.io.idl import readsav

            print('~ reading existing pstalk..')

            ps = readsav(join(sim.pc_datadir, 'tmp', 'pstalk.sav'))('pstalk')

            for key in set(ps.dtype.fields.keys()):
                if hasattr(self, key.lower()): continue
                setattr(self, key.lower(), ps[key][0].T)

        else:
            try:
                cwd = os.getcwd()
                from idlpy import IDL
                os.chdir(cwd)

                print('~ reading pstalk in IDL..')

                idl_call = ', '.join([
                    'pc_read_pstalk', 'obj=pstalk',
                    'datadir="' + datadir + '"', 'it0=' + str(tmin),
                    'it1=' + str(tmax), 'quiet=' + quiet,
                    'swap_endian=' + swap_endian, 'noutmax=' + str(noutmax)
                ])

                IDL.run(idl_call)

                print('~ parsing pstalk from IDL to python..')
                ps = IDL.pstalk

                for key in set(ps.keys()):
                    if hasattr(self, key.lower()): continue
                    setattr(self, key.lower(), ps[key].T)

            except:
                print(
                    '! ERROR: no idl<->python bridge found. Try whats written in pstalk-comment to fix that issue.'
                )
                print('! ')
                print(
                    '! Use something like: (ensure you have IDL 8.5.1 or larger)'
                )
                print(
                    '! export PYTHONPATH=$PYTHONPATH:$IDL_HOME/lib/bridges:$IDL_HOME/bin/bin.linux.x86_64'
                )
                print(
                    '! export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib64:$IDL_HOME/bin/bin.linux.x86_64'
                )
                print('! in your .bashrc')
                print('! ')
                print(
                    '! If you have it already installed, try: from idlpy import IDL and check for errors'
                )
                print('! ')
                print(
                    '~ BACKUP SOLUTION: reading pstalk via pidly, starting IDL..'
                )

                from pencil.backpack import pidly
                IDL = pidly.IDL(long_delay=0.05)  # start IDL engine
                from scipy.io.idl import readsav
                from pencil.io import mkdir

                ## read tstalk file
                print('## reading particle stalker file..')
                IDL('pc_read_pstalk, object=pstalk, datadir="' + sim.datadir +
                    '"' + ', quiet=' + quiet + ', it0=' + str(tmin) +
                    ', it1=' + str(tmax) + ', swap_endian=' + swap_endian)

                print('## transfering pstalk file from IDL to python..')
                mkdir(join(sim.pc_datadir, 'tmp'))
                IDL('save, pstalk, filename="' +
                    join(sim.pc_datadir, 'tmp', 'pstalk_' + str(tmin) + '_' +
                         str(tmax) + '.sav') + '"')
                ps = readsav(join(sim.pc_datadir, 'tmp',
                                  'pstalk.sav'))('pstalk')

                #from pc.io import debug_breakpoint; debug_breakpoint()

                for key in set(ps.dtype.fields.keys()):
                    if hasattr(self, key.lower()): continue
                    setattr(self, key.lower(), ps[key][0].T)