示例#1
0
def test_unchecked_remove():
    """Also tests mkdir_parent"""
    from pykern import pkunit
    from pykern import pkio

    with pkunit.save_chdir_work():
        fn = 'f1'
        # Should not throw an exception
        pkio.unchecked_remove(fn)
        pkio.write_text(fn, 'hello')
        pkio.unchecked_remove(fn)
        assert not os.path.exists(fn), \
            'When file removed, should be gone'
        for f in ('d1', 'd2/d3'):
            assert py.path.local(f) == pkio.mkdir_parent(f), \
                'When mkdir_parent is called, returns path passed in'
        assert os.path.exists('d1'), \
            'When single directory, should exist'
        assert os.path.exists('d2/d3'), \
            'When nested directory, should exist'
        with pytest.raises(AssertionError):
            pkio.unchecked_remove('.')
        with pytest.raises(AssertionError):
            pkio.unchecked_remove(os.getcwd())
        with pytest.raises(AssertionError):
            pkio.unchecked_remove('/')
示例#2
0
def _proprietary_codes():
    """Get proprietary files and put it in the proprietary code dir

    Args:
      uri (str): where to get file (file:// or http://)
    """
    import sirepo.feature_config
    import sirepo.sim_data
    import sirepo.srdb
    import urllib.error
    import urllib.request

    for s in sirepo.feature_config.cfg().proprietary_sim_types:
        f = sirepo.sim_data.get_class(s).proprietary_code_tarball()
        if not f:
            return
        r = pkio.mkdir_parent(sirepo.srdb.proprietary_code_dir(s), ).join(f)
        # POSIT: download/installers/flash-tarball/radiasoft-download.sh
        u = f'{cfg.proprietary_code_uri}/{s}-dev.tar.gz'
        try:
            urllib.request.urlretrieve(u, r)
        except urllib.error.URLError as e:
            if not isinstance(e.reason, FileNotFoundError):
                raise
            pkdlog('uri={} not found; mocking empty file={}', u, r)
            pkio.write_text(
                r,
                'mocked by sirepo.pkcli.setup_dev',
            )
示例#3
0
文件: pkjinja.py 项目: moellep/pykern
def render_file(filename,
                j2_ctx,
                output=None,
                strict_undefined=False,
                jinja_env=None):
    """Render filename as template with j2_ctx.

    Args:
        basename (str): name without jinja extension
        j2_ctx (dict): how to replace values in Jinja2 template
        output (str): file name of output; if None, return str
        strict_undefined (bool): set `jinja2.StrictUndefined` if True
        jinja_env (dict): add values to jinja2 environment

    Returns:
        str: rendered template
    """
    t = pkio.read_text(filename)
    kw = dict(
        trim_blocks=True,
        lstrip_blocks=True,
        keep_trailing_newline=True,
        extensions=['jinja2.ext.do'],
    )
    if strict_undefined:
        kw['undefined'] = jinja2.StrictUndefined
    if jinja_env:
        kw.update(jinja_env)
    je = jinja2.Environment(**kw)
    res = je.from_string(t).render(j2_ctx)
    if output:
        pkio.write_text(output, res)
    return res
示例#4
0
def test_importer(import_req):
    from pykern import pkcollections
    from pykern import pkjson
    from pykern.pkunit import pkeq
    from sirepo.template import zgoubi
    import sirepo.sim_data

    with pkunit.save_chdir_work() as w:
        for fn in pkio.sorted_glob(pkunit.data_dir().join('*.dat')):
            error = None
            try:
                data = zgoubi.import_file(import_req(fn), unit_test_mode=True)
                sirepo.sim_data.get_class('zgoubi').fixup_old_data(data)
                #TODO(pjm): easier way to convert nested dict to pkcollections.Dict?
                data = pkcollections.json_load_any(pkjson.dump_pretty(data))
            except Exception as e:
                pkdlog(pkdexc())
                error = e.message
            if error:
                actual = error
            else:
                actual = zgoubi.python_source_for_model(data)
            outfile = fn.basename + '.txt'
            pkio.write_text(outfile, actual)
            e = pkunit.data_dir().join(outfile)
            expect = pkio.read_text(e)
            pkeq(expect, actual, 'diff {} {}', e, w.join(outfile))
示例#5
0
def _proprietary_codes():
    """Get proprietary RPMs and put it in the proprietary code dir

    Args:
      uri (str): where to get RPM (file:// or http://)
    """
    import sirepo.feature_config
    import sirepo.sim_data
    import sirepo.srdb
    import urllib.error
    import urllib.request

    for s in sirepo.feature_config.cfg().proprietary_sim_types:
        r = pkio.mkdir_parent(sirepo.srdb.proprietary_code_dir(s), ).join(
            sirepo.sim_data.get_class(s).proprietary_code_rpm(), )
        # POSIT: download/installers/rpm-code/dev-build.sh
        u = f'{cfg.proprietary_code_uri}/rscode-{s}-dev.rpm'
        try:
            urllib.request.urlretrieve(u, r)
        except urllib.error.URLError as e:
            if not isinstance(e.reason, FileNotFoundError):
                raise
            pkdlog('uri={} not found; mocking empty rpm={}', u, r)
            pkio.write_text(
                r,
                'mocked by sirepo.pkcli.setup_dev',
            )
示例#6
0
    def write_files(self, data, source_path, dest_dir):
        """writes files for the simulation


        Returns:
            PKDict: structure of files written (debugging only)
        """
        class _G(_Generate):
            def _abspath(basename):
                return source_path.new(basename=basename)

            def _input_file(self, model_name, field, filename):
                return filename

            def _lattice_filename(self, value):
                return value

        g = _G(data)
        g.sim()
        v = g.jinja_env
        r = PKDict(
            commands=dest_dir.join(source_path.basename),
            lattice=self._lattice_path(dest_dir, data),
        )
        pkio.write_text(r.commands, v.commands)
        pkio.write_text(r.lattice, v.rpn_variables + v.lattice)
        for f in set(
                LatticeUtil(data, _SCHEMA).iterate_models(
                    lattice.InputFileIterator(_SIM_DATA)).result, ):
            f = _SIM_DATA.lib_file_name_without_type(f)
            dest_dir.join(f).mksymlinkto(source_path.new(basename=f),
                                         absolute=False)
        f = g.filename_map
        r.output_files = [f[k] for k in f.keys_in_order]
        return r
