<commit_before># -*- coding: utf-8 -*-
"""
Created on Mon Feb 15 11:28:53 2016

@author: ih3
"""

import numpy
from scipy.optimize import brentq
from scipy.integrate import odeint
from copy import deepcopy
from .state import State

def rarefaction_dwdp(w, p, q_known, wavenumber):
    r"""
    There is a tricky point here that needs investigation. If
    the input p is used here, rather than local_state.p, then they
    can diverge (when :math:`v_t` is significant) leading to overflows of g. By
    using local_state we avoid the overflow, but it may mean the final
    state is not very accurate.
    """
    lr_sign = wavenumber - 1
    dwdp = numpy.zeros_like(w)
    rho, v, eps = w
    vt = q_known.vt_from_known(rho, v, eps)
    local_state = State(rho, v, vt, eps, q_known.eos)
    cs = local_state.cs
    h = local_state.h
    W_lorentz = local_state.W_lorentz
    xi = local_state.wavespeed(wavenumber)
    # NOTE: what is g?
    g = vt**2 * (xi**2 - 1.0) / (1.0 - xi * v)**2
    dwdp[0] = 1.0 / (h * cs**2)
    dwdp[1] = lr_sign / (rho * h * W_lorentz**2 * cs) / numpy.sqrt(1.0 + g)
    dwdp[2] = local_state.p / (rho**2 * h * cs**2)
    return dwdp

class Wave(object):

    def __init__(self, q_known, unknown_value, wavenumber, unknown_eos=None,
                 t_i=None):
        """
        Initialize a wave.

        There are two possibilities: the wave is linear (wavenumber = 1),
        which is a contact, where the known value is the left state and the
        unknown value the right state.

        The second possibility is that the wave is nonlinear (wavenumber = 0,2)
        which is either a shock or a rarefaction, where the known value is the
        left/right state (wavenumber = 0,2 respectively) and the unknown value
        the pressure in the star state.
        """
        self.trivial = False # NOTE: what does this variable represent?
        assert(wavenumber in [0, 1, 2]), "wavenumber must be 0, 1, 2"
        self.wavenumber = wavenumber
        if self.wavenumber == 1:
            self.type = "Contact"
            assert(isinstance(unknown_value, State)), "unknown_value must " \
            "be a State when wavenumber is 1"
            self.q_l = q_known
            self.q_r = unknown_value
            assert(numpy.allclose(self.q_l.v, self.q_r.v)), "For a contact, "\
            "wavespeeds must match across the wave"
            assert(numpy.allclose(self.q_l.p, self.q_r.p)), "For a contact, "\
            "pressure must match across the wave"
            if numpy.allclose(self.q_l.state(), self.q_r.state()):
                self.trivial = True
            self.wave_speed = numpy.array([self.q_l.v, self.q_r.v])
            self.name = r"{\cal C}"
        elif self.wavenumber == 0:
            self.q_l = deepcopy(q_known)
            if (self.q_l.p < unknown_value):
                if unknown_eos is not None:
                    assert(t_i is not None)
                    self.solve_detonation(q_known, unknown_value, unknown_eos, t_i)
                else:
                    self.solve_shock(q_known, unknown_value)
            else:
                if unknown_eos is not None:
                    assert(t_i is not None)<commit_msg>Split out the WaveSection parts.<commit_after># -*- coding: utf-8 -*-
"""
Created on Mon Feb 15 11:28:53 2016

@author: ih3
"""

import numpy
from scipy.optimize import brentq
from scipy.integrate import odeint
from copy import deepcopy
from .state import State

