text:
# Copyright (c) 2020, Fabio Muratore, Honda Research Institute Europe GmbH, and
# Technical University of Darmstadt.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. Neither the name of Fabio Muratore, Honda Research Institute Europe GmbH,
#    or Technical University of Darmstadt, nor the names of its contributors may
#    be used to endorse or promote products derived from this software without
#    specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL FABIO MURATORE, HONDA RESEARCH INSTITUTE EUROPE GMBH,
# OR TECHNICAL UNIVERSITY OF DARMSTADT BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
# IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""
Train an agent to solve the Qube swing-up task using Bayesian Domain Randomization.
"""
import numpy as np

import pyrado
from pyrado.algorithms.episodic.power import PoWER
from pyrado.algorithms.meta.bayrn import BayRn
from pyrado.domain_randomization.default_randomizers import (
    create_default_domain_param_map_qq,
    create_zero_var_randomizer,
)
from pyrado.domain_randomization.utils import wrap_like_other_env
from pyrado.environment_wrappers.domain_randomization import DomainRandWrapperLive, MetaDomainRandWrapper
from pyrado.environments.pysim.quanser_qube import QQubeSwingUpSim
from pyrado.logger.experiment import save_dicts_to_yaml, setup_experiment
from pyrado.policies.special.environment_specific import QQubeSwingUpAndBalanceCtrl
from pyrado.spaces import BoxSpace
from pyrado.utils.argparser import get_argparser

if __name__ == "__main__":
    # Parse command line arguments
    args = get_argparser().parse_args()

    # Experiment (set seed before creating the modules)
    ex_dir = setup_experiment(
        QQubeSwingUpSim.name,
        f"{BayRn.name}-{PoWER.name}_{QQubeSwingUpAndBalanceCtrl.name}",
        f"sim2sim_rand-mass_pend_pole-mass_rot_pole_seed-{args.seed}",
    )

    # Set seed if desired
    pyrado.set_seed(args.seed, verbose=True)

    # Environments
    env_sim_hparams = dict(dt=1 / 100.0, max_steps=600)
    env_sim = QQubeSwingUpSim(**env_sim_hparams)
    env_sim = DomainRandWrapperLive(env_sim, create_zero_var_randomizer(env_sim))
    dp_map = create_default_domain_param_map_qq()
    env_sim = MetaDomainRandWrapper(env_sim, dp_map)

    env_real = QQubeSwingUpSim(**env_sim_hparams)
    env_real.domain_param = dict(
        mass_pend_pole=0.024 * 1.1,
        mass_rot_pole=0.095 * 1.1,
    )
    env_real_hparams = env_sim_hparams
    env_real = wrap_like_other_env(env_real, env_sim)

    # PoWER and energy-based controller setup
    policy_hparam = dict(energy_gain=0.587, ref_energy=0.827, acc_max=10.0)
    policy = QQubeSwingUpAndBalanceCtrl(env_sim.spec, **policy_hparam)
    subrtn_hparam = dict(
        max_iter=5,
        pop_size=50,
        num_init_states_per_domain=4,
        num_domains=10,
        num_is_samples=5,
        expl_std_init=2.0,
        expl_std_min=0.02,
        symm_sampling=False,
        num_workers=12,
    )
    subrtn = PoWER(ex_dir, env_sim, policy, **subrtn_hparam)

    # PoWER and linear policy setup
    # policy_hparam = dict(
    #     feats=FeatureStack(identity_feat, sign_feat, abs_feat, squared_feat,
    #                        MultFeat((2, 5)), MultFeat((3, 5)), MultFeat((4, 5)))
    # )
    # policy = LinearPolicy(spec=env_sim.spec, **policy_hparam)
    # subrtn_hparam = dict(
    #     max_iter=20,
    #     pop_size=200,
    #     num_init_states_per_domain=6,
    #     num_is_samples=10,
    #     expl_std_init=2.0,
    #     expl_std_min=0.02,
    #     symm_sampling=False,
    #     num_workers=32,
    # )
    # subrtn = PoWER(ex_dir, env_sim, policy, **subrtn_hparam)

    # Set the boundaries for the GP
    dp_nom = QQubeSwingUpSim.get_nominal_domain_param()
    ddp_space = BoxSpace(
        bound_lo=np.array([0.8 * dp_nom["mass_pend_pole"], 1e-8, 0.8 * dp_nom["mass_rot_pole"], 1e-8]),
        bound_up=np.array([1.2 * dp_nom["mass_pend_pole"], 1e-7, 1.2 * dp_nom["mass_rot_pole"], 1e-7]),
    )

    # Algorithm
    bayrn_hparam = dict(
        max_iter=15,
        acq_fc="UCB",
        acq_param=dict(beta=0.25),
        acq_restarts=500,
        acq_samples=1000,
        num_init_cand=4,
        warmstart=False,
        num_eval_rollouts_real=100,
        thold_succ_subrtn=300,
    )

    # Save the environments and the hyper-parameters (do it before the init routine of BayRn)
    save_dicts_to_yaml(
        dict(env_sim=env_sim_hparams, env_real=env_real_hparams, seed=args.seed),
        dict(policy=policy_hparam),
        dict(subrtn=subrtn_hparam, subrtn_name=PoWER.name),
        dict(algo=bayrn_hparam, algo_name=BayRn.name, dp_map=dp_map),
        save_dir=ex_dir,
    )

    algo = BayRn(ex_dir, env_sim, env_real, subrtn, ddp_space, **bayrn_hparam)

    # Jeeeha
    algo.train(snapshot_mode="latest", seed=args.seed)

max_stars_repo_name: theogruner/SimuRLacra

max_stars_repo_licenses: DOC Zlib BSD-3-Clause

max_issues_repo_name: theogruner/SimuRLacra

max_issues_repo_licenses: DOC Zlib BSD-3-Clause

max_forks_repo_name: theogruner/SimuRLacra

max_forks_repo_licenses: DOC Zlib BSD-3-Clause


text:
# -*- coding: utf-8 -*-
from sympy import Symbol, Matrix, Integral, log, Rational, Derivative, exp, \
        sqrt, pi, Function, sin, cos, pprint_use_unicode, oo, Eq, Le, \
        Gt, Ne, Limit, factorial, gamma, conjugate, I, Piecewise, S
from sympy.printing.pretty import pretty as xpretty

x = Symbol('x')
y = Symbol('y')
th  = Symbol('theta')
ph  = Symbol('phi')

def pretty(expr):
    # ascii-pretty by default
    return xpretty(expr, use_unicode=False)

def test_pretty_str():
    assert pretty( 'xxx' ) == "'xxx'"
    assert pretty( "xxx" ) == "'xxx'"
    assert pretty( 'xxx\'xxx' ) == '"xxx\'xxx"'
    assert pretty( 'xxx"xxx' ) == "'xxx\"xxx'"
    assert pretty( 'xxx\"xxx' ) == "'xxx\"xxx'"
    assert pretty( "xxx'xxx" ) == '"xxx\'xxx"'
    assert pretty( "xxx\'xxx" ) == '"xxx\'xxx"'
    assert pretty( "xxx\"xxx" ) == "'xxx\"xxx'"
    assert pretty( "xxx\"xxx\'xxx" ) == '\'xxx"xxx\\\'xxx\''
    assert pretty( "xxx\nxxx\'xxx" ) == '"xxx\\nxxx\'xxx"'

def test_pretty_unicode():
    assert pretty( u'xxx' ) == "u'xxx'"
    assert pretty( u"xxx" ) == "u'xxx'"
    assert pretty( u'xxx\'xxx' ) == 'u"xxx\'xxx"'
    assert pretty( u'xxx"xxx' ) == "u'xxx\"xxx'"
    assert pretty( u'xxx\"xxx' ) == "u'xxx\"xxx'"
    assert pretty( u"xxx'xxx" ) == 'u"xxx\'xxx"'
    assert pretty( u"xxx\'xxx" ) == 'u"xxx\'xxx"'
    assert pretty( u"xxx\"xxx" ) == "u'xxx\"xxx'"
    assert pretty( u"xxx\"xxx\'xxx" ) == 'u\'xxx"xxx\\\'xxx\''
    assert pretty( u"xxx\nxxx\'xxx" ) == 'u"xxx\\nxxx\'xxx"'

def test_pretty_basic():
    # Simple numbers/symbols
    assert pretty( -Rational(1)/2 ) == '-1/2'
    assert pretty( -Rational(13)/22 ) == '  13\n- --\n  22'
    assert pretty( oo ) == 'oo'

    # Powers
    assert pretty( (x**2) ) == ' 2\nx '
    assert pretty( 1/x ) == '1\n-\nx'
    assert pretty( y*x**-2 ) == 'y \n--\n 2\nx '
    assert pretty( x**Rational(-5,2) ) == ' 1  \n----\n 5/2\nx   '
    assert pretty( (-2)**x ) == '    x\n(-2) '

    # Sums of terms
    assert pretty( (x**2 + x + 1))  in [
                '         2\n1 + x + x ',
                '     2    \nx + x  + 1',
                ' 2        \nx  + x + 1']
    assert pretty( 1-x ) in ['1 - x', '-x + 1']
    assert pretty( 1-2*x ) in ['1 - 2*x', '-2*x + 1']
    assert pretty( 1-Rational(3,2)*y/x ) in [
            '    3*y\n1 - ---\n    2*x',
            '  3*y    \n- --- + 1\n  2*x    ']

    # Multiplication
    assert pretty( x/y ) == 'x\n-\ny'
    assert pretty( -x/y ) == '-x\n--\ny '
    assert pretty( (x+2)/y ) in ['2 + x\n-----\n  y  ', 'x + 2\n-----\n  y  ']
    assert pretty( (1+x)*y ) in ['(1 + x)*y', 'y*(1 + x)', 'y*(x + 1)']

    # Check for proper placement of negative sign
    assert pretty( -5*x/(x+10) ) == ' -5*x \n------\n10 + x'
    assert pretty( 1 - Rational(3,2)*(x+1) ) == '       3*x\n-1/2 - ---\n        2 '

def test_pretty_relational():
    assert pretty(Eq(x, y)) == 'x = y'
    assert pretty(Le(x, y)) == 'x <= y'
    assert pretty(Gt(x,  y)) == 'y < x'
    assert pretty(Ne(x/(y+1), y**2)) in [
            '  x       2\n----- != y \n1 + y      ',
            '  x       2\n----- != y \ny + 1      ']

def test_pretty_unicode():
    assert xpretty( oo, use_unicode=True ) == u'\u221e'
    assert xpretty( pi, use_unicode=True ) == u'\u03c0'
    assert xpretty( pi+2*x, use_unicode=True ) in [u'\u03c0 + 2⋅x', u'2⋅x + \u03c0']
    assert xpretty( pi**2+exp(x), use_unicode=True ) in [
            u' 2    x\n\u03c0  + \u212f ',
            u' x    2\n\u212f  + \u03c0 ']
    assert xpretty( Ne(x, y), use_unicode=True ) == u'x \u2260 y'
    assert xpretty( gamma(x), use_unicode=True ) == u'\u0393(x)'

def test_pretty_unicode_defaults():
    use_unicode = pprint_use_unicode(True)
    assert xpretty(Symbol('alpha')) == u'\u03b1'
    pprint_use_unicode(False)
    assert xpretty(Symbol('alpha')) == 'alpha'

    pprint_use_unicode(use_unicode)

def test_pretty_functions():
    f = Function('f')

    # Simple
    assert pretty( (2*x + exp(x)) ) in [' x      \ne  + 2*x', '       x\n2*x + e ']
    assert pretty(abs(x)) == '|x|'
    assert pretty(abs(x/(x**2+1))) in [
            '|  x   |\n|------|\n|     2|\n|1 + x |',
            '|  x   |\n|------|\n| 2    |\n|x  + 1|']
    assert pretty(conjugate(x)) == '_\nx'
    assert pretty(conjugate(f(x+1))) in [
            '________\nf(1 + x)',
            '________\nf(x + 1)']

    # Univariate/Multivariate functions
    assert pretty(f(x)) == 'f(x)'
    assert pretty(f(x, y)) == 'f(x, y)'
    assert pretty(f(x/(y+1), y)) in [
            ' /  x     \\\nf|-----, y|\n \\1 + y   /',
            ' /  x     \\\nf|-----, y|\n \\y + 1   /',
            ]

    # Nesting of square roots
    assert pretty( sqrt((sqrt(x+1))+1) ) in [
            '   _______________\n  /       _______ \n\\/  1 + \\/ 1 + x  ',
            '   _______________\n  /   _______     \n\\/  \\/ x + 1  + 1 ']
    # Function powers
    assert pretty( sin(x)**2 ) == '   2   \nsin (x)'

    # Conjugates
    a,b = map(Symbol, 'ab')
    #assert pretty( conjugate(a+b*I) ) == '_     _\na - I*b'
    #assert pretty( conjugate(exp(a+b*I)) ) == ' _     _\n a - I*b\ne       '

def test_pretty_factorial():
    n = Symbol('n', integer=True)
    assert pretty( factorial(n) ) == 'n!'
    assert pretty( factorial(2*n) ) == '(2*n)!'
    assert pretty( factorial(factorial(factorial(n))) ) == '((n!)!)!'
    assert pretty( factorial(n+1) ) in ['(1 + n)!', '(n + 1)!']

def test_pretty_sqrt():
    assert pretty( sqrt(2) ) == '  ___\n\\/ 2 '
    assert pretty( 2**Rational(1,3) ) == '3 ___\n\\/ 2 '
    assert pretty( sqrt(x**2+1) ) == '   ________\n  /      2 \n\\/  1 + x  '
    assert pretty( 2**Rational(1,100) ) == '100___\n \\/ 2 '
    assert pretty( (1+sqrt(5))**Rational(1,3) ) == \
        '   ___________\n3 /       ___ \n\\/  1 + \\/ 5  '
    assert pretty( 2**(1/x) ) == 'x ___\n\\/ 2 '
    assert pretty( sqrt(2+pi) ) == '  ________\n\\/ 2 + pi '
    #and now everything together...
    assert pretty( (2+(1+x**2)/(2+x))**Rational(1,4)+
        (1+x**Rational(1,1000))/sqrt(3+x**2) ) == '''\
                   ____________
    1000___       /          2 \n\
1 +   \/ x       /      1 + x  \n\
----------- + 4 /   2 + ------ \n\
   ________   \/        2 + x  \n\
  /      2                     \n\
\/  3 + x                      \
'''

def test_pretty_derivatives():
    # Simple
    f_1 = Derivative(log(x), x, evaluate=False)
    assert pretty(f_1) == 'd         \n--(log(x))\ndx        '

    f_2 = Derivative(log(x), x, evaluate=False) + x
    assert pretty(f_2) in [
            '    d         \nx + --(log(x))\n    dx        ',
            'd             \n--(log(x)) + x\ndx            ']

    # Multiple symbols
    f_3 = Derivative(log(x) + x**2, x, y, evaluate=False)
    assert pretty(f_3) in [
            '   2              \n  d  / 2         \\\n-----\\x  + log(x)/\ndy dx             ',
            '   2              \n  d  /          2\\\n-----\\log(x) + x /\ndy dx             ']

    f_4 = Derivative(2*x*y, y, x, evaluate=False) + x**2
    assert pretty(f_4) in [
            '        2        \n 2     d         \nx  + -----(2*x*y)\n     dx dy       ',
            '   2             \n  d             2\n-----(2*x*y) + x \ndx dy            ']

def test_pretty_integrals():
    # Simple
    f_1 = Integral(log(x), x)
    assert pretty(f_1) == '  /         \n |          \n | log(x) dx\n |          \n/           '

    f_2 = Integral(x**2, x)
    assert pretty(f_2) == '  /     \n |      \n |  2   \n | x  dx\n |      \n/       '
    # Double nesting of pow
    f_3 = Integral(x**(2**x), x)
    assert pretty(f_3) == '  /        \n |         \n |  / x\\   \n |  \\2 /   \n | x     dx\n |         \n/          '

    # Definite integrals
    f_4 = Integral(x**2, (x,1,2))
    assert pretty(f_4) == '  2      \n  /      \n |       \n |   2   \n |  x  dx\n |       \n/        \n1        '

    f_5 = Integral(x**2, (x,Rational(1,2),10))
    assert pretty(f_5) == ' 10      \n  /      \n |       \n |   2   \n |  x  dx\n |       \n/        \n1/2      '

    # Nested integrals
    f_6 = Integral(x**2*y**2, x,y)
    assert pretty(f_6) == '  /  /           \n |  |            \n |  |  2  2      \n |  | x *y  dx dy\n |  |            \n/  /             '

def test_pretty_matrix():
    p = pretty( Matrix([[x**2+1, 1], [y, x+y]]) )
    s1 = \
"""\
[     2       ]
[1 + x     1  ]
[             ]
[  y     x + y]\
"""
    s2 = \
"""\
[ 2           ]
[x  + 1    1  ]
[             ]
[  y     y + x]\
"""
    assert p in [s1, s2]

def test_pretty_Piecewise():
    p = pretty(Piecewise((x,x<1),(x**2,True)))
    s = \
"""\
/x   for x < 1
|             \n\
< 2           \n\
|x   otherwise\n\
\             \
"""
    assert p == s

def test_pretty_seq():
    assert pretty([]) == '[]'
    assert pretty(()) == '()'
    assert pretty({}) == '{}'

    e = [x**2, 1/x, x, y, sin(th)**2/cos(ph)**2]
    p = pretty(e)
    s = \
"""\
                 2        \n\
  2  1        sin (theta) \n\