示例#7
0
def render_file(filename, j2_ctx, output=None, strict_undefined=False):
    """Render filename as template with j2_ctx.

    Args:
        basename (str): name without jinja extension
        j2_ctx (dict): how to replace values in Jinja2 template
        output (str): file name of output; if None, return str
        strict_undefined (bool): set `jinja2.StrictUndefined` if True

    Returns:
        str: rendered template
    """
    t = pkio.read_text(filename)
    kw = dict(
        trim_blocks=True,
        lstrip_blocks=True,
        keep_trailing_newline=True,
    )
    if strict_undefined:
        kw['undefined'] = jinja2.StrictUndefined
    je = jinja2.Environment(**kw)
    res = je.from_string(t).render(j2_ctx)
    if output:
        pkio.write_text(output, res)
    return res
示例#8
0
def write_parameters(data, run_dir, is_parallel):
    _extract_zip(data, run_dir)
    pkio.write_text(
        #TODO: generate python instead
        run_dir.join('flash.par'),
        _generate_parameters_file(data),
    )
示例#9
0
文件: radia.py 项目: cchall/sirepo
def write_parameters(data, run_dir, is_parallel):
    # remove centrailzed geom files
    pkio.unchecked_remove(_geom_file(data.simulationId), _dmp_file(data.simulationId))
    pkio.write_text(
        run_dir.join(template_common.PARAMETERS_PYTHON_FILE),
        _generate_parameters_file(data),
    )
示例#10
0
def test_import():
    from pykern import pkjson
    from pykern.pkunit import pkeq
    from sirepo.template import flash_parser
    import re

    def _parse_config(fn):
        return flash_parser.ConfigParser().parse(pkio.read_text(fn))

    def _parse_par(fn):
        data_file = fn.basename.replace('-flash.par', '')
        return flash_parser.ParameterParser().parse(
            pkjson.load_any(
                pkio.read_text(
                    pkunit.data_dir().join(f'{data_file}-sirepo-data.json'))),
            pkio.read_text(fn),
        )

    with pkunit.save_chdir_work():
        for fn in pkio.sorted_glob(pkunit.data_dir().join('*')):
            if re.search(r'-Config$', fn.basename):
                parser = _parse_config
            elif re.search(r'flash.par$', fn.basename):
                parser = _parse_par
            else:
                continue
            try:
                actual = pkjson.dump_pretty(parser(fn))
            except Exception as e:
                pkdlog(pkdexc())
                actual = str(e)
            outfile = f'{fn.basename}.out'
            pkio.write_text(outfile, actual)
            expect = pkio.read_text(pkunit.data_dir().join(outfile))
            pkeq(expect, actual)
示例#11
0
def test_from_elegant_to_madx_and_back():
    from pykern import pkio
    from pykern.pkunit import pkeq
    from sirepo.template import elegant, madx, madx_converter, madx_parser

    with pkunit.save_chdir_work() as d:
        for name in ('SPEAR3', 'Compact Storage Ring',
                     'Los Alamos Proton Storage Ring'):
            data = _example_data(name)
            mad = madx_parser.parse_file(
                elegant.python_source_for_model(data, 'madx'))
            madx._fixup_madx(mad)
            outfile = name.lower().replace(' ', '-') + '.madx'
            actual = madx.python_source_for_model(mad, None)
            pkio.write_text(outfile, actual)
            e = pkunit.data_dir().join(outfile)
            expect = pkio.read_text(e)
            pkeq(expect, actual, 'diff {} {}', e, d.join(outfile))

            data = madx_parser.parse_file(actual)
            lattice = madx_converter.from_madx(elegant.SIM_TYPE, data)
            outfile = name.lower().replace(' ', '-') + '.lte'
            actual = elegant.python_source_for_model(lattice, None)
            pkio.write_text(outfile, actual)
            e = pkunit.data_dir().join(outfile)
            expect = pkio.read_text(e)
            pkeq(expect, actual, 'diff {} {}', e, d.join(outfile))
示例#12
0
def _run_elegant():
    run_dir = os.getcwd()
    with open('in.json') as f:
        data = json.load(f)
    exec(pkio.read_text('elegant_parameters.py'), locals(), locals())
    pkio.write_text('elegant.lte', lattice_file)
    pkio.write_text('elegant.ele', elegant_file)
    call(['elegant', 'elegant.ele'])

    index = 0
    if sdds.sddsdata.InitializeInput(index, 'elegant.bun') != 1:
        sdds.sddsdata.PrintErrors(1)
    column_names = sdds.sddsdata.GetColumnNames(index)
    errorCode = sdds.sddsdata.ReadPage(index)
    if errorCode != 1:
        sdds.sddsdata.PrintErrors(1)
    bunch = data['models'][data['report']]
    x = sdds.sddsdata.GetColumn(index, column_names.index(bunch['x']))
    y = sdds.sddsdata.GetColumn(index, column_names.index(bunch['y']))
    nbins = int(bunch['histogramBins'])
    hist, edges = np.histogramdd([x, y], nbins)

    info = {
        'x_range': [float(edges[0][0]), float(edges[0][-1]), len(hist)],
        'y_range': [float(edges[1][0]), float(edges[1][-1]), len(hist[0])],
        'x_label': _FIELD_LABEL[bunch['x']],
        'y_label': _FIELD_LABEL[bunch['y']],
        'title': _plot_title(bunch),
        'z_matrix': hist.T.tolist(),
    }
    with open('out.json', 'w') as outfile:
        json.dump(info, outfile)
