cool_H2O_polyexpress.py

import numpy as np
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score
import re
import sympy as sp

sizea = 29
sizeb = 21
sizec = 21
# Initialize the 3D array and 1D array
inputs = np.zeros((sizea, sizeb, sizec, 3))
outco = np.zeros(sizea * sizeb * sizec)

# Reading and processing the file
with open('datafiles/coolH2O.dat', "r") as f:
    for row in f:
        srow = row.strip()
        # Skip comments and blank lines
        if srow == "" or srow.startswith("#") or srow.startswith("!"):
            continue
        # Split the row by whitespace and convert to floats
        arow = srow.split()
        arow = [float(x) for x in arow]

        # Extract indices and values
        i, j, k = int(arow[0]) - 1, int(arow[1]) - 1, int(arow[2]) - 1
        inputs[i, j, k] = [arow[3], arow[4], arow[5]]
        outco[i * sizeb * sizec + j * sizec + k] = arow[6]

# Flatten the input arrays and the output array
x_flat = inputs[:, :, :, 0].flatten()
y_flat = inputs[:, :, :, 1].flatten()
z_flat = inputs[:, :, :, 2].flatten()
outco_flat = outco.flatten()

# Stack the flattened input arrays
X = np.vstack([x_flat, y_flat, z_flat]).T

# Fit a polynomial of degree NN
# Tried a range from 1 - 20: best results are for degree 10
# even the best results over/under estimate the CO cooling rate in the worst case by 1.2 and 1.5 orders of magnitude
NN = np.arange(9, 10, 1)

for i in range(0, len(NN)):
    print('Fitting polynomial of degree: ', NN[i], ' to H2O cooling table')
    poly = PolynomialFeatures(degree=NN[i])
    X_poly = poly.fit_transform(X)

    # Perform linear regression
    clf = LinearRegression()
    clf.fit(X_poly, outco_flat)

    # Get the coefficients and intercept
    coefficients = clf.coef_
    intercept = clf.intercept_

    #print("Coefficients:", coefficients)
    #print("Intercept:", intercept)
    gof = r2_score(outco_flat, clf.predict(X_poly))
    print("r2 (goodness of fit) score: ", gof)

    # Create a human-readable form of the polynomial equation
    terms = poly.get_feature_names_out(['x', 'y', 'z'])
    polynomial_expression = f"{intercept:.6f} "
    for coef, term in zip(coefficients, terms):
        polynomial_expression += f"+ ({coef:.13f} * {term}) "

    #print(polynomial_expression)

    def interpolatorr(v):
        x, y, z = v
        v_poly = poly.transform([[x, y, z]])
        return clf.predict(v_poly)[0]


    # Example usage
    #v = [0.477, -2, -18]
    #print("Interpolated value:", interpolatorr(v))
    #v = [4, 14, 25]
    #print("Interpolated value:", interpolatorr(v))

    #Debug info:
    #'''
    if gof > 0.95:
        bb = np.zeros(len(X))
        for j in range(0, len(X)):
            aa = interpolatorr(X[j])
            bb[j] = aa - outco[j]

        print('Min max mean of interpolated - actual values: ', np.min(bb), np.max(bb), np.mean(bb))
        print('5 50 95 percentiles of interpolated - actual values: ', np.percentile(bb, 5), np.percentile(bb, 50), np.percentile(bb, 95))
        print('-----------------')
    #'''


#function below will format the polynomial_expression obtained above to python
def add_multiplication_operators(s):
    # This regex matches sequences of variables (with optional exponents) separated by spaces
    pattern = re.compile(r'([a-zA-Z](?:\^\d+)?)(\s+([a-zA-Z](?:\^\d+)?))+')
    
    def replace_func(match):
        # Extract all variable parts including optional exponents
        parts = match.group().split()
        # Join them with ' * '
        return ' * '.join(parts)
    
    return pattern.sub(replace_func, s)

polynomial_expression = add_multiplication_operators(polynomial_expression)
#now replace all "^" by "**"
polynomial_expression = polynomial_expression.replace("^", "**")

#get the min/max values to pass to sympychemratecollection
coolH2Ox1min, coolH2Ox2min, coolH2Ox3min = np.min(X, axis=0)
coolH2Ox1max, coolH2Ox2max, coolH2Ox3max = np.max(X, axis=0)

#write this info to a .py file that can be read in the main code:
with open('cool_H2O.py', 'w') as file:
    file.write('coolH2Ox1min = ' + str(coolH2Ox1min) + '\n')
    file.write('coolH2Ox2min = ' + str(coolH2Ox2min) + '\n')
    file.write('coolH2Ox3min = ' + str(coolH2Ox3min) + '\n')
    file.write('coolH2Ox1max = ' + str(coolH2Ox1max) + '\n')
    file.write('coolH2Ox2max = ' + str(coolH2Ox2max) + '\n')
    file.write('coolH2Ox3max = ' + str(coolH2Ox3max) + '\n')
    file.write("polynomial_expression = " + "'" + polynomial_expression + "'" + '\n')


'''
#scipy grid interpolator output:
coolCOx1min = 0.47712125E+000
coolCOx1max = 0.40000000E+001
coolCOx2min = -0.20000000E+001
coolCOx2max = 0.14000000E+002
coolCOx3min = -0.18000000E+002
coolCOx3max = 0.25000000E+002

x = np.linspace(coolCOx1min, coolCOx1max, size)
y = np.linspace(coolCOx2min, coolCOx2max, size)
z = np.linspace(coolCOx3min, coolCOx3max, size)

# Reshape outco to match the 3D grid
outco_reshaped = outco.reshape((size, size, size))

# Create the interpolator
enterpolatorr = scipy.interpolate.RegularGridInterpolator((x, y, z), outco_reshaped)
'''