<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import itertools
import importlib
from collections import namedtuple

import numpy as np
from scipy.special import factorial
from scipy.constants import e as qe
from cobjects import CBuffer, CObject, CField
from .mad_helper import madseq_to_generator


class Drift(CObject):
    _typeid = 2
    length = CField(0, 'real', default=0.0, alignment=8)


class DriftExact(CObject):
    _typeid = 3
    length = CField(0, 'real', default=0.0, alignment=8)


class BeamMonitor(CObject):
    _typeid = 10
    num_stores = CField(0, 'int64', default=0, alignment=8)
    start = CField(1, 'int64', default=0, alignment=8)
    skip = CField(2, 'int64', default=1, alignment=8)
    out_address = CField(3, 'uint64', default=0, alignment=8)
    max_particle_id = CField(4, 'int64', default=0, alignment=8)
    min_particle_id = CField(5, 'int64', default=0, alignment=8)
    is_rolling = CField(6, 'int64', default=0, alignment=8)
    is_turn_ordered = CField(7, 'int64', default=1, alignment=8)


def append_beam_monitors_to_lattice(
        beam_elements_buffer, until_turn_elem_by_elem, until_turn_turn_by_turn,
        until_turn, skip_turns=1, min_particle_id=0, max_particle_id=0,
        initial_at_turn=0):
    num_beam_monitors_added = 0
    start_turn_by_turn = max(initial_at_turn, until_turn_elem_by_elem)

    if until_turn_turn_by_turn > start_turn_by_turn:
        bm_turn_by_turn = BeamMonitor(
            start=start_turn_by_turn,
            num_stores=(
                until_turn_turn_by_turn -
                start_turn_by_turn),
            skip=1,
            out_address=0,
            min_particle_id=min_particle_id,
            max_particle_id=max_particle_id,
            is_rolling=False,
            is_turn_ordered=True,
            cbuffer=beam_elements_buffer)
        num_beam_monitors_added += 1

    start_output_turn = max(start_turn_by_turn, until_turn_turn_by_turn)

    if until_turn > start_output_turn:
        if skip_turns <= 0:
            skip_turns = 1

        num_stores = until_turn - start_output_turn
        remainder = num_stores % skip_turns
        num_stores = num_stores // skip_turns

        if remainder > 0:
            num_stores += 1

        bm_output = BeamMonitor(
            start=start_output_turn,
            num_stores=num_stores,
            skip=skip_turns,
            out_address=0,
            min_particle_id=min_particle_id,
            max_particle_id=max_particle_id,
            is_rolling=True,
            is_turn_ordered=True,
            cbuffer=beam_elements_buffer)
        num_beam_monitors_added += 1

    return num_beam_monitors_added