def rarefaction_dwdp(w, p, q_known, wavenumber):
    r"""
    There is a tricky point here that needs investigation. If
    the input p is used here, rather than local_state.p, then they
    can diverge (when :math:`v_t` is significant) leading to overflows of g. By
    using local_state we avoid the overflow, but it may mean the final
    state is not very accurate.
    """
    lr_sign = wavenumber - 1
    dwdp = numpy.zeros_like(w)
    rho, v, eps = w
    vt = q_known.vt_from_known(rho, v, eps)
    local_state = State(rho, v, vt, eps, q_known.eos)
    cs = local_state.cs
    h = local_state.h
    W_lorentz = local_state.W_lorentz
    xi = local_state.wavespeed(wavenumber)
    # g quantifies the effect of tangential velocities: see the Living Review
    # and original Pons et al paper for details.
    g = vt**2 * (xi**2 - 1.0) / (1.0 - xi * v)**2
    dwdp[0] = 1.0 / (h * cs**2)
    dwdp[1] = lr_sign / (rho * h * W_lorentz**2 * cs) / numpy.sqrt(1.0 + g)
    dwdp[2] = local_state.p / (rho**2 * h * cs**2)
    return dwdp
    
def mass_flux_squared(q_start, p_end, unknown_eos=None):

    if unknown_eos is None:
        unknown_eos = q_start.eos

    def shock_root_rho(rho):
        h = unknown_eos['h_from_rho_p'](rho, p_end)
        return (h**2 - q_start.h**2) - \
        (h/rho + q_start.h/q_start.rho) * (p_end - q_start.p)

    if p_end >= q_start.p:
        # Shock
        min_rho = q_start.rho
        shock_root_min = shock_root_rho(min_rho)
        max_rho = numpy.sqrt(p_end/q_start.p) * q_start.rho
        shock_root_max = shock_root_rho(max_rho)
        while(shock_root_min * shock_root_max > 0.0):
            min_rho /= 1.1 # Not sure - could end up with unphysical root?
            max_rho *= 10.0
            shock_root_min = shock_root_rho(min_rho)
            shock_root_max = shock_root_rho(max_rho)
    else:
        # Deflagration
        max_rho = q_start.rho
        shock_root_max = shock_root_rho(max_rho)
        min_rho = numpy.sqrt(p_end/q_start.p) * q_start.rho
        shock_root_min = shock_root_rho(min_rho)
        while(shock_root_min * shock_root_max > 0.0):
            min_rho /= 10.0 # Not sure - could end up with unphysical root?
            max_rho *= 1.1
            shock_root_min = shock_root_rho(min_rho)
            shock_root_max = shock_root_rho(max_rho)
    rho = brentq(shock_root_rho, min_rho, max_rho)
    h = unknown_eos['h_from_rho_p'](rho, p_end)
    eps = h - 1.0 - p_end / rho
    dp = p_end - q_start.p
    dh2 = h**2 - q_start.h**2
    j2 = -dp / (dh2 / dp - 2.0 * q_start.h / q_start.rho)

    return j2, rho, eps, dp

class WaveSection(object):
    
    def __init__(self, q_start, p_end, wavenumber):
        """
        A part of a wave. For a single shock or rarefaction, this will be the
        complete wave. For a deflagration or detonation, it may form part of
        the full wave.
        """
        self.trivial = False
        assert(wavenumber in [0, 1, 2]), "wavenumber must be 0, 1, 2"
        self.wavenumber = wavenumber
        
    def latex_string(self):
        if self.trivial:
            return ""
        else:
            s = self.name
            s += r": \lambda^{{({})}}".format(self.wavenumber)
            if len(self.wavespeed) > 1:
                s += r"\in [{:.4f}, {:.4f}]".format(self.wavespeed[0],
                             self.wavespeed[-1])
            else:
                s += r"= {:.4f}".format(self.wavespeed[0])
            return s

    def _repr_latex_(self):
        s = r"$" + self.latex_string() + r"$"
        return s

class Contact(WaveSection):
    
    def __init__(self, q_start, q_end, wavenumber):
        """
        A contact.
        """
        
        self.trivial = False
        assert(wavenumber in [1]), "wavenumber for a Contact must be 1"
        self.type = "Contact"
        self.wavenumber = wavenumber
        self.q_start = deepcopy(q_start)
        self.q_end = deepcopy(q_end)

        self.name = r"{\<NAME>}"
        
        self.wavespeed = [q_start.v]
        
        assert(numpy.allclose(q_start.v, q_end.v)), "Velocities of states "\
        "must match for a contact"
        assert(numpy.allclose(q_start.p, q_end.p)), "Pressures of states "\
        "must match for a contact"
        assert(numpy.allclose(q_start.wavespeed(wavenumber),
                              q_end.wavespeed(wavenumber))), "Wavespeeds of "\
        "states must match for a contact"
        
    def plotting_data(self):
        
        data = numpy.vstack((self.q_start.state(), self.q_end.state()))
        xi = numpy.array([self.q_start.wavespeed(self.wavenumber),
                       self.q_start.wavespeed(self.wavenumber)])
        
        return xi, data
        