示例#13
0
def sbatch_script(path):
    """Write script to path

    Args:
        path (str): where to write file
    """
    pkio.write_text(path, _script())
示例#14
0
    def generate_input_file(self, kwarg_dict, directory):

        # Get strings for each file and fill in arguments for this job
        for key, val in self.setup['file_mapping'].items():
            local_file_instance = self.get_file_def_module().__getattribute__(
                key).format(**kwarg_dict)
            pkio.write_text(os.path.join(directory, val), local_file_instance)
示例#15
0
    def write_files(self, data, source_path, dest_dir):
        """writes files for the simulation

        Returns:
            PKDict: structure of files written (debugging only)
        """
        def _unescape(value):
            return re.sub(r'\\\\', r'\\', value)

        class _G(_Generate):
            def _abspath(self, basename):
                return source_path.new(basename=basename)

            def _input_file(self, model_name, field, filename):
                return filename

            def _lattice_filename(self, value):
                return value

        g = _G(data, update_output_filenames=False)
        g.sim()
        v = g.jinja_env
        r = PKDict(
            commands=dest_dir.join(source_path.basename),
            lattice=self._lattice_path(dest_dir, data),
        )
        pkio.write_text(r.commands, _unescape(v.commands))
        if not r.lattice.exists():
            pkio.write_text(r.lattice, v.rpn_variables + v.lattice)
        self._write_input_files(data, source_path, dest_dir)
        f = g.filename_map
        r.output_files = [f[k] for k in f.keys_in_order]
        return r
示例#16
0
文件: jspec.py 项目: e-carlin/sirepo
def _run_jspec(data):
    _elegant_to_madx(data['models']['ring'])
    exec(pkio.read_text(template_common.PARAMETERS_PYTHON_FILE), locals(), locals())
    jspec_filename = template.JSPEC_INPUT_FILENAME
    pkio.write_text(jspec_filename, jspec_file)
    pksubprocess.check_call_with_signals(['jspec', jspec_filename], msg=pkdlog, output=template.JSPEC_LOG_FILE)
    return pkio.read_text(template.JSPEC_LOG_FILE)
示例#17
0
文件: madx.py 项目: mkeilman/sirepo
def _generate_ptc_particles_file(run_dir, data, twiss):
    bunch = data.models.bunch
    p = _ptc_particles(
        PKDict(
            x=_bunch_twiss(bunch, 'x'),
            y=_bunch_twiss(bunch, 'y'),
        ),
        bunch.numberOfParticles,
        bunch.randomSeed,
    )
    v = PKDict(
        x=template.to_floats(p.x.pos),
        px=template.to_floats(p.x.p),
        y=template.to_floats(p.y.pos),
        py=template.to_floats(p.y.p),
        t=template.to_floats(p.t.pos),
        pt=template.to_floats(p.t.p),
    )
    if 'bunchReport' in data.report:
        v.summaryData = twiss
        simulation_db.write_json(run_dir.join(template.BUNCH_PARTICLES_FILE), v)
    r = ''
    for i in range(len(v.x)):
        r += 'ptc_start'
        for f in ('x', 'px', 'y', 'py', 't', 'pt'):
           r += f', {f}={v[f][i]}'
        r +=';\n'
    pkio.write_text(run_dir.join(template.PTC_PARTICLES_FILE), r)
示例#18
0
def twiss_to_madx(elegant_twiss_file, madx_twiss_file):
    outfile = 'sdds_output.txt'
    twiss_file = 'twiss-with-mu.sdds'
    # convert elegant psix to mad-x MU, rad --> rad / 2pi
    pksubprocess.check_call_with_signals([
        'sddsprocess',
        elegant_twiss_file,
        '-define=column,mux,psix 2 pi * /',
        '-define=column,muy,psiy 2 pi * /',
        twiss_file,
    ],
                                         output=outfile,
                                         env=elegant_common.subprocess_env())
    pksubprocess.check_call_with_signals([
        'sdds2stream',
        twiss_file,
        '-columns={}'.format(','.join(
            map(lambda x: x[0], _ELEGANT_TO_MADX_COLUMNS))),
    ],
                                         output=outfile,
                                         env=elegant_common.subprocess_env())
    lines = pkio.read_text(outfile).split('\n')
    header = '* {}\n$ \n'.format(' '.join(
        map(lambda x: x[1], _ELEGANT_TO_MADX_COLUMNS)))
    pkio.write_text(madx_twiss_file, header + '\n'.join(lines) + '\n')
示例#19
0
文件: flash.py 项目: mkeilman/sirepo
def write_parameters(data, run_dir, is_parallel):
    if data.report == 'setupAnimation':
        return
    pkio.write_text(
        run_dir.join(_FLASH_PAR_FILE),
        _generate_parameters_file(data, run_dir=run_dir),
    )
示例#20
0
文件: hellweg.py 项目: ahebnl/Sirepo
def _run_hellweg(cfg_dir):
    r = template_common.exec_parameters()
    pkio.write_text(template.HELLWEG_INPUT_FILE, r.input_file)
    pkio.write_text(template.HELLWEG_INI_FILE, r.ini_file)
    s = solver.BeamSolver(template.HELLWEG_INI_FILE, template.HELLWEG_INPUT_FILE)
    s.solve()
    s.save_output(template.HELLWEG_SUMMARY_FILE)
    s.dump_bin(template.HELLWEG_DUMP_FILE)
示例#21
0
def _run_tunes_report(cfg_dir, data):
    r = template_common.exec_parameters()
    pkio.write_text(template.TUNES_INPUT_FILE, r.tunes_file)
    #TODO(pjm): uses datafile from animation directory
    os.symlink('../animation/zgoubi.fai', 'zgoubi.fai')
    subprocess.call([_TUNES_PATH])
    template_common.write_sequential_result(
        template.extract_tunes_report(cfg_dir, data))
