<commit_before>#! python
# -*- coding: utf-8 -*-
"""
Created on Wed Jan 25 13:30:52 2017

@author: <NAME>, UNIVERSITY COLLEGE LONDON.

Tools for calculating the Stark effect in Rydberg helium
using the Numerov method.

Based on:

    Stark structure of the Rydberg states of alkali-metal atoms
        <NAME> et al. Phys. Rev. A, 20 2251 (1979)
        http://dx.doi.org/10.1103/PhysRevA.20.2251

    Rydberg atom diamagnetism
        <NAME>, PhD thesis, Massachusetts Institute of Technology.
        Dept. of Physics. (1988)
        http://hdl.handle.net/1721.1/14367

    Cold atoms and molecules by Zeeman deceleration and Rydberg-
    Stark deceleration
        <NAME>, Habilitationsschrift, ETH Zurich (2012)
        http://dx.doi.org/10.3929/ethz-a-007577485

"""
from __future__ import print_function, division
from math import ceil, log, exp
import numpy as np
from numba import jit
from tqdm import trange
from scipy.constants import h, hbar, c, alpha, m_e, e, epsilon_0, atomic_mass, pi
from .drake1999 import quantum_defects
import errno
import os
from datetime import datetime

# constants
Ry_inf = 10973731.568508
En_h = 2*Ry_inf*h*c # alpha**2.0 * m_e * c**2.0
a_0 = hbar/ (m_e * c * alpha)
mass_He = 4.002602 * atomic_mass
Z = 2
mu_me = (mass_He - m_e) / mass_He
mu_M = m_e / mass_He
En_h_He = En_h * mu_me
a_0_He = a_0 / mu_me

@jit
def get_nl_vals(nmin, nmax, m):
    """ n and l vals for each matrix column, using range n_min to n_max.
    """
    n_rng = np.arange(nmin, nmax + 1)
    n_vals = np.array([], dtype='int32')
    l_vals = np.array([], dtype='int32')
    for n in n_rng:
        l_rng = np.arange(np.abs(m), n)
        n_vals = np.append(n_vals, np.array(np.zeros_like(l_rng) + n))
        l_vals = np.append(l_vals, l_rng)
    return n_vals, l_vals

@jit
def get_nlm_vals(nmin, nmax):
    """ n, l and m vals for each matrix column, using range n_min to n_max.
    """
    n_rng = np.arange(nmin, nmax + 1)
    n_vals = np.array([], dtype='int32')
    l_vals = np.array([], dtype='int32')
    m_vals = np.array([], dtype='int32')
    for n in n_rng:
        l_rng = np.arange(0, n)
        for l in l_rng:
            m_rng = np.arange(-l, l+1)
            n_vals = np.append(n_vals, np.array(np.zeros_like(m_rng) + n))
            l_vals = np.append(l_vals, np.array(np.zeros_like(m_rng) + l))
            m_vals = np.append(m_vals, m_rng)
    return n_vals, l_vals, m_vals

@jit
def get_J_vals(S, L_vals, diff):
    """ J = L + diff; unless L == 0, in which case J = S.
    """
    J_vals = L_vals + diff
    J_vals[L_vals == 0] = S
    return J_vals

@jit
def get_triplet_nLJ(nmin, nmax, m):
    """ n and L and J vals for each matrix column, using range n_min to n_max.
    """
    n_rng = np.arange(nmin, nmax + 1, dtype='int32')
    n_vals = np.array([], dtype='int32')
    L_vals = np.array([], dtype='int32')
    J_vals = np.array([], dtype='int32')
    for n in n_rng:
        l_rng = np.arange(m, n, dtype='int32')
        for l in l_rng:
            if l == 0:
                n_vals = np.append(n_vals, n)
                L_vals = np.append(L_vals, 0)
                J_vals = np.append(J_vals, 1)
            else:
                n_vals = np.append(n_vals, np.zeros(3, dtype='int32') + n)
                L_vals = np.append(L_vals, np.zeros(3, dtype='int32') + l)
                J_vals = np.append(J_vals, np.arange(l - 1, l + 2, dtype='int32'))
    return n_vals, L_vals, J_vals

@jit
def get_qd(S, n_vals, L_vals, J_vals):
    """ Calculate quantum defects.
    """
    iterations = int(10)
    num_cols = len(n_vals)
    qd = np.zeros(num_cols)
    for i in range(num_cols):
        n = n_vals[i]
        L = L_vals[i]
        J = J_vals[i]
        if L in quantum_defects[S]:
            if J in quantum_defects[S][L]:
                delta = quantum_defects[S][L][J]
                # calculate quantum defects
                qd_i = delta[0]
                for rep in range(iterations):
                    # repeat to get convergence
                    m = n - qd_i
                    defect = delta[0]
                    for j, d in enumerate(delta[1:]):
                        defect = defect + d*m**(-2.0*(j + 1))
                    qd_i = defect
                    qd[i] = defect
            else:
                qd[i] = np.nan
        else:
            qd[i] = 0.0
    return qd