class Rarefaction(WaveSection):
    
    def __init__(self, q_start, p_end, wavenumber):
        """
        A rarefaction.
        """
        
        self.trivial = False
        assert(wavenumber in [0, 2]), "wavenumber for a Rarefaction "\
        "must be in 0, 2"
        assert(q_start.p >= p_end), "For a rarefaction, p_start >= p_end"
        self.type = "Rarefaction"
        self.wavenumber = wavenumber
        self.q_start = deepcopy(q_start)

        self.name = r"{\cal R}"
        if self.wavenumber == 0:
            label = r"\star_L"
            self.name += r"_{\leftarrow}"
        else:
            label = r"\star_R"
            self.name += r"_{\rightarrow}"

        v_known = q_start.wavespeed(self.wavenumber)

        if numpy.allclose(q_start.p, p_end):
            self.trivial = True
            self.q_end = State(q_start.rho, q_start.v, q_start.vt, q_start.eps,
            q_start.eos, label)
            v_unknown = v_known
        else:
            w_all = odeint(rarefaction_dwdp,
                           numpy.array([q_start.rho, q_start.v, q_start.eps]),
                           [q_start.p, p_end], rtol = 1e-12, atol = 1e-10,
                           args=((q_start, self.wavenumber)))
            self.q_end = State(w_all[-1, 0], w_all[-1, 1],
                              q_start.vt_from_known(w_all[-1, 0], w_all[-1, 1], w_all[-1, 2]),
                              w_all[-1, 2], q_start.eos, label)
            v_unknown = self.q_end.wavespeed(self.wavenumber)

        self.wavespeed = []
        if self.wavenumber == 0:
            self.wavespeed = numpy.array([v_known, v_unknown])
        else:
            self.wavespeed = numpy.array([v_unknown, v_known])
    
    def plotting_data(self):
        p = numpy.linspace(self.q_start.p, self.q_end.p)
        w_all = odeint(rarefaction_dwdp,
                       numpy.array([self.q_start.rho, self.q_start.v, self.q_start.eps]),
                           p, rtol = 1e-12, atol = 1e-10,
                           args=(self.q_start, self.wavenumber))
        data = numpy.zeros((len(p),8))
        xi = numpy.zeros_like(p)
        for i in range(len(p)):
            state = State(w_all[i,0], w_all[i,1],
                          self.q_start.vt_from_known(w_all[i,0], w_all[i,1], w_all[i,2]),
                          w_all[i, 2], self.q_start.eos)
            xi[i] = state.wavespeed(self.wavenumber)
            data[i,:] = state.state()
        
        return xi, data