示例#22
0
def write_status(status, run_dir):
    """Write status to simulation

    Args:
        status (str): pending, running, completed, canceled
        run_dir (py.path): where to write the file
    """
    pkio.write_text(run_dir.join(_STATUS_FILE), status)
示例#23
0
def write_parameters(data, run_dir, is_parallel):
    m = re.search('^user_alert=(.*)', data.models.dog.breed)
    if m:
        raise sirepo.util.UserAlert(m.group(1), 'log msg should not be sent')
    pkio.write_text(
        run_dir.join(template_common.PARAMETERS_PYTHON_FILE),
        _generate_parameters_file(data),
    )
示例#24
0
def write_parameters(data,
                     run_dir,
                     is_parallel,
                     python_file=template_common.PARAMETERS_PYTHON_FILE):
    pkio.write_text(
        run_dir.join(python_file),
        _generate_parameters_file(data),
    )
示例#25
0
def write_status(status, run_dir):
    """Write status to simulation

    Args:
        status (str): pending, running, completed, canceled
        run_dir (py.path): where to write the file
    """
    pkio.write_text(run_dir.join(_STATUS_FILE), status)
示例#26
0
文件: jspec.py 项目: ahebnl/Sirepo
def _run_jspec(data):
    _elegant_to_madx(data['models']['ring'])
    r = template_common.exec_parameters()
    f = template.JSPEC_INPUT_FILENAME
    pkio.write_text(f, r.jspec_file)
    pksubprocess.check_call_with_signals(['jspec', f],
                                         msg=pkdlog,
                                         output=template.JSPEC_LOG_FILE)
示例#27
0
def _add_host(j2_ctx, srv, host):
    from rsconf.component import docker_registry
    from rsconf.component import rsconf

    netrc = rsconf.host_init(j2_ctx, host)
    pkio.write_text(srv.join(host + '-netrc'), netrc)
    if host == j2_ctx.master:
        docker_registry.host_init(j2_ctx, host)
示例#28
0
def _run_hellweg(cfg_dir):
    with pkio.save_chdir(cfg_dir):
        exec(pkio.read_text(template_common.PARAMETERS_PYTHON_FILE), locals(), locals())
        pkio.write_text(template.HELLWEG_INPUT_FILE, input_file)
        pkio.write_text(template.HELLWEG_INI_FILE, ini_file)
        solver = BeamSolver(template.HELLWEG_INI_FILE, template.HELLWEG_INPUT_FILE)
        solver.solve()
        solver.save_output(template.HELLWEG_SUMMARY_FILE)
        solver.dump_bin(template.HELLWEG_DUMP_FILE)
示例#29
0
文件: jspec.py 项目: pir8aye/sirepo
def _run_jspec(run_dir):
    with pkio.save_chdir(run_dir):
        data = simulation_db.read_json(template_common.INPUT_BASE_NAME)
        _elegant_to_madx(data['models']['ring'])
        exec(pkio.read_text(template_common.PARAMETERS_PYTHON_FILE), locals(), locals())
        jspec_filename = template.JSPEC_INPUT_FILENAME
        pkio.write_text(jspec_filename, jspec_file)
        pksubprocess.check_call_with_signals(['jspec', jspec_filename], msg=pkdp, output=template.JSPEC_LOG_FILE)
        return pkio.read_text(template.JSPEC_LOG_FILE)
示例#30
0
文件: jspec.py 项目: yeeon/sirepo
def _run_jspec(data):
    _elegant_to_madx(data['models']['ring'])
    exec(pkio.read_text(template_common.PARAMETERS_PYTHON_FILE), locals(),
         locals())
    jspec_filename = template.JSPEC_INPUT_FILENAME
    pkio.write_text(jspec_filename, jspec_file)
    pksubprocess.check_call_with_signals(['jspec', jspec_filename],
                                         msg=pkdlog,
                                         output=template.JSPEC_LOG_FILE)
    return pkio.read_text(template.JSPEC_LOG_FILE)
示例#31
0
def _run_tunes_report(cfg_dir, data):
    with pkio.save_chdir(cfg_dir):
        exec(pkio.read_text(template_common.PARAMETERS_PYTHON_FILE), locals(),
             locals())
        pkio.write_text(template.TUNES_INPUT_FILE, tunes_file)
        #TODO(pjm): uses datafile from animation directory
        os.symlink('../animation/zgoubi.fai', 'zgoubi.fai')
        subprocess.call([_TUNES_PATH])
        simulation_db.write_result(template.extract_tunes_report(
            cfg_dir, data))
示例#32
0
def write_parameters(data, schema, run_dir, is_parallel):
    """Write the parameters file

    Args:
        data (dict): input
        schema (dict): to validate data
        run_dir (py.path): where to write
        is_parallel (bool): run in background?
    """
    pkio.write_text(run_dir.join(template_common.PARAMETERS_PYTHON_FILE), generate_parameters_file(data, is_parallel))
示例#33
0
def _run_hellweg(cfg_dir):
    with pkio.save_chdir(cfg_dir):
        exec(pkio.read_text(template_common.PARAMETERS_PYTHON_FILE), locals(),
             locals())
        pkio.write_text(template.HELLWEG_INPUT_FILE, input_file)
        pkio.write_text(template.HELLWEG_INI_FILE, ini_file)
        s = solver.BeamSolver(template.HELLWEG_INI_FILE,
                              template.HELLWEG_INPUT_FILE)
        s.solve()
        s.save_output(template.HELLWEG_SUMMARY_FILE)
        s.dump_bin(template.HELLWEG_DUMP_FILE)
