text:
"""writes the MPT/MPTS table"""
from __future__ import annotations
from copy import deepcopy
from collections import defaultdict
from struct import pack, Struct
from typing import List, Union, TYPE_CHECKING

import numpy as np

from .geom1_writer import write_geom_header, close_geom_table
from .geom4_writer import write_header, write_header_nvalues

integer_types = (int, np.int32, np.int64)
if TYPE_CHECKING:  # pragma: no cover
    from cpylog import SimpleLogger
    from pyNastran.bdf.cards.aero.aero import CAERO1, CAERO2, PAERO1, PAERO2, AESURF, AESURFS
    from pyNastran.bdf.cards.aero.static_loads import AEROS, AESTAT # , CSSCHD, DIVERG, TRIM, TRIM2
    from pyNastran.bdf.cards.aero.dynamic_loads import AERO, MKAERO1, FLUTTER # , FLFACT, MKAERO2
    from pyNastran.op2.op2_geom import OP2Geom, BDF

def write_edt(op2_file, op2_ascii, model: Union[BDF, OP2Geom], endian: bytes=b'<') -> None:
    """writes the EDT/EDTS table"""
    if not hasattr(model, 'loads'):  # OP2
        return
    card_types = [
        'MKAERO1', # 'MKAERO2',
        'AERO', 'AEROS',
        'CAERO1', 'PAERO1', 'SPLINE1',
        'CAERO2', 'PAERO2', 'SPLINE2',
        #'AELIST', 'AEFACT', 'AESURF', 'AESURFS'
        'AESURF', 'AESTAT',
        'TRIM', 'FLUTTER',
        'DEFORM',
        'FLFACT',
    ]

    cards_to_skip = [
    ]
    #card_types = [
        #'CAERO1', 'CAERO2',
        #'SPLINE1', 'SPLINE2',
        #'AEFACT', 'AELIST', 'AESTAT',
    #]
    out = defaultdict(list)
    # geometry
    for unused_load_id, loads in model.loads.items():
        for load in loads:
            if load.type in ['DEFORM', 'CLOAD']:
                out[load.type].append(load)

    for eid, caero in sorted(model.caeros.items()):
        out[caero.type].append(eid)
    for pid, paero in sorted(model.paeros.items()):
        out[paero.type].append(pid)
    for spline_id, spline in sorted(model.splines.items()):
        out[spline.type].append(spline_id)
    for aesurf_id, aesurf in sorted(model.aesurf.items()):
        out[aesurf.type].append(aesurf_id)
    for aesurfs_id, aesurfs in sorted(model.aesurfs.items()):
        out[aesurfs.type].append(aesurfs_id)
    for aelink_id, aelink in sorted(model.aelinks.items()):
        out[aelink.type].append(aelink_id)
    for name, aecomp in sorted(model.aecomps.items()):
        out[aecomp.type].append(name)
    #for name, aecompl in sorted(model.aecompl.items()):
        #out[aecompl.type].append(name)

    # loads
    for trim_id, trim in sorted(model.trims.items()):
        out[trim.type].append(trim_id)
    for aestat_id, aestat in sorted(model.aestats.items()):
        out[aestat.type].append(aestat_id)
    if model.aeros:
        out[model.aeros.type].append(model.aeros)

    # flutter
    for flutter_id, flutter in sorted(model.flutters.items()):
        out[flutter.type].append(flutter_id)
    for flfact_id, flfact in sorted(model.flfacts.items()):
        out[flfact.type].append(flfact_id)
    for mkaero in model.mkaeros:
        out[mkaero.type].append(mkaero)
    if model.aero:
        out[model.aero.type].append(model.aero)

    remove_unsupported_cards(out, card_types, model.log)
    # other
    if len(out) == 0:
        return

    write_geom_header(b'EDT', op2_file, op2_ascii, endian=endian)
    itable = -3

    for name, ids in sorted(out.items()):
        model.log.debug('EDT %s %s' % (name, ids))
        ncards = len(ids)
        assert ncards > 0, ncards
        if name in cards_to_skip:
            model.log.warning('skipping EDT-%s' % name)
            continue

        #if nmaterials == 0:
            #continue
        try:
            func = EDT_MAP[name]
        except KeyError:  # pragma: no cover
            raise NotImplementedError(name)

        nbytes = func(model, name, ids, ncards, op2_file, op2_ascii, endian)
        op2_file.write(pack('i', nbytes))
        itable -= 1
        data = [
            4, itable, 4,
            4, 1, 4,
            4, 0, 4]
        op2_file.write(pack('9i', *data))
        op2_ascii.write(str(data) + '\n')

    #-------------------------------------
    #print('itable', itable)
    close_geom_table(op2_file, op2_ascii, itable)
    #-------------------------------------

def remove_unsupported_cards(card_dict: Dict[str, Any],
                             card_types: List[str],
                             log: SimpleLogger):

    for card_type in list(card_dict):
        if card_type not in card_types:
            del card_dict[card_type]
            log.warning(f"removing {card_type} in OP2 writer because it's unsupported")

def _write_trim(model: Union[BDF, OP2Geom], name: str,
                trim_ids: List[int], ncards: int,
                op2_file, op2_ascii, endian: bytes) -> int:
    """
    (2402, 24, 342)
    MSC 2018.2

    Word Name Type Description
    1 ID           I
    2 MACH        RS
    3 Q           RS
    4 AEQRATIO    RS

    5 LABEL(2) CHAR4
    7 UX          RS
    Words 5 through 7 repeat until (-1,-1,-1) occurs

    """
    #struct1 = Struct(endian + b'i 3f')
    #struct2 = Struct(endian + b'8sf')
    #struct_end = Struct(endian + b'3i')
    key = (2402, 24, 342)
    #nfields = 16

    nlabels_total = 0
    for trim_id in trim_ids:
        trim = model.trims[trim_id]
        assert len(trim.labels) == len(trim.uxs), trim.get_stats()
        nlabels = len(trim.labels)
        #nuxs = len(trim.uxs)
        nlabels_total += nlabels # + nuxs

    # the *3 comes from:
    #  - label (2 fields of 4s)
    #  - ux
    nvalues = 7 * ncards + nlabels_total * 3
    nbytes = write_header_nvalues(name, nvalues, key, op2_file, op2_ascii)

    all_data = []
    for trim_id in trim_ids:
        trim = model.trims[trim_id]
        nlabels = len(trim.labels)
        data = [trim.sid, trim.mach, trim.q, trim.aeqr]
        for label, ux in zip(trim.labels, trim.uxs):
            label_bytes = b'%-8s' % label.encode('latin1')
            data.extend([label_bytes, ux])
        data += [-1, -1, -1]
        assert None not in data, data
        fmt = endian + b'ifff ' + b' 8sf' * nlabels + b' 3i'
        structi = Struct(fmt)
        op2_ascii.write(f'  TRIM data={data}\n')
        op2_file.write(structi.pack(*data))
        all_data += data
    return nbytes

def _write_caero1(model: Union[BDF, OP2Geom], name: str,
                  caero_ids: List[int], ncards: int,
                  op2_file, op2_ascii, endian: bytes) -> int:
    """
    MSC 2018.2

    Word Name Type Description
    1 EID    I
    2 PID    I
    3 CP     I
    4 NSPAN  I
    5 NCHORD I
    6 LSPAN  I
    7 LCHORD I
    8 IGID   I
    9 X1    RS
    10 Y1   RS
    11 Z1   RS
    12 X12  RS
    13 X4   RS
    14 Y4   RS
    15 Z4   RS
    16 X43  RS

    """
    key = (3002, 30, 263)
    nfields = 16
    structi = Struct(endian + b'8i 8f')
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for caero_id in caero_ids:
        caero = model.caeros[caero_id] # type: CAERO1
        x1, y1, z1 = caero.p1
        x4, y4, z4 = caero.p4
        #print(caero.get_stats())
        data = [caero.eid, caero.pid, caero.cp,
                caero.nspan, caero.nchord,
                caero.lspan, caero.lchord,
                caero.igroup, x1, y1, z1, caero.x12,
                x4, y4, z4, caero.x43]

        assert None not in data, data
        op2_ascii.write(f'  CAERO1 data={data}\n')
        op2_file.write(structi.pack(*data))
    return nbytes

def _write_caero2(model: Union[BDF, OP2Geom], name: str,
                  caero_ids: List[int], ncards: int,
                  op2_file, op2_ascii, endian: bytes) -> int:
    """
    MSC 2018.2

    Word Name Type Description
    1 EID          I
    2 PID          I
    3 CP           I
    4 NSB          I
    5 NINT         I
    6 LSB          I
    7 LINT         I
    8 IGID         I
    9 X1          RS
    10 Y1         RS
    11 Z1         RS
    12 X12        RS
    13 UNDEF(4) none

    data = (54000, 4020, 0, 8, 8, 0, 0, 1, -5.0, 0, 0, 40.0, 0, 0, 0, 0),
    """
    key = (4301, 43, 167)
    nfields = 16
    structi = Struct(endian + b'8i 8f')
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for caero_id in caero_ids:
        caero = model.caeros[caero_id] # type: CAERO2
        x1, y1, z1 = caero.p1
        #print(caero.get_stats())
        data = [caero.eid, caero.pid, caero.cp,
                caero.nsb, caero.nint, caero.lsb, caero.lint,
                caero.igroup, x1, y1, z1, caero.x12,
                0.0, 0.0, 0.0, 0.0]

        assert None not in data, data
        op2_ascii.write(f'  CAERO2 data={data}\n')
        op2_file.write(structi.pack(*data))
    return nbytes

def _write_paero1(model: Union[BDF, OP2Geom], name: str,
                  paero_ids: List[int], ncards: int,
                  op2_file, op2_ascii, endian: bytes) -> int:
    """
    (3102, 31, 264)
    MSC 2018.2

    Word Name Type Description
    1 PID      I
    2 B1       I
    3 B2       I
    4 B3       I
    5 B4       I
    6 B5       I
    7 B6       I
    8 UNDEF none

    PAERO1  100001

    """
    key = (3102, 31, 264)
    nfields = 8
    structi = Struct(endian + b'8i')
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for paero_id in paero_ids:
        paero = model.paeros[paero_id] # type: PAERO1
        n0s = 7 - len(paero.caero_body_ids)
        data = [paero.pid, ] + paero.caero_body_ids
        if n0s:
            data += [0] * n0s
        assert None not in data, data
        op2_ascii.write(f'  PAERO1 data={data}\n')
        op2_file.write(structi.pack(*data))
    return nbytes