[x , -, x, y, -----------]\n\
     x            2       \n\
               cos (phi)  \
"""
    assert p == s

    assert pretty((x,)) == '(x,)'

    p = pretty((1/x,))
    s = \
"""\
 1  \n\
(-,)\n\
 x  \
"""
    assert p == s

    e = tuple(e)
    p = pretty(e)
    s = \
"""\
                 2        \n\
  2  1        sin (theta) \n\
(x , -, x, y, -----------)\n\
     x            2       \n\
               cos (phi)  \
"""
    assert p == s

    e = {x: sin(x)}
    p = pretty(e)
    s = \
"""\
{x: sin(x)}\
"""
    assert p == s

    e = {1/x: 1/y, x: sin(x)**2}
    p = pretty(e)
    s1 = \
"""\
 1  1        2    \n\
{-: -, x: sin (x)}\n\
 x  y             \
"""
    s2 = \
"""\
       2     1  1 \n\
{x: sin (x), -: -}\n\
             x  y \
"""
    # NOTE dict does not guaranty ordering
    assert (p == s1) or (p == s2)

def test_pretty_limits():
    assert pretty( Limit(x, x, oo) ) == ' lim x\nx->oo '
    assert pretty( Limit(x**2, x, 0) ) == '     2\nlim x \nx->0  '
    assert pretty( Limit(1/x, x, 0) ) == '    1\nlim -\nx->0x'

def test_pretty_class():
    """test that printer dispatcher correctly handles classes"""
    class C: pass   # C has no .__class__ and this was causing problems
    class D(object): pass

    assert pretty( C ) == str( C )
    assert pretty( D ) == str( D )

def test_infinity():
    assert pretty(I*oo) == "oo*I"

def test_full_prec():
    assert xpretty(S("0.3"), full_prec=True) == "0.300000000000000"
    assert xpretty(S("0.3"), full_prec="auto") == "0.300000000000000"
    assert xpretty(S("0.3"), full_prec=False) == "0.3"
    assert xpretty(S("0.3")*x, full_prec=True, use_unicode=False) in [
            "0.300000000000000*x",
            "x*0.300000000000000"
            ]
    assert xpretty(S("0.3")*x, full_prec="auto", use_unicode=False) in [
            "0.3*x",
            "x*0.3"
            ]
    assert xpretty(S("0.3")*x, full_prec=False, use_unicode=False) in [
            "0.3*x",
            "x*0.3"
            ]

ext: py

max_stars_repo_name: gnulinooks/sympy

max_issues_repo_name: gnulinooks/sympy

max_forks_repo_name: gnulinooks/sympy

reason: from sympy


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Mar 26 17:26:02 2019

@author: Bing
"""

import numpy as np
from scipy.sparse import lil_matrix, csr_matrix, kron, identity, linalg, issparse
import scipy
import sys

from lime.units import au2fs, au2k, au2ev
from lime.phys import dag, coth, ket2dm, comm, anticomm, pauli, destroy,\
    basis, isherm, basis_transform
from lime.mol import Mol
from lime.oqs import Lindblad_solver

