05_CreatePatches.py

'''
This script creates the patches for the test dataset using the trained models.
'''
import numpy as np
import pandas as pd
import subprocess
import os
import argparse
import shutil
import time

from models.LinearRegression import LinearRegressionModel
from models.RBFInterpolator import RadialBasisRegressionModel
from models.GaussianProcessRegressor import GaussianProcessRegressionModel

from miscellaneous.DataPreparation import save_numpy_to_xml, get_matrix_size

# Degree Elevation - Command Line Input
parser = argparse.ArgumentParser(
        description="Script that reads the Elevation Degrees from CMD"
    )
parser.add_argument("--e", required=False, type=int)
parser.add_argument("--r", required=True, type=int)
parser.add_argument("--n", required=False, type=int)
parser.add_argument("--R_lower", required=False, type=int)
parser.add_argument("--R_upper", required=False, type=int)
parser.add_argument("--R_step", required=False, type=int)
parser.add_argument("--Architecture", required=False, type=str)
args = parser.parse_args()
if args.e is not None:
    DEGREE_ELEVATIONS = args.r
else:
    DEGREE_ELEVATIONS = 1
if args.r is not None:
    H_REFINEMENTS = args.r
else:
    H_REFINEMENTS = 0
if args.n is not None:
    N_RUNS = args.n
else:
    N_RUNS = 1
print('DEGREE_ELEVATIONS:', DEGREE_ELEVATIONS)
print('H_REFINEMENTS:', H_REFINEMENTS)
DataFolder = f'Data/H_REFINEMENTS_{H_REFINEMENTS}'
PatchesFolder = f'{DataFolder}/Patches'
if not os.path.exists(PatchesFolder):
    os.mkdir(PatchesFolder)

# Sizes of the microstructures
VELOCITY_SIZE = get_matrix_size(f'{DataFolder}/train/velocities/velocity_field_0.xml')
PRESSURE_SIZE = get_matrix_size(f'{DataFolder}/train/pressure/pressure_field_0.xml')

path_to_executable = 'gismo/stokes_recreation'
if args.Architecture is not None:
    if args.Architecture == 'MacOS':
        path_to_executable = 'gismo/stokes_recreation'
    elif args.Architecture == 'Linux':
        path_to_executable =  'gismo/stokes_recreation_linux'
    elif args.Architecture == 'Windows':
        path_to_executable =  'gismo/stokes_recreation_linux'
    else:
        raise ValueError('Architecture not recognized')
    
# Choose the models to create the patches
models = {
    'LinearRegression': LinearRegressionModel(),
    # 'GP_DotWhite': GaussianProcessRegressionModel(),
    # 'GP_RBF': GaussianProcessRegressionModel(),
    # 'GP_Matern': GaussianProcessRegressionModel(),
    'GP_RationalQuadratic': GaussianProcessRegressionModel(),
    # 'GP_ExpSineSquared': GaussianProcessRegressionModel(),
    'RBF_Linear': RadialBasisRegressionModel(),
    'RBF_thinplatespline': RadialBasisRegressionModel(),
    'RBF_cubic': RadialBasisRegressionModel(),
    # 'RBF_quintic': RadialBasisRegressionModel(),
    # 'RBF_multiquadric': RadialBasisRegressionModel(),
    # 'RBF_inversemultiquadric': RadialBasisRegressionModel(),
    # 'RBF_inversequadratic': RadialBasisRegressionModel(),
    # 'RBF_gaussian': RadialBasisRegressionModel(),
}

df = pd.read_excel(f'{DataFolder}/test/parameter_input.xlsx').transpose()
SAMPLES_TEST = df.shape[1]

for model_name in models:
    print(f'Model: \t\t\t{model_name}')
    ModelFolder = f'{DataFolder}/TrainedModels/{model_name}'
    # Create temp folder
    if not os.path.exists(f'temp_{model_name}'):
        os.mkdir(f'temp_{model_name}')

    # Read POD modes
    if args.R_lower is not None and args.R_upper is not None and args.R_step is not None:
        R = np.linspace(args.R_lower, args.R_upper, args.R_step, dtype=int)
    else:
        R = np.linspace(1, 21, 21, dtype=int)
    print(f'R: {R}')

    timing_vec = []
    # Iterate over POD models
    for R_ in R:
        timing_vec_runs = []
        # Iterate for different runtimes
        for n_run in range(N_RUNS):
            print(f'Run: {n_run+1}/{N_RUNS}')
            print(f'Singular value: \t{R_}')
            # Create folder for the specific model
            if not os.path.exists(f'{PatchesFolder}/{model_name}_{int(R_)}'):
                os.mkdir(f'{PatchesFolder}/{model_name}_{int(R_)}')
            model = models[model_name]

            # Load model
            model_velocity, model_pressure = model.load_model(f'{ModelFolder}/velocity_{int(R_)}.pkl', f'{ModelFolder}/pressure_{int(R_)}.pkl')

            # Predict values
            time_start = time.time()
            test_velocity_predict, test_pressure_predict = model.predict_test(model_velocity, model_pressure, rescale=True)
            time_end = time.time()
            timing_vec_runs.append(time_end - time_start)

            # Iterate over each testcase
            for sample in range(SAMPLES_TEST):
                save_numpy_to_xml(test_velocity_predict[:,sample], f'temp_{model_name}/predict_velocity_field_{sample}.xml', rows=VELOCITY_SIZE, cols='1')

                save_numpy_to_xml(test_pressure_predict[:,sample], f'temp_{model_name}/predict_pressure_field_{sample}.xml', rows=PRESSURE_SIZE, cols='1')

                # Get all the input files for the evaluation
                geometry_file = f'{DataFolder}/test/microstructure/index_{sample}.xml'
                velocity_solution = f'{DataFolder}/test/velocities/velocity_field_{sample}.xml'
                pressure_solution = f'{DataFolder}/test/pressure/pressure_field_{sample}.xml'
                velocity_predicted = f'temp_{model_name}/predict_velocity_field_{sample}.xml'
                pressure_predicted = f'temp_{model_name}/predict_pressure_field_{sample}.xml'

                # Create patches for the geometry with gismo script
                program_output = subprocess.Popen(
                    [f'./{path_to_executable}', f'-f {geometry_file}', f'-v {velocity_solution}', f'-p {pressure_solution}', f'-w {velocity_predicted}', f'-q {pressure_predicted}', f'-e {DEGREE_ELEVATIONS}', f'-r {H_REFINEMENTS}', '--no-plot'],
                    stdout=subprocess.PIPE,
                    text=True
                ).communicate()[0]

                # move the output files to the results folder
                # field patches
                os.system(f'mv pressure_field_patches.xml {PatchesFolder}/{model_name}_{int(R_)}/pressure_field_patches_{sample}.xml')
                os.system(f'mv velocity_field_patches.xml {PatchesFolder}/{model_name}_{int(R_)}/velocity_field_patches_{sample}.xml')
                # field reconstructed patches
                os.system(f'mv pressure_field_rec_patches.xml {PatchesFolder}/{model_name}_{int(R_)}/pressure_field_rec_patches_{sample}.xml')
                os.system(f'mv velocity_field_rec_patches.xml {PatchesFolder}/{model_name}_{int(R_)}/velocity_field_rec_patches_{sample}.xml')

            # timing_vec.append(timing_vec_runs)
            np.savez(f'{PatchesFolder}/{model_name}_{int(R_)}/timing.npz', timing_vec=timing_vec_runs)
    
    # remove temp folder
    shutil.rmtree(f'temp_{model_name}')