# Source code for fluids.jet_pump

    # -*- coding: utf-8 -*-
    '''Chemical Engineering Design Library (ChEDL). Utilities for process modeling.
    Copyright (C) 2018, Caleb Bell <Caleb.Andrew.Bell@gmail.com>

    Permission is hereby granted, free of charge, to any person obtaining a copy
    of this software and associated documentation files (the "Software"), to deal
    in the Software without restriction, including without limitation the rights
    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    copies of the Software, and to permit persons to whom the Software is
    furnished to do so, subject to the following conditions:

    The above copyright notice and this permission notice shall be included in all
    copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    SOFTWARE.'''

    from __future__ import division
    from math import log, pi, sqrt
    from random import uniform, seed
    from pprint import pprint
    from fluids.numerics import newton, brenth

    __all__ = ['liquid_jet_pump', 'liquid_jet_pump_ancillary']

    def liquid_jet_pump_ancillary(rhop, rhos, Kp, Ks, d_nozzle=None, d_mixing=None,
                                  Qp=None, Qs=None, P1=None, P2=None):
        r'''Calculates the remaining variable in a liquid jet pump when solving for
        one if the inlet variables only and the rest of them are known. The
        equation comes from conservation of energy and momentum in the mixing
        chamber.

        The variable to be solved for must be one of `d_nozzle`, `d_mixing`,
        `Qp`, `Qs`, `P1`, or `P2`.

        .. math::
            P_1 - P_2 = \frac{1}{2}\rho_pV_n^2(1+K_p)
            - \frac{1}{2}\rho_s V_3^2(1+K_s)

        Rearrange to express V3 in terms of Vn, and using the density ratio `C`,
        the expression becomes:

        .. math::
            P_1 - P_2 = \frac{1}{2}\rho_p V_n^2\left[(1+K_p) - C(1+K_s)
            \left(\frac{MR}{1-R}\right)^2\right]

        Using the primary nozzle area and flow rate:

        .. math::
            P_1 - P_2 = \frac{1}{2}\rho_p \left(\frac{Q_p}{A_n}\right)^2
            \left[(1+K_p) - C(1+K_s) \left(\frac{MR}{1-R}\right)^2\right]

        For `P`, `P2`, `Qs`, and `Qp`, the equation can be rearranged explicitly
        for them. For `d_mixing` and `d_nozzle`, a bounded solver is used searching
        between 1E-9 m and 20 times the other diameter which was specified.

        Parameters
        ----------
        rhop : float
            The density of the primary (motive) fluid, [kg/m^3]
        rhos : float
            The density of the secondary fluid (drawn from the vacuum chamber),
            [kg/m^3]
        Kp : float
            The primary nozzle loss coefficient, [-]
        Ks : float
            The secondary inlet loss coefficient, [-]
        d_nozzle : float, optional
            The inside diameter of the primary fluid's nozle, [m]
        d_mixing : float, optional
            The diameter of the mixing chamber, [m]
        Qp : float, optional
            The volumetric flow rate of the primary fluid, [m^3/s]
        Qs : float, optional
            The volumetric flow rate of the secondary fluid, [m^3/s]
        P1 : float, optional
            The pressure of the primary fluid entering its nozzle, [Pa]
        P2 : float, optional
            The pressure of the secondary fluid at the entry of the ejector, [Pa]

        Returns
        -------
        solution : float
            The parameter not specified (one of `d_nozzle`, `d_mixing`,
            `Qp`, `Qs`, `P1`, or `P2`), (units of `m`, `m`, `m^3/s`, `m^3/s`,
            `Pa`, or `Pa` respectively)

        Notes
        -----
        The following SymPy code was used to obtain the analytical formulas (
        they are not shown here due to their length):

        >>> from sympy import *
        >>> A_nozzle, A_mixing, Qs, Qp, P1, P2, rhos, rhop, Ks, Kp = symbols('A_nozzle, A_mixing, Qs, Qp, P1, P2, rhos, rhop, Ks, Kp')
        >>> R = A_nozzle/A_mixing
        >>> M = Qs/Qp
        >>> C = rhos/rhop
        >>> rhs = rhop/2*(Qp/A_nozzle)**2*((1+Kp) - C*(1 + Ks)*((M*R)/(1-R))**2 )
        >>> new = Eq(P1 - P2,  rhs)
        >>> #solve(new, Qp)
        >>> #solve(new, Qs)
        >>> #solve(new, P1)
        >>> #solve(new, P2)

        Examples
        --------
        Calculating primary fluid nozzle inlet pressure P1:

        >>> liquid_jet_pump_ancillary(rhop=998., rhos=1098., Ks=0.11, Kp=.04,
        ... P2=133600, Qp=0.01, Qs=0.01, d_mixing=0.045, d_nozzle=0.02238)
        426434.60314398084

        References
        ----------
        .. [1] Ejectors and Jet Pumps. Design and Performance for Incompressible
           Liquid Flow. 85032. ESDU International PLC, 1985.
        '''
        unknowns = sum(i is None for i in (d_nozzle, d_mixing, Qs, Qp, P1, P2))
        if unknowns > 1:
            raise Exception('Too many unknowns')
        elif unknowns < 1:
            raise Exception('Overspecified')
        C = rhos/rhop

        if Qp is not None and Qs is not None:
            M = Qs/Qp
        if d_nozzle is not None:
            A_nozzle = pi/4*d_nozzle*d_nozzle
            if d_mixing is not None:
                A_mixing = pi/4*d_mixing*d_mixing
                R = A_nozzle/A_mixing

        if P1 is None:
            return rhop/2*(Qp/A_nozzle)**2*((1+Kp) - C*(1 + Ks)*((M*R)/(1-R))**2 ) + P2
        elif P2 is None:
            return -rhop/2*(Qp/A_nozzle)**2*((1+Kp) - C*(1 + Ks)*((M*R)/(1-R))**2 ) + P1
        elif Qs is None:
            try:
                return ((-2*A_nozzle**2*P1 + 2*A_nozzle**2*P2 + Kp*Qp**2*rhop + Qp**2*rhop)/(C*rhop*(Ks + 1)))**0.5*(A_mixing - A_nozzle)/A_nozzle
            except ValueError:
                return -1j
        elif Qp is None:
            return A_nozzle*((2*A_mixing**2*P1 - 2*A_mixing**2*P2 - 4*A_mixing*A_nozzle*P1 + 4*A_mixing*A_nozzle*P2 + 2*A_nozzle**2*P1 - 2*A_nozzle**2*P2 + C*Ks*Qs**2*rhop + C*Qs**2*rhop)/(rhop*(Kp + 1)))**0.5/(A_mixing - A_nozzle)
        elif d_nozzle is None:
            def err(d_nozzle):
                return P1 - liquid_jet_pump_ancillary(rhop=rhop, rhos=rhos, Kp=Kp, Ks=Ks, d_nozzle=d_nozzle, d_mixing=d_mixing, Qp=Qp, Qs=Qs,
                                  P1=None, P2=P2)
            return brenth(err, 1E-9, d_mixing*20)
        elif d_mixing is None:
            def err(d_mixing):
                return P1 - liquid_jet_pump_ancillary(rhop=rhop, rhos=rhos, Kp=Kp, Ks=Ks, d_nozzle=d_nozzle, d_mixing=d_mixing, Qp=Qp, Qs=Qs,
                                  P1=None, P2=P2)
            try:
                return brenth(err, 1E-9, d_nozzle*20)
            except:
                return newton(err, d_nozzle*2)

    def liquid_jet_pump_pressure_ratio(rhop, rhos, Km, Kd, Ks, Kp,
                         d_nozzle=None, d_mixing=None, d_diffuser=None,
                         Qp=None, Qs=None, P1=None, P2=None, P5=None,
                         nozzle_retracted=True):
        C = rhos/rhop
        if nozzle_retracted:
            j = 0.0
        else:
            j = 1.0

        R = d_nozzle**2/d_mixing**2
        alpha = d_mixing**2/d_diffuser**2
        M = Qs/Qp

        M2, R2, alpha2 = M*M, R*R, alpha*alpha
        num = 2.0*R + 2*C*M2*R2/(1.0 - R)
        num -= R2*(1.0 + C*M)*(1.0 + M)*(1.0 + Km + Kd + alpha2)
        num -= C*M2*R2/(1.0 - R)**2*(1.0 + Ks)

        den = (1.0 + Kp) - 2.0*R - 2.0*C*M2*R2/(1.0 - R)
        den += R2*(1.0 + C*M)*(1.0 + M)*(1.0 + Km + Kd + alpha2)
        den += (1.0 - j)*(C*M2/((1.0 - R)/R)**2)*(1.0 - Ks)
        N = num/den
        if P1 is None:
            P1 = (-P2 + P5*N + P5)/N
        elif P2 is None:
            P2 = -P1*N + P5*N + P5
        elif P5 is None:
            P5 = (P1*N + P2)/(N + 1.0)
        else:
            return N - (P5 - P2)/(P1 - P5)

        solution = {}
        solution['P1'] = P1
        solution['P2'] = P2
        solution['P5'] = P5
    #    solution['d_nozzle'] = d_nozzle
    #    solution['d_mixing'] = d_mixing
    #    solution['d_diffuser'] = d_diffuser
    #    solution['Qs'] = Qs
    #    solution['Qp'] = Qp
    #    solution['N'] = N
    #    solution['M'] = M
    #    solution['R'] = R
    #    solution['alpha'] = alpha
    #    solution['efficiency'] = M*N
        return solution

    def liquid_jet_pump(rhop, rhos, Kp=0.0, Ks=0.1, Km=.15, Kd=0.1,
                        d_nozzle=None, d_mixing=None, d_diffuser=None,
                        Qp=None, Qs=None, P1=None, P2=None, P5=None,
                        nozzle_retracted=True, max_variations=100):
        r'''Calculate the remaining two variables in a liquid jet pump, using a
        model presented in [1]_ as well as [2]_, [3]_, and [4]_.

        .. math::
            N = \frac{2R + \frac{2 C M^2 R^2}{1-R} - R^2 (1+CM) (1+M) (1 + K_m
            + K_d + \alpha^2) - \frac{CM^2R^2}{(1-R)^2} (1+K_s)}
            {(1+K_p) - 2R - \frac{2CM^2R^2}{1-R} + R^2(1+CM)(1+M)(1+K_m + K_d
            + \alpha^2) + (1-j)\left(\frac{CM^2R^2}{({1-R})^2} \right)(1+K_s)}

        .. math::
            P_1 - P_2 = \frac{1}{2}\rho_p \left(\frac{Q_p}{A_n}\right)^2
            \left[(1+K_p) - C(1+K_s) \left(\frac{MR}{1-R}\right)^2\right]

        .. math::
            \text{Pressure ratio} = N = \frac{P_5 - P_2}{P_1 - P_5}

        .. math::
            \text{Volume flow ratio} = M =  \frac{Q_s}{Q_p}

        .. math::
            \text{Jet pump efficiency} = \eta = M\cdot N =
            \frac{Q_s(P_5-P_2)}{Q_p(P_1 - P_5)}

        .. math::
            R = \frac{A_n}{A_m}

        .. math::
            C = \frac{\rho_s}{\rho_p}

        There is no guarantee a solution will be found for the provided variable
        values, but every combination of two missing variables are supported.

        Parameters
        ----------
        rhop : float
            The density of the primary (motive) fluid, [kg/m^3]
        rhos : float
            The density of the secondary fluid (drawn from the vacuum chamber),
            [kg/m^3]
        Kp : float, optional
            The primary nozzle loss coefficient, [-]
        Ks : float, optional
            The secondary inlet loss coefficient, [-]
        Km : float, optional
            The mixing chamber loss coefficient, [-]
        Kd : float, optional
            The diffuser loss coefficient, [-]
        d_nozzle : float, optional
            The inside diameter of the primary fluid's nozle, [m]
        d_mixing : float, optional
            The diameter of the mixing chamber, [m]
        d_diffuser : float, optional
            The diameter of the diffuser at its exit, [m]
        Qp : float, optional
            The volumetric flow rate of the primary fluid, [m^3/s]
        Qs : float, optional
            The volumetric flow rate of the secondary fluid, [m^3/s]
        P1 : float, optional
            The pressure of the primary fluid entering its nozzle, [Pa]
        P2 : float, optional
            The pressure of the secondary fluid at the entry of the ejector, [Pa]
        P5 : float, optional
            The pressure at the exit of the diffuser, [Pa]
        nozzle_retracted : bool, optional
            Whether or not the primary nozzle's exit is before the mixing chamber,
            or somewhat inside it, [-]
        max_variations : int, optional
            When the initial guesses do not lead to a converged solution, try this
            many more guesses at converging the problem, [-]

        Returns
        -------
        solution : dict
            Dictionary of calculated parameters, [-]

        Notes
        -----

        The assumptions of the model are:

        * The flows are one dimensional except in the mixing chamber.
        * The mixing chamber has constant cross-sectional area.
        * The mixing happens entirely in the mixing chamber, prior to entry into
          the diffuser.
        * The primary nozzle is in a straight line with the middle of the mixing
          chamber.
        * Both fluids are incompressible, and have no excess volume on mixing.
        * Primary and secondary flows both enter the mixing throat with their
          own uniform velocity distribution; the mixed stream leaves with a uniform
          velocity profile.
        * If the secondary fluid is a gas, it undergoes isothermal compression in
          the throat and diffuser.
        * If the secondary fluid is a gas or contains a bubbly gas, it is
          homogeneously distributed in a continuous liquid phase.
        * Heat transfer between the fluids is negligible - there is no change in
          density due to temperature changes
        * The change in the solubility of a dissolved gas, if there is one, is
          negigibly changed by temperature or pressure changes.

        The model can be derived from the equations in
        :py:func:`~.liquid_jet_pump_ancillary` and the following:

        Conservation of energy at the primary nozzle, secondary inlet, and diffuser exit:
        .. math::
            P_1 = P_3 + \frac{1}{2}\rho_p V_n^2 + K_p\left(\frac{1}{2}\rho_p V_n^2\right)

        .. math::
            P_2 = P_3 + \frac{1}{2}\rho_s V_3^2 + K_s\left(\frac{1}{2}\rho_s V_3^2\right)

        .. math::
            P_5 = P_4 + \frac{1}{2}\rho_d V_4^2 - K_d\left(\frac{1}{2}\rho_d V_4^2\right)

        Continuity of the ejector:

        .. math::
            \rho_p Q_p + \rho_s Q_s = \rho_d Q_d

        Examples
        --------
        >>> ans = liquid_jet_pump(rhop=998., rhos=1098., Km=.186, Kd=0.12, Ks=0.11,
        ... Kp=0.04, d_mixing=0.045, Qs=0.01, Qp=.01, P2=133600,
        ... P5=200E3, nozzle_retracted=False, max_variations=10000)
        >>> s = []
        >>> for key, value in ans.items():
        ...     s.append('%s: %g' %(key, value))
        >>> sorted(s)
        ['M: 1', 'N: 0.293473', 'P1: 426256', 'P2: 133600', 'P5: 200000', 'Qp: 0.01', 'Qs: 0.01', 'R: 0.247404', 'alpha: 1e-06', 'd_diffuser: 45', 'd_mixing: 0.045', 'd_nozzle: 0.0223829', 'efficiency: 0.293473']

        References
        ----------
        .. [1] Karassik, Igor J., Joseph P. Messina, Paul Cooper, and Charles C.
           Heald. Pump Handbook. 4th edition. New York: McGraw-Hill Education, 2007.
        .. [2] Winoto S. H., Li H., and Shah D. A. "Efficiency of Jet Pumps."
           Journal of Hydraulic Engineering 126, no. 2 (February 1, 2000): 150-56.
           https://doi.org/10.1061/(ASCE)0733-9429(2000)126:2(150).
        .. [3] Elmore, Emily, Khalid Al-Mutairi, Bilal Hussain, and A. Sheriff
           El-Gizawy. "Development of Analytical Model for Predicting Dual-Phase
           Ejector Performance," November 11, 2016, V007T09A013.
        .. [4] Ejectors and Jet Pumps. Design and Performance for Incompressible
           Liquid Flow. 85032. ESDU International PLC, 1985.
        '''
        solution_vars = ['d_nozzle', 'd_mixing', 'Qp', 'Qs', 'P1', 'P2', 'P5']
        unknown_vars = []
        for i in solution_vars:
             if locals()[i] is None:
                 unknown_vars.append(i)

        if len(unknown_vars) > 2:
            raise Exception('Too many unknowns')
        elif len(unknown_vars) < 2:
            raise Exception('Overspecified')

        vals = {'d_nozzle': d_nozzle, 'd_mixing': d_mixing, 'Qp': Qp,
                'Qs': Qs, 'P1': P1, 'P2': P2, 'P5': P5}
        var_guesses = []
        # Initial guess algorithms for each variable here
        # No clever algorithms invented yet
        for v in unknown_vars:
            if v == 'd_nozzle':
                try:
                    var_guesses.append(d_mixing*0.4)
                except:
                    var_guesses.append(0.01)
            if v == 'd_mixing':
                try:
                    var_guesses.append(d_nozzle*2)
                except:
                    var_guesses.append(0.02)
            elif v == 'P1':
                try:
                    var_guesses.append(P2*5)
                except:
                    var_guesses.append(P5*5)
            elif v == 'P2':
                try:
                    var_guesses.append((P1 + P5)*0.5)
                except:
                    try:
                        var_guesses.append(P1/1.1)
                    except:
                        var_guesses.append(P5*1.25)
            elif v == 'P5':
                try:
                    var_guesses.append(P1*1.12)
                except:
                    var_guesses.append(P2*1.12)
            elif v == 'Qp':
                try:
                    var_guesses.append(Qs*1.04)
                except:
                    var_guesses.append(0.01)
            elif v == 'Qs':
                try:
                    var_guesses.append(Qp*0.5)
                except:
                    var_guesses.append(0.01)

        C = rhos/rhop
        if nozzle_retracted:
            j = 0.0
        else:
            j = 1.0
        # The diffuser diameter, if not specified, is set to a very large diameter
        # so as to not alter the results
        if d_diffuser is None:
            if d_mixing is not None:
                d_diffuser = d_mixing*1E3
            elif d_nozzle is not None:
                d_diffuser = d_nozzle*1E3
            else:
                d_diffuser = 1000.0
        vals['d_diffuser'] = d_diffuser

        def obj_err(val):
            # Use the dictionary `vals` to keep track of the currently iterating
            # variables
            for i, v in zip(unknown_vars, val):
                vals[i] = abs(float(v))

            # Keep the pressure limits sane
    #        if 'P1' in unknown_vars:
    #            if 'P5' not in unknown_vars:
    #                vals['P1'] = max(vals['P1'], 1.001*vals['P5'])
    #            elif 'P2' not in unknown_vars:
    #                vals['P1'] = max(vals['P1'], 1.001*vals['P2'])
    #        if 'P2' in unknown_vars:
    #            if 'P1' not in unknown_vars:
    #                vals['P2'] = min(vals['P2'], 0.999*vals['P1'])
    #            if 'P5' not in unknown_vars:
    #                vals['P2'] = max(vals['P2'], 1.001*vals['P2'])

            # Prelimary numbers
            A_nozzle = pi/4*vals['d_nozzle']**2
            alpha = vals['d_mixing']**2/d_diffuser**2
            R = vals['d_nozzle']**2/vals['d_mixing']**2
            M = vals['Qs']/vals['Qp']

            err1 = liquid_jet_pump_pressure_ratio(rhop=rhop, rhos=rhos, Km=Km, Kd=Kd,
                                                  Ks=Ks, Kp=Kp, d_nozzle=vals['d_nozzle'],
                                                  d_mixing=vals['d_mixing'],
                                                  Qs=vals['Qs'], Qp=vals['Qp'],
                                                  P2=vals['P2'], P1=vals['P1'],
                                                  P5=vals['P5'],
                                                  nozzle_retracted=nozzle_retracted,
                                                  d_diffuser=d_diffuser)

            rhs = rhop/2.0*(vals['Qp']/A_nozzle)**2*((1.0 + Kp) - C*(1.0 + Ks)*((M*R)/(1.0 - R))**2 )

            err2 = rhs  - (vals['P1'] - vals['P2'])

            vals['N'] = N = (vals['P5'] - vals['P2'])/(vals['P1']-vals['P5'])
            vals['M'] = M
            vals['R'] = R
            vals['alpha'] = alpha
            vals['efficiency'] = M*N

            if vals['efficiency'] < 0:
                if err1 < 0:
                    err1 -= abs(vals['efficiency'])
                else:
                    err1 += abs(vals['efficiency'])
                if err2 < 0:
                    err2 -= abs(vals['efficiency'])
                else:
                    err2 += abs(vals['efficiency'])

    #        elif vals['N'] < 0:
    #            err1, err2 =  abs(vals['N']) + err1,  abs(vals['N']) + err2
    #        print(err1, err2)
            return err1, err2

        # Only one unknown var
        if 'P5' in unknown_vars:
            ancillary = liquid_jet_pump_ancillary(rhop=rhop, rhos=rhos, Kp=Kp,
                                                  Ks=Ks, d_nozzle=d_nozzle,
                                                  d_mixing=d_mixing, Qp=Qp, Qs=Qs,
                                                  P1=P1, P2=P2)
            if unknown_vars[0] == 'P5':
                vals[unknown_vars[1]] = ancillary
            else:
                vals[unknown_vars[0]] = ancillary

            vals['P5'] = liquid_jet_pump_pressure_ratio(rhop=rhop, rhos=rhos, Km=Km, Kd=Kd, Ks=Ks, Kp=Kp, d_nozzle=vals['d_nozzle'],
                                   d_mixing=vals['d_mixing'], Qs=vals['Qs'], Qp=vals['Qp'], P2=vals['P2'],
                                   P1=vals['P1'], P5=None,
                                   nozzle_retracted=nozzle_retracted, d_diffuser=d_diffuser)['P5']
            # Compute the remaining parameters
            obj_err([vals[unknown_vars[0]], vals[unknown_vars[1]]])
            return vals

        from scipy.optimize import fsolve, root

        def solve_with_fsolve(var_guesses):
            res = fsolve(obj_err, var_guesses, full_output=True)
            if sum(abs(res[1]['fvec'])) > 1E-7:
                raise ValueError('Could not solve')

            for u, v in zip(unknown_vars, res[0].tolist()):
                vals[u] = abs(v)
            return vals

        try:
            return solve_with_fsolve(var_guesses)
        except:
            pass

        # Tying different guesses with fsolve is faster than trying different solvers
        for meth in ['hybr', 'lm', 'broyden1', 'broyden2']: #
            try:
                res = root(obj_err, var_guesses, method=meth, tol=1E-9)
                if sum(abs(res['fun'])) > 1E-7:
                    raise ValueError('Could not solve')

                for u, v in zip(unknown_vars, res['x'].tolist()):
                    vals[u] = abs(v)
                return vals
            except (ValueError, OverflowError):
                continue

        # Just do variations on this until it works
        for _ in range(int(max_variations/8)):
            for idx in [0, 1]:
                for r in [(1, 10), (0.1, 1)]:
                    i = uniform(*r)
                    try:
                        l = list(var_guesses)
                        l[idx] = l[idx]*i
                        return solve_with_fsolve(l)
                    except:
                        pass
        # Vary both parameters at once
        for _ in range(int(max_variations/8)):
            for r in [(1, 10), (0.1, 1)]:
                i = uniform(*r)
                for s in [(1, 10), (0.1, 1)]:
                    j = uniform(*s)
                    try:
                        l = list(var_guesses)
                        l[0] = l[0]*i
                        l[1] = l[1]*j
                        return solve_with_fsolve(l)
                    except:
                        pass
        raise ValueError('Could not solve')