class Composite(Mol):
    def __init__(self, A, B):
        """

        Parameters
        ----------
        A : mol object
            Quantum system.
        B : object
            Quantum system.

        Returns
        -------
        None.

        """

        self.A = A
        self.B = B
        self.idm = kron(A.idm, B.idm)  # identity matrix
        self.ida = A.idm
        self.idb = B.idm
        self.H = None
        self.nonhermH = None
        self.rdm_a = None
        self.rdm_b = None
        self.dim = A.dim * B.dim
        self.dims = [A.dim, B.dim]
        self.eigvals = None
        self.eigvecs = None

    def getH(self, a_ops=None, b_ops=None, g=0):
        """
        Compute the Hamiltonian for the composite A + B system.

        The interaction between A and B are
            V_AB = sum_i g[i] * a_ops[i] * b_ops[i]
        Parameters
        ----------
        a_ops: list of arrays
            coupling operators for subsystem A
        b_ops: list of arrays
            coupling operators for subsystem B
        g: coupling constants

        Returns
        -------

        """

        H = kron(self.A.H, self.idb) + kron(self.ida, self.B.H)

        if a_ops == None:

            print('Warning: there is no coupling between the two subsystems.')

        elif isinstance(a_ops, list):

            for i, a_op in enumerate(a_ops):
                b_op = b_ops[i]
                H += g[i] * kron(a_op, b_op)

        elif isinstance(a_ops, np.ndarray):

            H += g * a_ops @ b_ops

        self.H = H

        return H

    def get_nonhermH(self, a_ops=None, b_ops=None, g=0):
        """
        The interaction between A and B are
            V_AB = sum_i g[i] * a_ops[i] * b_ops[i]
        Parameters
        ----------
        a_ops: list of arrays
            coupling operators for subsystem A
        b_ops: list of arrays
            coupling operators for subsystem B
        g: coupling constants

        Returns
        -------

        """
        if self.A.nonhermH is None:
            raise ValueError('Call get_nonhermH() for subsystem A first.')

        if self.B.nonhermH is None:
            raise ValueError('Call get_nonhermH() for subsystem B first.')

        H = kron(self.A.nonhermH, self.idb) + kron(self.ida, self.B.nonhermH)

        if a_ops == None:

            print('Warning: there is no coupling between the two subsystems.')

        elif isinstance(a_ops, list):

            for i, a_op in enumerate(a_ops):
                b_op = b_ops[i]
                H += g[i] * kron(a_op, b_op)

        elif isinstance(a_ops, np.ndarray):

            H += g * a_ops @ b_ops

        self.nonhermH = H

        return H

    def promote(self, o, subspace='A'):
        """
        promote an operator in subspace to the full Hilbert space
        E.g. A = A \otimes I_B
        """
        if subspace == 'A':
            return kron(o, self.B.idm)

        elif subspace == 'B':
            return kron(self.A.idm, o)
        else:
            raise ValueError('The subspace option can only be A or B.')

    def promote_ops(self, ops, subspaces=None):
        if subspaces is None:
            subspaces = ['A'] * len(ops)

        new_ops = []
        for i, op in enumerate(ops):
            new_ops.append(self.promote(op, subspaces[i]))

        return new_ops

    def eigenstates(self, k=None):
        """
        compute the polaritonic spectrum

        Parameters
        ----------
        k : int, optional
            number of eigenstates. The default is 1.
        sparse : TYPE, optional
            if the Hamiltonian is sparse. The default is True.

        Returns
        -------
        evals : TYPE
            DESCRIPTION.
        evecs : TYPE
            DESCRIPTION.
        n_ph : TYPE
            photonic fractions in polaritons.

        """

        if self.H is None:
            sys.exit('Please call getH to compute the Hamiltonian first.')

        if k is None:

            # compute the full polariton states
            if issparse(self.H):
                evals, evecs = scipy.linalg.eigh(self.H.toarray())
            else:
                evals, evecs = scipy.linalg.eigh(self.H)

            self.eigvals = evals
            self.eigvecs = evecs
            return evals, evecs

        elif k < self.dim:

            if issparse(self.H):
                evals, evecs = linalg.eigsh(self.H, k, which='SA')
            else:
                raise TypeError('H is not sparse matrix.')

            self.eigvals = evals
            self.eigvecs = evecs
            return evals, evecs

        else:
            raise ValueError('k cannot exceed the size of H.')

    def spectrum(self):
        if self.H is None:
            sys.exit('Call getH() to compute the full Hamiltonian first.')
        else:
            eigvals, eigvecs = np.linalg.eigh(self.H.toarray())

            return eigvals, eigvecs

    def transform_basis(self, a):
        """
        transform the operator a from the direct product basis to polariton basis

        Parameters
        ----------
        a : TYPE
            DESCRIPTION.

        Returns
        -------
        2d array

        """
        if self.eigvecs is None:
            self.eigenstates()

        return basis_transform(a, self.eigvecs)

            # raise ValueError('Call eigenstates() to compute eigvecs first.')

    def rdm(self, psi, which='A'):
        """
        compute the reduced density matrix of A/B from a pure state

        Parameters
        ----------
        psi: array
            pure state
        which: str
            indicator of which rdm A or B. Default 'A'.

        Returns
        -------

        """
        na = self.A.dim
        nb = self.B.dim
        psi_reshaped = psi.reshape((na, nb))
        if which == 'A':

            rdm = psi_reshaped @ dag(psi_reshaped)
            return rdm

        elif which == 'B':
            rdm = psi_reshaped.T @ psi_reshaped.conj()
            return rdm
        else:
            raise ValueError('which option can only be A or B.')

    def purity(self, psi):
        rdm = self.rdm(psi)

        return np.trace(rdm @ rdm)

def ham_ho(freq, N, ZPE=False):
    """
    input:
        freq: fundamental frequency in units of Energy
        n : size of matrix
    output:
        h: hamiltonian of the harmonic oscillator
    """

    if ZPE:
        energy = np.arange(N + 0.5) * freq
    else:
        energy = np.arange(N) * freq

    H = lil_matrix((N, N))
    H.setdiag(energy)

    return H.tocsr()

def fft(t, x, freq=np.linspace(0, 0.1)):

    t = t/au2fs

    dt = (t[1] - t[0]).real

    sp = np.zeros(len(freq), dtype=np.complex128)

    for i in range(len(freq)):
        sp[i] = x.dot(np.exp(1j * freq[i] * t - 0.002*t)) * dt

    return sp

# def dag(H):
#     return H.conj().T

# def coth(x):
#     return 1./np.tanh(x)

# def ket2dm(psi):
#     return np.einsum("i, j -> ij", psi.conj(), psi)

# def obs(A, rho):
#     """
#     compute observables
#     """
#     return A.dot( rho).diagonal().sum()

def rk4_step(a, fun, dt, *args):

    dt2 = dt/2.0

    k1 = fun(a, *args)
    k2 = fun(a + k1*dt2, *args)
    k3 = fun(a + k2*dt2, *args)
    k4 = fun(a + k3*dt, *args)

    a += (k1 + 2*k2 + 2*k3 + k4)/6. * dt
    return a

class Pulse:
    def __init__(self, delay, sigma, omegac, amplitude=0.01, cep=0.):
        """
        Gaussian pulse A * exp(-(t-T)^2/2 / sigma^2)
        A: amplitude
        T: time delay
        sigma: duration
        """
        self.delay = delay
        self.sigma = sigma
        self.omegac = omegac  # central frequency
        self.unit = 'au'
        self.amplitude = amplitude
        self.cep = cep

    def envelop(self, t):
        return np.exp(-(t-self.delay)**2/2./self.sigma**2)

    def spectrum(self, omega):
        """
        Fourier transform of the Gaussian pulse
        """
        omegac = self.omegac
        sigma = self.sigma
        a = self.amplitude
        return a * sigma * np.sqrt(2.*np.pi) * np.exp(-(omega-omegac)**2 * sigma**2/2.)

    def field(self, t):
        '''
        electric field
        '''
        omegac = self.omegac
        delay = self.delay
        a = self.amplitude
        sigma = self.sigma
        return a * np.exp(-(t-delay)**2/2./sigma**2)*np.cos(omegac * (t-delay))

class Cavity():
    def __init__(self, freq, ncav, Q=None, polarization=None):
        self.frequency = self.freq = self.resonance = freq
        self.n_cav = ncav
        self.n = ncav
        self.dim = ncav # dimension of Fock space

        self.idm = identity(ncav)
        # self.create = self.get_create()
        # self.annihilate = destroy(ncav)
        self.H = self.getH()
        self.nonhermH = None

        self.quality_factor = Q
        self.decay = freq/2./Q
        self.polarization = polarization

#    @property
#    def hamiltonian(self):
#        return self._hamiltonian
#
#    @hamiltonian.setter
#    def hamiltonian(self):
#        self._hamiltonian = ham_ho(self.resonance, self.n)

    def get_ham(self, zpe=False):
        return self.getH(zpe)
    #
    # def H(self, ZPE=False):
    #     self._H = ham_ho(self.freq, self.dim, ZPE=ZPE)
    #     return self._H

    def getH(self, ZPE=False):
        self.H = ham_ho(self.freq, self.n_cav, ZPE=ZPE)
        return self.H

    def setQ(self, Q):
        self.quality_factor = Q

    def get_nonhermitianH(self):
        '''
        non-Hermitian Hamiltonian for the cavity mode

        Params:
            kappa: decay constant

        Returns
        -------
        TYPE
            DESCRIPTION.

        '''
        ncav = self.n_cav
        if self.quality_factor is not None:
            kappa = self.freq/2./self.quality_factor
        else:
            raise ValueError('The quality factor cannot be None.')

        self.nonhermH = self.H - 1j * kappa * np.identity(ncav)
        return self.nonhermH

    def get_nonhermH(self):
        return self.get_nonhermitianH()

    def ham(self, ZPE=False):
        return ham_ho(self.freq, self.n_cav, ZPE=ZPE)

    def get_create(self):
        n_cav = self.n_cav
        c = lil_matrix((n_cav, n_cav))
        c.setdiag(np.sqrt(np.arange(1, n_cav)), -1)
        return c.tocsr()

    def get_annihilate(self):
        n_cav = self.n_cav
        a = lil_matrix((n_cav, n_cav))
        a.setdiag(np.sqrt(np.arange(1, n_cav)), 1)

        return a.tocsr()

    def create(self):
        n_cav = self.n_cav
        c = lil_matrix((n_cav, n_cav))
        c.setdiag(np.sqrt(np.arange(1, n_cav)), -1)
        return c.tocsr()

    def annihilate(self):
        n_cav = self.n_cav
        a = lil_matrix((n_cav, n_cav))
        a.setdiag(np.sqrt(np.arange(1, n_cav)), 1)

        return a.tocsr()

    def vacuum(self, sparse=True):
        """
        get initial density matrix for cavity vacuum state
        """
        vac = np.zeros(self.n_cav)
        vac[0] = 1.
        if sparse:
            return csr_matrix(vac)
        else:
            return vac

    def vacuum_dm(self):
        """
        get initial density matrix for cavity vacuum state
        """
        vac = np.zeros(self.n_cav)
        vac[0] = 1.
        return ket2dm(vac)

    def get_num(self):
        """
        number operator
        """
        ncav = self.n_cav
        a = lil_matrix((ncav, ncav))
        a.setdiag(range(ncav), 0)

        return a.tocsr()

    def num(self):
        """
        number operator
        input:
            N: integer
                number of states
        """
        N = self.n_cav
        a = lil_matrix((N, N))
        a.setdiag(range(N), 0)
        return a.tocsr()