class Multipole(CObject):
    _typeid = 4
    order = CField(0, 'int64', default=0, const=True, alignment=8)
    length = CField(1, 'real', default=0.0, alignment=8)
    hxl = CField(2, 'real', default=0.0, alignment=8)
    hyl = CField(3, 'real', default=0.0, alignment=8)
    bal = CField(4, 'real', default=0.0, alignment=8, pointer=True,
                    length='2 * order + 2')

    def __init__(self, order=None, knl=None, ksl=None, bal=None, **kwargs):
        if bal is None and \
                (knl is not None or ksl is not None or order is not None):
            if knl is None:
                knl = []
            if ksl is None:
                ksl = []
            if order is None:
                order = 0

            n = max((order + 1), max(len(knl), len(ksl)))
            assert(n > 0)

            _knl = np.array(knl)
            nknl = np.zeros(n, dtype=_knl.dtype)
            nknl[:len(knl)] = knl
            knl = nknl
            del(_knl)
            assert(len(knl) == n)

            _ksl = np.array(ksl)
            nksl = np.zeros(n, dtype=_ksl.dtype)
            nksl[:len(ksl)] = ksl
            ksl = nksl
            del(_ksl)
            assert(len(ksl) == n)

            order = n - 1
            bal = np.zeros(2 * order + 2)

            idx = np.array([ii for ii in range(0, len(knl))])
            inv_factorial = 1.0 / factorial(idx, exact=True)
            bal[0::2] = knl * inv_factorial[idx // 2]
            bal[1::2] = ksl * inv_factorial[idx // 2]

            kwargs["bal"] = bal
            kwargs["order"] = order

        elif bal is not None and bal and len(bal) > 2 and ((len(bal) % 2) == 0):
            kwargs["bal"] = bal
            kwargs["order"] = (len(bal) - 2) / 2

        super().__init__(**kwargs)

    @property
    def knl(self):
        idx = np.array([ii for ii in range(0, len(self.bal), 2)])
        return self.bal[idx] * factorial(idx, exact=True)

    @property
    def ksl(self):
        idx = np.array([ii for ii in range(0, len(self.bal), 2)])
        return self.bal[idx + 1] * factorial(idx // 2, exact=True)

    def set_knl(self, value, order):
        assert order <= self.order
        self.bal[order * 2] = value / factorial(order, exact=True)

    def set_ksl(self, value, order):
        assert order <= self.order
        self.bal[order * 2 + 1] = value / factorial(order, exact=True)


class Cavity(CObject):
    _typeid = 5
    voltage = CField(0, 'real', default=0.0, alignment=8)
    frequency = CField(1, 'real', default=0.0, alignment=8)
    lag = CField(2, 'real', default=0.0, alignment=8)


class XYShift(CObject):
    _typeid = 6<commit_msg>python/sixtracklib: attempts to fix issues with knl and ksl calculation<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import itertools
import importlib
from collections import namedtuple

import numpy as np
from scipy.special import factorial
from scipy.constants import e as qe
from cobjects import CBuffer, CObject, CField
from .mad_helper import madseq_to_generator


class Drift(CObject):
    _typeid = 2
    length = CField(0, 'real', default=0.0, alignment=8)


class DriftExact(CObject):
    _typeid = 3
    length = CField(0, 'real', default=0.0, alignment=8)


class BeamMonitor(CObject):
    _typeid = 10
    num_stores = CField(0, 'int64', default=0, alignment=8)
    start = CField(1, 'int64', default=0, alignment=8)
    skip = CField(2, 'int64', default=1, alignment=8)
    out_address = CField(3, 'uint64', default=0, alignment=8)
    max_particle_id = CField(4, 'int64', default=0, alignment=8)
    min_particle_id = CField(5, 'int64', default=0, alignment=8)
    is_rolling = CField(6, 'int64', default=0, alignment=8)
    is_turn_ordered = CField(7, 'int64', default=1, alignment=8)


def append_beam_monitors_to_lattice(
        beam_elements_buffer, until_turn_elem_by_elem, until_turn_turn_by_turn,
        until_turn, skip_turns=1, min_particle_id=0, max_particle_id=0,
        initial_at_turn=0):
    num_beam_monitors_added = 0
    start_turn_by_turn = max(initial_at_turn, until_turn_elem_by_elem)

    if until_turn_turn_by_turn > start_turn_by_turn:
        bm_turn_by_turn = BeamMonitor(
            start=start_turn_by_turn,
            num_stores=(
                until_turn_turn_by_turn -
                start_turn_by_turn),
            skip=1,
            out_address=0,
            min_particle_id=min_particle_id,
            max_particle_id=max_particle_id,
            is_rolling=False,
            is_turn_ordered=True,
            cbuffer=beam_elements_buffer)
        num_beam_monitors_added += 1

    start_output_turn = max(start_turn_by_turn, until_turn_turn_by_turn)

    if until_turn > start_output_turn:
        if skip_turns <= 0:
            skip_turns = 1

        num_stores = until_turn - start_output_turn
        remainder = num_stores % skip_turns
        num_stores = num_stores // skip_turns

        if remainder > 0:
            num_stores += 1

        bm_output = BeamMonitor(
            start=start_output_turn,
            num_stores=num_stores,
            skip=skip_turns,
            out_address=0,
            min_particle_id=min_particle_id,
            max_particle_id=max_particle_id,
            is_rolling=True,
            is_turn_ordered=True,
            cbuffer=beam_elements_buffer)
        num_beam_monitors_added += 1

    return num_beam_monitors_added


class Multipole(CObject):
    _typeid = 4
    order = CField(0, 'int64', default=0, const=True, alignment=8)
    length = CField(1, 'real', default=0.0, alignment=8)
    hxl = CField(2, 'real', default=0.0, alignment=8)
    hyl = CField(3, 'real', default=0.0, alignment=8)
    bal = CField(4, 'real', default=0.0, alignment=8, pointer=True,
                    length='2 * order + 2')

    def __init__(self, order=None, knl=None, ksl=None, bal=None, **kwargs):
        if bal is None and \
                (knl is not None or ksl is not None or order is not None):
            if knl is None:
                knl = []
            if ksl is None:
                ksl = []
            if order is None:
                order = 0

            n = max((order + 1), max(len(knl), len(ksl)))
            assert(n > 0)

            _knl = np.array(knl)
            nknl = np.zeros(n, dtype=_knl.dtype)
            nknl[:len(knl)] = knl
            knl = nknl
            del(_knl)
            assert(len(knl) == n)

            _ksl = np.array(ksl)
            nksl = np.zeros(n, dtype=_ksl.dtype)
            nksl[:len(ksl)] = ksl
            ksl = nksl
            del(_ksl)
            assert(len(ksl) == n)

            order = n - 1
            bal = np.zeros(2 * order + 2)

            idx = np.array([ii for ii in range(0, len(knl))])
            inv_factorial = 1.0 / factorial(idx, exact=True)
            bal[0::2] = knl * inv_factorial
            bal[1::2] = ksl * inv_factorial

            kwargs["bal"] = bal
            kwargs["order"] = order

        elif bal is not None and bal and len(bal) > 2 and ((len(bal) % 2) == 0):
            kwargs["bal"] = bal
            kwargs["order"] = (len(bal) - 2) / 2

        super().__init__(**kwargs)

    @property
    def knl(self):
        idx = np.array([ii for ii in range(0, len(self.bal), 2)])
        return self.bal[idx] * factorial(idx // 2, exact=True)

    @property
    def ksl(self):
        idx = np.array([ii for ii in range(0, len(self.bal), 2)])
        return self.bal[idx + 1] * factorial(idx // 2, exact=True)

    def set_knl(self, value, order):
        assert order <= self.order
        self.bal[order * 2] = value / factorial(order, exact=True)

    def set_ksl(self, value, order):
        assert order <= self.order
        self.bal[order * 2 + 1] = value / factorial(order, exact=True)


class Cavity(CObject):
    _typeid = 5
    voltage = CField(0, 'real', default=0.0, alignment=8)
    frequency = CField(1, 'real', default=0.0, alignment=8)
    lag = CField(2, 'real', default=0.0, alignment=8)


class XYShift(CObject):
    _typeid = 6
{'fraction_non_alphanumeric': 0.11415226529659038, 'fraction_numerical': 0.02494161606725829, 'mean_word_length': 2.2992295839753467, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '1063009', 'n_tokens_mistral': 3840, 'n_tokens_neox': 3655, 'n_words': 825}