@jit
def En_0(neff):
    """ Field-free energy. Ignores extra correction terms.

        -- atomic units --
    """
    energy = np.array([])
    for n in neff:
        en = -0.5 * n**-2.0
        energy = np.append(energy, en)
    return energy

@jit
def W_n(S, n_vals, L_vals, J_vals):
    """ Field-free energy. Includes extra correction terms.

        -- atomic units --
    """
    neff = n_vals - get_qd(S, n_vals, L_vals, J_vals)
    energy = np.array([])
    for i, n in enumerate(n_vals):
        en = -0.5 * (neff[i]**-2.0 - 3.0 * alpha**2.0 / (4.0 * n**4.0) + \
                     mu_M**2.0 * ((1.0 + (5.0 / 6.0) * (alpha * Z)**2.0)/ n**2.0))
        energy = np.append(energy, en)
    return energy

@jit
def E_zeeman(m_vals, B_z):
    """ Energy shift due to the interaction of the orbital angular momentum of the Rydberg electron with the magnetic field.

        -- atomic units --
    """
    return m_vals * B_z * (1/2)

def chose_step(nmax, lmax, step_low, step_high, interp_type, interp_params):
    if interp_type == 'poly':
        return chose_step_poly(nmax, lmax, step_low, step_high, *interp_params)
    elif interp_type == 'sigmoid':
        return chose_step_sigmoid(nmax, lmax, step_low, step_high, *interp_params)
    
@jit
def chose_step_poly(nmax, lmax, step_low, step_high, exponent):
    return (lmax/nmax)**exponent * (step_high-step_low) + step_low

@jit
def chose_step_sigmoid(nmax, lmax, step_low, step_high, width, mid):
    return (( step_high-step_low )/( 1+np.exp( -width*( (lmax/nmax)-mid ) ) )) + step_low<commit_msg>Added rounding to the dynamic step<commit_after>#! python
# -*- coding: utf-8 -*-
"""
Created on Wed Jan 25 13:30:52 2017

@author: <NAME>, UNIVERSITY COLLEGE LONDON.

Tools for calculating the Stark effect in Rydberg helium
using the Numerov method.

Based on:

    Stark structure of the Rydberg states of alkali-metal atoms
        <NAME> al. Phys. Rev. A, 20 2251 (1979)
        http://dx.doi.org/10.1103/PhysRevA.20.2251

    Rydberg atom diamagnetism
        <NAME>, PhD thesis, Massachusetts Institute of Technology.
        Dept. of Physics. (1988)
        http://hdl.handle.net/1721.1/14367

    Cold atoms and molecules by Zeeman deceleration and Rydberg-
    Stark deceleration
        <NAME>, Habilitationsschrift, ETH Zurich (2012)
        http://dx.doi.org/10.3929/ethz-a-007577485

"""
from __future__ import print_function, division
from math import ceil, log, exp
import numpy as np
from numba import jit
from tqdm import trange
from scipy.constants import h, hbar, c, alpha, m_e, e, epsilon_0, atomic_mass, pi
from .drake1999 import quantum_defects
import errno
import os
from datetime import datetime

# constants
Ry_inf = 10973731.568508
En_h = 2*Ry_inf*h*c # alpha**2.0 * m_e * c**2.0
a_0 = hbar/ (m_e * c * alpha)
mass_He = 4.002602 * atomic_mass
Z = 2
mu_me = (mass_He - m_e) / mass_He
mu_M = m_e / mass_He
En_h_He = En_h * mu_me
a_0_He = a_0 / mu_me

@jit
def get_nl_vals(nmin, nmax, m):
    """ n and l vals for each matrix column, using range n_min to n_max.
    """
    n_rng = np.arange(nmin, nmax + 1)
    n_vals = np.array([], dtype='int32')
    l_vals = np.array([], dtype='int32')
    for n in n_rng:
        l_rng = np.arange(np.abs(m), n)
        n_vals = np.append(n_vals, np.array(np.zeros_like(l_rng) + n))
        l_vals = np.append(l_vals, l_rng)
    return n_vals, l_vals