class Polariton(Composite):
    def __init__(self, mol, cav):

        super(Polariton, self).__init__(mol, cav)

        self.mol = mol
        self.cav = cav
        self._ham = None
        self.dip = None
        self.cav_leak = None
        self.H = None
        self.dims = [mol.dim, cav.n_cav]
        self.dim = mol.dim * cav.n_cav
        #self.dm = kron(mol.dm, cav.get_dm())

    def getH(self, g, RWA=False):
        """

        Parameters
        ----------
        g : float
            single photon electric field strength sqrt(hbar * omegac / 2 episilon_0 V)
        RWA : TYPE, optional
            DESCRIPTION. The default is False.

        Returns
        -------
        TYPE
            Hamiltonian.

        """

        mol = self.mol
        cav = self.cav

        hmol = mol.getH()
        hcav = cav.getH()

        Icav = cav.idm
        Imol = mol.idm

        if RWA:

            hint = g * (kron(mol.raising, cav.annihilate()) +
                        kron(mol.lowering, cav.create()))

        else:

            hint = g * kron(mol.dip, cav.create() + cav.annihilate())

        self.H = kron(hmol, Icav) + kron(Imol, hcav) + hint

        return self.H

    def get_nonhermitianH(self, g, RWA=False):

        mol = self.mol
        cav = self.cav

        hmol = mol.get_nonhermitianH()
        hcav = cav.get_nonhermitianH()

        Icav = cav.idm
        Imol = mol.idm

        if RWA:

            hint = g * (kron(mol.raising(), cav.get_annihilate()) +
                        kron(mol.lowering(), cav.get_create()))

        else:

            hint = g * kron(mol.dip, cav.get_create() + cav.get_annihilate())

        H = kron(hmol, Icav) + kron(Imol, hcav) + hint

        return H

    def get_ham(self, RWA=False):
        return self.getH(RWA)

    def setH(self, h):
        self.H = h
        return

    def get_dip(self, basis='product'):
        '''
        transition dipole moment in the direct product basis

        Returns
        -------
        TYPE
            DESCRIPTION.

        '''
        return kron(self.mol.dip, self.cav.idm)

    def get_dm(self):
        return kron(self.mol.dm, self.cav.vacuum_dm())

    def get_cav_leak(self):
        """
        damp operator for the cavity mode
        """
        if self.cav_leak == None:
            self.cav_leak = kron(self.mol.idm, self.cav.annihilate)

        return self.cav_leak

    def eigenstates(self, k=None):
        """
        compute the polaritonic spectrum

        Parameters
        ----------
        k : int, optional
            number of eigenstates. The default is 1.
        sparse : TYPE, optional
            if the Hamiltonian is sparse. The default is True.

        Returns
        -------
        evals : TYPE
            DESCRIPTION.
        evecs : TYPE
            DESCRIPTION.
        n_ph : TYPE
            photonic fractions in polaritons.

        """

        if self.H == None:
            sys.exit('Please call getH to compute the Hamiltonian first.')

        if k is None:

            """
            compute the full polariton states with numpy
            """

            h = self.H.toarray()
            # evals, evecs = scipy.linalg.eigh(h, subset_by_index=[0, self.dim])
            evals, evecs = scipy.linalg.eigh(h)
            # number of photons in polariton states
            num_op = self.cav.num()
            num_op = kron(self.mol.idm, num_op)

            n_ph = np.zeros(self.dim)
            for j in range(self.dim):
                n_ph[j] = np.real(evecs[:, j].conj().dot(
                    num_op.dot(evecs[:, j])))

            self.eigvals = evals
            self.eigvecs = evecs

            return evals, evecs

        elif k < self.dim:

            evals, evecs = linalg.eigsh(self.H, k, which='SA')

            # number of photons in polariton states
            num_op = self.cav.num()
            num_op = kron(self.mol.idm, num_op)

            n_ph = np.zeros(k)
            for j in range(k):
                n_ph[j] = np.real(evecs[:, j].conj().dot(
                    num_op.dot(evecs[:, j])))

            return evals, evecs, n_ph

    def rdm_photon(self):
        """
        return the reduced density matrix for the photons
        """
    def transform_basis(self, a):
        """
        transform the operator a from the direct product basis to polariton basis

        Parameters
        ----------
        a : TYPE
            DESCRIPTION.

        Returns
        -------
        2d array

        """
        if self.eigvecs is None:
            self.eigenstates()

        return basis_transform(a, self.eigvecs)

def QRM(omega0, omegac, ncav=2):
    '''
    Quantum Rabi model / Jaynes-Cummings model

    Parameters
    ----------
    omega0 : float
        atomic transition frequency
    omegac : float
        cavity frequency
    g : float
        cavity-molecule coupling strength

    Returns
    -------
    rabi: object

    '''
    s0, sx, sy, sz = pauli()

    hmol = 0.5 * omega0 * (-sz + s0)

    mol = Mol(hmol, sx)  # atom
    cav = Cavity(omegac, ncav)  # cavity

    return Polariton(mol, cav)

class Dicke(Cavity):

    def __init__(self, onsite, nsites, omegac):
        self.nsites = nsites
        self.onsite = onsite

        self.omegac = omegac

        return

    def Hsubspace(self, g, nexc=1):
        '''
        single-polariton Hamiltonian for N two-level molecules coupled to
        a single cavity mode including the ground state

        input:
            g_cav: single cavity-molecule coupling strength
        '''
        nsites = self.nsites
        onsite = self.onsite
        omegac = self.omegac

        if nexc == 1:
            nstates = nsites + 1 + 1 # number of states in the system

            dip = np.zeros((nstates, nstates))
            # the states are respectively, |000>, |100>, |010>, |001>,
            # the position are ordered as

            H = np.diagflat([0.] + onsite + [omegac])
            H[1:nsites+1, -1] = H[-1, 1:nsites+1] = g

            dip[0,1:nsites+1] = dip[1:nsites+1, 0] =  1

            return csr_matrix(H), csr_matrix(dip)
        else:
            raise NotImplementedError()

    def Hfull(self):
        pass

if __name__ == '__main__':

    #mol = QRM(1, 1)
    from lime.phys import quadrature
    s0, sx, sy, sz = pauli()

    hmol = 0.5 * (-sz + s0)

    mol = Mol(hmol, sx)  # atom
    mol.set_decay_for_all(0.05)
    mol.get_nonhermH()

    cav = Cavity(1, 2, Q=100)  # cavity
    cav.get_nonhermH()

    pol = Polariton(mol, cav)
    H = pol.get_nonhermH(g=0.1)

    print(isherm(H.toarray()))

    # set up the initial state
    psi0 = basis(2, 0)
    rho0 = ket2dm(psi0)

    # mesolver = Lindblad_solver(H, c_ops=[0.2 * sx], e_ops = [sz])
    Nt = 200
    dt = 0.05

    # results = mesolver.evolve(rho0, dt=dt, Nt=Nt)
    # pol.getH([sx], [cav.annihilate() + cav.create()], [0.1])
    # evals, evecs = pol.eigenstates()
    # print(pol.purity(evecs[:,2]))

    # fig, ax = subplots()
    # ax.plot(times, )

reason: import numpy,import scipy,from scipy


text:
# Convert set of 3D points to a cylindrical mesh
# to represent the surface

import numpy as np
import matplotlib.pyplot as plt
from math import pi
import pandas as pd

identity_matrix = np.array([[1,0,0],[0,1,0],[0,0,1]])

def cylindrical2cartesian(array):
    new = pd.DataFrame()
    r = array[:,0]
    new['x'] = r * np.cos(array[:,1])
    new['y'] = r * np.sin(array[:,1])
    new['z'] = array[:,2]
    return new.values

def init_cylinder(radius, N):
    height = np.linspace(0., 1., num=int(np.sqrt(N)))
    angle = np.linspace(0., 2.*pi, num=int(np.sqrt(N)))
    array = np.ndarray((len(height) * len(angle), 3))
    count = 0
    for i in range(len(height)):
        for j in range(len(angle)):
            array[count,0]=radius
            array[count,1]=angle[j]
            array[count,2]=height[i]
            count+=1
    points = cylindrical2cartesian(array)
    return points

def unit(vector):
    return vector/np.linalg.norm(vector)

def rotate(vector_1, vector_2):
    vec_1=unit(vector_1)
    vec_2=unit(vector_2)
    cross = np.cross(vec_1, vec_2)
    dot = np.dot(vec_1, vec_2)
    skew_matrix = np.array([[0        , -cross[2], cross[1]],
                            [cross[2] , 0        , -cross[0]],
                            [-cross[1], cross[0] , 0       ]])

    rotation = identity_matrix + skew_matrix + (1./(1.+dot))*np.dot(skew_matrix, skew_matrix)
    return rotation

def transform(init, atom_1, atom_2):

    # rotate
    target = atom_2 - atom_1
    unit_vector = np.array([0,0,1])
    R = rotate(unit_vector, target)

    # translate
    T = atom_1

    # scale length
    L = np.array([[1,0,0,0],
                  [0,1,0,0],
                  [0,0,np.linalg.norm(target),0],
                  [0,0,0,1]])

    transformation_matrix = np.array([[R[0,0], R[0,1],   R[0,2]  ,  T[0]],
                                      [R[1,0], R[1,1],   R[1,2]  ,  T[1]],
                                      [R[2,0], R[2,1],   R[2,2],    T[2]],
                                      [0     , 0     ,       0   ,   1  ]])

    transformation_matrix = np.matmul(transformation_matrix, L)

    temp_points = np.ones((init.shape[0], init.shape[1]+1))
    temp_points[:,:-1] = init
    points = np.matmul(transformation_matrix, temp_points.transpose())
    points = points.transpose()[:,:-1]

    return points

def cylinder(atom_1, atom_2, radius, N):
    init = init_cylinder(radius, N)
    final = transform(init, atom_1, atom_2)
    return final

class Cylinder():

    def __init__(self, pair, radius, N):
        """
        Constructs a cylinder between a pair of atoms

        Input:

        pair : np.array([[x1,y1,z1],[x2,y2,z2]]) - positions of atoms
        radius : float - radius of cylinder
        N : int - number of points overall

        Attributes:

        points : np.ndarray((N, 3)) evenly distributed grid of points
        """
        atom_1 = pair[0]
        atom_2 = pair[1]
        self.points = cylinder(atom_1, atom_2, radius, N)

class Molecule():
    def __init__(self, array, sigma=1.1, N=100):
        self.atoms = array
        self.radius = 2. * sigma
        points = np.array([[0,0,0]])
        for i in range(len(array)-1):
            pair = array[i:i+2, :]
            grid = Cylinder(pair, self.radius, N).points
            points = np.concatenate((points, grid))
        self.points = points[1:,:]

    def plot(self):
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.plot(self.points[:, 0],
                self.points[:, 1],
                self.points[:, 2], 'bo', markersize=1)
        ax.plot(self.atoms[:, 0],
                self.atoms[:, 1],
                self.atoms[:, 2], 'ko-', markersize=10)
        plt.show()

max_stars_repo_name: jakublala/starpolymers

max_issues_repo_name: jakublala/starpolymers

max_forks_repo_name: jakublala/starpolymers

reason: import numpy


text:
from __future__ import division
import math
import os
import sys
import time
import numpy as np
import scipy.stats
from scipy.optimize import curve_fit
from math import log10
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from scipy import pi as nombrepi

data = str(input("Numero de plot : "))

junk, nchan, nbin, I, Q, U, V = np.loadtxt('data/16_puppi_57772_C0531+33_0007_2695.dm559.72.calibP.RM.DD.ASCII', delimiter=' ', unpack=True)

inputfile = open("input.dat","w")

n = len(junk)

frequencemax = int(nchan[n-1])+1

binmax = int(nbin[n-1])+1

frequence = np.zeros(frequencemax)

intensite = np.zeros(frequencemax)

pas = 1.5625 # given by the author

## here you will choose the range of time where the signal is (you can choose larger) after
## seeing it with the ecriture-matrice-intensite.py

# burst # 16
tmin = 800
tmax = 1100

# burst #1
# tmin = 1500
# tmax = 2000

# nothing to touch here, juste calculating the noise, and summing on time #

intensitebruit1 = np.zeros(tmin-1)
intensitebruit2 = np.zeros(binmax-tmax)
variance = np.zeros(binmax)

