moved script with scipy minimize script for x0

import sympy as sy

import pandas as pd

import os

# import scipy as sp

import numpy as np

import math

import itertools

import random

import matplotlib.pyplot as plt

import copy

# local contex

# from numerical_labelling.labelling import *

# constants for pressure loss, membrane properties, pump curves, etc.

data_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")

pump_curve_path = os.path.join(data_path, "ro_pump_curve_data.csv")

pump_curve = pd.read_csv(pump_curve_path)

pump_curve.columns = ["flow_rate_gpm", "pressure_psi"]

# unit conversion

pump_curve["flow_rate_m3_s"] = pump_curve["flow_rate_gpm"] * 6.30902e-5  # Convert GPM to m³/s

pump_curve["pressure_pa"] = pump_curve["pressure_psi"] * 6894.76  # Convert PSI to Pa"]

pump_curve["pressure_Mpa"] = pump_curve["pressure_pa"]*1e-6  # Convert Pa to MPa

m_pump, n_pump, b_pump = np.polyfit(pump_curve['flow_rate_m3_s'], pump_curve['pressure_pa'], 2)

x_poly = np.linspace(min(pump_curve['flow_rate_m3_s']), max(pump_curve['flow_rate_m3_s']), 100)

y_poly = m_pump * x_poly**2 + n_pump * x_poly + b_pump

def add_constraint(constr_dict, eqn):

    """

    add constraint to constraint dict

    """

    indx = len(constr_dict)

    constr_dict[indx] = eqn

    return constr_dict

def solve_constr(constr, var:str, subs = {}):

    """

    solve an equation for a single variable

    constr: the constraint equation (sympy equation)

    var: the variable to be solved for

    subs: dictionary of numeric values to be substituted for symbolic

    """

    if var in subs.keys():

        subs_new = copy.deepcopy(subs)

        subs_new.pop(var)

    else:

        subs_new = subs

    constr_sub = constr.subs(subs_new)

    eq = sy.Eq(constr_sub, 0)

    ans = sy.solve(eq, var)

    if len(ans) == 1: 

        return ans[0]

    else: 

        raise ValueError(f"multiple answers detected: {ans}")

# def constr_mem(M_in, M_out, x_in, x_out, p_in, p_out, simple=False):

#     """

#     M_in: mass flow rate in

#     M_out: mass flow rate out (permeate)

#     x_in: mass fraction salt in

#     x_out: mass fraction salt out (permeate)

#     p_in: pressure at membrane feed

#     p_out: pressure at membrane exit (permeate)

#     returns: the two membrane constraints 

#     """

#     c_in = get_concentration(x_in, simple=simple)

#     c_out = get_concentration(x_out, simple=simple)

#     pi_in = c_in*R*T

#     pi_out = c_out*R*T

#     dPi = pi_in - pi_out

#     dP = p_in - p_out

#     constr_1 = M_in - c_perm*(dP - dPi)

#     constr_2 = x_in*M_in - c_conc*(c_in-c_out)

#     return constr_1, constr_2

def constr_mem(m_in, m_out, x_in, x_out, p_in, p_out, simple=False):

    # m_in, m2, m3, x1, x2, x3, p1, p2, p3 = sy.symbols("m1 m2 m3 x1 x2 x3 p1 p2 p3")

    A, B, C, D = sy.symbols("A, B, C, D")

    varb_symbs = [m1, m2, m3, x1, x2, x3, p1, p2, p3, A, B, C, D]

    labels = ["m1", "m2", "m3", "x1", "x2", "x3", "p1", "p2", "p3", "A", "B", "C", "D"]

    varbs = dict(zip(labels, varb_symbs))

    constr = {}

    constr[0] = m1 - m2  - m3

    constr[1] = x1*m1 -x2*m2 -x3*m3

    constr[2] = m2 - A*((p1 - p2) - C*(x1-x2))

    constr[3] = m2*x2 - B*C*(x1-x2) # x2 may be very small, but for integer example maybe it's better to keep

    constr[4] = p1 - p3 - D*(m2+m3)**2

def constr_pump(m_in, x_in, m_out, p_in, p_out, simple=False, get_dp = False):

    # pump1 pressure increase (streams 1 and 2)

    # need volume flow rate

    rho_pump_1 = get_density(x_in, simple=simple)

    # rho_pump_1 = rho_w 

    vf = m_in/rho_pump_1

    pump_dp_1 = (p_out - p_in)

    poly_fit_1 = m_pump*vf**2 + n_pump*vf + b_pump

    constr = pump_dp_1 - poly_fit_1

    return constr