def _write_paero2(model: Union[BDF, OP2Geom], name: str,
                  paero_ids: List[int], ncards: int,
                  op2_file, op2_ascii, endian: bytes) -> int:
    """
    MSC 2018.2

    Word Name Type Description
    1 PID        I
    2 ORIENT CHAR4
    3 UNDEF   none (orient carryover)
    4 WIDTH     RS
    5 AR        RS
    6 LRSB       I
    7 LRIB       I
    8 LTH1       I
    9 LTH2       I
    10 THI1      I
    11 THN1      I
    12 THI2      I
    13 THN2      I
    14 THI3      I
    15 THN3      I
    PAERO2  100001

    data = (100001, 100001, 0, 10, 0, 0, 24, 1,
            99.3, 21.45, -11.65, 42.86, 101.8387, 122.62, -2.69, 32.71)

    """
    key = (4601, 46, 170)
    nfields = 15
    structi = Struct(endian + b'i 4s 3f 10i')
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for paero_id in paero_ids:
        paero = model.paeros[paero_id] # type: PAERO2
        orient = b'%-4s' % paero.orient.encode('ascii')
        lrsb = 0 if paero.lrsb is None else paero.lrsb
        lrib = 0 if paero.lrib is None else paero.lrib
        lth1 = 0 if paero.lth[0] is None else paero.lth[0]
        lth2 = 0 if paero.lth[1] is None else paero.lth[1]

        data = [
            paero.pid, orient, 0, paero.width, paero.AR,
            lrsb, lrib, lth1, lth2,
        ]
        # we need to copy this so we don't modify the original
        thi = deepcopy(paero.thi)
        nthi = len(thi)
        if nthi < 3:
            thi.extend([0]*(3-nthi))

        thn = deepcopy(paero.thn)
        nthn = len(paero.thn)
        if nthn < 3:
            thn.extend([0]*(3-nthn))

        #assert len(paero.thi) == 3, paero.thi
        #assert len(paero.ltn) == 3, paero.ltn
        for thii, thni in zip(thi, thn):
            assert isinstance(thii, integer_types), f'i={i} thi={thii}'
            assert isinstance(thni, integer_types), f'i={i} thn={thni}'
            data.extend([thii, thni])
        assert None not in data, data
        op2_ascii.write(f'  PAERO2 data={data}\n')
        #print('npaero2', len(data), data)
        op2_file.write(structi.pack(*data))
    return nbytes

def _write_spline1(model: Union[BDF, OP2Geom], name: str,
                   spline_ids: List[int], ncards: int,
                   op2_file, op2_ascii, endian: bytes) -> int:
    """
    Word Name Type Description
    1 EID   I
    2 CAERO I
    3 BOX1  I
    4 BOX2  I
    5 SETG  I
    6 DZ    RS
    7 METHOD(2) CHAR4 Method: IPS|TPS|FPS
    9 USAGE(2) CHAR4 Usage flag: FORCE|DISP|BOTH
    11 NELEM I Number of elements for FPS on x-axis
    12 MELEM I Number of elements for FPS on y-axis

    """
    key = (3302, 33, 266)
    nfields = 12
    structi = Struct(endian + b'5if 8s 8s 2i')
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for spline_id in spline_ids:
        spline = model.splines[spline_id] # type: SPLINE1

        method = b'%-8s' % spline.method.encode('ascii')
        usage = b'%-8s' % spline.usage.encode('ascii')
        data = [spline.eid, spline.CAero(), spline.box1, spline.box2,
                spline.Set(), spline.dz, method, usage, spline.nelements, spline.melements]
        #print(spline.get_stats())
        assert None not in data, data
        op2_ascii.write(f'  SPLINE1 data={data}\n')
        op2_file.write(structi.pack(*data))
    return nbytes

def _write_spline2(model: Union[BDF, OP2Geom], name: str,
                   spline_ids: List[int], ncards: int,
                   op2_file, op2_ascii, endian: bytes) -> int:
    """
    Writes the SPLINE2 card

    Word Name Type Description
    1 EID   I
    2 CAERO I
    3 ID1   I
    4 ID2   I
    5 SETG  I
    6 DZ    RS
    7 DTOR  RS
    8 CID   I
    9 DTHX  RS
    10 DTHY RS
    11 USAGE(2) CHAR4 Usage flag: FORCE|DISP|BOTH

    """
    key = (3402, 34, 267)
    nfields = 12
    structi = Struct(endian + b'5i 2f i 2f 8s')
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for spline_id in spline_ids:
        spline = model.splines[spline_id] # type: SPLINE1

        usage_bytes = b'%-8s' % spline.usage.encode('ascii')
        dthx = 0.0 if spline.dthx is None else spline.dthx
        dthy = 0.0 if spline.dthx is None else spline.dthx
        #print(spline.get_stats())
        data = [spline.eid, spline.caero, spline.box1, spline.box2,
                spline.setg, spline.dz, spline.dtor, spline.cid,
                dthx, dthy, usage_bytes]
        assert None not in data, data
        op2_ascii.write(f'  SPLINE2 data={data}\n')
        op2_file.write(structi.pack(*data))
    return nbytes

def _write_aesurf(model: Union[BDF, OP2Geom], name: str,
                  aesurf_ids: List[int], ncards: int,
                  op2_file, op2_ascii, endian: bytes) -> int:
    """
    MSC 2018.2

    Word Name Type Description
    1 ID           I Identification of an aerodynamic trim variable degree of freedom >0, no default
    2 LABEL(2) CHAR4 Control Surface (CS) name, no default
    4 CID1         I IDentification of a rectangular Coordinate system with y-axis that defines the hinge line of the CS component
    5 ALID1        I IDentification of an AELIST bulk data entry that identifies all aerodynamic elements that make up the CS comp
    6 CID2         I IDentification of a rectangular Coordinate system with y-axis that defines the hinge line of the CS component
    7 ALID2        I IDentification of an AELIST bulk data entry that identifies all aerodynamic elements that make up the CS comp
    8 EFF         RS Control surface EFFectiveness, default=1.0
    9 LDW          I =0 create a linear down wash, >0 no linear downwash
    10 CREFC      RS REFerence Chord Length for the CS > 0.0, default = 1.0
    11 CREFS      RS REFerence area for the CS > 0.0, default = 1.0
    12 PLLIM      RS Lower deflection   Limit for the control surface in radians, default=no limit
    13 PULIM      RS Upper deflection   Limit for the control surface in radians, default=no limit
    14 HMLLIM     RS Lower Hinge Moment Limit for the control surface in force-length units, default=no limit
    15 HMULIM     RS Upper Hinge Moment Limit for the control surface in force-length units, default=no limit
    16 TQLLIM     RS Lower deflection   Limit for the control surface as fct(q), >0, default=no limit
    17 TQULIM     RS Upper deflection   Limit for the control surface as fct(q), >0, default=no limit

    """
    key = (2002, 20, 338)
    nfields = 17
    structi = Struct(endian + b'i 8s 4i fi 2f')

    structf = Struct(endian + b'f')
    structs = Struct(endian + b'4s')
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for aesurf_id in aesurf_ids:
        aesurf = model.aesurf[aesurf_id] # type: AESURF

        label_bytes = b'%-8s' % aesurf.label.encode('ascii')

        if aesurf.ldw == 'LDW':
            ldw_int = 0
        elif aesurf.ldw == 'NOLDW':
            ldw_int = 1
        else:
            raise NotImplementedError(aesurf.ldw)
        cid2 = 0 if aesurf.cid2 is None else aesurf.cid2
        alid2 = 0 if aesurf.alid2 is None else aesurf.alid2

        #print(aesurf.get_stats())
        data = [
            aesurf.aesid, label_bytes,
            aesurf.cid1, aesurf.alid1,
            cid2, alid2,
            aesurf.eff, ldw_int,
            aesurf.crefc, aesurf.crefs,
        ]
        assert None not in data, data
        op2_file.write(structi.pack(*data))

        float_strs = [
            aesurf.pllim, aesurf.pulim,
            aesurf.hmllim, aesurf.hmulim,
            aesurf.tqllim, aesurf.tqulim]
        for val in float_strs:
            if val is None:
                op2_file.write(structs.pack(b'    '))
            else:
                op2_file.write(structf.pack(val))
        op2_ascii.write(f'  AESURF data={data}; float_strs={float_strs}\n')
    return nbytes

def _write_aesurfs(model: Union[BDF, OP2Geom], name: str,
                   aesurfs_ids: List[int], ncards: int,
                   op2_file, op2_ascii, endian: bytes) -> int:
    """
    Word Name Type Description
    1 ID       I     Identification of an aerodynamic trim variable degree
                     of freedom >0, no default
    2 LABEL(2) CHAR4 Control Surface (CS) name, no default
    4 LIST1    I     Identification of a SET1 that contains the grids ids
                     associated with this control surface
    5 LIST2    I     Identification of a SET1 that contains the grids ids
                     associated with this control surface
    """
    key = (7701, 77, 581)
    nfields = 5
    structi = Struct(endian + b'i 8s 2i')
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for aesurfs_id in aesurfs_ids:
        aesurfs = model.aesurfs[aesurfs_id] # type: AESURFS
        label_bytes = b'%-8s' % aesurfs.label.encode('ascii')

        #print(aesurfs.get_stats())
        data = [
            aesurfs.aesid, label_bytes,
            aesurfs.list1, aesurfs.list2]
        assert None not in data, data
        op2_ascii.write(f'  AESURFS data={data}\n')
        op2_file.write(structi.pack(*data))
    return nbytes

def _write_aestat(model: Union[BDF, OP2Geom], name: str,
                  aestat_ids: List[int], ncards: int,
                  op2_file, op2_ascii, endian: bytes) -> int:
    """
    MSC 2018.2

    Word Name Type Description
    1 ID I
    2 LABEL(2) CHAR4

    """
    key = (2102, 21, 339)
    nfields = 3
    structi = Struct(endian + b'i 8s')
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for aestat_id in aestat_ids:
        aestat = model.aestats[aestat_id] # type: AESTAT

        label_bytes = b'%-8s' % aestat.label.encode('ascii')

        #print(aestat.get_stats())
        data = [aestat.aestat_id, label_bytes]
        assert None not in data, data
        op2_ascii.write(f'  AESTAT data={data}\n')
        op2_file.write(structi.pack(*data))
    return nbytes

def _write_flutter(model: Union[BDF, OP2Geom], name: str,
                   flutter_ids: List[int], ncards: int,
                   op2_file, op2_ascii, endian: bytes) -> int:
    """
    (3902, 39, 272)
    MSC 2018.2

    Word Name Type Description
    1 SID           I
    2 METHOD(2) CHAR4
    4 DENS          I
    5 MACH          I
    6 RFREQ         I
    7 IMETH(2)  CHAR4
    SFLG=0 (std)
      9 NEIGN  I  nvalue
      10 EPR  RS
      11 SFLG  I SWEEP FLAG
    SFLG=1 (sweep)
      9 FMAX  RS maximum frequency
      10 EPR  RS
      11 SFLG  I SWEEP FLAG
    End SFLG
    Words 1 through max repeat until End of Record

    NX:
      data = (30, PK, 1, 2, 3, L, 3, 0.001, -1)
    """
    key = (3902, 39, 272)
    nfields = 11
    structi = Struct(endian + b'i 8s 3i 8s ifi')
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for flutter_id in flutter_ids:
        flutter = model.flutters[flutter_id]  # type: FLUTTER
        #print(flutter.get_stats())
        assert isinstance(flutter.method, str), f'flutter.method={flutter.method}; type={type(flutter.method)}'
        assert isinstance(flutter.imethod, str), f'flutter.imethod={flutter.imethod}; type={type(flutter.imethod)}'

        #if isinstance(flutter.method, str):
        method = b'%-8s' % flutter.method.encode('latin1')
        imethod = b'%-8s' % flutter.imethod.encode('latin1')
        #if 0:
            #try:
                #method = b'-%8s' % flutter.method.decode('latin1')
                #imethod = b'-%8s' % flutter.imethod.decode('latin1')
            #except AttributeError:
                #print(flutter.get_stats())
                #raise
        nvalue = flutter.nvalue
        if nvalue is None:
            nvalue = 0
        #assert isinstance(flutter.nvalue, int), f'flutter.nvalue={flutter.nvalue}; type={type(flutter.nvalue)}'

        data = [flutter.sid, method,
                flutter.density, flutter.mach, flutter.reduced_freq_velocity,
                imethod, nvalue, flutter.epsilon, -1]
        assert len(data) == 9, data
        assert None not in data, data
        op2_ascii.write(f'  FLUTTER data={data}\n')
        op2_file.write(structi.pack(*data))
    return nbytes