for i in range(1,frequencemax+1) :

	frequence[i-1] = (i-1)

	### Calcul des incertitudes sur chaque frequence

	for j in range(1,binmax+1) :

		variance[j-1] =  I[j-1 + binmax*(frequencemax-i)]

	### Calcul de l'intensite du signal ###

	for j in range(tmin,tmax+1) :

		intensite[i-1] = intensite[i-1] + (I[j-1 + binmax*(frequencemax-i)])/(tmax+1-tmin) # frequences rangees en ordre inverse

	### Calcul du bruit et soustraction ###

	for j in range(1,tmin) :

		intensitebruit1[j-1] = (I[j-1 + binmax*(frequencemax-i)])/(tmin-1)

	for j in range(tmax+1,binmax+1) :

		intensitebruit2[j-1-tmax] = (I[j-1 + binmax*(frequencemax-i)])/(binmax-tmax)

	intensite[i-1] = intensite[i-1] - (sum(intensitebruit1)+sum(intensitebruit2))/2

	inputfile.write(str(frequence[i-1]))
	inputfile.write(" ")

	inputfile.write(str(intensite[i-1]))
	inputfile.write(" ")

	variance_freq = np.var(variance)
	inputfile.write(str(variance_freq))
	inputfile.write(" ")

	moyenne_freq = np.mean(variance)
	inputfile.write(str(moyenne_freq))

	inputfile.write("\n")

### Now you have the file "input(numberoftheburst).dat" ready to be used.

inputfile.close()


text:
import time
import copy
import os
from multiprocessing import Pool

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib.animation as animation

import flowrect
from flowrect.simulations.util import calculate_age, calculate_mt, eta_SRM, moving_average
from flowrect.simulations import (
    particle_population_nomemory,
    ASA1,
    quasi_renewal_pde,
)

def _simulate_QR(args):
    I_ext, params, a_cutoff = args
    t = time.time()
    params["I_ext"] = I_ext
    ts_QR, _, _, _, _, A_QR = quasi_renewal_pde(a_cutoff=a_cutoff, **params)
    print(f"I_ext = {I_ext:.2f}, QR done in {time.time() -t:.2f}s")
    Ainf_QR = A_QR[-1]

    return Ainf_QR

def _simulate_pdes(args):
    I_ext, params, a_cutoff = args
    t = time.time()
    params["I_ext"] = I_ext
    ts_PDE, _, _, _, _, _, A_PDE = ASA1(a_cutoff=a_cutoff, **params)
    print(f"I_ext = {I_ext:.2f}, PDE done in {time.time() -t:.2f}s")
    Ainf_PDE = A_PDE[-1]

    return Ainf_PDE

def _simulate_particle(args):
    I_ext, params, N, w = args
    t = time.time()
    cparams = copy.copy(params)
    cparams["I_ext"] = I_ext
    ts_P, A_P, _, _ = particle_population_nomemory(N=N, **cparams)
    dt = ts_P[1] - ts_P[0]

    new_A = A_P
    if w:
        new_A = moving_average(A_P, w)

    # Take the last seconds in activity
    last_seconds_idx = len(new_A) - 1 - int(1 / dt * 5)
    last_A_P = new_A[last_seconds_idx:]
    Ainf_P = np.mean(last_A_P)
    Ainf_std_P = np.std(last_A_P)
    print(f"I_ext = {I_ext:.2f}, Particle done in {time.time() -t:.2f}s")
    return Ainf_P, Ainf_std_P

def plot_A_inf(
    params,
    N,
    pool,
    num_sim=30,
    num_p_sim=10,
    I_end=5,
    w=None,
    cache_path="cache",
    cache_suppl_name="",
    QR_params=None,
    p_params=None,
    a_cutoff=7,
    savepath="",
    savename="",
    save=False,
    usetex=False,
    figsize=(8, 8),
    dpi=None,
    title=None,
    font_family="serif",
    font_size="12",
    noshow=False,
):
    if usetex:
        plt.rc("text", usetex=True)
        plt.rc("font", family=font_family, size=font_size)

    # External input range
    I_vec = np.linspace(0, I_end, num_sim)
    I_vec_p = np.linspace(0, I_end, num_p_sim)

    # Parameters
    if not p_params:
        p_params = copy.deepcopy(params)
    if not QR_params:
        QR_params = copy.deepcopy(params)

    # Check for cache

    t = time.time()
    params_multi_pde = [(I_vec[i], params, a_cutoff) for i in range(len(I_vec))]
    pde_res = pool.map(_simulate_pdes, params_multi_pde)

    params_multi_QR = [(I_vec[i], QR_params, a_cutoff) for i in range(len(I_vec))]
    QR_res = pool.map(_simulate_QR, params_multi_QR)

    params_multi_p = [(I_vec_p[i], p_params, N, w) for i in range(len(I_vec_p))]
    p_res = pool.map(_simulate_particle, params_multi_p)

    Ainf_PDE = []
    Ainf_QR = []
    Ainf_P = []
    Ainf_std_P = []

    # Fetch results
    for A_PDE in pde_res:
        Ainf_PDE.append(A_PDE)
    for A_QR in QR_res:
        Ainf_QR.append(A_QR)
    for A_P, A_std_P in p_res:
        if A_P:
            Ainf_P.append(A_P)
            Ainf_std_P.append(A_std_P)

    plt.figure()
    if title:
        plt.title(title)

    plt.plot(I_vec, Ainf_PDE, "-r", label="ASMA")
    plt.plot(I_vec, Ainf_QR, "-b", label="QR")
    plt.errorbar(I_vec_p, Ainf_P, Ainf_std_P, fmt=".k", capsize=2.0, label=r"$25\cdot10^3$ neurons")
    plt.xlabel(r"$I_1$ (A)")
    plt.ylabel(r"$A_{\infty}$ (Hz)")
    plt.legend()
    plt.tight_layout()

    if save:
        if dpi:
            plt.savefig(os.path.join(savepath, savename), dpi=dpi, transparent=True)
        else:
            plt.savefig(os.path.join(savepath, savename), transparent=True)
    if not noshow:
        plt.show()

max_stars_repo_name: jokteur/ASMA

max_issues_repo_name: jokteur/ASMA

max_forks_repo_name: jokteur/ASMA


text:
# -*- coding: utf-8 -*-
# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
    this file contain modifications of the original code, in hope new developers,
    can have better understanding how gpt2 Language Modeling is constructed from scratch.