示例#34
0
def _save_field_csv(field_type, vectors, scipy_rotation, path):
    # reserve first line for a header
    data = [f'x,y,z,{field_type}x,{field_type}y,{field_type}z']
    pts, mags, dirs = _rotate_fields(vectors, scipy_rotation, True)
    for i in range(len(mags)):
        j = 3 * i
        r = pts[j:j + 3]
        r = numpy.append(r, mags[i] * dirs[j:j + 3])
        data.append(','.join(map(str, r)))
    pkio.write_text(path, '\n'.join(data))
    return path
示例#35
0
def write_parameters(data, run_dir, is_parallel):
    pkio.write_text(
        run_dir.join(OPAL_INPUT_FILE),
        _generate_parameters_file(data),
    )
    if is_parallel:
        pkio.write_text(
            run_dir.join(OPAL_POSITION_FILE),
            'import os\n' \
            + 'os.system("python data/opal_ElementPositions.py --export-vtk")\n',
        )
示例#36
0
def test_generate_python():
    from pykern import pkio
    from pykern.pkunit import pkeq
    from sirepo.template import shadow

    with pkunit.save_chdir_work():
        for name in ('Complete Beamline', 'Wiggler'):
            data = _example_data(name)
            data['report'] = 'watchpointReport{}'.format(data.models.beamline[-1].id)
            actual = shadow._generate_parameters_file(data)
            outfile = data.models.simulation.simulationId + '.txt'
            pkio.write_text(outfile, actual)
            expect = pkio.read_text(pkunit.data_dir().join(outfile))
            pkeq(expect, actual)
示例#37
0
def run_background(cfg_dir):
    with pkio.save_chdir(cfg_dir):
        fn = 'run_background.py'
        cmd = [sys.executable or 'python', fn]
        script = pkio.read_text('srw_parameters.py')
        p = dict(pkcollections.map_items(cfg))
        if cfg.slave_processes > 1:
            cmd[0:0] = [
                'mpiexec',
                '-n',
                # SRW includes a master process so 2 really needs 3 processes
                str(cfg.slave_processes + 1),
            ]
            script += '''
from mpi4py import MPI
if MPI.COMM_WORLD.Get_rank():
    import signal
    signal.signal(signal.SIGTERM, lambda x, y: MPI.COMM_WORLD.Abort(1))
'''
        else:
            # In interactive (dev) mode, output as frequently as possible
            p['particles_per_slave'] = 1
        script += '''
import srwl_bl
v = srwl_bl.srwl_uti_parse_options(get_srw_params(), use_sys_argv=False)
source_type, mag = setup_source(v)
v.wm = True
v.wm_nm = {total_particles}
v.wm_na = {particles_per_slave}
# Number of "iterations" per save is best set to num processes
v.wm_ns = {slave_processes}
op = get_beamline_optics()
srwl_bl.SRWLBeamline(_name=v.name).calc_all(v, op)
'''.format(**p)
        pkio.write_text(fn, script)
        try:
            p = subprocess.Popen(
                cmd,
                stdin=open(os.devnull),
                stdout=open('run_background.out', 'w'),
                stderr=subprocess.STDOUT,
            )
            signal.signal(signal.SIGTERM, lambda x, y: p.terminate())
            rc = p.wait()
            if rc != 0:
                p = None
                raise RuntimeError('child terminated: retcode={}'.format(rc))
        finally:
            if not p is None:
                p.terminate()
示例#38
0
def run_background(cfg_dir):
    with pkio.save_chdir(cfg_dir):
        fn = 'run_background.py'
        cmd = [sys.executable or 'python', fn]
        script = pkio.read_text('srw_parameters.py')
        p = dict(pkcollections.map_items(cfg))
        if cfg.slave_processes > 1:
            cmd[0:0] = [
                'mpiexec',
                '-n',
                # SRW includes a master process so 2 really needs 3 processes
                str(cfg.slave_processes + 1),
            ]
            script += '''
from mpi4py import MPI
if MPI.COMM_WORLD.Get_rank():
    import signal
    signal.signal(signal.SIGTERM, lambda x, y: MPI.COMM_WORLD.Abort(1))
'''
        else:
            # In interactive (dev) mode, output as frequently as possible
            p['particles_per_slave'] = 1
        script += '''
import srwl_bl
v = srwl_bl.srwl_uti_parse_options(get_srw_params(), use_sys_argv=False)
source_type, mag = setup_source(v)
v.wm = True
v.wm_nm = {total_particles}
v.wm_na = {particles_per_slave}
# Number of "iterations" per save is best set to num processes
v.wm_ns = {slave_processes}
op = get_beamline_optics()
srwl_bl.SRWLBeamline(_name=v.name).calc_all(v, op)
'''.format(**p)
        pkio.write_text(fn, script)
        try:
            p = subprocess.Popen(
                cmd,
                stdin=open(os.devnull),
                stdout=open('run_background.out', 'w'),
                stderr=subprocess.STDOUT,
            )
            signal.signal(signal.SIGTERM, lambda x, y: p.terminate())
            rc = p.wait()
            if rc != 0:
                p = None
                raise RuntimeError('child terminated: retcode={}'.format(rc))
        finally:
            if not p is None:
                p.terminate()
示例#39
0
def assert_object_with_json(basename, actual):
    """Converts actual to JSON and compares with data_dir/basename.json

    Reads data_dir/basename.json and compares with actual
    converted to json. Trailing newline is managed properly. The
    keys are sorted and indentation is 4. actual written to work_dir.

    Args:
        expected_basename (str): file to be found in data_dir with json suffix
        actual (object): to be serialized as json
    """
    actual = pkdpretty(actual)
    fn = '{}.json'.format(basename)
    pkio.write_text(work_dir().join(fn), actual)
    expect = pkio.read_text(data_dir().join(fn))
    assert expect == actual, \
        '{}: unexpected result'.format(basename)
