<commit_before># -*- coding: utf-8 -*-
u"""simulation data operations

:copyright: Copyright (c) 2020 RadiaSoft LLC.  All Rights Reserved.
:license: http://www.apache.org/licenses/LICENSE-2.0.html
"""
from pykern.pkcollections import PKDict
from pykern.pkdebug import pkdc, pkdlog, pkdp
from sirepo.template.lattice import LatticeUtil
from sirepo.template.template_common import ParticleEnergy
import math
import numpy
import re
import sirepo.sim_data
import sirepo.simulation_db


class SimData(sirepo.sim_data.SimDataBase):

    @classmethod
    def add_ptc_track_commands(cls, data):

        def _set_ptc_ids(ptc_commands, data):
            m = LatticeUtil.max_id(data) + 1
            for i,  c in enumerate(ptc_commands):
                c._id = m + i
            return ptc_commands
        data.models.bunch.beamDefinition = 'gamma'
        data.models.commands.extend(_set_ptc_ids(
            [
                PKDict(_type='ptc_create_universe'),
                PKDict(_type='ptc_create_layout'),
                PKDict(_type='ptc_track', file='1', icase='6'),
                PKDict(_type='ptc_track_end'),
                PKDict(_type='ptc_end'),
            ],
            data,
        ))

    @classmethod
    def beamline_elements(cls, madx):
        elmap = PKDict({e._id: e for e in madx.elements})
        for el_id in madx.beamlines[0]['items']:
            yield elmap[el_id]

    @classmethod
    def controls_madx_dir(cls):
        return sirepo.simulation_db.simulation_dir('madx')

    @classmethod
    def current_field(cls, kick_field):
        return 'current_{}'.format(kick_field)

    @classmethod
    def default_optimizer_settings(cls, madx):
        targets = []
        for el in cls.beamline_elements(madx):
            if el.type in ('MONITOR', 'HMONITOR', 'VMONITOR'):
                item = cls.model_defaults('optimizerTarget')
                item.name = el.name
                if el.type == 'HMONITOR':
                    del item['y']
                elif el.type == 'VMONITOR':
                    del item['x']
                targets.append(item)
        opts = cls.model_defaults('optimizerSettings').pkupdate(PKDict(
            targets=targets,
        ))
        cls.init_optimizer_inputs(opts, madx)
        return opts

    @classmethod
    def fixup_old_data(cls, data):
        dm = data.models
        cls._init_models(
            dm,
            (
                'beamPositionAnimation',
                'bunch',
                'command_beam',
                'dataFile',
                'deviceServerMonitor',
                'initialMonitorPositionsReport',
                'instrumentAnimationAll',
                'instrumentAnimationTwiss',
            ),
        )
        if 'externalLattice' in dm:
            sirepo.sim_data.get_class('madx').fixup_old_data(dm.externalLattice)
            if 'optimizerSettings' not in dm:
                dm.optimizerSettings = cls.default_optimizer_settings(dm.externalLattice.models)
            if 'controlSettings' not in dm:
                cls.init_process_variables(dm)
                cls.init_currents(dm.command_beam, dm.externalLattice.models)
            cls._init_models(dm, ('controlSettings', 'optimizerSettings'))
            if 'inputs' not in dm.optimizerSettings:
                cls.init_optimizer_inputs(dm.optimizerSettings, dm.externalLattice.models)
            cls._remove_old_command(dm.externalLattice.models)
        if dm.command_beam.gamma == 0 and 'pc' in dm.command_beam and dm.command_beam.pc > 0:
            cls.update_beam_gamma(dm.command_beam)
            dm.command_beam.pc = 0
        if 'command_twiss' in dm:
            for f in dm.command_twiss:
                if f in dm.bunch:
                    dm.bunch[f] = dm.command_twiss[f]
            del dm['command_twiss']
            cls.add_ptc_track_commands(dm.externalLattice)


    @classmethod
    def init_optimizer_inputs(cls, optimizerSettings, madx):
        optimizerSettings.inputs = PKDict(
                    kickers=PKDict(),
                    quads=PKDict()
                )
        for el in cls.beamline_elements(madx):
            if el.type == 'QUADRUPOLE':
                optimizerSettings.inputs.quads[str(el._id)] = False
            elif 'KICKER' in el.type:
                optimizerSettings.inputs.kickers[str(el._id)] = True

    @classmethod
    def init_currents(cls, beam, models):
        def is_kick_field(field):
            return re.search(r'^(.?kick|k1)$', field)
        ac = AmpConverter(beam)
        for el in cls.beamline_elements(models):
            for f in list(el.keys()):
                if is_kick_field(f):
                    el[cls.current_field(f)] = ac.kick_to_current(el[f])

    @classmethod
    def init_process_variables(cls, models):
        pvs = []
        def _add_pv(elId, dim, write='0'):
            pvs.append(PKDict(
                elId=elId,<commit_msg>fix #4248 controls fixup (#4252)<commit_after># -*- coding: utf-8 -*-
u"""simulation data operations

:copyright: Copyright (c) 2020 RadiaSoft LLC.  All Rights Reserved.
:license: http://www.apache.org/licenses/LICENSE-2.0.html
"""
from pykern.pkcollections import PKDict
from pykern.pkdebug import pkdc, pkdlog, pkdp
from sirepo.template.lattice import LatticeUtil
from sirepo.template.template_common import ParticleEnergy
import math
import numpy
import re
import sirepo.sim_data
import sirepo.simulation_db


class SimData(sirepo.sim_data.SimDataBase):

    @classmethod
    def add_ptc_track_commands(cls, data):

        def _set_ptc_ids(ptc_commands, data):
            m = LatticeUtil.max_id(data) + 1
            for i,  c in enumerate(ptc_commands):
                c._id = m + i
            return ptc_commands
        data.models.bunch.beamDefinition = 'gamma'
        data.models.commands.extend(_set_ptc_ids(
            [
                PKDict(_type='ptc_create_universe'),
                PKDict(_type='ptc_create_layout'),
                PKDict(_type='ptc_track', file='1', icase='6'),
                PKDict(_type='ptc_track_end'),
                PKDict(_type='ptc_end'),
            ],
            data,
        ))

    @classmethod
    def beamline_elements(cls, madx):
        elmap = PKDict({e._id: e for e in madx.elements})
        for el_id in madx.beamlines[0]['items']:
            yield elmap[el_id]

    @classmethod
    def controls_madx_dir(cls):
        return sirepo.simulation_db.simulation_dir('madx')

    @classmethod
    def current_field(cls, kick_field):
        return 'current_{}'.format(kick_field)

    @classmethod
    def default_optimizer_settings(cls, madx):
        targets = []
        for el in cls.beamline_elements(madx):
            if el.type in ('MONITOR', 'HMONITOR', 'VMONITOR'):
                item = cls.model_defaults('optimizerTarget')
                item.name = el.name
                if el.type == 'HMONITOR':
                    del item['y']
                elif el.type == 'VMONITOR':
                    del item['x']
                targets.append(item)
        opts = cls.model_defaults('optimizerSettings').pkupdate(PKDict(
            targets=targets,
        ))
        cls.init_optimizer_inputs(opts, madx)
        return opts

    @classmethod
    def fixup_old_data(cls, data):
        dm = data.models
        cls._init_models(
            dm,
            (
                'beamPositionAnimation',
                'bunch',
                'command_beam',
                'dataFile',
                'deviceServerMonitor',
                'initialMonitorPositionsReport',
                'instrumentAnimationAll',
                'instrumentAnimationTwiss',
            ),
        )
        if 'externalLattice' in dm:
            sirepo.sim_data.get_class('madx').fixup_old_data(dm.externalLattice)
            if 'optimizerSettings' not in dm:
                dm.optimizerSettings = cls.default_optimizer_settings(dm.externalLattice.models)
            if 'controlSettings' not in dm:
                cls.init_process_variables(dm)
                cls.init_currents(dm.command_beam, dm.externalLattice.models)
            cls._init_models(dm, ('controlSettings', 'optimizerSettings'))
            if 'inputs' not in dm.optimizerSettings:
                cls.init_optimizer_inputs(dm.optimizerSettings, dm.externalLattice.models)
            cls._remove_old_command(dm.externalLattice.models)
        if dm.command_beam.gamma == 0 and 'pc' in dm.command_beam and dm.command_beam.pc > 0:
            cls.update_beam_gamma(dm.command_beam)
            dm.command_beam.pc = 0
        if 'command_twiss' in dm:
            for f in dm.command_twiss:
                if f in dm.bunch:
                    dm.bunch[f] = dm.command_twiss[f]
            del dm['command_twiss']
            if 'externalLattice' in dm:
                cls.add_ptc_track_commands(dm.externalLattice)


    @classmethod
    def init_optimizer_inputs(cls, optimizerSettings, madx):
        optimizerSettings.inputs = PKDict(
                    kickers=PKDict(),
                    quads=PKDict()
                )
        for el in cls.beamline_elements(madx):
            if el.type == 'QUADRUPOLE':
                optimizerSettings.inputs.quads[str(el._id)] = False
            elif 'KICKER' in el.type:
                optimizerSettings.inputs.kickers[str(el._id)] = True

    @classmethod
    def init_currents(cls, beam, models):
        def is_kick_field(field):
            return re.search(r'^(.?kick|k1)$', field)
        ac = AmpConverter(beam)
        for el in cls.beamline_elements(models):
            for f in list(el.keys()):
                if is_kick_field(f):
                    el[cls.current_field(f)] = ac.kick_to_current(el[f])

    @classmethod
    def init_process_variables(cls, models):
        pvs = []
        def _add_pv(elId, dim, write='0'):
            pvs.append(PKDict(
                elId=elId,
{'fraction_non_alphanumeric': 0.10214782696379954, 'fraction_numerical': 0.0040334778662902086, 'mean_word_length': 1.9882494727327509, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4411393', 'n_tokens_mistral': 3153, 'n_tokens_neox': 2956, 'n_words': 565}