"""

"""PyTorch OpenAI GPT-2 model."""

import sys
import os
import math
import random
import numpy as np
import pickle
import time

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.nn import CrossEntropyLoss, MSELoss

def gelu(x):
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))

class GPT2Config(nn.Module):
    """Configuration class to store the configuration of a `GPT2Model`.

    Args:
        vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `GPT2Model` or a configuration json file.
        n_positions: Number of positional embeddings.
        n_ctx: Size of the causal mask (usually same as n_positions).
        n_embd: Dimensionality of the embeddings and hidden states.
        n_layer: Number of hidden layers in the Transformer encoder.
        n_head: Number of attention heads for each attention layer in
            the Transformer encoder.
        layer_norm_epsilon: epsilon to use in the layer norm layers
        resid_pdrop: The dropout probabilitiy for all fully connected
            layers in the embeddings, encoder, and pooler.
        attn_pdrop: The dropout ratio for the attention
            probabilities.
        embd_pdrop: The dropout ratio for the embeddings.
        initializer_range: The sttdev of the truncated_normal_initializer for
            initializing all weight matrices.
    """

    def __init__(
                self,
        vocab_size_or_config_json_file=50257,
        n_positions=1024, #1024, 512 768
        n_ctx=1024, #1024, 512 768
        n_embd=600, #768, 256, 512, 1024 --> make sure n_embd%n_head==0
        n_layer=12, # 6, 32
        n_head=10, #12, 8, 16
        resid_pdrop=0.1,
        embd_pdrop=0.1,
        attn_pdrop=0.1,
        layer_norm_epsilon=1e-9,
        initializer_range=0.02,

        num_labels=1,
        summary_type='cls_index',
        summary_use_proj=True,
        summary_activation=None,
        summary_proj_to_labels=True,
        summary_first_dropout=0.1,
        **kwargs
    ):
        """Constructs GPT2Config.

        Args:
            vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `GPT2Model` or a configuration json file.
            n_positions: Number of positional embeddings.
            n_ctx: Size of the causal mask (usually same as n_positions).
            n_embd: Dimensionality of the embeddings and hidden states.
            n_layer: Number of hidden layers in the Transformer encoder.
            n_head: Number of attention heads for each attention layer in
                the Transformer encoder.
            layer_norm_epsilon: epsilon to use in the layer norm layers
            resid_pdrop: The dropout probabilitiy for all fully connected
                layers in the embeddings, encoder, and pooler.
            attn_pdrop: The dropout ratio for the attention
                probabilities.
            embd_pdrop: The dropout ratio for the embeddings.
            initializer_range: The sttdev of the truncated_normal_initializer for
                initializing all weight matrices.
        """
        super(GPT2Config, self).__init__(**kwargs)

        if isinstance(vocab_size_or_config_json_file, str) or (sys.version_info[0] == 2
                        and isinstance(vocab_size_or_config_json_file, unicode)):
            with open(vocab_size_or_config_json_file, "r", encoding="utf-8") as reader:
                json_config = json.loads(reader.read())
            for key, value in json_config.items():
                self.__dict__[key] = value
        elif isinstance(vocab_size_or_config_json_file, int):
            self.vocab_size = vocab_size_or_config_json_file
            self.n_ctx = n_ctx
            self.n_positions = n_positions
            self.n_embd = n_embd
            self.n_layer = n_layer
            self.n_head = n_head
            self.resid_pdrop = resid_pdrop
            self.embd_pdrop = embd_pdrop
            self.attn_pdrop = attn_pdrop
            self.layer_norm_epsilon = layer_norm_epsilon
            self.initializer_range = initializer_range

            self.num_labels = num_labels
            self.summary_type = summary_type
            self.summary_use_proj = summary_use_proj
            self.summary_activation = summary_activation
            self.summary_first_dropout = summary_first_dropout
            self.summary_proj_to_labels = summary_proj_to_labels

            ###
            self.output_attentions = False
            self.output_hidden_states = False
        else:
            raise ValueError(
                "First argument must be either a vocabulary size (int)"
                "or the path to a pretrained model config file (str)"
            )

    @property
    def max_position_embeddings(self):
        return self.n_positions

    @property
    def hidden_size(self):
        return self.n_embd

    @property
    def num_attention_heads(self):
        return self.n_head

    @property
    def num_hidden_layers(self):
        return self.n_layer

class Conv1D(nn.Module):
    def __init__(self, nf, nx):
        """ Conv1D layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2)
            Basically works like a Linear layer but the weights are transposed
        """
        super(Conv1D, self).__init__()
        self.nf = nf
        w = torch.empty(nx, nf)
        nn.init.normal_(w, std=0.02)
        self.weight = nn.Parameter(w)
        self.bias = nn.Parameter(torch.zeros(nf))

    def forward(self, x):
        size_out = x.size()[:-1] + (self.nf,)
        x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
        x = x.view(*size_out)
        return x

def prune_conv1d_layer(layer, index, dim=1):
    """ Prune a Conv1D layer (a model parameters) to keep only entries in index.
        A Conv1D work as a Linear layer (see e.g. BERT) but the weights are transposed.
        Return the pruned layer as a new layer with requires_grad=True.
        Used to remove heads.
    """
    index = index.to(layer.weight.device)
    W = layer.weight.index_select(dim, index).clone().detach()
    if dim == 0:
        b = layer.bias.clone().detach()
    else:
        b = layer.bias[index].clone().detach()
    new_size = list(layer.weight.size())
    new_size[dim] = len(index)
    new_layer = Conv1D(new_size[1], new_size[0]).to(layer.weight.device)
    new_layer.weight.requires_grad = False
    new_layer.weight.copy_(W.contiguous())
    new_layer.weight.requires_grad = True
    new_layer.bias.requires_grad = False
    new_layer.bias.copy_(b.contiguous())
    new_layer.bias.requires_grad = True
    return new_layer

class Attention(nn.Module):
    def __init__(self, nx, n_ctx, config, scale=False):
        super(Attention, self).__init__()
        self.output_attentions = config.output_attentions

        n_state = nx  # in Attention: n_state=768 (nx=n_embd)
        # [switch nx => n_state from Block to Attention to keep identical to TF implem]
        assert n_state % config.n_head == 0
        self.register_buffer("bias", torch.tril(torch.ones(n_ctx, n_ctx)).view(1, 1, n_ctx, n_ctx))
        self.n_head = config.n_head
        self.split_size = n_state
        self.scale = scale

        self.c_attn = Conv1D(n_state * 3, nx)
        self.c_proj = Conv1D(n_state, nx)
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        mask = torch.ones(self.n_head, self.split_size // self.n_head)
        heads = set(heads) - self.pruned_heads  # Convert to set and emove already pruned heads
        for head in heads:
            # Compute how many pruned heads are before the head and move the index accordingly
            head = head - sum(1 if h < head else 0 for h in self.pruned_heads)
            mask[head] = 0
        mask = mask.view(-1).contiguous().eq(1)
        index = torch.arange(len(mask))[mask].long()
        index_attn = torch.cat([index, index + self.split_size, index + (2*self.split_size)])

        # Prune conv1d layers
        self.c_attn = prune_conv1d_layer(self.c_attn, index_attn, dim=1)
        self.c_proj = prune_conv1d_layer(self.c_proj, index, dim=0)

        # Update hyper params
        self.split_size = (self.split_size // self.n_head) * (self.n_head - len(heads))
        self.n_head = self.n_head - len(heads)
        self.pruned_heads = self.pruned_heads.union(heads)

    def _attn(self, q, k, v, attention_mask=None, head_mask=None):
        w = torch.matmul(q, k)
        if self.scale:
            w = w / math.sqrt(v.size(-1))
        nd, ns = w.size(-2), w.size(-1)
        b = self.bias[:, :, ns-nd:ns, :ns]
        w = w * b - 1e4 * (1 - b)

        if attention_mask is not None:
            # Apply the attention mask
            w = w + attention_mask

        w = nn.Softmax(dim=-1)(w)
        w = self.attn_dropout(w)

        # Mask heads if we want to
        if head_mask is not None:
            w = w * head_mask

        outputs = [torch.matmul(w, v)]
        if self.output_attentions:
            outputs.append(w)
        return outputs

    def merge_heads(self, x):
        x = x.permute(0, 2, 1, 3).contiguous()
        new_x_shape = x.size()[:-2] + (x.size(-2) * x.size(-1),)
        return x.view(*new_x_shape)  # in Tensorflow implem: fct merge_states

    def split_heads(self, x, k=False):
        new_x_shape = x.size()[:-1] + (self.n_head, x.size(-1) // self.n_head)
        x = x.view(*new_x_shape)  # in Tensorflow implem: fct split_states
        if k:
            return x.permute(0, 2, 3, 1)  # (batch, head, head_features, seq_length)
        else:
            return x.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)

    def forward(self, x, layer_past=None, attention_mask=None, head_mask=None):
        x = self.c_attn(x)
        query, key, value = x.split(self.split_size, dim=2)
        query = self.split_heads(query)
        key = self.split_heads(key, k=True)
        value = self.split_heads(value)
        if layer_past is not None:
            past_key, past_value = layer_past[0].transpose(-2, -1), layer_past[1]  # transpose back cf below
            key = torch.cat((past_key, key), dim=-1)
            value = torch.cat((past_value, value), dim=-2)
        present = torch.stack((key.transpose(-2, -1), value))  # transpose to have same shapes for stacking

        attn_outputs = self._attn(query, key, value, attention_mask, head_mask)
        a = attn_outputs[0]

        a = self.merge_heads(a)
        a = self.c_proj(a)
        a = self.resid_dropout(a)

        outputs = [a, present] + attn_outputs[1:]
        return outputs  # a, present, (attentions)

class MLP(nn.Module):
    def __init__(self, n_state, config):  # in MLP: n_state=3072 (4 * n_embd)
        super(MLP, self).__init__()
        nx = config.n_embd
        self.c_fc = Conv1D(n_state, nx)
        self.c_proj = Conv1D(nx, n_state)
        self.act = gelu
        self.dropout = nn.Dropout(config.resid_pdrop)

    def forward(self, x):
        h = self.act(self.c_fc(x))
        h2 = self.c_proj(h)
        return self.dropout(h2)

class Block(nn.Module):
    def __init__(self, n_ctx, config, scale=False):
        super(Block, self).__init__()

        nx = config.n_embd
        self.ln_1 = nn.LayerNorm(nx, eps=config.layer_norm_epsilon)
        self.attn = Attention(nx, n_ctx, config, scale)
        self.ln_2 = nn.LayerNorm(nx, eps=config.layer_norm_epsilon)
        self.mlp = MLP(4 * nx, config)

    def forward(self, x, layer_past=None, attention_mask=None, head_mask=None):
        output_attn = self.attn(self.ln_1(x),
                                layer_past=layer_past,
                                attention_mask=attention_mask,
                                head_mask=head_mask)
        a = output_attn[0]  # output_attn: a, present, (attentions)

        x = x + a
        m = self.mlp(self.ln_2(x))
        x = x + m

        outputs = [x] + output_attn[1:]
        return outputs  # x, present, (attentions)

class GPT2Model(nn.Module):
    r"""
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``
            Sequence of hidden-states at the last layer of the model.
        **past**:
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            that contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `past` input) to speed up sequential decoding.
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
    """
    def __init__(self, config):
        super(GPT2Model, self).__init__()
        self.output_hidden_states = config.output_hidden_states
        self.output_attentions = config.output_attentions

        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
        self.wpe = nn.Embedding(config.n_positions, config.n_embd)
        self.drop = nn.Dropout(config.embd_pdrop)
        self.h = nn.ModuleList([Block(config.n_ctx, config, scale=True) for _ in range(config.n_layer)])
        self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)

        self.config = config
        self.init_weights()

    def init_weights(self):
        """ Initialize and prunes weights if needed. """
        # Initialize weights
        self.apply(self._init_weights)

    def _init_weights(self, module):
        """ Initialize the weights.
        """
        if isinstance(module, (nn.Linear, nn.Embedding, Conv1D)):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if isinstance(module, (nn.Linear, Conv1D)) and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def _resize_token_embeddings(self, new_num_tokens):
        self.wte = self._get_resized_embeddings(self.wte, new_num_tokens)
        return self.wte

    def _prune_heads(self, heads_to_prune):
        """ Prunes heads of the model.
            heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
        """
        for layer, heads in heads_to_prune.items():
            self.h[layer].attn.prune_heads(heads)

    def forward(self, input_ids, past=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None):
        input_shape = input_ids.size()
        input_ids = input_ids.view(-1, input_shape[-1])
        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, input_shape[-1])
        if position_ids is not None:
            position_ids = position_ids.view(-1, input_shape[-1])

        if past is None:
            past_length = 0
            past = [None] * len(self.h)
        else:
            past_length = past[0][0].size(-2)
        if position_ids is None:
            position_ids = torch.arange(past_length, input_ids.size(-1) + past_length, dtype=torch.long, device=input_ids.device)
            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)

        # Attention mask.
        if attention_mask is not None:
            attention_mask = attention_mask.view(-1, input_shape[-1])
            # We create a 3D attention mask from a 2D tensor mask.
            # Sizes are [batch_size, 1, 1, to_seq_length]
            # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
            # this attention mask is more simple than the triangular masking of causal attention
            # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)

            # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
            # masked positions, this operation will create a tensor which is 0.0 for
            # positions we want to attend and -10000.0 for masked positions.
            # Since we are adding it to the raw scores before the softmax, this is
            # effectively the same as removing these entirely.
            attention_mask = attention_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
            attention_mask = (1.0 - attention_mask) * -10000.0

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # head_mask has shape n_layer x batch x n_heads x N x N
        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
                head_mask = head_mask.expand(self.config.n_layer, -1, -1, -1, -1)
            elif head_mask.dim() == 2:
                head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
            head_mask = head_mask.to(dtype=next(self.parameters()).dtype) # switch to fload if need + fp16 compatibility
        else:
            head_mask = [None] * self.config.n_layer

        inputs_embeds = self.wte(input_ids)
        position_embeds = self.wpe(position_ids)
        if token_type_ids is not None:
            token_type_embeds = self.wte(token_type_ids)
        else:
            token_type_embeds = 0
        hidden_states = inputs_embeds + position_embeds + token_type_embeds
        hidden_states = self.drop(hidden_states)

        output_shape = input_shape + (hidden_states.size(-1),)

        presents = ()
        all_attentions = []
        all_hidden_states = ()
        for i, (block, layer_past) in enumerate(zip(self.h, past)):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states.view(*output_shape),)

            outputs = block(hidden_states,
                            layer_past=layer_past,
                            attention_mask=attention_mask,
                            head_mask=head_mask[i])

            hidden_states, present = outputs[:2]
            presents = presents + (present,)

            if self.output_attentions:
                all_attentions.append(outputs[2])

        hidden_states = self.ln_f(hidden_states)

        hidden_states = hidden_states.view(*output_shape)
        # Add last hidden state
        if self.output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        outputs = (hidden_states, presents)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
        if self.output_attentions:
            # let the number of heads free (-1) so we can extract attention even after head pruning
            attention_output_shape = input_shape[:-1] + (-1,) + all_attentions[0].shape[-2:]
            all_attentions = tuple(t.view(*attention_output_shape) for t in all_attentions)
            outputs = outputs + (all_attentions,)
        return outputs  # last hidden state, presents, (all hidden_states), (attentions)

class GPT2LMHeadModel(nn.Module):
    r"""
        **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Labels for language modeling.
            Note that the labels **are shifted** inside the model, i.e. you can set ``lm_labels = input_ids``
            Indices are selected in ``[-1, 0, ..., config.vocab_size]``
            All labels set to ``-1`` are ignored (masked), the loss is only
            computed for labels in ``[0, ..., config.vocab_size]``

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Language modeling loss.
        **prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.vocab_size)``
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        **past**:
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            that contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `past` input) to speed up sequential decoding.
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    """
    def __init__(self, config):
        super(GPT2LMHeadModel, self).__init__()
        self.transformer = GPT2Model(config)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=True)

        self.config = config

        self.init_weights()
        self.tie_weights()

    def init_weights(self):
        """ Initialize and prunes weights if needed. """
        # Initialize weights
        self.apply(self._init_weights)

    def _init_weights(self, module):
        """ Initialize the weights.
        """
        if isinstance(module, (nn.Linear, nn.Embedding, Conv1D)):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if isinstance(module, (nn.Linear, Conv1D)) and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def tie_weights(self):
        """ Make sure we are sharing the input and output embeddings.
            Export to TorchScript can't handle parameter sharing so we are cloning them instead.
        """
        self._tie_or_clone_weights(self.lm_head,
                                   self.transformer.wte)

    def _tie_or_clone_weights(self, first_module, second_module):
        """ Tie or clone module weights depending of weither we are using TorchScript or not
        """
        first_module.weight = second_module.weight

        if hasattr(first_module, 'bias') and first_module.bias is not None:
            first_module.bias.data = torch.nn.functional.pad(
                first_module.bias.data,
                (0, first_module.weight.shape[0] - first_module.bias.shape[0]),
                'constant',
                0
            )

    def forward(self, input_ids, past=None, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None,
                labels=None):
        transformer_outputs = self.transformer(input_ids,
                                               past=past,
                                               attention_mask=attention_mask,
                                               token_type_ids=token_type_ids,
                                               position_ids=position_ids,
                                               head_mask=head_mask)
        hidden_states = transformer_outputs[0]

        lm_logits = self.lm_head(hidden_states)

        outputs = (lm_logits,) + transformer_outputs[1:]
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()

            # Flatten the tokens
            loss_fct = CrossEntropyLoss(ignore_index=-1)
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)),
                            shift_labels.view(-1))
            outputs = (loss,) + outputs

        return outputs  # (loss), lm_logits, presents, (all hidden_states), (attentions)


text:
# Galileo E5b-Q code construction
#
# Copyright 2014 Peter Monta

import numpy as np

chip_rate = 10230000
code_length = 10230

# secondary code table from Table 20 (page 20) of Galileo ICD SIS (2014)
# index is PRN

secondary_code = {
  1: 'CFF914EE3C6126A49FD5E5C94',   2: 'FC317C9A9BF8C6038B5CADAB3',
  3: 'A2EAD74B6F9866E414393F239',   4: '72F2B1180FA6B802CB84DF997',
  5: '13E3AE93BC52391D09E84A982',   6: '77C04202B91B22C6D3469768E',
  7: 'FEBC592DD7C69AB103D0BB29C',   8: '0B494077E7C66FB6C51942A77',
  9: 'DD0E321837A3D52169B7B577C',  10: '43DEA90EA6C483E7990C3223F',
 11: '0366AB33F0167B6FA979DAE18',  12: '99CCBBFAB1242CBE31E1BD52D',
 13: 'A3466923CEFDF451EC0FCED22',  14: '1A5271F22A6F9A8D76E79B7F0',
 15: '3204A6BB91B49D1A2D3857960',  16: '32F83ADD43B599CBFB8628E5B',
 17: '3871FB0D89DB77553EB613CC1',  18: '6A3CBDFF2D64D17E02773C645',
 19: '2BCD09889A1D7FC219F2EDE3B',  20: '3E49467F4D4280B9942CD6F8C',
 21: '658E336DCFD9809F86D54A501',  22: 'ED4284F345170CF77268C8584',
 23: '29ECCE910D832CAF15E3DF5D1',  24: '456CCF7FE9353D50E87A708FA',
 25: 'FB757CC9E18CBC02BF1B84B9A',  26: '5686229A8D98224BC426BC7FC',
 27: '700A2D325EA14C4B7B7AA8338',  28: '1210A330B4D3B507D854CBA3F',
 29: '438EE410BD2F7DBCDD85565BA',  30: '4B9764CC455AE1F61F7DA432B',
 31: 'BF1F45FDDA3594ACF3C4CC806',  32: 'DA425440FE8F6E2C11B8EC1A4',
 33: 'EE2C8057A7C16999AFA33FED1',  34: '2C8BD7D8395C61DFA96243491',
 35: '391E4BB6BC43E98150CDDCADA',  36: '399F72A9EADB42C90C3ECF7F0',
 37: '93031FDEA588F88E83951270C',  38: 'BA8061462D873705E95D5CB37',
 39: 'D24188F88544EB121E963FD34',  40: 'D5F6A8BB081D8F383825A4DCA',
 41: '0FA4A205F0D76088D08EAF267',  42: '272E909FAEBC65215E263E258',
 43: '3370F35A674922828465FC816',  44: '54EF96116D4A0C8DB0E07101F',
 45: 'DE347C7B27FADC48EF1826A2B',  46: '01B16ECA6FC343AE08C5B8944',
 47: '1854DB743500EE94D8FC768ED',  48: '28E40C684C87370CD0597FAB4',
 49: '5E42C19717093353BCAAF4033',  50: '64310BAD8EB5B36E38646AF01'
}

# parse the hex string, return a 100-bit secondary-code array

def secondary_seq(s):
  x = np.zeros(100)
  for i in range(100):
    nib = i//4
    bit = 3-(i%4)
    x[i] = (int(s[nib],16)>>bit)&1
    x[i] = 1.0 - 2.0*x[i]
  return x

# transform the secondary-code table entries to sequences

for i in range(1,51):
  secondary_code[i] = secondary_seq(secondary_code[i])

# initial-state table from Table 16 (page 15) of Galileo ICD SIS (2014)
# index is PRN

e5bq_init = {
   1: 0o03331,    2: 0o06143,    3: 0o25322,    4: 0o23371,
   5: 0o00413,    6: 0o36235,    7: 0o17750,    8: 0o04745,
   9: 0o13005,   10: 0o37140,   11: 0o30155,   12: 0o20237,
  13: 0o03461,   14: 0o31662,   15: 0o27146,   16: 0o05547,
  17: 0o02456,   18: 0o30013,   19: 0o00322,   20: 0o10761,
  21: 0o26767,   22: 0o36004,   23: 0o30713,   24: 0o07662,
  25: 0o21610,   26: 0o20134,   27: 0o11262,   28: 0o10706,
  29: 0o34143,   30: 0o11051,   31: 0o25460,   32: 0o17665,
  33: 0o32354,   34: 0o21230,   35: 0o20146,   36: 0o11362,
  37: 0o37246,   38: 0o16344,   39: 0o15034,   40: 0o25471,
  41: 0o25646,   42: 0o22157,   43: 0o04336,   44: 0o16356,
  45: 0o04075,   46: 0o02626,   47: 0o11706,   48: 0o37011,
  49: 0o27041,   50: 0o31024
}

def e5bq_reg1_shift(x):
  return [x[13]^x[12]^x[10]^x[3]] + x[0:13]

def e5bq_reg2_shift(x):
  return [x[13]^x[9]^x[8]^x[5]^x[4]^x[0]] + x[0:13]

def make_e5bq_reg1():
  x = [1,1,1,1,1,1,1,1,1,1,1,1,1,1]
  y = np.zeros(code_length)
  for i in range(code_length):
    y[i] = x[13]
    x = e5bq_reg1_shift(x)
  return y

def make_e5bq_reg2(start):
  x = start
  y = np.zeros(code_length)
  for i in range(code_length):
    y[i] = x[13]
    x = e5bq_reg2_shift(x)
  return y

r1 = make_e5bq_reg1()

def seq(a):
  s = []
  for i in range(14):
    s = s + [(a>>i)&1]
  return s

codes = {}

def make_e5bq(prn):
  start = seq(e5bq_init[prn])
  r2 = make_e5bq_reg2(start)
  return np.logical_xor(r1,r2)

def e5bq_code(prn):
  if prn not in codes:
    codes[prn] = make_e5bq(prn)
  return codes[prn]

def code(prn,chips,frac,incr,n):
  c = e5bq_code(prn)
  idx = (chips%code_length) + frac + incr*np.arange(n)
  idx = np.floor(idx).astype('int')
  idx = np.mod(idx,code_length)
  x = c[idx]
  return 1.0 - 2.0*x

try:
  from numba import jit
except:
  def jit(**kwargs):
    return lambda x: x

@jit(nopython=True)
def correlate(x,prn,chips,frac,incr,c):
  n = len(x)
  p = 0.0j
  cp = (chips+frac)%code_length
  for i in range(n):
    p += x[i]*(1.0-2.0*c[int(cp)])
    cp = (cp+incr)%code_length
  return p

# test vectors in Galileo ICD

if __name__=='__main__':
  for prn in [1,2,3,4]:
    x = e5bq_code(prn)
    print(x[0:24].astype('int'))

max_stars_repo_name: wumouyan/GNSS-SDR-Python

max_issues_repo_name: wumouyan/GNSS-SDR-Python

max_forks_repo_name: wumouyan/GNSS-SDR-Python

reason: import numpy,from numba


text:
import numpy as np
from astropy.io import fits
#import matplotlib.pyplot as plt
import os
from scipy.ndimage import zoom
import glob

def rot8(im,nrot):
    """
    Rotation of a 2d image with the 8 possible geometries. Rotation 0-3
    do not flip the image, 4-7 perform a flip

    nrot = 0 -> same as input
    nrot = 1 -> 90deg counter-clock-wise
    nrot = 2 -> 180deg
    nrot = 3 -> 90deg clock-wise
    nrot = 4 -> flip top-bottom
    nrot = 5 -> flip top-bottom and rotate 90 deg counter-clock-wise
    nrot = 6 -> flip top-bottom and rotate 180 deg
    nrot = 7 -> flip top-bottom and rotate 90 deg clock-wise
    nrot >=8 -> performs a modulo 8 anyway

    :param im: input image
    :param nrot: integer between 0 and 7
    :return: rotated and/or flipped image
    """
    nrot = int(nrot % 8)
    return np.rot90( im[::1-2*(nrot//4)],nrot % 4 )

def med32(im):
    # input image MUST be 4096x4096
    # and have 32 amplifiers with a mirror
    # odd/even symmetry. We fold the image
    # into a 32x4096x128 cube where odd amplifiers
    # are flipped left/right. The median
    # amplifier structure is then re-padded
    # in an image that maps amp x-talk.
    # We expect the orders to be masked with NaNs
    # and the image to be high-passed so that only
    # high-frequency structures are left here.

    # cube to contain ordres in an easily
    # managed form
    cube = np.zeros([32, 4096, 128])
    for i in range(32):
        if (i % 2) == 0: # for left/right flipping
            i1 = i * 128
            i2 = i * 128 + 128
            sig = 1
        else:
            i1 = i * 128 + 127
            i2 = i * 128 -1
            sig = -1

        cube[i, :, :] = im[:, i1:i2:sig]
    # derive median amplifier structure
    med = np.nanmedian(cube, axis=0)

    # pad back onto the output image
    im2 = np.zeros_like(im)
    for i in range(32): # unflip
        if (i % 2) == 0: # for left/right flipping
            i1 = i * 128
            i2 = i * 128 + 128
            sig = 1
        else:
            i1 = i * 128 + 127
            i2 = i * 128 -1
            sig = -1

        im2[:, i1:i2:sig] = med
    return im2

def medbin(im,bx,by):
    # median-bin an image to a given size through
    # some funny np.reshape. To be used for low-pass
    # filterning of an image.
    sz = np.shape(im)
    out = np.nanmedian(np.nanmedian(im.reshape([bx,sz[0]//bx,by, sz[1]//bx]), axis=1), axis=2)
    return out

def mk_mask():
    # creation of the mask image
    im_flat = '20200401164256_ramp_020_HA_FLAT_FLAT.fits'
    im = fits.getdata(im_flat)
    # find pixel that are more than 10 absolute deviations
    # from the image median
    im -= np.nanmedian(im)

    sig = np.nanmedian(np.abs(im))

    # generate a first estimate of the mask
    mask = im>10*sig
    # now apply a proper filtering of the image
    im2 = np.array(im)
    im2[mask] = np.nan

    ################################################
    # Same code as for an individual science frame #
    ################################################

    # median-bin and expand back to original size
    binsize = 32
    lowf = zoom(medbin(im2, binsize, binsize), 4096 // binsize)

    # subtract low-frequency from masked image
    im2 -= lowf
    # find the amplifier x-talk map
    xtalk = med32(im2)

    # subtract both low-frequency and x-talk from input image
    im -= (lowf + xtalk)

    # generate a better estimate of the mask
    mask = im>10*sig

    fits.writeto('mask.fits',np.array(mask,dtype = int), overwrite = True)

    return []

def nirps_pp(files):

    ref_hdr = fits.getheader('ref_hdr.fits')

    if type(files) == str:
        files = glob.glob(files)

    for file in files:
        outname = '_pp.'.join(file.split('.'))

        if '_pp.' in file:
            print(file+' is a _pp file')
            continue

        if os.path.isfile(outname):
            print('File : '+outname +' exists')
            continue
        else:
            print('We pre-process '+file)

            hdr = fits.getheader(file)
            im = fits.getdata(file)

            mask = np.array(fits.getdata('mask.fits'),dtype = bool)

            im2 = np.array(im)
            im2[mask] = np.nan

            # we find the low level frequencies
            # we bin in regions of 32x32 pixels. This CANNOT be
            # smaller than the order footprint on the array
            # as it would lead to a set of NaNs in the downsized
            # image and chaos afterward
            binsize = 32 # pixels

            # median-bin and expand back to original size
            lowf = zoom(medbin(im2,binsize,binsize),4096//binsize)

            # subtract low-frequency from masked image
            im2 -= lowf

            # find the amplifier x-talk map
            xtalk = med32(im2)
            im2 -= xtalk

            # subtract both low-frequency and x-talk from input image
            im -= (lowf+xtalk)

            tmp = np.nanmedian(im2,axis=0)

            im -= np.tile(tmp, 4096).reshape(4096, 4096)

            # rotates the image so that it matches the order geometry of SPIRou and HARPS
            # redder orders at the bottom and redder wavelength within each order on the left

            # NIRPS = 5
            # SPIROU = 3
            im = rot8(im,5)

            #DPRTYPE
            """
            MJDMID  =    58875.10336167315 / Mid Observation time [mjd]
            BERVOBSM= 'header  '           / BERV method used to calc observation time
            DPRTYPE = 'FP_FP   '           / The type of file (from pre-process)
            PVERSION= '0.6.029 '           / DRS Pre-Processing version
            DRSVDATE= '2020-01-27'         / DRS Release date
            DRSPDATE= '2020-01-30 22:16:00.344' / DRS Processed date
            DRSPID  = 'PID-00015804225603440424-JKBM' / The process ID that outputted this f
            INF1000 = '2466774a.fits'      / Input file used to create output infile=0
            QCC001N = 'snr_hotpix'         / All quality control passed
            QCC001V =    876.2474157597072 / All quality control passed
            QCC001L = 'snr_hotpix < 1.00000e+01' / All quality control passed
            QCC001P =                    1 / All quality control passed
            QCC002N = 'max(rms_list)'      / All quality control passed
            QCC002V = 0.002373232122258537 / All quality control passed
            QCC002L = 'max(rms_list) > 1.5000e-01' / All quality control passed
            QCC002P =                    1 / All quality control passed
            QCC_ALL =                    T
            DETOFFDX=                    0 / Pixel offset in x from readout lag
            DETOFFDY=                    0 / Pixel offset in y from readout lag

            """

            if 'MJDEND' not in hdr:
                hdr['MJDEND'] = 0.00
                hdr['EXPTIME'] = 5.57*len(hdr['INTT*'])

            hdr['MJDMID'] = hdr['MJDEND'] - hdr['EXPTIME']/2.0/86400.0

            hdr['INF1000'] = file
            DPRTYPES = ['DARK_DARK','DARK_FP','FLAT_FLAT','DARK_FLAT',
                        'FLAT_DARK','HC_FP','FP_HC','FP_FP','OBJ_DARK','OBJ_FP','HC_DARK','DARK_HC','HC_HC']

            if 'STAR_DARK' in file:
                hdr['DPRTYPE'] = 'OBJ_DARK'

            if 'STAR_FP' in file:
                hdr['DPRTYPE'] = 'OBJ_FP'

            for DPRTYPE in DPRTYPES:
                if DPRTYPE in file:
                    if DPRTYPE == 'DARK_DARK':
                        hdr['DPRTYPE'] = 'DARK_DARK_TEL'
                    elif DPRTYPE == 'HC_HC':
                        hdr['DPRTYPE'] = 'HCONE_HCONE'
                    elif DPRTYPE == 'FP_HC':
                        hdr['DPRTYPE'] = 'FP_HCONE'
                    elif DPRTYPE == 'HC_FP':
                        hdr['DPRTYPE'] = 'HCONE_FP'
                    elif DPRTYPE == 'DARK_HC':
                        hdr['DPRTYPE'] = 'DARK_HCONE'
                    elif DPRTYPE == 'HC_DARK':
                        hdr['DPRTYPE'] = 'HCONE_DARK'
                    elif DPRTYPE == 'FP_DARK':
                        hdr['DPRTYPE'] = 'FP_DARK'
                    elif DPRTYPE == 'DARK_FP':
                        hdr['DPRTYPE'] = 'DARK_FP'
                    else:
                        hdr['DPRTYPE '] = DPRTYPE

            if 'DPRTYPE' not in hdr:
                print('error, with DPRTYPE for ',file)
                return

            if 'OBJECT' not in hdr:
                hdr['OBJECT'] = 'none'

            if 'RDNOISE' not in hdr:
                hdr['RDNOISE']= 10.0,'rdnoise *not* provided, added by _pp'

            if 'GAIN' not in hdr:
                hdr['GAIN']= 1.000,'gain *not* provided, added by _pp'

            if 'SATURATE' not in hdr:
                hdr['SATURATE']= 60000,'saturate *not* provided, added by _pp'

            if 'PVERSION' not in hdr:
                hdr['PVERSION'] = 'NIRPS_SIMU_PP'

            if 'OBSTYPE' not in hdr:
                if hdr['DPRTYPE'][0:4] == 'FLAT':
                    hdr['OBSTYPE'] = 'FLAT'

                if hdr['DPRTYPE'][0:4] == 'DARK':
                    hdr['OBSTYPE'] = 'DARK'

                if hdr['DPRTYPE'][0:2] == 'FP':
                    hdr['OBSTYPE'] = 'ALIGN'

                if hdr['DPRTYPE'][0:2] == 'HC':
                    hdr['OBSTYPE'] = 'COMPARISON'

                if hdr['DPRTYPE'][0:3] == 'OBJ':
                    hdr['OBSTYPE'] = 'OBJECT'

            if hdr['DPRTYPE'][0:3] == 'OBJ':
                hdr['TRG_TYPE'] = 'TARGET'
            else:
                hdr['TRG_TYPE'] = ''

            necessary_kwrd = ['OBSTYPE','TRG_TYPE','OBJECT','OBJRA','OBJDEC','OBJECT','OBJEQUIN','OBJRAPM','OBJDECPM','AIRMASS','RELHUMID','OBJTEMP','GAIA_ID','OBJPLX','OBSRV','GAIN','RDNOISE','FRMTIME','EXPTIME','PI_NAME','CMPLTEXP','NEXP','MJDATE','MJDEND','SBCREF_P','SBCCAS_P','SBCALI_P','SBCDEN_P','DATE-OBS','UTC-OBS','SATURATE','TEMPERAT','SB_POL_T']

            missing = False
            for key in necessary_kwrd:
                if key not in hdr:
                    print('missing keyword : {0}'.format(key))
                    missing = True

                    if key in ref_hdr:
                        hdr[key] = ref_hdr[key]

            b = fits.getdata(file,ext = 2)
            errslope = fits.getdata(file,ext = 3)
            n = fits.getdata(file,ext = 4)

            b = rot8(b,5)
            errslope = rot8(errslope,5)
            n = rot8(n,5)

            hdu1 = fits.PrimaryHDU()
            hdu1.header = hdr
            hdu1.header['NEXTEND'] = 4
            hdu2 = fits.ImageHDU(im)
            hdu2.header['UNITS'] = ('ADU/S', 'Slope of fit, flux vs time')
            hdu2.header['EXTNAME'] = ('slope', 'Slope of fit, flux vs time')

            hdu3 = fits.ImageHDU(b)
            hdu3.header['UNITS'] = ('ADU', 'Intercept of the pixel/time fit.')
            hdu3.header['EXTNAME'] = ('intercept', 'Intercept of the pixel/time fit.')

            hdu4 = fits.ImageHDU(errslope)
            hdu4.header['UNITS'] = ('ADU/S', 'Formal error on slope fit')
            hdu4.header['EXTNAME'] = ('errslope', 'Formal error on slope fit')

            hdu5 = fits.ImageHDU(n)
            hdu5.header['UNITS'] = ('Nimages', 'N readouts below saturation')
            hdu5.header['EXTNAME'] = ('count', 'N readouts below saturation')

            new_hdul = fits.HDUList([hdu1, hdu2, hdu3, hdu4, hdu5])

            # just to avoid an error message with writeto
            if os.path.isfile(outname):
                print('file : ' + outname + ' exists, we are overwriting it')
                os.system('rm ' + outname + '')

            new_hdul.writeto(outname, overwrite=True)

    return []

reason: import numpy,from scipy,from astropy


text:
import gym
import numpy as np
import random

class Exponential(gym.core.Env):
    action_space = gym.spaces.Box(-10 * np.ones([1]), np.zeros([1]))

    observation_space = gym.spaces.Box(np.zeros([1]), np.ones([1]))
    max_turns = 100

    @property
    def useless_obs(self):
        return np.random.rand(1, 1)

    def render(self, mode='human'):
        pass

    def __init__(self):
        self.done = False
        self.ctr = 0
        self.optimal_log = -4 * random.random() - 2

    def reset(self):
        self.done = False
        self.ctr = 0
        return self.useless_obs

    def step(self, action):
        if self.done:
            raise Exception("reset the env!")

        self.ctr += 1
        if self.ctr >= self.max_turns:
            self.done = True

        action_float =  np.squeeze(action)
        log_miss = abs( self.optimal_log - action_float )
        reward = 10 ** (-log_miss/ 10)

        return self.useless_obs, reward, self.done, {}

    def __str__(self):
        return "Exponential range environment"