def _makero_temp(data, i: int, nloops: int):
    if i == (nloops - 1):
        data_temp = data[i*8:]
    else:
        #machs_temp = [-1] * 8
        data_temp = data[i*8:i*8+8]
    return data_temp

def _write_mkaero1(model: Union[BDF, OP2Geom], name: str,
                   mkaero1s: List[MKAERO1], ncards: int,
                   op2_file, op2_ascii, endian: bytes) -> int:
    """writes the MKAERO1

    data = (1.3, -1, -1, -1, -1, -1, -1, -1,
           0.03, 0.04, 0.05, -1, -1, -1, -1, -1)
    """
    key = (3802, 38, 271)

    makero1s_temp = []
    makero1s_final = []
    for mkaero in mkaero1s:
        nmachs = len(mkaero.machs)
        nkfreqs = len(mkaero.reduced_freqs)
        assert nmachs > 0, mkaero
        assert nkfreqs > 0, mkaero

        if nmachs <= 8 and nkfreqs <= 8:
            # no splitting required
            makero1s_final.append((mkaero.machs, mkaero.reduced_freqs))
        elif nmachs <= 8 or nkfreqs <= 8:
            # one of machs or kfreqs < 8
            makero1s_temp.append((mkaero.machs, mkaero.reduced_freqs))
        else:
            # both machs and kfreqs > 8
            nloops_mach = int(np.ceil(nmachs/8))
            for i in range(nloops_mach):
                machs_temp = _makero_temp(mkaero.machs, i, nloops_mach)
                assert len(machs_temp) > 0, (i, nloops_mach, machs_temp)
                makero1s_temp.append((machs_temp, mkaero.reduced_freqs))

    for (machs, reduced_freqs) in makero1s_temp:
        nmachs = len(machs)
        nkfreqs = len(reduced_freqs)
        assert nmachs > 0, nmachs
        assert nkfreqs > 0, nkfreqs

        if nmachs <= 8 and nkfreqs <= 8:  # pragma: no cover
            raise RuntimeError(f'this should never happen...nmachs={nmachs} knfreqs={nkfreqs}')
        if nmachs <= 8:
            # nkfreqs > 8
            nloops = int(np.ceil(nkfreqs/8))
            for i in range(nloops):
                reduced_freqs_temp = _makero_temp(reduced_freqs, i, nloops)
                makero1s_final.append((machs, reduced_freqs_temp))
        elif nkfreqs <= 8:
            # nmachs > 8
            nloops = int(np.ceil(nmachs/8))
            for i in range(nloops):
                machs_temp = _makero_temp(machs, i, nloops)
                assert len(machs_temp) > 0, (i, nloops_mach, machs_temp)
                makero1s_final.append((machs_temp, reduced_freqs))
        else:  # pragma: no cover
            raise RuntimeError(f'this should never happen...nmachs={nmachs} knfreqs={nkfreqs}')
            #raise RuntimeError((nmachs, nkfreqs))

    ncards = len(makero1s_final)
    nfields = 16
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for machs, reduced_freqs in makero1s_final:
        data = []
        nmachs = len(machs)
        nkfreqs = len(reduced_freqs)
        assert nmachs > 0, machs
        assert nkfreqs > 0, reduced_freqs

        nint_mach = 8 - nmachs
        nint_kfreq = 8 - nkfreqs
        fmt1 = b'%if' % nmachs + b'i' * nint_mach
        fmt2 = b'%if' % nkfreqs + b'i' * nint_kfreq
        spack = Struct(endian + fmt1 + fmt2)

        data.extend(machs.tolist())
        assert nint_mach < 8, nint_mach
        if nint_mach:
            data.extend([-1]*nint_mach)
        data.extend(reduced_freqs.tolist())
        if nint_kfreq:
            data.extend([-1]*nint_kfreq)
        op2_ascii.write(f'  mkaero1 data={data}\n')
        op2_file.write(spack.pack(*data))
    return nbytes

def _write_aero(model: Union[BDF, OP2Geom], name: str,
                aero: List[AERO], ncards: int,
                op2_file, op2_ascii, endian: bytes) -> int:
    """
    Word Name Type Description
    1 ACSID     I
    2 VELOCITY RS
    3 REFC     RS
    4 RHOREF   RS
    5 SYMXZ     I
    6 SYMXY     I

    """
    aeroi = aero[0]
    spack = Struct(endian + b'i 3f 2i')

    key = (3202, 32, 265)
    nfields = 6
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    # cref   : 0.3048
    # is_anti_symmetric_xy : False
    # is_anti_symmetric_xz : False
    # is_symmetric_xy : False
    # is_symmetric_xz : False
    # rho_ref : 1.225
    # sym_xy : 0
    # sym_xz : 0
    # type   : 'AERO'
    # velocity : 0.0
    velocity = 0.0
    if aeroi.velocity is not None:
        velocity = aeroi.velocity
    data = [aeroi.acsid, velocity, aeroi.cref, aeroi.rho_ref, aeroi.sym_xz, aeroi.sym_xy]
    op2_ascii.write(f'  AERO data={data}\n')
    op2_file.write(spack.pack(*data))
    return nbytes

def _write_aeros(model: Union[BDF, OP2Geom], name: str,
                 aeros: List[AEROS], ncards: int,
                 op2_file, op2_ascii, endian: bytes) -> int:
    """
    AEROS(2202, 22, 340)

    AEROS   0       100     36.     360.    12960.
    data = (0, 100, 36.0, 360.0, 12960.0, 0, 0)

    """
    aeroi = aeros[0]
    #print(aeroi.get_stats())
    spack = Struct(endian + b'2i 3f 2i')

    key = (2202, 22, 340)
    nfields = 7
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    data = [aeroi.acsid, aeroi.rcsid,
            aeroi.cref, aeroi.bref, aeroi.sref,
            aeroi.sym_xz, aeroi.sym_xy]
    #print(data)
    assert None not in data, data
    op2_ascii.write(f'  AEROS data={data}\n')
    op2_file.write(spack.pack(*data))
    return nbytes

def _write_deform(model: Union[BDF, OP2Geom], name: str,
                  loads: List[AEROS], ncards: int,
                  op2_file, op2_ascii, endian: bytes) -> int:
    """
    (104, 1, 81)
    NX 2019.2

    Word Name Type Description
    1 SID I Deformation set identification number
    2 EID I Element number
    3 D RS Deformation

    """
    key = (104, 1, 81)
    nfields = 3
    structi = Struct(endian + b'iif')
    nbytes = write_header(name, nfields, ncards, key, op2_file, op2_ascii)

    for load in loads:
        data = [load.sid, load.eid, load.deformation]
        #flutter = model.loads[flutter_id]  # type: FLUTTER
        #print(flutter.get_stats())
        assert None not in data, data
        op2_ascii.write(f'  DEFORM data={data}\n')
        op2_file.write(structi.pack(*data))
    return nbytes

def _write_flfact(model: Union[BDF, OP2Geom], name: str,
                  flfact_ids, ncards: int,
                  op2_file, op2_ascii, endian: bytes) -> int:
    """
    (4102, 41, 274)
    NX 2019.2

    data = (1, 0.206, -1,
            2, 1.3, -1,
            3, 14400.0, 15600.0, 16800.0, 18000.0, 19200.0, 20400.0, -1)

    """
    key = (4102, 41, 274)
    #nfields = 3

    fmt = ''
    data = []
    for flfact_id in flfact_ids:
        flfact = model.flfacts[flfact_id]
        #print(flfact.get_stats())
        nfactors = len(flfact.factors)
        fmt += 'i %if i' % nfactors
        if isinstance(flfact.factors, np.ndarray):
            factors = flfact.factors.tolist()
        else:
            factors = flfact.factors
        datai = [flfact.sid] + factors + [-1]
        data.extend(datai)
        #flutter = model.loads[flutter_id]  # type: FLUTTER
        assert None not in datai, datai
        op2_ascii.write(f'  FLFACT data={data}\n')

    nvalues = len(data)
    nbytes = write_header_nvalues(name, nvalues, key, op2_file, op2_ascii)
    structi = Struct(fmt)
    op2_file.write(structi.pack(*data))
    return nbytes

EDT_MAP = {
    'CAERO1': _write_caero1,
    'CAERO2': _write_caero2,
    'PAERO1': _write_paero1,
    'PAERO2': _write_paero2,
    'MKAERO1': _write_mkaero1,
    'AERO': _write_aero,
    'AEROS': _write_aeros,
    'TRIM': _write_trim,
    'FLUTTER': _write_flutter,
    'DEFORM': _write_deform,
    'FLFACT': _write_flfact,
    'SPLINE1': _write_spline1,
    'SPLINE2': _write_spline2,
    'AESURF': _write_aesurf,
    #'AESURFS': _write_aesurfs,
    'AESTAT': _write_aestat,
}


text:
# coding: utf-8
# Copyright (c) Scanlon Materials Theory Group
# Distributed under the terms of the MIT License.

"""
Module containing helper functions for dealing with band structures.

todo:
  * Extend get projections to allow specifying lm orbitals and atomic sites.
"""

import numpy as np
import itertools as it
from copy import deepcopy

from collections import defaultdict

from pymatgen.electronic_structure.core import Spin