示例#40
0
def test_init_tree():
    """Normal case"""
    with pkunit.save_chdir_work():
        name = 'proj1'
        pkio.mkdir_parent(name)
        with pkio.save_chdir(name):
            subprocess.check_call(['git', 'init', '.'])
            subprocess.check_call(['git', 'config', 'user.email', '*****@*****.**'])
            subprocess.check_call(['git', 'config', 'user.name', 'pykern'])
            projex.init_tree(
                name=name,
                author='zauthor',
                author_email='*****@*****.**',
                description='some python project',
                license='MIT',
                url='http://example.com',
            )
            pkio.write_text('tests/test_1.py', 'def test_1(): pass')
            for expect_fn, expect_re in (
                ('.gitignore', 'MANIFEST.in'),
                ('LICENSE', 'The MIT License'),
                ('README.md', 'licenses/MIT'),
                ('docs/_static/.gitignore', ''),
                ('docs/_templates/.gitignore', ''),
                ('docs/index.rst', name),
                ('requirements.txt', 'pykern'),
                ('setup.py', "author='zauthor'"),
                ('setup.py', r':copyright:.*zauthor\.'),
                ('tests/.gitignore', '_work'),
                (name + '/__init__.py', ''),
                (name + '/package_data/.gitignore', ''),
                (
                    '{}/{}_console.py'.format(name, name),
                    r"main\('{}'\)".format(name),
                ),
            ):
                assert re.search(expect_re, pkio.read_text(expect_fn)), \
                    '{} should exist and match "{}"'.format(expect_fn, expect_re)
            subprocess.check_call(['git', 'commit', '-m', 'initial'])
            # Do not install from PyPI
            pkio.write_text(
                'requirements.txt',
                '-e ' + str(py.path.local(__file__).dirpath().dirpath().dirpath()),
            );
            subprocess.check_call(['python', 'setup.py', 'test'])
            subprocess.check_call(['python', 'setup.py', 'tox'])
示例#41
0
def test_importer():
    from pykern import pkcollections
    from pykern import pkio
    from pykern.pkunit import pkeq
    from sirepo.template import elegant

    with pkunit.save_chdir_work():
        for fn in pkio.sorted_glob(pkunit.data_dir().join('*')):
            if not pkio.has_file_extension(fn, ('ele', 'lte')) \
                or fn.basename.endswith('ele.lte'):
                continue
            error = None
            try:
                data = elegant.import_file(FlaskRequest(fn))
            except Exception as e:
                pkdlog(pkdexc())
                error = e.message
            if error:
                actual = error
            else:
                if pkio.has_file_extension(fn, 'lte'):
                    data['models']['commands'] = []
                    actual = '{}{}'.format(
                        elegant._generate_variables(data),
                        elegant.generate_lattice(
                            data,
                            elegant._build_filename_map(data),
                            elegant._build_beamline_map(data),
                            pkcollections.Dict(),
                        ),
                    )
                else:
                    data2 = elegant.import_file(FlaskRequest('{}.lte'.format(fn)), test_data=data)
                    actual = elegant._generate_commands(
                        data2,
                        elegant._build_filename_map(data2),
                        elegant._build_beamline_map(data2),
                        pkcollections.Dict(),
                    )
            outfile = fn.basename + '.txt'
            pkio.write_text(outfile, actual)
            expect = pkio.read_text(pkunit.data_dir().join(outfile))
            #TODO(pjm): this takes too long if there are a lot of diffs
            #assert expect == actual
            pkeq(expect, actual)
示例#42
0
def _run_elegant(bunch_report=False, with_mpi=False):
    exec(pkio.read_text(template_common.PARAMETERS_PYTHON_FILE), locals(), locals())
    pkio.write_text('elegant.lte', lattice_file)
    ele = 'elegant.ele'
    pkio.write_text(ele, elegant_file)
    kwargs = {
        'output': ELEGANT_LOG_FILE,
        'env': elegant_common.subprocess_env(),
    }
    try:
        #TODO(robnagler) Need to handle this specially, b/c different binary
        if execution_mode == 'parallel' and with_mpi and mpi.cfg.cores > 1:
            mpi.run_program(['Pelegant', ele], **kwargs)
        else:
            pksubprocess.check_call_with_signals(['elegant', ele], msg=pkdlog, **kwargs)
    except Exception as e:
        # ignore elegant failures - errors will be parsed from the log
        pass
示例#43
0
def hack_nfs_write_status(status, run_dir):
    """Verify status file exists before writing.

    NFS doesn't propagate files immediately so there
    is a race condition when the celery worker starts.
    This file handles this case.

    Args:
        status (str): pending, running, completed, canceled
        run_dir (py.path): where to write the file
    """
    fn = run_dir.join(_STATUS_FILE)
    for i in range(cfg.nfs_tries):
        if fn.check(file=True):
            break
        time.sleep(cfg.nfs_sleep)
    # Try once always
    pkio.write_text(fn, status)
示例#44
0
def write_parameters(data, run_dir, is_parallel):
    """Write the parameters file

    Args:
        data (dict): input
        run_dir (py.path): where to write
        is_parallel (bool): run in background?
    """
    pkio.write_text(
        run_dir.join(template_common.PARAMETERS_PYTHON_FILE),
        generate_parameters_file(
            data,
            is_parallel,
        ),
    )
    for f in _simulation_files(data):
        if re.search(r'SCRIPT-commandFile', f):
            os.chmod(str(run_dir.join(f)), stat.S_IRUSR | stat.S_IXUSR)
