<commit_before># coding: utf-8

from __future__ import division, print_function

__author__ = "adrn <<EMAIL>>"

# Standard library
import os
import sys

# Third-party
from astropy import log as logger
import numpy as np

# Project
import gary.potential as gp
from gary.units import galactic
from gary.util import get_pool

from streammorphology.initialconditions import loop_grid, box_grid
from streammorphology.util import worker
from streammorphology.util import _shape

base_path = os.path.split(os.path.split(os.path.abspath(__file__))[0])[0]

def main(E, loopbox, mpi=False, overwrite=False, ngrid=None, disk=False):
    np.random.seed(42)

    if disk:
        potential = gp.OblateMWPotential()
    else:
        potential = gp.LeeSutoTriaxialNFWPotential(v_c=0.239225, r_s=30.,
                                                   a=1., b=0.75, c=0.55,
                                                   units=galactic)

    # get a pool object for multiprocessing / MPI
    pool = get_pool(mpi=mpi)
    if mpi:
        logger.info("Using MPI")

    path = os.path.join(base_path, 'output',
                        'E{:.3f}_{}_{}'.format(E, potential.__class__.__name__, loopbox))
    logger.info("Caching to: {}".format(path))
    allfreqs_filename = os.path.join(path, "allfreqs.dat")
    if not os.path.exists(path):
        os.mkdir(path)

    # initial conditions
    if loopbox == 'loop':
        w0 = loop_grid(E, potential, Naxis=ngrid)
    else:
        w0 = box_grid(E, potential, Ntotal=ngrid*ngrid)

    norbits = len(w0)
    logger.info("Number of orbits: {}".format(norbits))

    if os.path.exists(allfreqs_filename):
        if overwrite:
            os.remove(allfreqs_filename)
        else:
            pool.close()
            return

    allfreqs_shape = (norbits,) + _shape
    d = np.memmap(allfreqs_filename, mode='w+', dtype='float64', shape=allfreqs_shape)

    # save the initial conditions
    w0_filename = os.path.join(path, 'w0.npy')
    np.save(w0_filename, w0)

    tasks = [dict(index=i, w0_filename=w0_filename,
                  allfreqs_filename=allfreqs_filename,
                  potential=potential) for i in range(norbits)]
    pool.map(worker, tasks)

    pool.close()

if __name__ == '__main__':
    from argparse import ArgumentParser
    import logging

    # Define parser object
    parser = ArgumentParser(description="")
    parser.add_argument("-v", "--verbose", action="store_true", dest="verbose",
                        default=False, help="Be chatty! (default = False)")
    parser.add_argument("-q", "--quiet", action="store_true", dest="quiet",
                        default=False, help="Be quiet! (default = False)")
    parser.add_argument("-o", "--overwrite", action="store_true", dest="overwrite",
                        default=False, help="DESTROY. DESTROY. (default = False)")

    parser.add_argument("-E", "--energy", dest="energy", type=float, required=True,
                        help="Energy of the orbits.")<commit_msg>only do ones that aren't done<commit_after># coding: utf-8

from __future__ import division, print_function

__author__ = "adrn <<EMAIL>>"

# Standard library
import os
import sys

# Third-party
from astropy import log as logger
import numpy as np

# Project
import gary.potential as gp
from gary.units import galactic
from gary.util import get_pool

from streammorphology.initialconditions import loop_grid, box_grid
from streammorphology.util import worker
from streammorphology.util import _shape

base_path = os.path.split(os.path.split(os.path.abspath(__file__))[0])[0]

def main(E, loopbox, mpi=False, overwrite=False, ngrid=None, disk=False):
    np.random.seed(42)

    if disk:
        potential = gp.OblateMWPotential()
    else:
        potential = gp.LeeSutoTriaxialNFWPotential(v_c=0.239225, r_s=30.,
                                                   a=1., b=0.75, c=0.55,
                                                   units=galactic)

    # get a pool object for multiprocessing / MPI
    pool = get_pool(mpi=mpi)
    if mpi:
        logger.info("Using MPI")

    path = os.path.join(base_path, 'output',
                        'E{:.3f}_{}_{}'.format(E, potential.__class__.__name__, loopbox))
    logger.info("Caching to: {}".format(path))
    allfreqs_filename = os.path.join(path, "allfreqs.dat")
    if not os.path.exists(path):
        os.mkdir(path)

    # initial conditions
    if loopbox == 'loop':
        w0 = loop_grid(E, potential, Naxis=ngrid)
    else:
        w0 = box_grid(E, potential, Ntotal=ngrid*ngrid)

    norbits = len(w0)
    logger.info("Number of orbits: {}".format(norbits))

    # save the initial conditions
    w0_filename = os.path.join(path, 'w0.npy')
    np.save(w0_filename, w0)

    if os.path.exists(allfreqs_filename) and overwrite:
        os.remove(allfreqs_filename)

    allfreqs_shape = (norbits,) + _shape
    if not os.path.exists(allfreqs_filename):
        d = np.memmap(allfreqs_filename, mode='w+', dtype='float64', shape=allfreqs_shape)
        tasks = [dict(index=i, w0_filename=w0_filename,
                      allfreqs_filename=allfreqs_filename,
                      potential=potential) for i in range(norbits)]

    else:
        d = np.memmap(allfreqs_filename, mode='r', dtype='float64', shape=allfreqs_shape)
        not_done = np.where(d[:,0,7] != 1)[0]
        tasks = [dict(index=i, w0_filename=w0_filename,
                      allfreqs_filename=allfreqs_filename,
                      potential=potential) for i in not_done]

    pool.map(worker, tasks)
    pool.close()

if __name__ == '__main__':
    from argparse import ArgumentParser
    import logging

    # Define parser object
    parser = ArgumentParser(description="")
    parser.add_argument("-v", "--verbose", action="store_true", dest="verbose",
                        default=False, help="Be chatty! (default = False)")
    parser.add_argument("-q", "--quiet", action="store_true", dest="quiet",
                        default=False, help="Be quiet! (default = False)")
    parser.add_argument("-o", "--overwrite", action="store_true", dest="overwrite",
                        default=False, help="DESTROY. DESTROY. (default = False)")

    parser.add_argument("-E", "--energy", dest="energy", type=float, required=True,
                        help="Energy of the orbits.")
{'fraction_non_alphanumeric': 0.12521711669035213, 'fraction_numerical': 0.01168482551713248, 'mean_word_length': 2.6932944606413995, 'pattern_counts': {'":': 0, '<': 7, '<?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': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1259450', 'n_tokens_mistral': 2067, 'n_tokens_neox': 1965, 'n_words': 446}