def get_projections_by_branches(bs, selection, normalise=None):
    """Returns orbital projections for each branch in a band structure.

    Args:
        bs (:obj:`~pymatgen.electronic_structure.bandstructure.BandStructureSymmLine`):
            The band structure.
        selection (list): A list of :obj:`tuple` or :obj:`string`
            identifying which projections to return. Projections can be
            specified by both element and orbital, for example::

                [('Sn', 's'), ('Bi', 'p'), ('S', 'p')]

            If just the element is specified then all the orbitals of that
            element are combined. For example, the following will combine
            all the S orbitals into a single projection::

                [('Bi', 's'), ('Bi', 'p'), 'S']

            Particular orbitals can also be combined, for example::

                [('Bi', 's'), ('Bi', 'p'), ('S', ('s', 'p', 'd'))]

        normalise (:obj:`str`, optional): Normalisation the projections.
            Options are:

              * ``'all'``: Projections normalised against the sum of all
                   other projections.
              * ``'select'``: Projections normalised against the sum of the
                   selected projections.
              * ``None``: No normalisation performed.

            Defaults to ``None``.

    Returns:
        list: A ``list`` of orbital projections for each branch of the band
        structure, in the same order as specified in ``selection``, with
        the format::

            [ [ {spin: projections} ], [ {spin: projections} ], ... ]

        Where spin is a :obj:`pymatgen.electronic_structure.core.Spin`
        object and projections is a :obj:`numpy.array` of::

            projections[band_index][kpoint_index]

        If there are no projections in the band structure, then an array of
        zeros is returned for each spin.
    """
    spins = bs.bands.keys()
    projections = get_projections(bs, selection, normalise=normalise)

    branches = []
    for b in bs.branches:
        s = b['start_index']
        e = b['end_index'] + 1

        branch_proj = deepcopy(projections)
        for spin, i in it.product(spins, range(len(projections))):
            branch_proj[i][spin] = projections[i][spin][:, s:e]

        branches.append(branch_proj)
    return branches

def get_projections(bs, selection, normalise=None):
    """Returns orbital projections from a band structure.

    Args:
        bs (:obj:`~pymatgen.electronic_structure.bandstructure.BandStructureSymmLine`):
            The band structure.
        selection (list): A list of :obj:`tuple` or :obj:`string`
            identifying which projections to return. Projections can be
            specified by both element and orbital, for example::

                [('Bi', 's'), ('Bi', 'p'), ('S', 'p')]

            If just the element is specified then all the orbitals of
            that element are combined. For example, the following will combine
            all the S orbitals into a single projection::

                [('Bi', 's'), ('Bi', 'p'), 'S']

            Particular orbitals can also be combined, for example::

                [('Bi', 's'), ('Bi', 'p'), ('S', ('s', 'p', 'd'))]

        normalise (:obj:`str`, optional): Normalisation the projections.
            Options are:

              * ``'all'``: Projections normalised against the sum of all
                   other projections.
              * ``'select'``: Projections normalised against the sum of the
                   selected projections.
              * ``None``: No normalisation performed.

            Defaults to ``None``.

    Returns:
        list: A ``list`` of orbital projections, in the same order as specified
        in ``selection``, with the format::

            [ {spin: projections}, {spin: projections} ... ]

        Where spin is a :obj:`pymatgen.electronic_structure.core.Spin`
        object and projections is a :obj:`numpy.array` of::

            projections[band_index][kpoint_index]

        If there are no projections in the band structure, then an array of
        zeros is returned for each spin.
    """
    spins = bs.bands.keys()
    nbands = bs.nb_bands
    nkpts = len(bs.kpoints)

    # if we are to normalise the data later we need access to all projections
    elements = bs.structure.symbol_set
    all_orbitals = ['s', 'p', 'd', 'f']

    # dictio has the form: {'el1': [s, p, d, f], 'el2': [s, p, d, f]...}
    dictio = dict(zip(elements, [all_orbitals] * len(elements)))

    # bs.get_projection_on_elements_and_orbitals() returns the data in a
    # really frustrating format, namely:
    #     {spin: [band_index][kpoint_index]{element: {orbital: projection}}}
    all_proj = bs.get_projections_on_elements_and_orbitals(dictio)

    # Make a defaultdict of defaultdicts
    dict_proj = defaultdict(lambda: defaultdict(dict))
    sum_proj = {s: np.zeros((nbands, nkpts)) for s in spins}

    # store the projections for all elements and orbitals in a useable format
    for spin, element, orbital in it.product(spins, elements, all_orbitals):

        # convert data to [nb][nk]
        el_orb_proj = [[all_proj[spin][nb][nk][element][orbital] for nk in range(nkpts)]
                       for nb in range(nbands)]

        dict_proj[element][orbital][spin] = np.array(el_orb_proj)

        if normalise == 'all':
            sum_proj[spin] += el_orb_proj

    # now go through the selected orbitals and extract what's needed
    spec_proj = []
    for spec in selection:

        if isinstance(spec, str):
            # spec is just an element type, therefore sum all orbitals
            element = spec
            orbitals = all_orbitals
        else:
            element, orbitals = spec
            # even if there is only one orbital, make sure we can loop over it
            orbitals = tuple(orbitals)

        proj = {s: np.zeros((nbands, nkpts)) for s in spins}
        for spin, orbital in it.product(spins, orbitals):
            proj[spin] += dict_proj[element][orbital][spin]

            if normalise == 'select':
                sum_proj[spin] += dict_proj[element][orbital][spin]

        spec_proj.append(proj)

    if normalise:
        # to prevent warnings/errors relating to divide by zero,
        # catch warnings and surround divide with np.nan_to_num
        with np.errstate(divide='ignore', invalid='ignore'):
            for spin, i in it.product(spins, range(len(spec_proj))):
                spec_proj[i][spin] = np.nan_to_num(spec_proj[i][spin] / sum_proj[spin])
    return spec_proj

def get_reconstructed_band_structure(list_bs, efermi=None):
    """Combine a list of band structures into a single band structure.

    This is typically very useful when you split non self consistent
    band structure runs in several independent jobs and want to merge back
    the results.

    This method will also ensure that any BandStructure objects will contain
    branches.

    Args:
        list_bs (:obj:`list` of \
        :obj:`~pymatgen.electronic_structure.bandstructure.BandStructure` \
        or :obj:`~pymatgen.electronic_structure.bandstructure.BandStructureSymmLine`):
            The band structures.
        efermi (:obj:`float`, optional): The Fermi energy of the reconstructed
            band structure. If `None`, an average of all the Fermi energies
            across all band structures is used.

    Returns:
        :obj:`pymatgen.electronic_structure.bandstructure.BandStructure` or \
        :obj:`pymatgen.electronic_structure.bandstructureBandStructureSymmLine`:
        A band structure object. The type depends on the type of the band
        structures in ``list_bs``.
    """
    if efermi is None:
        efermi = sum([b.efermi for b in list_bs]) / len(list_bs)

    kpoints = np.concatenate([[k.frac_coords for k in bs.kpoints] for bs in list_bs])
    nb_bands = min([list_bs[i].nb_bands for i in range(len(list_bs))])
    dicts = [bs.labels_dict for bs in list_bs]
    labels_dict = {k: v.frac_coords for d in dicts for k, v in d.items()}

    eigenvals = {}
    projections = {}
    for spin in list_bs[0].bands.keys():
        spin_bands = [bs.bands[spin][:nb_bands] for bs in list_bs]
        eigenvals[spin] = np.concatenate(spin_bands, axis=1)

        if len(list_bs[0].projections) != 0:
            spin_projs = [bs.projections[Spin.up][:nb_bands] for bs in list_bs]
            projections[Spin.up] = np.concatenate(spin_projs, axis=1)

    bs = type(list_bs[0])(
        kpoints,
        eigenvals,
        list_bs[0].lattice_rec,
        efermi,
        labels_dict,
        structure=list_bs[0].structure,
        projections=projections
    )
    return force_branches(bs)

def force_branches(bandstructure):
    """Force a linemode band structure to contain branches.

    Branches give a specific portion of the path from one high-symmetry point
    to another. Branches are required for the plotting methods to function correctly.
    Unfortunately, due to the pymatgen BandStructure implementation they require
    duplicate k-points in the band structure path. To avoid this unnecessary
    computational expense, this function can reconstruct branches in band structures
    without the duplicate k-points.

    Args:
        bandstructure: A band structure object.

    Returns:
        A band structure with brnaches.
    """
    kpoints = np.array([k.frac_coords for k in bandstructure.kpoints])
    labels_dict = {k: v.frac_coords for k, v in bandstructure.labels_dict.items()}

    # pymatgen band structure objects support branches. These are formed when
    # two kpoints with the same label are next to each other. This bit of code
    # will ensure that the band structure will contain branches, if it doesn't
    # already.
    dup_ids = []
    high_sym_kpoints = tuple(map(tuple, labels_dict.values()))
    for i, k in enumerate(kpoints):
        dup_ids.append(i)
        if (tuple(k) in high_sym_kpoints and i != 0 and i != len(kpoints) - 1 and (
                not np.array_equal(kpoints[i + 1], k) or not np.array_equal(kpoints[i - 1], k))):
            dup_ids.append(i)

    kpoints = kpoints[dup_ids]

    eigenvals = {}
    projections = {}
    for spin, spin_energies in bandstructure.bands.items():
        eigenvals[spin] = spin_energies[:, dup_ids]
        if len(bandstructure.projections) != 0:
            projections[spin] = bandstructure.projections[spin][:, dup_ids]

    return type(bandstructure)(
        kpoints,
        eigenvals,
        bandstructure.lattice_rec,
        bandstructure.efermi,
        labels_dict,
        structure=bandstructure.structure,
        projections=projections
    )

def string_to_spin(spin_string):
    """Function to convert 'spin' cli argument to pymatgen Spin object"""
    if spin_string in ['up', 'Up', '1', '+1']:
        return Spin.up

    elif spin_string in ['down', 'Down', '-1']:
        return Spin.down

    elif spin_string is None:
        return None

    else:
        raise ValueError("Unable to parse 'spin' argument")

ext: py

max_stars_repo_name: mkrack/sumo

max_issues_repo_name: mkrack/sumo

max_forks_repo_name: mkrack/sumo

reason: import numpy


text:
import numpy as np
from gym import Env
from .contextual_point_mass import ContextualPointMass

class ContextualPointMass2D(Env):

    def __init__(self, context=np.array([0., 2.])):
        self.env = ContextualPointMass(np.concatenate((context, [0.])))
        self.action_space = self.env.action_space
        self.observation_space = self.env.observation_space

    def set_context(self, context):
        self.env.context = np.concatenate((context, [0.]))

    def get_context(self):
        return self.env.context.copy()

    context = property(get_context, set_context)

    def reset(self):
        return self.env.reset()

    def step(self, action):
        return self.env.step(action)

    def render(self, mode='human'):
        self.env.render(mode)

max_stars_repo_name: psclklnk/spdl

max_issues_repo_name: psclklnk/spdl

max_forks_repo_name: psclklnk/spdl


text:
import argparse

import cv2
import numpy as np
from keras.applications import vgg19
from scipy.optimize import fmin_l_bfgs_b
import keras.backend as K

# Weights for content & style
style_weight = 1
content_weight = 0.25
total_variation_weight = 1
image_width = 224
image_height = 224
iterations = 20

def deprocess_image(x):
    x = x.reshape((3, image_width, image_height))
    x = x.transpose((1, 2, 0))

    # Remove zero-center by mean pixel
    x[:, :, 0] += 103.939
    x[:, :, 1] += 116.779
    x[:, :, 2] += 123.68
    # 'BGR'->'RGB'
    x = x[:, :, ::-1]
    x = np.clip(x, 0, 255).astype('uint8')
    return x