class Shock(WaveSection):
    
    def __init__(self, q_start, p_end, wavenumber):
        """
        A shock.
        """

        self.trivial = False
        assert(wavenumber in [0, 2]), "wavenumber for a Shock "\
        "must be in 0, 2"
        assert(q_start.p <= p_end), "For a shock, p_start <= p_end"
        self.type = "Shock"
        self.wavenumber = wavenumber
        lr_sign = self.wavenumber - 1
        self.q_start = deepcopy(q_start)

        self.name = r"{\cal S}"
        if self.wavenumber == 0:
            label = r"\star_L"
            self.name += r"_{\leftarrow}"
        else:
            label = r"\star_R"
            self.name += r"_{\rightarrow}"

        if numpy.allclose(q_start.p, p_end):
            self.trivial = True
            self.q_end = State(q_start.rho, q_start.v, q_start.vt, q_start.eps,
            q_start.eos, label)
            v_shock = q_start.wavespeed(self.wavenumber)
        else:
            j2, rho, eps, dp = mass_flux_squared(q_start, p_end, 
                                                      q_start.eos)
            j = numpy.sqrt(j2)
            v_shock = (q_start.rho**2 * q_start.W_lorentz**2 * q_start.v + \
                lr_sign * j**2 * \
                numpy.sqrt(1.0 + q_start.rho**2 * q_start.W_lorentz**2 * (1.0 - q_start.v**2) / j**2)) / \
                (q_start.rho**2 * q_start.W_lorentz**2 + j**2)
            W_lorentz_shock = 1.0 / numpy.sqrt(1.0 - v_shock**2)
            v = (q_start.h * q_start.W_lorentz * q_start.v + lr_sign * dp * W_lorentz_shock / j) / \
                (q_start.h * q_start.W_lorentz + dp * (1.0 / q_start.rho / q_start.W_lorentz + \
                lr_sign * q_start.v * W_lorentz_shock / j))
            vt = q_start.vt_from_known(rho, v, eps)
            self.q_end = State(rho, v, vt, eps, q_start.eos, label)

        self.wavespeed = numpy.array([v_shock, v_shock])
        
    def plotting_data(self):
        
        data = numpy.vstack((self.q_start.state(), self.q_end.state()))
        xi = numpy.array([self.q_start.wavespeed(self.wavenumber),
                       self.q_start.wavespeed(self.wavenumber)])
        
        return xi, data
        
def build_inert_wave_section(q_known, unknown_value, wavenumber):
    """
    Object factory for the WaveSection; non-reactive case
    """
    
    if wavenumber == 1:
        return Contact(q_known, unknown_value, wavenumber)
    elif q_known.p < unknown_value:
        return Shock(q_known, unknown_value, wavenumber)
    else:
        return Rarefaction(q_known, unknown_value, wavenumber)

class Wave(object):

    def __init__(self, q_known, unknown_value, wavenumber, unknown_eos=None,
                 t_i=None):
        """
        Initialize a wave.

        There are two possibilities: the wave is linear (wavenumber = 1),
        which is a contact, where the known value is the left state and the
        unknown value the right state.

        The second possibility is that the wave is nonlinear (wavenumber = 0,2)
        which is either a shock or a rarefaction, where the known value is the
        left/right state (wavenumber = 0,2 respectively) and the unknown value
        the pressure in the star state.
        """
        # The wave is trivial if no quantites jump across it
        # This should be used to simplify output, especially for 
        # Riemann Problems with eg only a single shock.
        self.trivial = False 
        assert(wavenumber in [0, 1, 2]), "wavenumber must be 0, 1, 2"
        self.wavenumber = wavenumber
        if self.wavenumber == 1:
            self.type = "Contact"
            assert(isinstance(unknown_value, State)), "unknown_value must " \
            "be a State when wavenumber is 1"
            self.q_l = q_known
            self.q_r = unknown_value
            assert(numpy.allclose(self.q_l.v, self.q_r.v)), "For a contact, "\
            "wavespeeds must match across the wave"
            assert(numpy.allclose(self.q_l.p, self.q_r.p)), "For a contact, "\
            "pressure must match across the wave"
            if numpy.allclose(self.q_l.state(), self.q_r.state()):
                self.trivial = True
            self.wave_speed = numpy.array([self.q_l.v, self.q_r.v])
            self.name = r"{\<NAME>}"
        elif self.wavenumber == 0:
            self.q_l = deepcopy(q_known)
            if (self.q_l.p < unknown_value):
                if unknown_eos is not None:
                    assert(t_i is not None)
                    self.solve_detonation(q_known, unknown_value, unknown_eos, t_i)
                else:
                    self.solve_shock(q_known, unknown_value)
            else:
                if unknown_eos is not None:
                    assert(t_i is not None)
{'fraction_non_alphanumeric': 0.11215437415010036, 'fraction_numerical': 0.013080360033672214, 'mean_word_length': 1.907926944078328, 'pattern_counts': {'":': 1, '<': 10, '<?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': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2520411', 'n_tokens_mistral': 5442, 'n_tokens_neox': 5192, 'n_words': 1297}