def constr_fitting(m_in, m_out, x_in, p_in, p_out, K, simple=False):

    # friction loss over pipe fitting

    mavg = ((m_in + m_out)/2)

    rho = get_density(x_in)

    constr = (p_in - p_out) - K*(mavg**2/2*rho**2)

    return constr

def constr_fitting_mem(m_in, m_out, x_in, x_out, p_in, p_out, K_mem, simple=False):

    """

    need to take average

    """

    # friction loss over fitting

    mavg = ((m_in + m_out)/2)

    if simple:

        rhoavg = get_density(x_in, simple=simple)

    else:

        rho_in = get_density(x_in, simple=simple)

        rho_out = get_density(x_out, simple=simple)

        rhoavg = ((rho_in + rho_out)/2)

    constr = (p_in - p_out) - K_mem*(mavg**2/2*rhoavg**2)

    return constr

# generate constraints

def constr_mem_nl(simple=False, subf = False):

    """

    constraint equations for membrane example

    """

    m1, m2, m3, m4, m5, m6, m7, m8 = sy.symbols('m1 m2 m3 m4 m5 m6 m7 m8') # overall mass flow rates

    f1, f2, f3, f4, f5, f6, f7, f8 = sy.symbols('f1 f2 f3 f4 f5 f6 f7 f8') # component flow rates

    x1, x2, x3, x4, x5, x6, x7, x8 = sy.symbols('x1 x2 x3 x4 x5 x6 x7 x8') # component mass fractions

    p1, p2, p3, p4, p5, p6, p7, p8 = sy.symbols('p1 p2 p3 p4 p5 p6 p7 p8') # pressures

    varbs = [

        m1, m2, m3, m4, m5, m6, m7, m8,

        x1, x2, x3, x4, x5, x6, x7, x8,

        f1, f2, f3, f4, f5, f6, f7, f8,

        p1, p2, p3, p4, p5, p6, p7, p8,

        ]

    var_labels = [

        'M1', 'M2', 'M3', 'M4', 'M5', 'M6', 'M7', 'M8',

        'X1', 'X2', 'X3', 'X4', 'X5', 'X6', 'X7', 'X8',

        'F1', 'F2', 'F3', 'F4', 'F5', 'F6', 'F7', 'F8',

        'P1', 'P2', 'P3', 'P4', 'P5', 'P6', 'P7', 'P8',

        ]

    # Create a dictionary to store all variables and their values

    var_dict = dict(zip(var_labels, varbs))

    constr = {}

    # mass balance

    constr = add_constraint(constr, m1 - m2)

    constr = add_constraint(constr, m2 - m3 + m8)

    constr = add_constraint(constr, m3 - m4 - m5)

    constr = add_constraint(constr, m5 - m6 - m7)

    constr = add_constraint(constr, m7 - m8)

    # component balance

    constr = add_constraint(constr, f1 - f2)

    constr = add_constraint(constr, f2 - f3 + f8)

    constr = add_constraint(constr, f3 - f4 - f5)

    constr = add_constraint(constr, f5 - f6 - f7)

    constr = add_constraint(constr, f7 - f8)

    # solve for f

    constr = add_constraint(constr, f1 - m1 * x1)

    constr = add_constraint(constr, f2 - m2 * x2)

    constr = add_constraint(constr, f3 - m3 * x3)

    constr = add_constraint(constr, f4 - m4 * x4)

    constr = add_constraint(constr, f5 - m5 * x5)

    constr = add_constraint(constr, f6 - m6 * x6)

    constr = add_constraint(constr, f7 - m7 * x7)

    constr = add_constraint(constr, f8 - m8 * x8)

    # equal concentration

    constr = add_constraint(constr, x5 - x6)

    constr = add_constraint(constr, x5 - x7)

    mem_constrs = constr_mem(

        M_in = m3, M_out = m4, 

        x_in = x3, x_out = x4, 

        p_in = p3, p_out = p4,

        simple = simple,

        )

    for constr_eq in mem_constrs: 

        constr = add_constraint(constr, constr_eq)

    # pump1 constraints (s1 in, s2 out)

    pump_constr_1 = constr_pump(

        m_in=m1, x_in=x1, 

        m_out=m2, 

        p_in=p1, p_out=p2, 

        simple=simple

    )

    constr = add_constraint(constr, pump_constr_1)

    # pump2 constraints (s7 in, s8 out)

    pump_constr_2 = constr_pump(

        m_in=m7, x_in=x7, 

        m_out=m8, 

        p_in=p7, p_out=p8, 

        simple=simple

    )

    constr = add_constraint(constr, pump_constr_2)

    # friction loss over membrane (streams 3-5)

    mem_fric_constr = constr_fitting_mem(

        m_in=m3, m_out=m5, 

        x_in=m3, x_out=m5, 

        p_in=p3, p_out=p5, 

        K_mem=K_mem,

        simple=simple,

        )

    constr = add_constraint(constr, mem_fric_constr)

    # remaining friction losses are equivalent (add friction loss coeff later)

    constr = add_constraint(constr, p2 - p3)

    constr = add_constraint(constr, p8 - p2)

    constr = add_constraint(constr, p5 - p6)

    constr = add_constraint(constr, p5 - p7)

    if subf:

        f_dict = {

            "f1":x1*m1, 

            "f2":x2*m2,

            "f3":x3*m3,

            "f4":x4*m4,

            "f5":x5*m5,

            "f6":x6*m6,

            "f7":x7*m7,

            "f8":x8*m8,

        }

        constr_subf = constr.copy()

        for i, iconstr in constr.items():

            constr[i] = iconstr.subs(f_dict) 

        varbs = [

            m1, m2, m3, m4, m5, m6, m7, m8,

            x1, x2, x3, x4, x5, x6, x7, x8,

            p1, p2, p3, p4, p5, p6, p7, p8,

            ]

        var_labels = [

            'M1', 'M2', 'M3', 'M4', 'M5', 'M6', 'M7', 'M8',

            'X1', 'X2', 'X3', 'X4', 'X5', 'X6', 'X7', 'X8',

            'P1', 'P2', 'P3', 'P4', 'P5', 'P6', 'P7', 'P8',

            ]

    # Create a dictionary to store all variables and their values

    var_dict = dict(zip(var_labels, varbs))

    return constr, var_dict