示例#45
0
def _run_elegant(bunch_report=False, with_mpi=False):
    exec(pkio.read_text(template_common.PARAMETERS_PYTHON_FILE), locals(), locals())
    if bunch_report and re.search('\&sdds_beam\s', elegant_file):
        return
    pkio.write_text('elegant.lte', lattice_file)
    ele = 'elegant.ele'
    pkio.write_text(ele, elegant_file)
    # TODO(robnagler) Need to handle this specially, b/c different binary
    env = copy.deepcopy(os.environ)
    env['RPN_DEFNS'] = pkresource.filename('defns.rpn')
    if with_mpi and mpi.cfg.cores > 1:
        return mpi.run_program(['Pelegant', ele], output=ELEGANT_LOG_FILE, env=env)
    pksubprocess.check_call_with_signals(
        ['elegant', ele],
        output=ELEGANT_LOG_FILE,
        env=env,
        msg=pkdp,
    )
示例#46
0
def twiss_to_madx(elegant_twiss_file, madx_twiss_file):
    outfile = 'sdds_output.txt'
    twiss_file = 'twiss-with-mu.sdds'
    # convert elegant psix to mad-x MU, rad --> rad / 2pi
    pksubprocess.check_call_with_signals([
        'sddsprocess',
        elegant_twiss_file,
        '-define=column,mux,psix 2 pi * /',
        '-define=column,muy,psiy 2 pi * /',
        twiss_file,
    ], output=outfile, env=elegant_common.subprocess_env())
    pksubprocess.check_call_with_signals([
        'sdds2stream',
        twiss_file,
        '-columns={}'.format(','.join(map(lambda x: x[0], _ELEGANT_TO_MADX_COLUMNS))),
    ], output=outfile, env=elegant_common.subprocess_env())
    lines = pkio.read_text(outfile).split('\n')
    header = '* {}\n$ \n'.format(' '.join(map(lambda x: x[1], _ELEGANT_TO_MADX_COLUMNS)))
    pkio.write_text(madx_twiss_file, header + '\n'.join(lines) + '\n')
示例#47
0
def test_walk_tree_and_sorted_glob():
    """Looks in work_dir"""
    from pykern import pkunit
    from pykern import pkio

    with pkunit.save_chdir_work() as pwd:
        for f in ('d1/d7', 'd2/d3', 'd4/d5/d6'):
            pkio.mkdir_parent(f)
        expect = []
        for f in ['d1/d7/f1', 'd4/d5/f2', 'd2/d3/f3']:
            pkio.write_text(f, '')
            expect.append(py.path.local(f))
        assert sorted(expect) == list(pkio.walk_tree('.')), \
            'When walking tree, should only return files'
        assert [expect[2]] == list(pkio.walk_tree('.', 'f3')), \
            'When walking tree with file_re, should only return matching files'
        assert [expect[0]] == list(pkio.walk_tree('.', '^d1')), \
            'When walking tree with file_re, file to match does not include dir being searched'
        assert pkio.sorted_glob('*[42]') == [py.path.local(f) for f in ('d2', 'd4')]
示例#48
0
def _start_simulation(data, run_async=False):
    """Setup and start the simulation.

    Args:
        data (dict): app data
        run_async (bool): run-background or run

    Returns:
        object: _Command or daemon instance
    """
    run_dir = simulation_db.simulation_run_dir(data, remove_dir=True)
    pkio.mkdir_parent(run_dir)
    #TODO(robnagler) create a lock_dir -- what node/pid/thread to use?
    #   probably can only do with celery.
    simulation_type = data['simulationType']
    sid = simulation_db.parse_sid(data)
    data = simulation_db.fixup_old_data(simulation_type, data)
    assert simulation_type in simulation_db.APP_NAMES, \
        '{}: invalid simulation type'.format(simulation_type)
    template = sirepo.template.import_module(simulation_type)
    simulation_db.save_simulation_json(simulation_type, data)
    for d in simulation_db.simulation_dir(simulation_type, sid), simulation_db.simulation_lib_dir(simulation_type):
        for f in glob.glob(str(d.join('*.*'))):
            if os.path.isfile(f):
                py.path.local(f).copy(run_dir)
    with open(str(run_dir.join('in{}'.format(simulation_db.JSON_SUFFIX))), 'w') as outfile:
        json.dump(data, outfile)
    pkio.write_text(
        run_dir.join(simulation_type + '_parameters.py'),
        template.generate_parameters_file(
            data,
            _schema_cache(simulation_type),
            run_dir=run_dir,
            run_async=run_async,
        )
    )

    cmd = [_ROOT_CMD, simulation_type] \
        + ['run-background' if run_async else 'run'] + [str(run_dir)]
    if run_async:
        return cfg.job_queue(sid, run_dir, cmd)
    return _Command(cmd, cfg.foreground_time_limit)
示例#49
0
文件: mpi.py 项目: mrakitin/sirepo
def run_script(script):
    """Execute python script with mpi.

    Args:
        script (str): python text
    """
    abort = '''

from mpi4py import MPI
if MPI.COMM_WORLD.Get_rank():
    import signal
    signal.signal(signal.SIGTERM, lambda x, y: MPI.COMM_WORLD.Abort(1))

'''
    n = re.sub(r'^from __future.*', abort, script, count=1, flags=re.MULTILINE)
    script = abort + script if n == script else n
    fn = 'mpi_run.py'
    pkio.write_text(fn, script)
    p = None
    return run_program([sys.executable or 'python', fn])
示例#50
0
def render_file(filename, values, output=None):
    """Render filename as template with values.

    Args:
        basename (str): name without jinja extension
        values (dict): how to replace values
        output (str): file name of output; if None, return str

    Returns:
        str: rendered template
    """
    t = pkio.read_text(filename)
    je = jinja2.Environment(
        trim_blocks=True,
        lstrip_blocks=True,
        keep_trailing_newline=True,
    )
    res = je.from_string(t).render(values)
    if output:
        pkio.write_text(output, res)
    return res