def load_image_to_tensor(image_path):
    im = cv2.resize(cv2.imread(image_path), (224, 224)).astype(np.float32)
    im[:, :, 0] -= 103.939
    im[:, :, 1] -= 116.779
    im[:, :, 2] -= 123.68
    im = im.transpose((2, 0, 1))
    im = np.expand_dims(im, axis=0)
    return vgg19.preprocess_input(im)

# Content loss is created using mean square error
def content_loss(content_image, result_image):
    return K.sum(K.square(result_image - content_image))

# Create gram matrix of input array
def gram_matrix(x):
    matrix = K.batch_flatten(x)
    return K.dot(matrix, K.transpose(matrix))

# Style loss is based on Gram-Matrices
def style_loss(style_image, result_image):
    style_gram_matrix = gram_matrix(style_image)
    result_gram_matrix = gram_matrix(result_image)
    return K.sum(K.square(style_gram_matrix - result_gram_matrix)) / (4. * 9 * ((image_width * image_height) ** 2))

# Creates a dictionary of layer outputs based on Content and Style images
def prepare_model_layers(style_path, content_path):
    style_image = K.variable(load_image_to_tensor(style_path))
    content_image = K.variable(load_image_to_tensor(content_path))
    result_image = K.placeholder((1, 3, image_width, image_height))
    tensor_mix = K.concatenate([content_image, style_image, result_image], axis=0)
    model = vgg19.VGG19(input_tensor=tensor_mix, weights='imagenet', include_top=False)
    return dict([(layer.name, layer.output) for layer in model.layers]), result_image

def total_variation_loss(x):
    a = K.square(x[:, :, :image_height - 1, :image_width - 1] - x[:, :, 1:, :image_width - 1])
    b = K.square(x[:, :, :image_height - 1, :image_width - 1] - x[:, :, :image_height - 1, 1:])
    return K.sum(K.pow(a + b, 1.25))

def loss_compute(layer_values, result_image):
    loss = K.variable(0.)
    layer_features = layer_values['block4_conv2']
    base_image_features = layer_features[0, :, :, :]
    combination_features = layer_features[2, :, :, :]
    loss += content_weight * content_loss(base_image_features,combination_features)

    feature_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1']
    for layer_name in feature_layers:
        layer_features = layer_values[layer_name]
        style_reference_features = layer_features[1, :, :, :]
        combination_features = layer_features[2, :, :, :]
        sl = style_loss(style_reference_features, combination_features)
        loss += (style_weight / len(feature_layers)) * sl
    loss += total_variation_weight * total_variation_loss(result_image)

    # get the gradients of the generated image wrt the loss
    grads = K.gradients(loss, result_image)

    outputs = [loss]
    if isinstance(grads, (list, tuple)):
        outputs += grads
    else:
        outputs.append(grads)

    f_outputs = K.function([result_image], outputs)
    return f_outputs

def eval_loss_and_grads(x, outputs_function):
    x = x.reshape((1, 3, image_width, image_height))
    outs = outputs_function([x])
    loss_value = outs[0]
    if len(outs[1:]) == 1:
        grad_values = outs[1].flatten().astype('float64')
    else:
        grad_values = np.array(outs[1:]).flatten().astype('float64')
    return loss_value, grad_values

class Evaluator(object):

    def __init__(self, output_function):
        self.loss_value = None
        self.grads_values = None
        self.output_function = output_function

    def loss(self, x):
        assert self.loss_value is None
        loss_value, grad_values = eval_loss_and_grads(x, self.output_function)
        self.loss_value = loss_value
        self.grad_values = grad_values
        return self.loss_value

    def grads(self, x):
        assert self.loss_value is not None
        grad_values = np.copy(self.grad_values)
        self.loss_value = None
        self.grad_values = None
        return grad_values

def create_argparser():
    parser = argparse.ArgumentParser(description='Runs neural noodle')
    parser.add_argument('style_path', metavar='style_path', help='path to style image')
    parser.add_argument('content_path', metavar='content_path', help='path to content image')
    return parser

def main(style_path, content_path):
    output_image = "results/" + content_path[:content_path.rfind(".")].replace("/", "_") + "_" + style_path[:style_path.rfind(".")].replace("/","_") + "_iter{}.jpg"
    layers, result_image_tensor = prepare_model_layers(style_path, content_path)
    evaluator = Evaluator(loss_compute(layers, result_image_tensor))
    random_image = np.random.uniform(0, 255, (1, 3, image_width, image_height)) - 128.
    for i in range(iterations):
        print('Start of iteration', i)
        random_image, min_val, info = fmin_l_bfgs_b(evaluator.loss, random_image.flatten(),
                                                    fprime=evaluator.grads, maxfun=20)
        img = deprocess_image(random_image.copy())
        fname = output_image.format(i)
        cv2.imwrite(fname, img)
        print('Image saved as', fname)

if __name__ == '__main__':
    parser = create_argparser()
    args = parser.parse_args()
    main(args.style_path, args.content_path)

max_stars_repo_name: Azkel/NeuralNoodle

max_issues_repo_name: Azkel/NeuralNoodle

max_forks_repo_name: Azkel/NeuralNoodle

reason: import numpy,from scipy


text:
# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import absolute_import, division, print_function, unicode_literals
from astropy.coordinates import Angle
from astropy.coordinates import SkyCoord
from ...utils.testing import assert_quantity_allclose
from ...utils.testing import requires_data, requires_dependency
from ...utils.energy import EnergyBounds, Energy
from ...data import DataStore, ObservationList

@requires_dependency('scipy')
@requires_data('gammapy-extra')
def test_make_psftable():
    position = SkyCoord(83.63, 22.01, unit='deg')
    data_store = DataStore.from_dir('$GAMMAPY_EXTRA/datasets/hess-crab4-hd-hap-prod2')
    obs1 = data_store.obs(23523)
    obs2 = data_store.obs(23526)
    energy = EnergyBounds.equal_log_spacing(1, 10, 100, "TeV")
    energy_band = Energy([energy[0].value, energy[-1].value], energy.unit)

    psf1 = obs1.make_psf(position=position, energy=energy, rad=None)
    psf2 = obs2.make_psf(position=position, energy=energy, rad=None)
    psf1_int = psf1.table_psf_in_energy_band(energy_band, spectral_index=2.3)
    psf2_int = psf2.table_psf_in_energy_band(energy_band, spectral_index=2.3)
    obslist = ObservationList([obs1, obs2])
    psf_tot = obslist.make_mean_psf(position=position, energy=energy)
    psf_tot_int = psf_tot.table_psf_in_energy_band(energy_band, spectral_index=2.3)

    # Check that the mean PSF is consistent with the individual PSFs
    # (in this case the R68 of the mean PSF is in between the R68 of the individual PSFs)
    assert_quantity_allclose(psf1_int.containment_radius(0.68), Angle(0.1050259592154517, 'deg'))
    assert_quantity_allclose(psf2_int.containment_radius(0.68), Angle(0.09173224724288895, 'deg'))
    assert_quantity_allclose(psf_tot_int.containment_radius(0.68), Angle(0.09838901174312292, 'deg'))

max_stars_repo_name: gabemery/gammapy

max_issues_repo_name: gabemery/gammapy

max_forks_repo_name: gabemery/gammapy

reason: from astropy


text:
# !/usr/bin/python
# -*- coding: utf-8 -*-
# @time    : 2021/4/15 21:59
# @author  : Mo
# @function: encode of bert-whiteing

from __future__ import print_function, division, absolute_import, division, print_function

# 适配linux
import sys
import os
path_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "./."))
sys.path.append(path_root)
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
print(path_root)

from bertWhiteTools import txt_read, txt_write, save_json, load_json

from bert4keras.models import build_transformer_model
from bert4keras.snippets import sequence_padding
from bert4keras.tokenizers import Tokenizer
from bert4keras.backend import keras, K
from keras.models import Model

from argparse import Namespace
# from tqdm import tqdm
import pandas as pd
import numpy as np
import json
import time