def get_x0_nl(m1 = 0.2, x1 = 0.2, p1 = 1e5, simple=False):

    # nonlinear operating point: 

    # for failure case make recycle really large, other streams very small

    constr, _ = constr_mem_nl(simple=simple)

    M1 = m1

    M8 = 1.1

    M2 = M1

    M3 = M2 + M8

    M4 = 0.1

    M5 = M3 - M4

    M6 = 0.1

    M7 = M8

    X1 = x1

    X2 = x1

    X3 = 0.01

    X4 = 0.01

    X5 = 0.01

    X6 = 0.01

    X7 = 0.01

    X8 = 0.01

    F1 = M1*X1

    F2 = M2*X2

    F3 = M3*X3

    F4 = M4*X4

    F5 = M5*X5

    F6 = M6*X6

    F7 = M7*X7

    F8 = M8*X8

    x0_nl = {

        'm1': M1, 

        'm2': M2, 

        'm3': M3, 

        'm4': M4, 

        'm5': M5, 

        'm6': M6, 

        'm7': M7, 

        'm8': M8,

        'f1': F1, 

        'f2': F2, 

        'f3': F3, 

        'f4': F4, 

        'f5': F5, 

        'f6': F6, 

        'f7': F7, 

        'f8': F8,

        'x1': X1, 

        'x2': X2, 

        'x3': X3, 

        'x4': X4, 

        'x5': X5, 

        'x6': X6, 

        'x7': X7,

        'x8': X8, 

    }

    # pressure constraints

    P1 = p1

    x0_nl["p1"] = P1

    # NOTE: constraint numbering is hardcoded. picking which constraint 

    # corresponds to which is not addressed for the operating point

    # pump 1 constraint eq

    P2 = solve_constr(constr = constr[22], var="p2", subs=x0_nl )

    P2 = float(P2)

    x0_nl["p2"] = P2

    P3 = P2

    x0_nl["p3"] = P3

    P4 = solve_constr(constr = constr[20], var="p4", subs=x0_nl)

    P4 = float(P4)

    x0_nl["p4"] = P4

    P5 = solve_constr(constr = constr[24], var="p5", subs=x0_nl)

    x0_nl["p5"] = P5

    P6 = P5

    x0_nl["p6"] = P6

    P7 = P5

    x0_nl["p7"] = P7

    P8 = solve_constr(constr = constr[23], var="p8", subs=x0_nl)

    x0_nl["p8"] = P8

    # # resolve the mass flow out, concentration out

    # x0_nl_x4 = copy.deepcopy(x0_nl)

    # x0_nl_x4.pop("x4")

    M4 = solve_constr(constr[21], var="m4", subs=x0_nl)

    # # resolve the mass flow out, concentration out

    # x0_nl_x4 = copy.deepcopy(x0_nl)

    # x0_nl_x4.pop("x4")

    X4 = solve_constr(constr[21], var="x4", subs=x0_nl)

    return x0_nl