示例#51
0
def test_init_and_run(monkeypatch):
    from pykern import pkio
    from pykern import pkunit
    from pykern.pkcli import sim
    from pykern.pkcli import rsmanifest
    import netrc
    import os
    import os.path
    import re
    import subprocess

    cfg = pkunit.cfg.aux.get('sim_test', None)
    if not cfg:
        # No testing if there's no auth config
        return
    u, p = cfg.split(' ')
    monkeypatch.setattr(netrc, 'netrc', _netrc)
    _netrc.result = (u, None, p)
    with pkunit.save_chdir_work(is_pkunit_prefix=True):
        f = 'out/log'
        expect_code = pkunit.random_alpha()
        pkio.write_text('run.sh', 'echo {}>{}'.format(expect_code, f))
        rsmanifest.pkunit_setup()
        sim._cmd_init()
        sim._cmd_run()
        x = subprocess.check_output(['git', 'remote', '-v']),
        m = re.search(r'/(sim-sim_work-\d+-\d+)\.git', x[0])
        repo = m.group(1)
        pkunit.pkok(m, 'git remote: failed: {}', x)
        pkunit.pkeq(expect_code, pkio.read_text('out/log').rstrip())
        os.remove('run.sh')
        sim._cmd_pip('djson')
        pkio.write_text('run.py', 'import djson'.format(expect_code, f))
        sim._cmd_run()
        sim._git_api_request(
            'delete',
            'repositories/{user}/{repo}',
            dict(repo=repo),
        )
示例#52
0
def test_importer():
    from sirepo.template import elegant
    with pkunit.save_chdir_work():
        for filename in _FILES:
            error, data = elegant.import_file(TestFlaskRequest(filename))
            outfile = '{}.txt'.format(filename)
            if error:
                actual = error
            else:
                if '.lte' in filename:
                    data['models']['commands'] = []
                    actual = '{}{}'.format(
                        elegant._generate_variables(data),
                        elegant.generate_lattice(data, elegant._build_filename_map(data), elegant._build_beamline_map(data), {}))
                else:
                    err2, data2 = elegant.import_file(TestFlaskRequest('{}.lte'.format(filename)), test_data=data)
                    actual = elegant._generate_commands(data2, elegant._build_filename_map(data2), elegant._build_beamline_map(data2), {})
            pkio.write_text(outfile, actual)
            expect = pkio.read_text(pkunit.data_dir().join(outfile))
            #TODO(pjm): this takes too long if there are a lot of diffs
            #assert expect == actual
            if expect != actual:
                assert False
示例#53
0
def test_write_text():
    """Also tests read_text"""
    from pykern import pkunit
    from pykern import pkio

    d = pkunit.empty_work_dir()
    expect_res = d.join('anything')
    expect_content = 'something'
    res = pkio.write_text(str(expect_res), expect_content)
    assert expect_res == res, \
        'Verify result is file path as py.path.Local'
    with open(str(expect_res)) as f:
        assert expect_content == f.read(), \
            'When write_text is called, it should write "something"'
    assert expect_content == pkio.read_text(str(expect_res)), \
        'When read_text, it should read "something"'
示例#54
0
def import_python(code, tmp_dir, lib_dir, user_filename=None, arguments=None):
    """Converts script_text into json and stores as new simulation.

    Avoids too much data back to the user in the event of an error.
    This could be a potential security issue, because the script
    could be used to probe the system.

    Args:
        simulation_type (str): always "srw", but used to find lib dir
        code (str): Python code that runs SRW
        user_filename (str): uploaded file name for log
        arguments (str): argv to be passed to script

    Returns:
        dict: simulation data
    """
    script = None

    # Patch for the mirror profile for the exported .py file from Sirepo:
    code = _patch_mirror_profile(code, lib_dir)

    try:
        with pkio.save_chdir(tmp_dir):
            # This string won't show up anywhere
            script = pkio.write_text('in.py', code)
            o = SRWParser(
                script,
                lib_dir=py.path.local(lib_dir),
                user_filename=user_filename,
                arguments=arguments,
            )
            return o.data
    except Exception as e:
        lineno = script and _find_line_in_trace(script)
        # Avoid
        pkdlog(
            'Error: {}; exception={}; script={}; filename={}; stack:\n{}',
            e.message,
            e,
            script,
            user_filename,
            pkdexc(),
        )
        e = str(e)[:50]
        raise ValueError(
            'Error on line {}: {}'.format(lineno, e) if lineno
            else 'Error: {}'.format(e))
示例#55
0
def import_python(code, tmp_dir, lib_dir, user_filename=None):
    """Converts script_text into json and stores as new simulation.

    Avoids too much data back to the user in the event of an error.
    This could be a potential security issue, because the script
    could be used to probe the system.

    Args:
        simulation_type (str): always "srw", but used to find lib dir
        code (str): Python code that runs SRW
        user_filename (str): uploaded file name for log

    Returns:
        error: string containing error or None
        dict: simulation data
    """
    error = 'Import failed: error unknown'
    script = None
    try:
        with pkio.save_chdir(tmp_dir):
            # This string won't show up anywhere
            script = pkio.write_text('in.py', code)
            o = SRWParser(
                script,
                lib_dir=py.path.local(lib_dir),
                user_filename=user_filename,
            )
            return None, o.data
    except Exception as e:
        lineno = _find_line_in_trace(script) if script else None
        # Avoid
        pkdp(
            'Error: {}; exception={}; script={}; filename={}; stack:\n{}',
            error,
            e,
            script,
            user_filename,
            traceback.format_exc(),
        )
        error = 'Error on line {}: {}'.format(lineno or '?', str(e)[:50])
    return error, None
示例#56
0
def write_parameters(data, run_dir, is_parallel):
    pkio.write_text(
        run_dir.join(template_common.PARAMETERS_PYTHON_FILE),
        _generate_parameters_file(data),
    )
示例#57
0
def test_1():
    with pkunit.save_chdir_work():
        pkio.write_text('do_not_include_in_sdist.py', 'some text')