class NonMaskingLayer(keras.layers.Layer):
    """
    fix convolutional 1D can"t receive masked input, detail: https://github.com/keras-team/keras/issues/4978
    thanks for https://github.com/jacoxu
    """

    def __init__(self, **kwargs):
        self.supports_masking = True
        super(NonMaskingLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        pass

    def compute_mask(self, input, input_mask=None):
        # do not pass the mask to the next layers
        return None

    def call(self, x, mask=None):
        return x

    def get_output_shape_for(self, input_shape):
        return input_shape

class BertWhiteModel:
    def __init__(self, config=None):
        """ 初始化超参数、加载预训练模型等 """
        self.config = Namespace(**config)
        self.load_pretrain_model()
        self.eps = 1e-8

    def transform_and_normalize(self, vecs, kernel=None, bias=None):
        """应用变换，然后标准化
        """
        if not (kernel is None or bias is None):
            vecs = (vecs + bias).dot(kernel)
        norms = (vecs ** 2).sum(axis=1, keepdims=True) ** 0.5
        return vecs / np.clip(norms, self.eps, np.inf)

    def compute_kernel_bias(self, vecs):
        """计算kernel和bias
        最后的变换：y = (x + bias).dot(kernel)
        """
        mu = vecs.mean(axis=0, keepdims=True)
        cov = np.cov(vecs.T)
        u, s, vh = np.linalg.svd(cov)
        W = np.dot(u, np.diag(1 / np.sqrt(s)))
        return W[:, :self.config.n_components], -mu

    def convert_to_vecs(self, texts):
        """转换文本数据为向量形式
        """

        token_ids = self.convert_to_ids(texts)
        vecs = self.bert_white_encoder.predict(x=[token_ids, np.zeros_like(token_ids)],
                                               batch_size=self.config.batch_size, verbose=self.config.verbose)
        return vecs

    def convert_to_ids(self, texts):
        """转换文本数据为id形式
        """
        token_ids = []
        for text in texts:
            # token_id = self.tokenizer.encode(text, maxlen=self.config.maxlen)[0]
            token_id = self.tokenizer.encode(text, max_length=self.config.maxlen)[0]
            token_ids.append(token_id)
        token_ids = sequence_padding(token_ids)
        return token_ids

    def load_pretrain_model(self):
        """ 加载预训练模型, 和tokenizer """
        self.tokenizer = Tokenizer(os.path.join(self.config.bert_dir, self.config.dict_path), do_lower_case=True)
        # bert-load
        if self.config.pooling == "pooler":
            bert = build_transformer_model(os.path.join(self.config.bert_dir, self.config.config_path),
                                           os.path.join(self.config.bert_dir, self.config.checkpoint_path),
                                           model=self.config.model, with_pool="linear")
        else:
            bert = build_transformer_model(os.path.join(self.config.bert_dir, self.config.config_path),
                                           os.path.join(self.config.bert_dir, self.config.checkpoint_path),
                                           model=self.config.model)
        # output-layers
        outputs, count = [], 0
        while True:
            try:
                output = bert.get_layer("Transformer-%d-FeedForward-Norm" % count).output
                outputs.append(output)
                count += 1
            except:
                break
        # pooling
        if self.config.pooling == "first-last-avg":
            outputs = [NonMaskingLayer()(output_i) for output_i in [outputs[0], outputs[-1]]]
            outputs = [keras.layers.GlobalAveragePooling1D()(fs) for fs in outputs]
            output = keras.layers.Average()(outputs)
        elif self.config.pooling == "first-last-max":
            outputs = [NonMaskingLayer()(output_i) for output_i in [outputs[0], outputs[-1]]]
            outputs = [keras.layers.GlobalMaxPooling1D()(fs) for fs in outputs]
            output = keras.layers.Average()(outputs)
        elif self.config.pooling == "cls-max-avg":
            outputs = [NonMaskingLayer()(output_i) for output_i in [outputs[0], outputs[-1]]]
            outputs_cls = [keras.layers.Lambda(lambda x: x[:, 0])(fs) for fs in outputs]
            outputs_max = [keras.layers.GlobalMaxPooling1D()(fs) for fs in outputs]
            outputs_avg = [keras.layers.GlobalAveragePooling1D()(fs) for fs in outputs]
            output = keras.layers.Concatenate()(outputs_cls + outputs_avg)
        elif self.config.pooling == "last-avg":
            output = keras.layers.GlobalAveragePooling1D()(outputs[-1])
        elif self.config.pooling == "cls-3":
            outputs = [keras.layers.Lambda(lambda x: x[:, 0])(fs) for fs in [outputs[0], outputs[-1], outputs[-2]]]
            output = keras.layers.Concatenate()(outputs)
        elif self.config.pooling == "cls-2":
            outputs = [keras.layers.Lambda(lambda x: x[:, 0])(fs) for fs in [outputs[0], outputs[-1]]]
            output = keras.layers.Concatenate()(outputs)
        elif self.config.pooling == "cls-1":
            output = keras.layers.Lambda(lambda x: x[:, 0])(outputs[-1])
        elif self.config.pooling == "pooler":
            output = bert.output
        # 最后的编码器
        self.bert_white_encoder = Model(bert.inputs, output)
        print("load bert_white_encoder success!" )

    def train(self, texts):
        """
        训练
        """
        print("读取文本数:".format(len(texts)))
        print(texts[:3])
        # 文本转成向量vecs
        vecs = self.convert_to_vecs(texts)
        # 训练, 计算变换矩阵和偏置项
        self.config.kernel, self.config.bias = self.compute_kernel_bias(vecs)
        if self.config.ues_white:
            # 生成白化后的句子, 即qa对中的q
            vecs = self.transform_and_normalize(vecs, self.config.kernel, self.config.bias)
        return vecs

    def prob(self, texts):
        """
        编码、白化后的向量
        """
        vecs_encode = self.convert_to_vecs(texts)
        if self.config.ues_white:
            vecs_encode = self.transform_and_normalize(vecs=vecs_encode, kernel=self.config.kernel, bias=self.config.bias)
        return vecs_encode

class BertWhiteFit:
    def __init__(self, config):
        # 训练
        self.bert_white_model = BertWhiteModel(config)
        self.config = Namespace(**config)
        self.docs = []

    def load_bert_white_model(self, path_config):
        """ 模型, 超参数加载 """
        # 超参数加载
        config = load_json(path_config)
        # bert等加载
        self.bert_white_model = BertWhiteModel(config)
        self.config = Namespace(**config)
        # 白化超参数初始化
        self.bert_white_model.config.kernel = np.array(self.bert_white_model.config.kernel)
        self.bert_white_model.config.bias = np.array(self.bert_white_model.config.bias)
        # 加载qa文本数据
        self.answers_dict = load_json(os.path.join(self.config.save_dir, self.config.path_answers))
        self.docs_dict = load_json(os.path.join(self.config.save_dir, self.config.path_docs))
        self.qa_idx = load_json(os.path.join(self.config.save_dir, self.config.path_qa_idx))

        # 加载问题question预训练语言模型bert编码、白化后的encode向量
        self.docs_encode = np.loadtxt(os.path.join(self.config.save_dir, self.config.path_docs_encode))
        # index of vector
        if self.config.use_annoy or self.config.use_faiss:
            from indexAnnoy import AnnoySearch
            self.annoy_model = AnnoySearch(dim=self.config.n_components, n_cluster=self.config.n_cluster)
            self.annoy_model.load(os.path.join(self.config.save_dir, self.config.path_index))
        else:
            self.docs_encode_norm = np.linalg.norm(self.docs_encode, axis=1)
        print("load_bert_white_model success!")

    def read_qa_from_csv(self, sep="\t"):
        """
        从csv文件读取QA对
        """
        # ques_answer = txt_read(os.path.join(self.config.save_dir, self.config.path_qa)) # common qa, sep="\t"
        ques_answer = txt_read(self.config.path_qa)
        self.answers_dict = {}
        self.docs_dict = {}
        self.qa_idx = {}
        count = 0
        for i in range(len(ques_answer)):
            count += 1
            if count > 320:
                break
            ques_answer_sp = ques_answer[i].strip().split(sep)
            if len(ques_answer_sp) != 2:
                print(ques_answer[i])
                continue
            question = ques_answer_sp[0]
            answer = ques_answer_sp[1]
            self.qa_idx[str(i)] = i
            self.docs_dict[str(i)] = question.replace("\n", "").strip()
            self.answers_dict[str(i)] = answer.replace("\n", "").strip()
        self.bert_white_model.config.qa_idx = self.qa_idx

    def build_index(self, vectors):
        """ 构建索引, annoy 或者 faiss """
        if self.config.use_annoy:
            from indexAnnoy import AnnoySearch as IndexSearch
        elif self.config.use_faiss:
            from indexFaiss import FaissSearch as IndexSearch
        self.index_model= IndexSearch(dim=self.config.n_components, n_cluster=self.config.n_cluster)
        self.index_model.fit(vectors)
        self.index_model.save(os.path.join(self.config.save_dir, self.config.path_index))
        print("build index")

    def load_index(self):
        """ 加载索引, annoy 或者 faiss """
        if self.config.use_annoy:
            from indexAnnoy import AnnoySearch as IndexSearch
        elif self.config.use_faiss:
            from indexFaiss import FaissSearch as IndexSearch
        self.index_model = IndexSearch(dim=self.config.n_components, n_cluster=self.config.n_cluster)
        self.index_model.load(self.config.path_index)

    def remove_index(self, ids):
        self.index_model.remove(np.array(ids))

    def predict_with_mmr(self, texts, topk=12):
        """ 维护匹配问题的多样性 """
        from mmr import MMRSum

        res = bwf.predict(texts, topk)
        mmr_model = MMRSum()
        result = []
        for r in res:
            # 维护一个 sim:dict字典存储
            r_dict = {ri.get("sim"):ri for ri in r}
            r_mmr = mmr_model.summarize(text=[ri.get("sim") for ri in r], num=8, alpha=0.6)
            r_dict_mmr = [r_dict[rm[1]] for rm in r_mmr]
            result.append(r_dict_mmr)
        return result

    def predict(self, texts, topk=12):
        """ 预训练模型bert等编码，白化, 获取这一批数据的kernel和bias"""
        texts_encode = self.bert_white_model.prob(texts)
        result = []
        if self.config.use_annoy or self.config.use_faiss:
            index_tops = self.index_model.k_neighbors(vectors=texts_encode, k=topk)
            if self.config.use_annoy:
                for i, index_top in enumerate(index_tops):
                    [dist, idx] = index_top
                    res = []
                    for j, id in enumerate(idx):
                        score = float((2 - (dist[j] ** 2)) / 2)
                        res_i = {"score": score, "text": texts[i], "sim": self.docs_dict[str(id)],
                                 "answer": self.answers_dict[str(id)]}
                        res.append(res_i)
                    result.append(res)
            else:
                distances, indexs = index_tops
                for i in range(len(distances)):
                    res = []
                    for j in range(len(distances[i])):
                        score = distances[i][j]
                        id = indexs[i][j]
                        id = id if id != -1 else len(self.docs_dict) - 1
                        res_i = {"score": score, "text": texts[i], "sim": self.docs_dict[str(id)],
                                 "answer": self.answers_dict[str(id)]}
                        res.append(res_i)
                    result.append(res)
        else:
            for i, te in enumerate(texts_encode):
                # scores = np.matmul(texts_encode, self.docs_encode_reshape)
                score = np.sum(te*self.docs_encode, axis=1) / (np.linalg.norm(te) * self.docs_encode_norm + 1e-9)
                idxs = np.argsort(score)[::-1]
                res = []
                for j in idxs[:topk]:
                    res_i = {"score": float(score[j]), "text": texts[i], "sim": self.docs_dict[str(j)],
                             "answer": self.answers_dict[str(j)]}
                    res.append(res_i)
                result.append(res)
        return result

    def trainer(self):
        """ 预训练模型bert等编码，白化, 获取这一批数据的kernel和bias """
        # 加载数据
        self.read_qa_from_csv()
        # bert编码、训练
        self.docs_encode = self.bert_white_model.train([self.docs_dict.get(str(i), "") for i in range(len(self.docs_dict))])
        self.bert_white_model.config.kernel = self.bert_white_model.config.kernel.tolist()
        self.bert_white_model.config.bias = self.bert_white_model.config.bias.tolist()
        # 存储qa文本数据
        save_json(self.bert_white_model.config.qa_idx, os.path.join(self.config.save_dir, self.config.path_qa_idx))
        save_json(self.answers_dict, os.path.join(self.config.save_dir, self.config.path_answers))
        save_json(self.docs_dict, os.path.join(self.config.save_dir, self.config.path_docs))
        # 包括超参数等
        save_json(vars(self.bert_white_model.config), os.path.join(self.config.save_dir, self.config.path_config))
        # 存储问题question预训练语言模型bert编码、白化后的encode向量
        np.savetxt(os.path.join(self.config.save_dir, self.config.path_docs_encode), self.docs_encode)
        # 索引 或者 正则化
        if self.config.use_annoy or self.config.use_faiss:
            self.build_index(self.docs_encode.astype(np.float32))
        else:
            self.docs_encode_norm = np.linalg.norm(self.docs_encode, axis=1)
        print(" bert-white-trainer success! ")

if __name__ == '__main__':
    # 训练并存储
    from bertWhiteConf import bert_white_config
    bwf = BertWhiteFit(config=bert_white_config)
    bwf.trainer()

    texts = ["小姜机器人", "你叫什么名字"]
    res = bwf.predict(texts)
    print(res)
    res_mmr = bwf.predict_with_mmr(texts)
    print(res_mmr)

    # bwf.index_model.remove([i for i in range(100)])

    while True:
        print("请输入:")
        ques = input()
        res_mmr = bwf.predict_with_mmr(texts)
        print(res_mmr)
        res = bwf.predict([ques])
        print(res)

# python bertWhiteTrain.py


text:
import math
import numpy as np

# useful functions
aval = [0.5, 1, 2, 4]
bval = [150, 120, 80, 50]

def floor_freq(l, EI, m):
    f1 = (math.pi/2/l**2)*math.sqrt(EI/m)
    return f1

def vel(f1, b, l, m, EIl, EIt):
    n40 = math.pow(((40/f1)**2-1.0)*((b/l)**4)*(EIl/EIt), 0.25)
    print(f"n40 = {n40}")
    v = 4*(0.4+0.6*n40)/(m*b*l+200)
    return v

def b_from_a(a):
    return np.interp(a, aval, bval)

def a_from_b(b):
    return np.interp(b, np.flip(bval), aval)

def vlim(f1, b, damp):
    return math.pow(b, f1*damp-1.0)

max_stars_repo_name: pcachim/eurocodepy

max_issues_repo_name: pcachim/eurocodepy

max_forks_repo_name: pcachim/eurocodepy


text:
# Project: finmarketpy project
# Filename: learn_network_structure
# Objective: compute a network graph for a group of asset return time series
# Created: 2019-11-02 12:05
# Version: 0.0
# Author: FS

# importing packages
import numpy as np
from sklearn import cluster, covariance, manifold

def learn_network_structure(ts_returns_data, names, alphas=4, cv=5, mode='cd',
							assume_centered = False,
							n_components=2, n_neighbors=5,
							eigen_solver="dense", method='standard',
							neighbors_algorithm="auto",
							random_state = None, n_jobs=None,
							standardise=False):
	"""

	Parameters
	----------
	ts_returns_data : array-like of shape [n_samples, n_instruments]
	                  time series matrix of returns

	names : array-like of shape [n_samples, 1]
	        Individual names of the financial instrument

	alphas : int or positive float, optional
	         Number of points on the grids to be used

	cv : int, optional
	     Number of folds for cross-validation splitting strategy

	mode : str, optional
	       Solver to use to compute the graph

	assume_centered : bool, optional
                      Centre the data if False.

	n_components : int
	               Number of components for the manifold

	n_neighbors: int
                 Number of neighbours to consider for each point

	eigen_solver : str
	               Algorithm to compute eigenvalues

	method : str
             Algorithm to use for local linear embedding
	neighbors_algorithm : str
	                      Algorithm to use for nearest neighbours search

	random_state : int, RandomState instance or None, optional
	               If int, random_state is the seed used by the random number generator.
	               If RandomState instance, random_state is the random number generator.
	               If None, the random number generator is the RandomState instance used by np.random.
	               Used when eigen_solver == ‘arpack’

	n_jobs : int or None, optional
	         number of parallel jobs to run

	standardise : bool
	              standardise data if True

	Returns : sklearn.covariance.graph_lasso_.GraphicalLassoCV

              sklearn.manifold.locally_linear.LocallyLinearEmbedding

              array-like of shape [n_components, n_instruments]
              Transformed embedding vectors

              array-like of shape [n_instruments, 1]
              numeric identifier of each cluster

	-------
	"""

	if not isinstance(ts_returns_data, (np.ndarray, np.generic)):
		raise TypeError("ts_returns_data must be of class ndarray")

	# learn graphical structure
	edge_model = covariance.GraphicalLassoCV(alphas=alphas, cv=cv, mode=mode,
											 assume_centered=assume_centered)
	edge_model.fit(ts_returns_data)

	# cluster using affinity propagation
	_, labels = cluster.affinity_propagation(edge_model.covariance_)
	n_labels = labels.max()
	for i in range(n_labels + 1):
		print('Cluster %i: %s' % ((i + 1), ', '.join(names[labels == i])))

	# find low-dimension embedding - useful for 2D plane visualisation
	node_position_model = manifold.LocallyLinearEmbedding(
			n_components=n_components, eigen_solver=eigen_solver,
			n_neighbors=n_neighbors, method=method,
			neighbors_algorithm=neighbors_algorithm,
			random_state=random_state, n_jobs=n_jobs)
	embedding = node_position_model.fit_transform(ts_returns_data.T).T

	if standardise:
		# standardise returns
		standard_ret = ts_returns_data.copy()
		standard_ret /= ts_returns_data.std(axis=0)

		# learn graph model
		edge_model.fit(standard_ret)

		# cluster using affinity propagation
		_, labels = cluster.affinity_propagation(edge_model.covariance_)
		n_labels = labels.max()
		for i in range(n_labels + 1):
			print('Cluster %i: %s' % ((i + 1), ', '.join(names[labels == i])))

		# find low-dimension embedding - useful for 2D plane visualisation
		node_position_model = manifold.LocallyLinearEmbedding(
				n_components=n_components, eigen_solver=eigen_solver,
				n_neighbors=n_neighbors, method=method,
				neighbors_algorithm=neighbors_algorithm,
				random_state=random_state, n_jobs=n_jobs)
		embedding = node_position_model.fit_transform(ts_returns_data.T).T

	return edge_model, node_position_model, embedding, labels

max_stars_repo_name: khosroheidari/finmarketpy

max_issues_repo_name: khosroheidari/finmarketpy

max_forks_repo_name: khosroheidari/finmarketpy


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Adaptive rejection sampling from a given distribution."""
import logging

import numpy as np
from scipy import optimize
from scipy.interpolate import RegularGridInterpolator

from .utils import get_centres

class Sampler(object):
    def __init__(self, pdf, dimensions, blocks=100, bounds=None, seed=-1):
        """Initialise sampler.

        Parameters
        ----------
        pdf : callable
            Pdf to sample from.
        dimensions : {1, 2}
            Used to define array of bounds, eg. np.array([[-1], [1]]) in 1D (from 0
            to 1) or np.array([[-2, -2], [2, 2]]) for a 2D square boundary with
            corners at (0, 0) and (1, 1).
        blocks : int
            Number of blocks to use for the adaptive sampling algorithm.
        bounds : array
            Can be used to override the automatic bounds as mentioned for
            `dimensions`.
        seed : int
            Random number generator seed. If -1, no seed will be used and results will
            vary from run to run. Pass `seed` > -1 for repeatable results.

        """
        self.logger = logging.getLogger(__name__)
        self.logger.debug("starting init")
        self.rng = np.random.default_rng(
            seed if seed is not None and seed > -1 else None
        )
        self.pdf = pdf
        if bounds is None:
            if dimensions == 1:
                self.bounds = np.array([-1, 1]).reshape(2, 1)
            elif dimensions == 2:
                self.bounds = np.array([[-2, -2], [2, 2]])
        else:
            self.bounds = bounds
        self.dims = dimensions

        # Sample the input pdf at `blocks` positions.
        self.pdf_values = np.zeros([blocks] * self.dims, dtype=np.float64)
        self.centres_list = []  # the block (bin) centres in each dimension
        self.edges_list = []  # the block (bin) edges in each dimension
        for lower, upper in self.bounds.T:
            # lower and upper bounds for each dimension
            sample_edges = np.linspace(lower, upper, blocks + 1)
            self.edges_list.append(sample_edges)
            sample_centres = get_centres(sample_edges)
            self.centres_list.append(sample_centres)
        # Now evaluate the pdf at each of the positions given in the centres_lists,
        # filling in the resulting values into self.pdf_values.
        for indices in np.ndindex(*self.pdf_values.shape):
            values = [
                centres[index] for centres, index in zip(self.centres_list, indices)
            ]
            self.pdf_values[indices] = self.pdf(np.array(values))

        # Within each grid domain, find the maximum pdf value, which will form the
        # basis for getting the 'g' function below.
        # Get an envelope of the pdf, `g` using `self.blocks` uniform pdfs that across
        # their specific intervals always have a value that is at least that of the
        # pdf in the same interval.
        self.logger.debug("starting fn max finding")
        self.max_box_values = np.zeros_like(self.pdf_values)

        def inv_pdf(x):
            return -self.pdf(np.array(x))

        for indices in np.ndindex(*self.pdf_values.shape):
            min_edges = [edges[index] for edges, index in zip(self.edges_list, indices)]
            max_edges = [
                edges[index + 1] for edges, index in zip(self.edges_list, indices)
            ]
            centres = [
                centres[index] for centres, index in zip(self.centres_list, indices)
            ]

            edges_array = np.array((min_edges, max_edges))
            max_args = []
            max_values = []
            target_value = self.pdf_values[indices]
            for x0_indices in np.ndindex(edges_array.squeeze().shape):
                x0 = [edges_array[j, k] for k, j in enumerate(x0_indices)]
                # Now perform the function minimization from `min_edges` to
                # `max_edges`, in order to find the maximum in that range. This is
                # achieved by minimizing -1 * pdf.
                self.logger.debug("calling at:{:}".format(x0))
                x_max = optimize.fmin_tnc(
                    func=inv_pdf,
                    x0=x0,
                    bounds=[(l, u) for l, u in zip(min_edges, max_edges)],
                    approx_grad=True,
                    disp=0,
                )[0]
                max_args.append(x_max)
                max_values.append(self.pdf(np.array(x_max)))
                if max_values[-1] > target_value:
                    break
            max_value = np.max(max_values)
            if np.isclose(max_value, 0):
                # The minimisation has been completed successfully.
                self.logger.debug("calling centre:{:}".format(centres))
                x_max = optimize.fmin_tnc(
                    func=inv_pdf,
                    x0=centres,
                    bounds=[(l, u) for l, u in zip(min_edges, max_edges)],
                    approx_grad=True,
                    disp=0,
                )[0]
            self.logger.debug(
                "max value:{:} at {:}".format(self.pdf(np.array(x_max)), x_max)
            )
            self.max_box_values[indices] = self.pdf(np.array(x_max))

        self.max_box_values += 1e-7
        # Prevent minuscule differences from getting buried in the floating point
        # limit.

        diffs = self.max_box_values - self.pdf_values
        invalid_indices = np.where(diffs < 0)
        self.max_box_values[invalid_indices] = np.max(self.max_box_values)
        diffs2 = self.max_box_values - self.pdf_values
        assert np.min(diffs2) >= 0, "g(x) needs to be > f(x)"
        # Trim the boundaries by reducing the bounds such that only non-zero parts of
        # the pdf within the sampling region. This is important for a very narrow
        # tophat distribution, for example.
        non_zero_mask = self.max_box_values > 1e-6
        non_zero_indices = np.where(non_zero_mask)

        bounds = np.zeros((2, dimensions), dtype=np.float64)
        for i, axes_indices in enumerate(non_zero_indices):
            # `axes_indices` contains the indices for one axis, which
            # contribute to the bounds array in one column.
            min_i, max_i = np.min(axes_indices), np.max(axes_indices)
            bounds[0, i] = self.edges_list[i][min_i]
            bounds[1, i] = self.edges_list[i][max_i + 1]
        if not np.all(non_zero_mask) and not np.all(self.bounds == bounds):
            # If there are some zero elements left, AND if the bounds have changed
            # compared to the input bounds - last condition is necessary to avoid an
            # endless loop in the 2D case.
            self.logger.debug(
                "calling init again with reduced bounds:\n{:}".format(bounds)
            )
            self.__init__(pdf, dimensions, blocks, bounds)

        self.logger.debug("starting lin_interp")
        self.lin_interp_cdf()

    def lin_interp_cdf(self):
        """Get linear interpolation for every block in terms of the inverted CDF, ie.
        an interpolation of probability vs position. This should be done separately
        for each dimension.

        These interpolators will be used in order to get from a randomly sampled
        probability to the corresponding position, which is then distributed according
        the to discrete pdf given by max_box_values. A comparison with the actual pdf
        at that point along with another randomly sampled number then completes to
        sampling process.

        What was described above produces samples across the entire region described
        by the given `bounds`. This is akin to sampling step sizes for a square
        domain if the walker is at the centre of the box. In order to further restrict
        the possible outputs of the sampling, it would be required to restrict the
        possible values returned by the linear interpolators. This can be done by
        defining inverse interpolators, that, given the bounds of the step sizes, give
        the bounds of the probabilities, which can then be used to restrict the
        probabilities with which the positions are sampled.

        """
        self.interpolators = []
        self.inverse_interpolators = []
        nr_axes = self.dims
        # Discrete cdf along each axis.
        first_discrete_cdf = np.cumsum(
            np.sum(
                self.max_box_values, axis=(tuple([i for i in range(nr_axes) if i != 0]))
            )
        )
        # Now we want the inverse of this to get the position along the first axis.
        # Also rescale this such that the maximum value is 1.
        first_discrete_cdf /= np.max(first_discrete_cdf)  # rescale to [0, 1]
        first_edges = self.edges_list[0]
        # Now we know the cdf and coordinate of each bin centre we need to linearly
        # interpolate such that the probabilities are on the x-axis, and the positions
        # on the y-axis.

        first_probs = np.hstack(
            (np.array([0]).reshape(1, 1), first_discrete_cdf.reshape(1, -1))
        )
        self.first_interpolator = RegularGridInterpolator(
            # Add the first 0 explicitly - this corresponds to the lowest coordinate
            # possible, ie. the first edge.
            first_probs.reshape(1, -1),
            first_edges.reshape(-1),
        )
        self.first_inv_interpolator = RegularGridInterpolator(
            first_edges.reshape(1, -1), first_probs.reshape(-1)
        )
        # The first interpolator in this list will be used in order to find the second
        # interpolator to use and so on, one for each dimension.
        self.interpolators.append(self.first_interpolator)
        self.inverse_interpolators.append(self.first_inv_interpolator)
        if nr_axes == 2:
            second_discrete_cdf = np.cumsum(
                np.sum(
                    self.max_box_values,
                    axis=(tuple([i for i in range(nr_axes) if i not in (0, 1)])),
                ),
                axis=1
                # Such that the cumulative sum increases in the y-direction and one
                # cumulative sum is carried out per entry in the x-axis.
            )
            filled_discrete_cdf = np.zeros(
                (second_discrete_cdf.shape[0], second_discrete_cdf.shape[1] + 1),
                dtype=np.float64,
            )
            filled_discrete_cdf[:, 1:] = second_discrete_cdf
            edges = self.edges_list[1]
            self.second_interpolators = []
            self.second_inv_interpolators = []
            for discrete_cdf_series in filled_discrete_cdf:
                discrete_cdf_series /= np.max(discrete_cdf_series)
                self.second_interpolators.append(
                    RegularGridInterpolator(
                        discrete_cdf_series.reshape(1, -1), edges.reshape(-1)
                    )
                )
                self.second_inv_interpolators.append(
                    RegularGridInterpolator(
                        edges.reshape(1, -1), discrete_cdf_series.reshape(-1)
                    )
                )
            self.interpolators.append(self.second_interpolators)
            self.inverse_interpolators.append(self.second_inv_interpolators)
        if nr_axes > 2:
            raise NotImplementedError("Higher Dimensions not Implemented")

    def sample(self, position):
        """Generate a sample from the pdf given a position.

        Parameters
        ----------
        position : array
            Position with respect to which to sample from the pdf.

        Returns
        -------
        sampled : array
            Sample.

        """
        output = []
        centre_indices = []
        if self.dims == 1:
            axes_step_bounds = np.array((-position, 1 - position)).reshape(1, 2)
        elif self.dims == 2:
            position = position.reshape(2, 1)
            axes_step_bounds = np.array([-1, 1]).reshape(1, 2) - position
        axes_step_bounds = np.clip(
            axes_step_bounds, np.min(self.bounds, axis=0), np.max(self.bounds, axis=0)
        )

        for (i, (interpolators, edges, inv_interpolators, step_bounds)) in enumerate(
            zip(
                self.interpolators,
                self.edges_list,
                self.inverse_interpolators,
                axes_step_bounds,
            )
        ):
            if i == 0:
                # Only a single interpolator.
                interpolator = interpolators
                inv_interpolator = inv_interpolators
            else:
                interpolator = interpolators[centre_indices[-1]]
                inv_interpolator = inv_interpolators[centre_indices[-1]]

            # Use the inv_interpolator in order to get the probability bounds which
            # will only return valid step sizes using the interpolator.
            min_prob, max_prob = inv_interpolator(step_bounds)

            prob = self.rng.uniform(min_prob, max_prob)

            coord = interpolator([prob])
            output.append(coord)

            interpolator_index = np.where(coord >= edges)[0][-1]
            centre_indices.append(interpolator_index)

        # Check that the probabilities in `probs` are indeed lower than those returned
        # by the original pdf.
        pdf_val = self.pdf(np.array(output))
        max_box_val = self.max_box_values[tuple(centre_indices)]
        ratio = pdf_val / max_box_val
        prob = self.rng.uniform(0, 1)

        # This comparison causes the output to be an approximation to the true pdf, as
        # opposed to simply the max_box_values representation of the pdf.
        if prob < ratio:
            return np.array(output)
        else:
            self.logger.debug("{:} more than {:}, calling again".format(prob, ratio))
            return self.sample(position=position)

max_stars_repo_name: akuhnregnier/bounded-rand-walkers

max_issues_repo_name: akuhnregnier/bounded-rand-walkers

max_forks_repo_name: akuhnregnier/bounded-rand-walkers


text:
"""
#################################################################################################
# Copyright Info :    Copyright (c) Davar Lab @ Hikvision Research Institute. All rights reserved.
# Filename       :    show_masks.py
# Abstract       :    Script for visualization

# Current Version:    1.0.0
# Date           :    2021-06-24
#################################################################################################
"""
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import cv2
import os
import time
import imageio
import json
import torch
import pyclipper
import mmcv

colorlist=[(255,0,0), (0, 255, 0), (0,0, 255), (255,255,0), (255,0,255), (0,255,255)]

def add_text(img, text, left, top, textColor, textSize):
    if isinstance(img, np.ndarray):
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img)
    fontStyle = ImageFont.truetype("./simsun.ttc", textSize, encoding="utf-8")
    draw.text((max(left,0), max(top,0)), text, textColor, font=fontStyle)
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)

def show_segmentation(img, mask_result, bboxes_result=None, alpha=0.5, num_grid=40, out_prefix="./"):

    img_mask = np.zeros(img.shape, np.uint8)
    img_mask[:, :, 2] = (mask_result * 255).astype(np.uint8)
    img_mask = cv2.addWeighted(img_mask, alpha, img, 1 - alpha, 0)

    H, W, _ = img.shape
    step_h = H / num_grid
    step_w = W / num_grid

    for i in range(num_grid):
        cv2.line(img_mask, (0, int(i * step_h)), (W, int(i * step_h)), (0, 0, 0))
        img_mask = add_text(img_mask, str(i), 0, int(i * step_h), (255, 255, 255),
                            textSize=int(min(step_h, step_w)))
        cv2.line(img_mask, (int(i * step_w), 0), (int(i * step_w), H), (0, 0, 0))
        img_mask = add_text(img_mask, str(i), int(i * step_w), 0, (255, 255, 255),
                            textSize=int(min(step_h, step_w)))
    cv2.imwrite(out_prefix + '_seg.jpg', img_mask)

def show_mask_att(img, mask_att_results, alpha=0.5, out_prefix="./"):
    img_mask_att = []

    for c_id in range(10):
        mask_att_c = mask_att_results[c_id, :, :]
        mask_att_c = mask_att_c - np.min(mask_att_c)
        mask_att_c = mask_att_c / np.max(mask_att_c)
        heatmap = cv2.applyColorMap(np.uint8(255 * mask_att_c), cv2.COLORMAP_JET)
        cam = (heatmap + np.float32(img)) / 255
        cam = cam / np.max(cam)
        img_mask_att_c = np.uint8(255 * cam)
        img_mask_att_c = cv2.putText(img_mask_att_c, 'Frame:' + str(c_id), (0, img.shape[0]), cv2.FONT_HERSHEY_SIMPLEX, 1,
                           (0, 0, 0), 2)
        img_mask_att.append(img_mask_att_c[:,:,::-1])
    imageio.mimsave(out_prefix + '_cma.gif', img_mask_att, 'GIF', duration=1)

def show_text(img, texts_result, bboxes_result, out_prefix="./"):
    img_draw = img.copy()
    for j, box in enumerate(bboxes_result):
        for i in range(0, len(box), 2):
            cv2.line(img_draw, (int(box[i]), int(box[i+1])),(int(box[(i+2)%len(box)]),  int(box[(i+3)%len(box)])), (255,0,0),2)
        img_draw = add_text(img_draw, str(texts_result[j]), int(box[i]-5), int(box[i+1]-5), (0, 0, 0), textSize=15)
    cv2.imwrite(out_prefix + '_text.jpg', img_draw)

def show_cate(img, cate_preds, resize_shape=(1000,1000), pad_size_divisor=128, grid_num=40, out_prefix="./"):
    def PadResize(img, resize_shape=(1000, 1000), size_divisor=128):
        img = mmcv.imrescale(img, resize_shape)
        padded_img = mmcv.impad_to_multiple(img, size_divisor)
        return padded_img

    img_cate = img.copy()
    img_cate = PadResize(img_cate, resize_shape=resize_shape, size_divisor=pad_size_divisor )
    H, W, _ = img_cate.shape
    step_h = H / grid_num
    step_w = W / grid_num

    for i in range(grid_num):
        cv2.line(img_cate, (0, int(i * step_h)), (W, int(i * step_h)), (0, 0, 0))
        cv2.line(img_cate, (int(i * step_w), 0), (int(i * step_w), H), (0, 0, 0))
        for j in range(grid_num):
            if cate_preds[i*grid_num + j] > 0:
                cv2.rectangle(img_cate, (int(j*step_w),int(i * step_h)), (int((j+1)*step_w), int((i+1) * step_h)), colorlist[cate_preds[i*grid_num + j]%6],-1)

    cv2.imwrite(out_prefix + '_cate.jpg', img_cate)

max_stars_repo_name: hikopensource/DAVAR-Lab-OCR

max_issues_repo_name: hikopensource/DAVAR-Lab-OCR

max_forks_repo_name: hikopensource/DAVAR-Lab-OCR