@jit
def get_nlm_vals(nmin, nmax):
    """ n, l and m vals for each matrix column, using range n_min to n_max.
    """
    n_rng = np.arange(nmin, nmax + 1)
    n_vals = np.array([], dtype='int32')
    l_vals = np.array([], dtype='int32')
    m_vals = np.array([], dtype='int32')
    for n in n_rng:
        l_rng = np.arange(0, n)
        for l in l_rng:
            m_rng = np.arange(-l, l+1)
            n_vals = np.append(n_vals, np.array(np.zeros_like(m_rng) + n))
            l_vals = np.append(l_vals, np.array(np.zeros_like(m_rng) + l))
            m_vals = np.append(m_vals, m_rng)
    return n_vals, l_vals, m_vals

@jit
def get_J_vals(S, L_vals, diff):
    """ J = L + diff; unless L == 0, in which case J = S.
    """
    J_vals = L_vals + diff
    J_vals[L_vals == 0] = S
    return J_vals

@jit
def get_triplet_nLJ(nmin, nmax, m):
    """ n and L and J vals for each matrix column, using range n_min to n_max.
    """
    n_rng = np.arange(nmin, nmax + 1, dtype='int32')
    n_vals = np.array([], dtype='int32')
    L_vals = np.array([], dtype='int32')
    J_vals = np.array([], dtype='int32')
    for n in n_rng:
        l_rng = np.arange(m, n, dtype='int32')
        for l in l_rng:
            if l == 0:
                n_vals = np.append(n_vals, n)
                L_vals = np.append(L_vals, 0)
                J_vals = np.append(J_vals, 1)
            else:
                n_vals = np.append(n_vals, np.zeros(3, dtype='int32') + n)
                L_vals = np.append(L_vals, np.zeros(3, dtype='int32') + l)
                J_vals = np.append(J_vals, np.arange(l - 1, l + 2, dtype='int32'))
    return n_vals, L_vals, J_vals

@jit
def get_qd(S, n_vals, L_vals, J_vals):
    """ Calculate quantum defects.
    """
    iterations = int(10)
    num_cols = len(n_vals)
    qd = np.zeros(num_cols)
    for i in range(num_cols):
        n = n_vals[i]
        L = L_vals[i]
        J = J_vals[i]
        if L in quantum_defects[S]:
            if J in quantum_defects[S][L]:
                delta = quantum_defects[S][L][J]
                # calculate quantum defects
                qd_i = delta[0]
                for rep in range(iterations):
                    # repeat to get convergence
                    m = n - qd_i
                    defect = delta[0]
                    for j, d in enumerate(delta[1:]):
                        defect = defect + d*m**(-2.0*(j + 1))
                    qd_i = defect
                    qd[i] = defect
            else:
                qd[i] = np.nan
        else:
            qd[i] = 0.0
    return qd

@jit
def En_0(neff):
    """ Field-free energy. Ignores extra correction terms.

        -- atomic units --
    """
    energy = np.array([])
    for n in neff:
        en = -0.5 * n**-2.0
        energy = np.append(energy, en)
    return energy

@jit
def W_n(S, n_vals, L_vals, J_vals):
    """ Field-free energy. Includes extra correction terms.

        -- atomic units --
    """
    neff = n_vals - get_qd(S, n_vals, L_vals, J_vals)
    energy = np.array([])
    for i, n in enumerate(n_vals):
        en = -0.5 * (neff[i]**-2.0 - 3.0 * alpha**2.0 / (4.0 * n**4.0) + \
                     mu_M**2.0 * ((1.0 + (5.0 / 6.0) * (alpha * Z)**2.0)/ n**2.0))
        energy = np.append(energy, en)
    return energy

@jit
def E_zeeman(m_vals, B_z):
    """ Energy shift due to the interaction of the orbital angular momentum of the Rydberg electron with the magnetic field.

        -- atomic units --
    """
    return m_vals * B_z * (1/2)

def chose_step(nmax, lmax, step_low, step_high, interp_type, interp_params, round_step=10):
    if interp_type == 'poly':
        return chose_step_poly(nmax, lmax, step_low, step_high, round_step=round_step, *interp_params)
    elif interp_type == 'sigmoid':
        return chose_step_sigmoid(nmax, lmax, step_low, step_high, round_step=round_step, *interp_params)
    
@jit
def chose_step_poly(nmax, lmax, step_low, step_high, exponent, round_step=10):
    return np.round( (lmax/nmax)**exponent * (step_high-step_low) + step_low, round_step)

@jit
def chose_step_sigmoid(nmax, lmax, step_low, step_high, width, mid, round_step=10):
    return np.round( (( step_high-step_low )/( 1+np.exp( -width*( (lmax/nmax)-mid ) ) )) + step_low, round_step)
{'fraction_non_alphanumeric': 0.1319195123987643, 'fraction_numerical': 0.0313935042164148, 'mean_word_length': 2.2744669218151996, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2315574', 'n_tokens_mistral': 4833, 'n_tokens_neox': 4536, 'n_words': 1198}